python中的双循环优化

Question

我正在尝试优化以下循环：

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计算产品总和可以获得更好的性能：

^{pr2}$

这导致

Time numpy : 4.1595

Time numba1 : 0.6993

Time numba2 : 1.0135

使用numba版本的sum函数（sum\u opt）性能非常好。但是我想知道为什么numba版本的双循环函数（numba2）会导致执行速度变慢。我试图使用jit而不是autojit，指定参数类型，但情况更糟。在

我还注意到，先在最小的循环上循环比先在最大的循环上循环慢。有什么解释吗？在

不管怎样，我确信这个双循环函数可以改进很多向量化问题（比如this）或使用其他方法（map？）但我对这些方法有点困惑。在

在我代码的其他部分，我使用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)`

您的解决方案非常优雅，但我必须在代码中大量使用此函数。因此，对于1000次迭代，这将导致

Time numba1 : 3.2487

Time numba2 : 3.7012

Time numba3 : 3.2088

Time convol : 22.7696

autojit和{}非常接近。 但是，在使用jit时，指定所有参数类型似乎很重要。在

当函数有多个输出时，我不知道是否有方法在jit修饰符中指定参数类型。有人吗？在

现在我没有找到其他的解决办法，除了使用numba。欢迎有新想法！在

Answer 1

你没有充分利用纽比的能力。解决问题的方法如下：

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数组的中心非零部分：

^{pr2}$

对bb和{}也是类似的。在

在我的系统中，使用示例输入，这段代码比numpy函数快300倍以上。根据你的时间安排，这将比numba快一到两个数量级。在

Answer 2

您基本上在那里执行2D卷积，只做了一个小的修改，即内核不会像通常的^{}操作那样反转。 所以，基本上，我们需要做两件事来使用^{}来解决我们的案子-

• 将输入数组rho切片，以选择在原始循环版本代码中使用的部分。这将是卷积的输入数据。在
• 反转内核c，并将其与切片数据一起提供给signal.convolve2d。在

请注意，这些操作将分别用于a和{}的计算。在

这是实现-

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的输出，这将进一步提高性能。在

运行时测试

^{pr2}$

因此，对于实际的输入数据大小，所提出的基于卷积的方法比循环代码快得多，比最快的基于numba的方法numba1快。在

Answer 3

Numba在^{} mode中非常快，但是在您的代码中，它必须回到object模式，这要慢得多。如果将nopython=True传递给jit装饰器，就可以看到这种情况。在

如果您将a和{}作为参数传递，则它确实在nopython模式下编译（至少在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被弃用，取而代之的是{}。在

---

显然你还不满意。那么基于stride_tricks的解决方案怎么样？在

^{pr2}$

此外，由于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