numpy中矩形子阵列的快速顺序加法

时间:2016-11-03 13:31:04

标签: python performance numpy vectorization

我遇到的问题是以矢量化形式重写一段代码。下面显示的代码是初始问题的简化说明

K = 20
h, w = 15, 20 
H, W = 1000-h, 2000-w 
q = np.random.randint(0, 20, size=(H, W, K))  # random just for illustration 
Q = np.zeros((H+h, W+w, K))
for n in range(H):
    for m in range(W):
        Q[n:n+h, m:m+w, :] += q[n, m, :]

这段代码需要很长时间才能执行,在我看来,允许矢量化实现相当简单。

我知道numpy的s_函数允许构建切片,这反过来可以帮助进行代码矢量化。但是因为Q中的每一个元素都是q中多个后续元素添加的结果,我发现很难以这种简单的方式进行。

我想np.add.at可能对处理顺序添加很有用。但我花了很多时间试图让这两个功能对我有用,并决定寻求帮助,因为我经常得到一个

IndexError: failed to coerce slice entry of type numpy.ndarray to integer

我做的任何尝试。

也许还有一些我不知道的神奇魔法,它可以帮助我完成任务,但谷歌似乎非常难以实现。

1 个答案:

答案 0 :(得分:3)

你基本上是沿着第一和第二轴的滑动窗口求和,在信号处理域中称为convolution。对于两个轴2D convolution。现在,Scipy将其实现为convolve2d,并且可以用于沿第三轴的每个切片。

因此,我们将实现它,就像这样 -

from scipy.signal import convolve2d

kernel = np.ones((h,w),dtype=int)
m,n,r = q.shape[0]+h-1, q.shape[1]+w-1, q.shape[2]
out = np.empty((m,n,r),dtype=q.dtype)
for i in range(r):
    out[...,i] = convolve2d(q[...,i],kernel)

事实证明,我们可以使用同一个repo中的fftconvolve,它允许我们使用更高维数组。这将以完全矢量化的方式获得输出,如此 -

from scipy.signal import fftconvolve

out = fftconvolve(q,np.ones((h,w,1),dtype=int))

运行时测试

功能定义 -

def original_app(q,K,h,w,H,W):
    Q = np.zeros((H+h-1, W+w-1, K))
    for n in range(H):
        for m in range(W):
            Q[n:n+h, m:m+w, :] += q[n, m, :]
    return Q        

def convolve2d_app(q,K,h,w,H,W):
    kernel = np.ones((h,w),dtype=int)
    m,n,r = q.shape[0]+h-1, q.shape[1]+w-1, q.shape[2]
    out = np.empty((m,n,r),dtype=q.dtype)
    for i in range(r):
        out[...,i] = convolve2d(q[...,i],kernel)
    return out

def fftconvolve_app(q,K,h,w,H,W):
    return fftconvolve(q,np.ones((h,w,1),dtype=int))

计时和验证 -

In [128]: # Setup inputs
     ...: K = 20
     ...: h, w = 15, 20 
     ...: H, W = 200-h, 400-w
     ...: q = np.random.randint(0, 20, size=(H, W, K))
     ...: 

In [129]: %timeit original_app(q,K,h,w,H,W)
1 loops, best of 3: 2.05 s per loop

In [130]: %timeit convolve2d_app(q,K,h,w,H,W)
1 loops, best of 3: 2.05 s per loop

In [131]: %timeit fftconvolve_app(q,K,h,w,H,W)
1 loops, best of 3: 233 ms per loop

In [132]: np.allclose(original_app(q,K,h,w,H,W),convolve2d_app(q,K,h,w,H,W))
Out[132]: True

In [133]: np.allclose(original_app(q,K,h,w,H,W),fftconvolve_app(q,K,h,w,H,W))
Out[133]: True

所以,似乎基于fftconvolve的方法在那里做得非常好!