优化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计算产品的和导致更好的性能：

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)`

这导致

时间点: 4.1595 时代numba1 : 0.6993 时代numba2 : 1.0135

使用和函数(sum_opt)的numba版本执行得非常好。但是我想知道为什么numba版本的双循环函数(numba2)会导致较慢的执行时间。我尝试使用jit而不是autojit，指定参数类型，但情况更糟。

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

不管是这样，我确信这个双循环函数可以得到很大的改进(比如这)，或者使用另一种方法(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)`

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

时代numba1 : 3.2487 时代numba2 : 3.7012 时代numba3 : 3.2088 时间会议: 22.7696

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

当函数有多个输出时，我不知道是否有一种方法可以在jit装饰器中指定参数类型。有人吗？

目前，除了使用numba之外，我没有找到其他解决方案。欢迎新想法！

Answer 1

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

如果将a和b作为参数传递，它确实在a模式下进行编译(至少在NumbaVersion0.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

请注意，在发布说明中，有人提到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

Answer 2

您基本上是在那里执行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更好。

Answer 3

你没有充分利用numpy的能力。处理你的问题的方法应该是：

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倍。根据你的时间表，这将比南巴快一两个数量级。