具有一侧填充的两个三维阵列的卷积太慢

Question

在我目前的项目中,我需要以一种略微不同寻常的方式" 卷积 "两个三维数组:

假设我们有两个三维数组A和B,其尺寸为dimA和dimB(每个轴都相同).现在我们要为每个轴创建一个尺寸为dimA + dimB的第三个数组C.

C的条目计算如下:

c_{x1+x2,y1+y2,z1+z2} += a_{x1,y1,z1} * b_{x2,y2,z2}

我目前的版本很简单:

dimA = A.shape[0]
dimB = B.shape[0]
dimC = dimA+dimB

C = np.zeros((dimC,dimC,dimC))
for x1 in range(dimA):
    for x2 in range(dimB):
        for y1 in range(dimA):
            for y2 in range(dimB):
                for z1 in range(dimA):
                    for z2 in range(dimB):
                        x = x1+x2
                        y = y1+y2
                        z = z1+z2
                        C[x,y,z] += A[x1,y1,z1] * B[x2,y2,z2] 

不幸的是,这个版本非常慢,无法使用.

我的第二个版本是:

C = scipy.signal.fftconvolve(A,B,mode="full")

但这只计算元素 max(dimA,dimB)

谁有更好的主意？

Answer 1

您是否尝试过使用Numba？它是一个包，允许您包装通常使用 JIT 编译器速度较慢的 Python 代码。我使用 Numba 快速解决了您的问题，并获得了显着的加速。使用 IPython 的魔法timeit魔法函数，该custom_convolution函数花费了大约 18 秒，而 Numba 的优化函数花费了 10.4 毫秒。这是一个超过 1700的加速。

以下是 Numba 的实现方式。

import numpy as np
from numba import jit, double

s = 15
array_a = np.random.rand(s ** 3).reshape(s, s, s)
array_b = np.random.rand(s ** 3).reshape(s, s, s)

# Original code
def custom_convolution(A, B):

    dimA = A.shape[0]
    dimB = B.shape[0]
    dimC = dimA + dimB

    C = np.zeros((dimC, dimC, dimC))
    for x1 in range(dimA):
        for x2 in range(dimB):
            for y1 in range(dimA):
                for y2 in range(dimB):
                    for z1 in range(dimA):
                        for z2 in range(dimB):
                            x = x1 + x2
                            y = y1 + y2
                            z = z1 + z2
                            C[x, y, z] += A[x1, y1, z1] * B[x2, y2, z2]
    return C

# Numba'ing the function with the JIT compiler
fast_convolution = jit(double[:, :, :](double[:, :, :],
                        double[:, :, :]))(custom_convolution)

如果您计算两个函数结果之间的残差，您将得到零。这意味着 JIT 实现工作没有任何问题。

slow_result = custom_convolution(array_a, array_b) 
fast_result = fast_convolution(array_a, array_b)

print np.max(np.abs(slow_result - fast_result))

我为此得到的输出是0.0.

您可以将 Numba 安装到当前的 Python 设置中，也可以使用来自 continuum.io的AnacondaCE包快速尝试。

最后但并非最不重要的是，Numba 的函数比scipy.signal.fftconvolve函数快几倍。

注意：我使用的是 Anaconda 而不是 AnacondaCE。两个包在 Numba 的性能上存在一些差异，但我认为不会有太大差异。