numpy fftn 对于多个图像的 2d fft 效率非常低

Question

我想计算多幅图像的傅立叶变换。因此，我将 numpy 的 fft.fftn 与暴力 for 循环进行了基准测试。

这是我用来对两种方法进行基准测试的代码（在 jupyter notebook 中）：

import numpy as np

x = np.random.rand(32, 256, 256)

def iterate_fft(arr):
    k = np.empty_like(arr, dtype=np.complex64)
    for i, a in enumerate(arr):
        k[i] = np.fft.fft2(a)
    return k

k_it = iterate_fft(x)
k_np = np.fft.fftn(x, axes=(1, 2))
np.testing.assert_allclose(k_it.real, k_np.real)
np.testing.assert_allclose(k_it.imag, k_np.imag)

%%timeit
k_it = iterate_fft(x)

输出：63.6 ms ± 1.23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
k_np = np.fft.fftn(x, axes=(1, 2))

输出：122 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

为什么会有这么大的差异？

Answer 1

numpy 中的这些例程目前似乎假定最后一个维度将始终是最小的。当这是真的时 fftn 更快，有时快很多。

也就是说，我发现这两种方法在性能上的差异比你的小得多（Python 3.7.4，numpy 1.17.2）。对于您的示例，iterate_fft 需要 46 毫秒，而 ffn 需要 50 毫秒。但是如果我翻转轴，到 (256, 256, 32)，我分别得到 55 毫秒和 40 毫秒。用 (256, 256, 2) 的形状进一步推动我分别得到 21ms 和 4ms。

请注意，如果性能确实是个问题，在某些情况下 perform better 还可以使用其他 FFT 库。此外，scipy 中的完整 fftpack 与 numpy 中更有限的代码相比，其性能可能大不相同。

请注意 your usage of fftn 基本上是：

x = np.random.rand(32, 256, 256)

a = np.fft.fft(x, n=256, axis=2)
a = np.fft.fft(a, n=256, axis=1)

np.testing.assert_allclose(np.fft.fftn(x, axes=(1, 2)), a)

Answer 2

所以参与 numpy fft 开发的人 has answered 关于 GitHub 的深层问题，结果表明减速很可能来自使用的一些多维数组重排通过 pocketfft.

当 numpy 切换到 scipy 1.4 实现时，这将成为一个回忆，使用我的基准测试可以证明它没有这些缺点。

numpy fftn 对于多个图像的 2d fft 效率非常低

numpy fftn very inefficient for 2d fft of several images

python

profiling

numpy

fft