Python 大型二维数组中的高效求和

Question

我的任务相当简单：我有一个大的二维矩阵，只包含零和一。对于这个矩阵中的每个位置，我想对这个位置周围 window 中的所有像素求和。问题是矩阵的形状为 (166667, 17668)，window 大小范围为 (333, 333) 到 (5333, 5333)。到目前为止，我只尝试了一部分数据。我得到的代码：

out_arr = np.array( in_arr.shape )
in_arr = np.pad(in_arr, windowsize//2, mode='reflect')
for y in range(out_arr.shape[0]):
    for x in range(out_arr.shape[1]):
        out_arr[y, x] = np.sum(in_arr[y:y+windowsize, x:x+windowsize])

显然，这需要很长时间。但就我而言，它比使用 numpy.stride_tricks.as_strided 的滚动 window 方法更快，如所述。我试过用cython编译，没有效果。

除了并行化之外，您对加快速度有何建议？
我手边有一台 Nvidia Titan X。有没有办法从中受益？（例如使用 cupy）

Answer 1

@Divakar 在评论中指出您可以使用 conv2d，他是对的。这是一个例子：

import numpy as np
from scipy import signal

data = np.random.rand(5,5) # you original data that you want to sum
kernel = np.ones((2,2)) # square matrix of your dimensions, filled with ones
output = signal.convolve2d(data,kernel,mode='same') # the convolution

Answer 2

对于 windowed 求和卷积实际上是矫枉过正，因为存在简单的 O(n) 解决方案：

import numpy as np
from scipy.signal import convolve

def winsum(in_arr, windowsize):
    in_arr = np.pad(in_arr, windowsize//2+1, mode='reflect')[:-1, :-1]
    in_arr[0] = 0
    in_arr[:, 0] = 0
    ps = in_arr.cumsum(0).cumsum(1)
    return ps[windowsize:, windowsize:] + ps[:-windowsize, :-windowsize] \
           -  ps[windowsize:, :-windowsize] - ps[:-windowsize, windowsize:]

这已经很快了，但您可以节省更多，因为 ps 为最大 window 尺寸计算一次可以重复用于所有较小的 window 尺寸。

但是，有一个潜在的缺点，那就是像这样对所有内容求和可能会产生非常大的数字。一个数字上更合理的版本通过首先考虑差异来消除这个问题。缺点：通过共享 ps 的额外节省不再可用。

def winsum_safe(in_arr, windowsize):
    in_arr = np.pad(in_arr, windowsize//2, mode='reflect')
    in_arr[windowsize:] -= in_arr[:-windowsize]
    in_arr[:, windowsize:] -= in_arr[:, :-windowsize]
    return in_arr.cumsum(0)[windowsize-1:].cumsum(1)[:, windowsize-1:]

作为参考，这是最接近的竞争对手，它是基于 fft 的卷积。您需要 scipy 的 up-to-date 版本才能有效地工作。在旧版本上使用 fftconvolve 而不是 convolve.

def winsumc(in_arr, windowsize):
    in_arr = np.pad(in_arr, windowsize//2, mode='reflect')
    kernel = np.ones((windowsize, windowsize), in_arr.dtype)
    return convolve(in_arr, kernel, 'valid')

下一个是模拟 scipy 的旧的——而且极其缓慢的——行为。

def winsum_nofft(in_arr, windowsize):
    in_arr = np.pad(in_arr, windowsize//2, mode='reflect')
    kernel = np.ones((windowsize, windowsize), in_arr.dtype)
    return convolve(in_arr, kernel, 'valid', method='direct')

测试和基准测试：

data = np.random.random((1000, 1000))

assert np.allclose(winsum(data, 333), winsumc(data, 333))
assert np.allclose(winsum(data, 333), winsum_safe(data, 333))

kwds = dict(globals=globals(), number=10)

from timeit import timeit
from time import perf_counter

print('data 100x1000, window 333x333')
print('cumsum:      ', timeit('winsum(data, 333)', **kwds)*100, 'ms')
print('cumsum safe: ', timeit('winsum_safe(data, 333)', **kwds)*100, 'ms')
print('fftconv:     ', timeit('winsumc(data, 333)', **kwds)*100, 'ms')


t = perf_counter()
res = winsum_nofft(data, 99) # 333 just takes too long
t = perf_counter() - t

assert np.allclose(winsum(data, 99), res)

print('data 100x1000, window 99x99')
print('conv:        ', t*1000, 'ms')

示例输出：

data 100x1000, window 333x333
cumsum:       70.33260859316215 ms
cumsum safe:  59.98647050000727 ms
fftconv:      298.60571819590405 ms
data 100x1000, window 99x99
conv:         135224.8261970235 ms

Python 大型二维数组中的高效求和

Python efficient summation in large 2D array

python

numpy

convolution