在我当前的项目中,我需要以一种稍微不寻常的方式“卷积”两个三维数组:
假设我们有两个三维数组 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)
谁有更好的主意?
最佳答案
您是否尝试过使用 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 的性能,这两个包之间存在一些差异,但我认为差异不会太大。
关于python - 一侧填充太慢的两个三维数组的卷积,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/14786920/