比较慢的python numpy 3D傅里叶变换

Question

对于我的工作，我需要对大图像执行离散傅立叶变换 (DFT)。在当前示例中，我需要 1921 x 512 x 512 图像的 3D FT（以及 512 x 512 图像的 2D FFT）。现在，我正在使用 numpy 包和相关的函数np.fft.fftn()。下面的代码 sn-p 以以下方式示例性地显示了相同大小/稍小的 2D/3D 随机数生成网格上的 2D 和 3D FFT 时间：

import sys
import numpy as np
import time

tas = time.time()
a = np.random.rand(512, 512)
tab = time.time()
b = np.random.rand(100, 512, 512)

tbfa = time.time()

fa = np.fft.fft2(a)
tfafb = time.time()
fb = np.fft.fftn(b)
tfbe = time.time()

print "initializing 512 x 512 grid:", tab - tas
print "initializing 100 x 512 x 512 grid:", tbfa - tab
print "2D FFT on 512 x 512 grid:", tfafb - tbfa
print "3D FFT on 100 x 512 x 512 grid:", tfbe - tfafb

输出：

initializing 512 x 512 grid: 0.00305700302124
initializing 100 x 512 x 512 grid: 0.301637887955
2D FFT on 512 x 512 grid: 0.0122730731964
3D FFT on 100 x 512 x 512 grid: 3.88418793678

我遇到的问题是我经常需要这个过程，所以每张图片花费的时间应该很短。在我自己的计算机上进行测试时（中段笔记本电脑，2GB RAM 分配给虚拟机（--> 因此较小的测试网格）），如您所见，3D FFT 需要约 5 秒（数量级）。现在，在工作中，机器要好得多，集群/网格架构系统和 FFT 更快。在这两种情况下，2D 都准瞬间完成。

但是对于 1921x512x512，np.fft.fftn() 需要大约 5 分钟。由于我猜 scipy 的实现速度并没有快多少，并且考虑到在 MATLAB FFT 上相同大小的网格在大约 5 秒内完成，我的问题是是否有一种方法可以将这个过程加速到或几乎达到 MATLAB 时间。我对 FFT 的了解有限，但显然 MATLAB 使用了 FFTW 算法，而 python 没有。使用一些 pyFFTW 包我得到类似时间的任何合理机会？此外，1921 似乎是一个不幸的选择，只有 2 个质因数（17、113），所以我认为这也起作用。另一方面，512 是非常适合的 2 的幂。如果可能的话，是否也可以在不填充零到 2048 的情况下实现类似 MATLAB 的时间？

我问是因为我将不得不大量使用 FFT（在一定程度上这种差异会产生巨大影响！），如果没有可能减少 python 中的计算时间，我必须切换到其他更快的实现。

Answer 1

是的，与numpy.fft 或scipy.fftpack 相比，通过接口pyfftw 使用FFTW 有可能减少您的计算时间。这些 DFT 算法实现的性能可以在基准测试中进行比较，例如 this one：Improving FFT performance in Python 报告了一些有趣的结果

我建议使用以下代码进行测试：

import pyfftw
import numpy
import time
import scipy

f = pyfftw.n_byte_align_empty((127,512,512),16, dtype='complex128')
#f = pyfftw.empty_aligned((33,128,128), dtype='complex128', n=16)
f[:] = numpy.random.randn(*f.shape)

# first call requires more time for plan creation
# by default, pyfftw use FFTW_MEASURE for the plan creation, which means that many 3D dft are computed so as to choose the fastest algorithm.
fftf=pyfftw.interfaces.numpy_fft.fftn(f)

#help(pyfftw.interfaces)
tas = time.time()
fftf=pyfftw.interfaces.numpy_fft.fftn(f) # here the plan is applied, nothing else.
tas = time.time()-tas
print "3D FFT, pyfftw:", tas

f = pyfftw.n_byte_align_empty((127,512,512),16, dtype='complex128')
#f = pyfftw.empty_aligned((33,128,128), dtype='complex128', n=16)
f[:] = numpy.random.randn(*f.shape)


tas = time.time()
fftf=numpy.fft.fftn(f)
tas = time.time()-tas
print "3D FFT, numpy:", tas

tas = time.time()
fftf=scipy.fftpack.fftn(f)
tas = time.time()-tas
print "3D FFT, scipy/fftpack:", tas

# first call requires more time for plan creation
# by default, pyfftw use FFTW_MEASURE for the plan creation, which means that many 3D dft are computed so as to choose the fastest algorithm.
f = pyfftw.n_byte_align_empty((128,512,512),16, dtype='complex128')
fftf=pyfftw.interfaces.numpy_fft.fftn(f)

tas = time.time()
fftf=pyfftw.interfaces.numpy_fft.fftn(f) # here the plan is applied, nothing else.
tas = time.time()-tas
print "3D padded FFT, pyfftw:", tas

对于 127*512*512 的尺寸，在我普通的电脑上，我得到了：

3D FFT, pyfftw: 3.94130897522
3D FFT, numpy: 16.0487070084
3D FFT, scipy/fftpack: 19.001199007
3D padded FFT, pyfftw: 2.55221295357

所以pyfftw 明显快于numpy.fft 和scipy.fftpack。使用填充甚至更快，但计算的东西不同。

最后，pyfftw 根据documentation 使用标志FFTW_MEASURE，在第一次运行时可能看起来更慢。当且仅当连续计算许多相同大小的 DFT 时，这是一件好事。

Answer 2

您可以尝试来自 Intel MKL（数学内核库）的 FFT，它是 faster，而不是 FFTW。 英特尔为 Python 提供了mkl-fft，它取代了 numpy.fft。您需要做的就是输入：

pip install mkl-fft

然后再次运行您的程序，不做任何更改。

另外，numpy 1.17（即将发布）将有新的 FFT 实现：

pocketfft 库替换基于 fftpack 的 FFT 模块

两种实现都有相同的祖先（Paul 的 Fortran77 FFTPACK N. Swarztrauber），但 pocketfft 包含额外的修改 在某些情况下可以提高准确性和性能。为了 FFT 长度包含大的素因子，pocketfft 使用 Bluestein 的 算法，它保持 O(N log N) 的运行时间复杂度，而不是 对于素数长度，向 O(N*N) 恶化。此外，准确度为 具有接近质数长度的实值 FFT 得到了改进并达到了同等水平 具有复值 FFT。