python - 一侧填充太慢的两个三维数组的卷积

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 设置中，或者使用 AnacondaCE 快速试用它来自 continuum.io 的包。

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

注意:我使用的是 Anaconda 而不是 AnacondaCE。对于 Numba 的性能，这两个包之间存在一些差异，但我认为差异不会太大。