python - 加速 numpy 3D 数组的卷积循环？

Question

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

Answer 1

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

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