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

Question

对于我的工作，我需要对大图像执行离散傅立叶变换 (DFT)。在当前示例中，我需要 1921 x 512 x 512 图像的 3D FT（以及 512 x 512 图像的 2D FFT）。现在，我正在使用 numpy 包和相关的函数np.fft.fftn()。下面的代码片段示例性地显示了在相同大小/略小的 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 次幂。如果可能的话，是否也可以实现类似 MATLAB 的时间而不用零填充到 2048？

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

Answer 1

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

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

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由于它FFTW_MEASURE根据文档使用该标志，因此第一次运行时可能看起来较慢。当且仅当连续计算许多相同大小的 DFT 时，这才是一件好事。