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

Question

我想计算几个图像的傅里叶变换。 因此，我将numpy的fft.fftn与暴力for循环进行了对比。你知道吗

这是我用来测试这两种方法的代码（在jupyter笔记本中）：

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会更快，有时会快很多。你知道吗

也就是说，我得到的这两种方法在性能上的差别要比你小得多（使用python3.7.4，numpy1.17.2）。例如，iterate_fft需要46毫秒，ffn需要50毫秒。但是如果我把轴翻转到(256, 256, 32)，我分别得到55毫秒和40毫秒。用(256, 256, 2)的形状进一步推，我分别得到21ms和4ms。你知道吗

注意，如果性能确实是一个问题，那么在某些情况下，还有其他可用的FFT库perform better。此外，scipy中的完整fftpack的性能可能与numpy中更有限的代码有很大不同。你知道吗

注意your usage of ^{}基本上是：

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)