为numpy 3D数组加速卷积循环？

Posted 2023-02-23

【中文标题】为numpy 3D数组加速卷积循环？

【英文标题】：Speed up for loop in convolution for numpy 3D array?

【发布时间】：2015-11-08 19:50:01

【问题描述】：

沿 3d numpy 数组的 Z 向量执行卷积，然后对结果进行其他操作，但它现在实现起来很慢。是 for 循环在这里减慢了我的速度还是卷积？我尝试整形为一维向量并在 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 方面的经验有限，会以最简单的解决方案为目标。

【问题讨论】：

什么是Gauss？某种一维高斯核？如果是这样，相对于z_depth 的大小是多少？

高斯核在循环之前生成一次。数据是一维向量（z_depth），通常长度约为 1535 个元素，长度通常为 79 的一维高斯核。我在 fftconvolve 中清理了一堆开销，基本上只是直接进入 irfftn(rfftn(in1, fshape) * rfftn(in2 , fshape), fshape)[fslice].copy()

【参考方案1】：

您使用的fftconvolve 函数大概来自SciPy。如果是这样，请注意它需要 N 维数组。因此，一种更快的卷积方法是生成 3d 内核，对应于在 x 和 y 维度中什么都不做，并在 z. 中进行 1d 高斯卷积

一些代码和计时结果如下。在我的机器上和一些玩具数据，这导致了 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

【讨论】：

尝试生成长度为 64 的数据，看看这是否会使其更快。 FFT 通常在 2 次方上效率更高。

实现了一个 3D 版本，它比我的慢：时间卷积：5851.7 毫秒（新 3D 版本）时间卷积：4093.4 毫秒（旧版本）

每个人的数据的确切形状是什么？ 3-D fftconvolve 对我来说在 256 x 256 x 256 大小的数组中保持更快，尽管比我上面发布的代码的 10 倍略低于 2 倍。

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

【参考方案2】：

我想你已经找到fftconvolve函数的source code了。通常，对于实际输入，它使用计算 N 维变换的 numpy.fft.rfftn 和 .irfftn 函数。由于目标是进行多次一维转换，因此您基本上可以像这样（简化）重写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 变体的输出格式非常尴尬（参见docs），这使得乘法变得困难。

另一个可能的后端是通过 pyFFTW 包装器实现的 FFTW。缺点是转换之前有一个缓慢的“规划阶段”，并且输入必须是 16 字节对齐才能获得最佳性能。这在pyFFTW tutorial 中有很好的解释。生成的代码可能是例如：

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 属性中的内存地址来对齐内存对齐 can be easily checked。

【讨论】：

将 rfftn 替换为 rfft 将性能提高了约 30%。但是 pyfftw 方法没有帮助： pyFFTW: 6.3 sec numpy rfft: 4.6 sec pyFFTW: 86.1 sec numpy rfft: 62.4 sec