在python中优化Double Loop [英] Optimize Double loop in python
问题描述
我正在尝试优化以下循环:
I am trying to optimize the following loop :
def numpy(nx, nz, c, rho):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum(c*rho[ix-1:ix+3, iz])
b[ix, iz] = sum(c*rho[ix-2:ix+2, iz])
return a, b
我尝试了不同的解决方案，并发现使用numba来计算乘积之和会带来更好的性能:
I tried different solutions and found using numba to calculate the sum of the product leads to better performances:
import numpy as np
import numba as nb
import time
@nb.autojit
def sum_opt(arr1, arr2):
s = arr1[0]*arr2[0]
for i in range(1, len(arr1)):
s+=arr1[i]*arr2[i]
return s
def numba1(nx, nz, c, rho):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
@nb.autojit
def numba2(nx, nz, c, rho):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
nx = 1024
nz = 256
rho = np.random.rand(nx, nz)
c = np.random.rand(4)
a = np.zeros((nx, nz))
b = np.zeros((nx, nz))
ti = time.clock()
a, b = numpy(nx, nz, c, rho)
print 'Time numpy : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
a, b = numba1(nx, nz, c, rho)
print 'Time numba1 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
a, b = numba2(nx, nz, c, rho)
print 'Time numba2 : ' + `round(time.clock() - ti, 4)`
这导致
时间numpy:4.1595
Time numpy : 4.1595
时间numba1:0.6993
Time numba1 : 0.6993
时间numba2:1.0135
Time numba2 : 1.0135
使用sum函数(sum_opt)的numba版本效果很好.但是我想知道为什么双循环函数的numba版本(numba2)会导致执行时间变慢.我尝试使用jit代替autojit来指定参数类型，但情况更糟.
Using the numba version of the sum function (sum_opt) performs very well. But I am wondering why the numba version of the double loop function (numba2) leads to slower execution times. I tried to use jit instead of autojit, specifying the argument types, but it was worse.
我还注意到，在最小的循环上先循环比在最大的循环上先循环慢.有什么解释吗?
I also noticed that looping first on the smallest loop is slower than looping first on the biggest loop. Is there any explanation ?
无论是哪种情况，我都可以肯定可以改善此双循环功能，从而使问题向量化(例如
Whether it is, I am sure this double loop function can be improved a lot vectorizing the problem (like this) or using another method (map ?) but I am a little bit confused about these methods.
在代码的其他部分，我使用numba和numpy切片方法来替换所有显式循环，但是在这种特殊情况下，我不设置它.
In the other parts of my code, I used numba and numpy slicing methods to replace all explicit loops but in this particular case, I don't how to set it up.
有什么想法吗?
编辑
感谢您的所有评论.我在这个问题上做了一些工作:
Thanks for all your comments. I worked a little on this problem:
import numba as nb
import numpy as np
from scipy import signal
import time
@nb.jit(['float64(float64[:], float64[:])'], nopython=True)
def sum_opt(arr1, arr2):
s = arr1[0]*arr2[0]
for i in xrange(1, len(arr1)):
s+=arr1[i]*arr2[i]
return s
@nb.autojit
def numba1(nx, nz, c, rho, a, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
@nb.jit(nopython=True)
def numba2(nx, nz, c, rho, a, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)
def numba3a(nx, nz, c, rho, a):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
return a
@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)
def numba3b(nx, nz, c, rho, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return b
def convol(nx, nz, c, aa, bb):
s1 = rho[1:nx-1,2:nz-3]
s2 = rho[0:nx-2,2:nz-3]
kernel = c[:,None][::-1]
aa[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid')
bb[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid')
return aa, bb
nx = 1024
nz = 256
rho = np.random.rand(nx, nz)
c = np.random.rand(4)
a = np.zeros((nx, nz))
b = np.zeros((nx, nz))
ti = time.clock()
for i in range(1000):
a, b = numba1(nx, nz, c, rho, a, b)
print 'Time numba1 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a, b = numba2(nx, nz, c, rho, a, b)
print 'Time numba2 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a = numba3a(nx, nz, c, rho, a)
b = numba3b(nx, nz, c, rho, b)
print 'Time numba3 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a, b = convol(nx, nz, c, a, b)
print 'Time convol : ' + `round(time.clock() - ti, 4)`
您的解决方案是非常优雅的Divakar，但是我必须在代码中大量使用此功能.因此，对于1000次迭代，这将导致
Your solution is very elegant Divakar, but I have to use this function a large number of time in my code. So, for 1000 iterations, this lead to
时间numba1:3.2487
Time numba1 : 3.2487
时间numba2:3.7012
Time numba2 : 3.7012
时间numba3:3.2088
Time numba3 : 3.2088
时间卷积:22.7696
Time convol : 22.7696
autojit和jit非常接近.
但是，使用jit时，指定所有参数类型似乎很重要.
autojit and jit are very close.
However, when using jit, it seems important to specify all argument types.
当函数具有多个输出时，我不知道是否有一种方法可以在jit装饰器中指定参数类型.有人吗
I do not know if there is a way to specify argument types in the jit decorator when the function has multiple outputs. Someone ?
目前，除了使用numba之外，我没有找到其他解决方案.欢迎新想法！
For now I did not find other solution than using numba. New ideas are welcomed !
推荐答案
Numba is very fast in nopython mode but with your code it has to fall back to object mode, which is a lot slower. You can see this happening if you pass nopython=True to the jit decorator.
如果您将a和b作为参数传递，它将以nopython模式编译(至少在Numba版本0.18.2中):
It does compile in nopython mode (at least in Numba version 0.18.2) if you pass a and b as arguments:
import numba as nb
@nb.jit(nopython=True)
def sum_opt(arr1, arr2):
s = arr1[0]*arr2[0]
for i in range(1, len(arr1)):
s+=arr1[i]*arr2[i]
return s
@nb.jit(nopython=True)
def numba2(nx, nz, c, rho, a, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
请注意，在发行说明中提到过以jit弃用autojit.
Note that in the release notes there is mention of autojit being deprecated in favor of jit.
显然您还不满意.那么基于stride_tricks的解决方案呢?
Apparently you're not satisfied yet. So how about a solution based on stride_tricks?
from numpy.lib.stride_tricks import as_strided
def stridetrick_einsum(c, rho, out):
ws = len(c)
nx, nz = rho.shape
shape = (nx-ws+1, ws, nz)
strides = (rho.strides[0],) + rho.strides
rho_windowed = as_strided(rho, shape, strides)
np.einsum('j,ijk->ik', c, rho_windowed, out=out)
a = np.zeros((nx, nz))
stridetrick_einsum(c, rho[1:-1,2:-3], a[2:-3,2:-3])
b = np.zeros((nx, nz))
stridetrick_einsum(c, rho[0:-2,2:-3], b[2:-3,2:-3])
此外，由于a和b显然几乎完全相同，因此您可以一次性计算它们，然后将值复制到以下位置:
What's more, since a and b are obviously almost exactly the same, you can calculate them in one go and then copy the values over:
a = np.zeros((nx, nz))
stridetrick_einsum(c, rho[:-1,2:-3], a[1:-3,2:-3])
b = np.zeros((nx, nz))
b[2:-3,2:-3] = a[1:-4,2:-3]
a[1,2:-3] = 0.0
这篇关于在python中优化Double Loop的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！
