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

【问题标题】：Python efficient summation in large 2D array大型二维数组中的 Python 高效求和
【发布时间】：2018-03-27 12:57:32
【问题描述】：

我的任务相当简单：我有一个大的二维矩阵，只包含零和一。对于该矩阵中的每个位置，我想对该位置周围的窗口中的所有像素求和。问题是矩阵的形状为 (166667, 17668)，窗口大小范围从 (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 的滚动窗口方法更快，如 here 所述。我尝试使用 cython 编译它，但没有效果。

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

【问题讨论】：

试试二维卷积？
试试 numba。
我不明白这怎么能比滚动窗口方法更快。
@kazemakase 更加优雅且易于阅读。
@mr_mo 我提到了 OP 的循环方法。卷积很好:)

标签： python numpy convolution

【解决方案1】：

对于加窗求和卷积实际上是多余的，因为存在一个简单的 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 为最大窗口大小计算一次，可以重复用于所有较小的窗口大小。

但是，有一个潜在的缺点，那就是像这样对所有内容求和可能会产生非常大的数字。数字上更合理的版本通过首先获取差异来消除这个问题。缺点：不再提供通过分享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 才能有效地工作。在旧版本上使用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

【讨论】：

今天我有时间测试你的解决方案。我选择了 cumsum_safe，因为我只有几个不同的窗口大小需要计算。它粉碎了所有其他人-谢谢！每天都变得更聪明 :) 是否有类似的圆形过滤器内核解决方案，或者为此必须退回到真正的卷积？
只要内核只是一个矩形，您应该能够简单地使用pad 和mode='wrap' 而不是'reflect'。

【解决方案2】：

@Divakar 在 cmets 中指出您可以使用 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

【讨论】：