我正在尝试优化以下循环：
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计算产品总和可以获得更好的性能：
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
时间numba1：0.6993
时间numba2：1.0135
使用sum函数的numba版本（sum_opt）表现非常好。但我想知道为什么双循环函数（numba2）的numba版本会导致执行时间变慢。我尝试使用jit而不是autojit，指定参数类型，但情况更糟。
我还注意到，在最小循环上首先循环比在最大循环上首先循环慢。有没有解释？
无论是否，我确信这个双循环函数可以很好地改进问题（如this）或使用其他方法（map？），但我对这些方法有点困惑。 / p>
在我的代码的其他部分，我使用numba和numpy切片方法来替换所有显式循环，但在这种特殊情况下，我不知道如何设置它。
有什么想法吗？
修改
感谢您的所有评论。我在这个问题上做了一些工作：
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次迭代，这导致
时间numba1：3.2487
时间numba2：3.7012
时间numba3：3.2088
时间缠绕：22.7696
autojit和jit非常接近。
但是，使用jit时，指定所有参数类型似乎很重要。
当函数有多个输出时，我不知道是否有办法在jit装饰器中指定参数类型。有人吗？
目前我没有找到其他解决方案而不是使用numba。欢迎新的想法！
答案 0 :(得分：7)
你基本上在那里进行2D卷积，只需要进行一些小的修改，你的内核就不会像通常的convolution操作那样反转。
所以，基本上，我们需要做两件事来使用signal.convolve2d来解决我们的案例 -
rho以选择在原始循环版本的代码中使用的部分。这将是卷积的输入数据。c并将其与切片数据一起提供给signal.convolve2d。请注意，要分别计算a和b这些内容。
这是实施 -
import numpy as np
from scipy import signal
# Slices for convolutions to get a and b respectively
s1 = rho[1:nx-1,2:nz-3]
s2 = rho[0:nx-2,2:nz-3]
kernel = c[:,None][::-1] # convolution kernel
# Setup output arrays and fill them with convolution results
a = np.zeros((nx, nz))
b = np.zeros((nx, nz))
a[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid')
b[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid')
如果您不需要在输出数组边界附加额外的零，则可以简单地使用signal.convolve2d的输出，这必须进一步提高性能。
运行时测试
In [532]: %timeit loop_based(nx, nz, c, rho)
1 loops, best of 3: 1.52 s per loop
In [533]: %timeit numba1(nx, nz, c, rho)
1 loops, best of 3: 282 ms per loop
In [534]: %timeit numba2(nx, nz, c, rho)
1 loops, best of 3: 509 ms per loop
In [535]: %timeit conv_based(nx, nz, c, rho)
10 loops, best of 3: 15.5 ms per loop
因此，对于实际的输入数据，建议的基于卷积的方法比循环代码快 100x ，而 20x 更好而不是基于numba的最快方法numba1。
答案 1 :(得分：6)
您没有充分利用numpy的功能。 numpythonic 处理问题的方式如下：
cs = np.zeros((nx+1, nz))
np.cumsum(c*rho, axis=0, out=cs[1:])
aa = cs[5:, 2:-3] - cs[1:-4, 2:-3]
bb = cs[4:-1, 2:-3] - cs[:-5, 2:-3]
aa现在将保留a数组的中心非零部分：
>>> a[:5, :5]
array([[ 0. , 0. , 0. , 0. , 0. ],
[ 0. , 0. , 0. , 0. , 0. ],
[ 0. , 0. , 2.31296595, 2.15743042, 2.5853117 ],
[ 0. , 0. , 2.02697233, 2.83191016, 2.58819583],
[ 0. , 0. , 2.4086584 , 2.45175615, 2.19628507]])
>>>aa[:3, :3]
array([[ 2.31296595, 2.15743042, 2.5853117 ],
[ 2.02697233, 2.83191016, 2.58819583],
[ 2.4086584 , 2.45175615, 2.19628507]])
，同样适用于bb和b。
在我的系统上，通过示例输入，此代码的运行速度比numpy函数快300倍。根据你的时间安排，这比numba快一到两个数量级。
答案 2 :(得分：2)
正如连续体博客上的performance section所述，autojit及时编译，而jit提前编译：
Numba可以使用autojit装饰器及时编译 时间与jit装饰者
这意味着在许多情况下，autojit意味着编译器可以对其编译的代码进行更有根据的猜测，然后进行优化。我知道准时编译听起来很矛盾，但是嘿。
但我想知道为什么双循环函数的numba版本 （numba2）导致执行时间变慢
Numba不会增加任意函数调用的性能。虽然我不能肯定地说，我的猜测是JIT编译的开销超过了这样做的好处（如果有任何好处的话）。
我还注意到，在最小的循环上首先循环比慢 首先在最大的循环上循环。有没有解释？
这可能是由于cache miss。二维数组被分配为大小为rows * columns的连续内存块。获取到高速缓存的内容基于时间（最近使用的内容）和空间（内存中与所使用的内容相近）的位置的组合，即被认为接下来使用的内容。
首先迭代行时，按照数据在内存中出现的顺序进行迭代。首先迭代列时，每次都“跳过”内存中行的宽度，使得访问的内存位置不太可能被提取到缓存中。
2D array: [[1,2,3],[4,5,6],[7,8,9]]
In memory: 1 2 3 4 5 6 7 8 9
让我们假设一个过度简化的，愚蠢的缓存提取算法，该算法获取3个后续内存位置。
迭代行优先：
In memory: 1 2 3 | 4 5 6 | 7 8 9
Accessed: 1 2 3 | 4 5 6 | 7 8 9
Cache miss: - - - | - - - | - - -
迭代列优先：
In memory: 1 2 3 | 4 5 6 | 7 8 9
Accessed: 1 4 7 | 2 5 8 | 3 6 9
Cache miss: - - - | x x x | x x x
答案 3 :(得分：1)
Numba在nopython mode中速度非常快，但是使用您的代码时，它必须回退到object模式，这种速度要慢得多。如果您将nopython=True传递给jit装饰器，就会发现这种情况。
如果您将nopython和a作为参数传递，它会在b模式下编译（至少在Numba版本0.18.2中）：
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
请注意，release notes中提及autojit已弃用jit。
显然你还不满意。那么基于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显然几乎完全相同，您可以一次性计算它们，然后复制值：
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