如何加速处理numpy 3D数组卷积中的for循环？

Question

在一个3D的Numpy数组中进行Z向量卷积，然后对结果执行其他操作，但是由于当前实现方式较慢。是否是for循环导致我的代码变慢了，或者是卷积本身？我尝试将其重新整形为1D向量，并在1次传递中执行卷积（就像在Matlab中所做的那样），而不使用for循环，但这并没有提高性能。我的Matlab版本比我在Python中能想到的任何版本都快大约50％。以下是相关代码片段：

convolved=np.zeros((y_lines,x_lines,z_depth))
for i in range(0, y_lines):
    for j in range(0, x_lines):
        convolved[i,j,:]= fftconvolve(data[i,j,:], Gauss) #80% of time here
        result[i,j,:]= other_calculations(convolved[i,j,:]) #20% of time here

有没有比for循环更好的方法？听说过Cython，但我现在对Python的经验有限，所以会尽力寻找最简单的解决方案。

- user4547612

---

什么是Gauss？是一种1-D高斯核吗？如果是，相对于z_depth的大小是多少？ - Curt F.

高斯核在循环之前生成一次。 数据是1D向量（z_depth），通常约有1535个元素，长度为79的1D高斯核。 我清理了fftconvolve中的大量开销，基本上直接转到 irfftn（rfftn（in1，fshape）* rfftn（in2，fshape），fshape）[fslice] .copy（）。 - user4547612

Answer 1

你正在使用的fftconvolve函数可能来自SciPy。如果是这样，请注意它需要N维数组。因此，执行卷积的更快方法是生成对应于在x和y维度上不进行任何操作，在z维度上进行一维高斯卷积的3D卷积核。

一些代码和时间结果如下。在我的机器上和一些玩具数据上，这导致了10倍的加速，如下所示：

import numpy as np
from scipy.signal import fftconvolve
from scipy.ndimage.filters import gaussian_filter

# use scipy filtering functions designed to apply kernels to isolate a 1d gaussian kernel
kernel_base = np.ones(shape=(5))
kernel_1d = gaussian_filter(kernel_base, sigma=1, mode='constant')
kernel_1d = kernel_1d / np.sum(kernel_1d)

# make the 3d kernel that does gaussian convolution in z axis only
kernel_3d = np.zeros(shape=(1, 1, 5,))
kernel_3d[0, 0, :] = kernel_1d

# generate random data
data = np.random.random(size=(50, 50, 50))

# define a function for loop based convolution for easy timeit invocation
def convolve_with_loops(data):
    nx, ny, nz = data.shape
    convolved=np.zeros((nx, ny, nz))
    for i in range(0, nx):
        for j in range(0, ny):
            convolved[i,j,:]= fftconvolve(data[i, j, :], kernel_1d, mode='same') 
    return convolved

# compute the convolution two diff. ways: with loops (first) or as a 3d convolution (2nd)
convolved = convolve_with_loops(data)
convolved_2 = fftconvolve(data, kernel_3d, mode='same')

# raise an error unless the two computations return equivalent results
assert np.all(np.isclose(convolved, convolved_2))

# time the two routes of the computation
%timeit convolved = convolve_with_loops(data)
%timeit convolved_2 = fftconvolve(data, kernel_3d, mode='same')

---

timeit 的结果：

10 loops, best of 3: 198 ms per loop
100 loops, best of 3: 18.1 ms per loop

- Curt F.

---

尝试生成长度为64的数据，看看是否可以加快速度。在二的幂次方上，FFT通常更加高效。 - cxrodgers

实现了一个3D版本，但速度比我之前的版本慢： 时间卷积：5851.7毫秒（新3D版本） 时间卷积：4093.4毫秒（旧版本） - user4547612

cxrodgers: fftconvolve 正在使用 def _next_regular(target): 来寻找数据的最佳大小（这里是 1620，用零填充一个 1535 元素向量）。 - user4547612

每个人的数据确切形状是什么？对于我来说，3D fftconvolve 在大小为256 x 256 x 256的数组上仍然更快，尽管速度只比我上面发布的代码快了不到2倍，而不是10倍。 - Curt F.

数据大小为最小200x200x1535，高斯核为1x79。 正如我所提到的，我从fftconvolve中删除了所有开销（if语句等），因此它直接进入卷积，减少了未修改的fftconvolve约30％的时间。 - user4547612

有趣的是，我尝试了一个200x200x1535的卷积，使用了一个仅为1x5形状的3D高斯核，但我的解决方案比for循环差了约两倍。我想这意味着最佳解决方案是使用for循环遍历数据块，而不仅仅是一个像素，然后在每个块中使用我的“3D”方法。 - Curt F.

Answer 2

我想您已经找到了fftconvolve函数的源代码。通常情况下，对于实数输入，它使用numpy.fft.rfftn和.irfftn函数来计算N维变换。由于目标是进行多个1-D变换，因此您可以将fftconvolve简化为以下形式：

from scipy.signal.signaltools import _next_regular

def fftconvolve_1d(in1, in2):

    outlen = in1.shape[-1] + in2.shape[-1] - 1
    n = _next_regular(outlen)

    tr1 = np.fft.rfft(in1, n)
    tr2 = np.fft.rfft(in2, n)
    out = np.fft.irfft(tr1 * tr2, n)

    return out[..., :outlen].copy()

并计算所需结果：

result = fftconvolve_1d(data, Gauss)

这段代码之所以有效，是因为numpy.fft.rfft和.irfft（注意名称中缺少n）在输入数组的单个轴上进行变换（默认情况下是最后一个轴）。在我的系统上，这比OP代码大约快了40%。

进一步提高速度的方法是使用不同的FFT后端。

首先，scipy.fftpack中的函数似乎比它们的Numpy等效函数更快。然而，Scipy变量的输出格式非常笨拙（请参见文档），这使得很难进行乘法运算。

另一个可能的后端是通过pyFFTW包装器使用FFTW。缺点是变换之前需要缓慢的“规划阶段”，并且输入必须对齐到16字节才能实现最佳性能。这在pyFFTW教程中解释得非常好。例如，生成的代码可以是：

from scipy.signal.signaltools import _next_regular
import pyfftw
pyfftw.interfaces.cache.enable()  # Cache for the "planning"
pyfftw.interfaces.cache.set_keepalive_time(1.0)

def fftwconvolve_1d(in1, in2):

    outlen = in1.shape[-1] + in2.shape[-1] - 1
    n = _next_regular(outlen)

    tr1 = pyfftw.interfaces.numpy_fft.rfft(in1, n)
    tr2 = pyfftw.interfaces.numpy_fft.rfft(in2, n)

    sh = np.broadcast(tr1, tr2).shape
    dt = np.common_type(tr1, tr2)
    pr = pyfftw.n_byte_align_empty(sh, 16, dt)
    np.multiply(tr1, tr2, out=pr)
    out = pyfftw.interfaces.numpy_fft.irfft(pr, n)

    return out[..., :outlen].copy()

通过对齐输入和缓存“规划”，我看到与OP中的代码相比加速了近3倍。内存对齐可以通过查看Numpy数组的ctypes.data属性中的内存地址来轻松检查，详情请见。

- user2379410

---

用rfft替换rfftn后，性能提高了约30％。但是pyfftw方法并没有起到帮助： pyFFTW：6.3秒 numpy rfft：4.6秒pyFFTW：86.1秒 numpy rfft：62.4秒 - user4547612