将嵌套循环计算转换为Numpy以加速

时间:2015-10-24 16:15:23

标签: python performance numpy optimization floating-point

我的Python程序的一部分包含以下代码,其中包含一个新网格 是根据旧网格中的数据计算的。

网格是浮动的二维列表。该代码使用三个for循环:

for t in xrange(0, t, step):
    for h in xrange(1, height-1):
        for w in xrange(1, width-1):
            new_gr[h][w] = gr[h][w] + gr[h][w-1] + gr[h-1][w] + t * gr[h+1][w-1]-2 * (gr[h][w-1] + t * gr[h-1][w])
    gr = new_gr

return gr

对于大型网格和大量时间 t ,代码极其缓慢。

我试图通过替换内循环来使用Numpy来加速这段代码 用:

J = np.arange(1, width-1)
new_gr[h][J] = gr[h][J] + gr[h][J-1] ...

但产生的结果(阵列中的浮子)比10%小 他们的列表计算对应物。

  • 使用 np.array(pylist)将浮动列表转换为Numpy浮点数组然后进行计算时,会出现什么样的精度损失?

  • 我应该如何将三重for循环转换为漂亮而快速的Numpy代码? (或者是否有其他建议可以显着加快代码的速度?)

2 个答案:

答案 0 :(得分:4)

如果gr是浮点数列表,那么如果您希望使用NumPy进行向量化,则第一步是将gr转换为带有np.array()的NumPy数组。

接下来,我假设您已new_gr使用形状为(height,width)的零进行初始化。在两个最里面的循环中执行的计算基本上代表2D convolution。因此,您可以将signal.convolve2d与适当的kernel一起使用。为了决定kernel,我们需要查看缩放因子并从中创建3 x 3内核并取消它们以模拟我们在每次迭代时所做的计算。因此,您将获得一个矢量化解决方案,其中两个最内层的循环被移除以获得更好的性能,如此 -

import numpy as np
from scipy import signal

# Get the scaling factors and negate them to get kernel
kernel = -np.array([[0,1-2*t,0],[-1,1,0,],[t,0,0]])

# Initialize output array and run 2D convolution and set values into it
out = np.zeros((height,width))
out[1:-1,1:-1] = signal.convolve2d(gr, kernel, mode='same')[1:-1,:-2]

验证输出和运行时测试

定义功能:

def org_app(gr,t):
    new_gr = np.zeros((height,width))
    for h in xrange(1, height-1):
        for w in xrange(1, width-1):
            new_gr[h][w] = gr[h][w] + gr[h][w-1] + gr[h-1][w] + t * gr[h+1][w-1]-2 * (gr[h][w-1] + t * gr[h-1][w]) 
    return new_gr

def proposed_app(gr,t):
    kernel = -np.array([[0,1-2*t,0],[-1,1,0,],[t,0,0]])
    out = np.zeros((height,width))
    out[1:-1,1:-1] = signal.convolve2d(gr, kernel, mode='same')[1:-1,:-2]
    return out

验证 -

In [244]: # Inputs
     ...: gr = np.random.rand(40,50)
     ...: height,width = gr.shape
     ...: t = 1
     ...: 

In [245]: np.allclose(org_app(gr,t),proposed_app(gr,t))
Out[245]: True

Timings -

In [246]: # Inputs
     ...: gr = np.random.rand(400,500)
     ...: height,width = gr.shape
     ...: t = 1
     ...: 

In [247]: %timeit org_app(gr,t)
1 loops, best of 3: 2.13 s per loop

In [248]: %timeit proposed_app(gr,t)
10 loops, best of 3: 19.4 ms per loop

答案 1 :(得分:2)

@Divakar,我在org_app上尝试了几种变体。完全矢量化的版本是:

def org_app4(gr,t):
    new_gr = np.zeros((height,width))
    I = np.arange(1,height-1)[:,None]
    J = np.arange(1,width-1)
    new_gr[I,J] = gr[I,J] + gr[I,J-1] + gr[I-1,J] + t * gr[I+1,J-1]-2 * (gr[I,J-1] + t * gr[I-1,J])
    return new_gr

虽然proposed_app的速度只有原来的一半,但它与原版的风格相近。因此可以帮助理解嵌套循环如何被矢量化。

重要的一步是将I转换为列数组,以便I,J一起索引一个值块。