numpy fftn 对于几个图像的 2d fft 非常低效答案

【问题标题】：numpy fftn very inefficient for 2d fft of several imagesnumpy fftn 对于几个图像的 2d fft 非常低效
【发布时间】：2020-01-20 09:56:14
【问题描述】：

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

这是我用来对这两种方法进行基准测试的代码（在 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)

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

【问题讨论】：

我在numpy 中打开了一个 GitHub issue 以了解他们对此的看法。

标签： python numpy profiling fft

【解决方案1】：

numpy 中的这些例程目前似乎假设最后一个维度总是最小的。如果这是真的，fftn 会更快，有时会快很多。

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

请注意，如果性能确实是个问题，在某些情况下还有其他可用的 FFT 库perform better。此外，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)

【讨论】：

感谢您的回答。我的numpy 版本是 1.16.4（不是一个强要求，我会更改看看会发生什么），我的 Python 版本是 3.6.8（我想保持这种状态）。我不完全确定我在阅读文档时理解您的第一点来自哪里，您能否提供更多详细信息？此外，我对批量大小为 2 的结果感到惊讶。批量大小为 1 会发生什么？似乎fftn 在 FFT2 中的执行速度比fft2 快。感谢您推荐其他库，性能对我来说并不是真正的问题，我只是感到惊讶。
numpy 1.17 似乎引入了pocketfft 的使用，这会稍微改变性能。请注意，我的形状更改只是展示了当前代码似乎已优化的情况，它正在做与您想做的不同的事情。鉴于当前的优化，迭代图像似乎更适合您的用例
是的，我刚刚尝试使用 numpy 1.17，并且首先使用批量维度进行迭代时效果也更好。我将尝试最后使用批量维度进行分析，尤其是批量大小为 1 的情况。是的，完全理解你在做什么，我只是不明白那些优化是什么使它在批量大小的最后一种情况下工作得更好，批量大小很小。
所以当批量维度是最后一个时，我也有fftn 以小幅度击败迭代版本。然而，对于 2 的批量大小，无论批量维度在哪里（最后一个或第一个），结果都是相同的，并且如下所示：迭代 2.9 毫秒，fftn 2.5 毫秒。因此，利润几乎没有你的那么大。我将上传一个要点来说明这一点。当批量大小为 1 时，我猜我们会支付 for 循环的成本，因为它是 2.2 毫秒与 1.2 毫秒。另外，我在文档中注意到了这一点：> fft2 只是 fftn，轴的默认值不同。
gist.github.com/zaccharieramzi/f05930697b40b74b412e14b420316972 这是我使用的代码（你可以看到批量大小 1 的结果，但它很容易适应 2 或 32）。无论如何，我认为这可能表明fftn 效率低下，对吧？

【解决方案2】：

所以涉及numpy fft开发的人has answered在github上的深刻问题，事实证明，减速可能来自pocketfft的一些多维数组重新排列。

当numpy 987654325 @ 1.4实现时，全部是一个存储器，可以使用电台显示，以没有这些缺点。

【讨论】：