相对较慢的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

输出：

^{pr2}$

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

但是对于1921x512x512，np.fft.fftn()需要大约5分钟。因为我猜scipy的实现速度不是很快，而且考虑到在MATLAB上相同大小网格的fft在~5s内完成，我的问题是是否有一种方法可以将处理速度提高到或几乎达到MATLAB的时间。我对fft的了解有限，但显然MATLAB使用FFTW算法，python没有。有没有合理的机会让我得到类似的时间？而且，1921年似乎是一个不幸的选择，只有2个基本因子（17113），所以我认为这也起了作用。另一方面，512是一个非常适合的二次幂。如果可能的话，类似MATLAB的时间是否也可以实现，而不需要在2048年之前填充零？在

我之所以这么问，是因为我将不得不大量使用fft（这种差异将产生巨大影响！）如果在python中不可能减少计算时间，我就不得不切换到其他更快的实现。在

Answer 1

你可以试试英特尔MKL（数学内核库）的FFT，它比FFTW faster。 英特尔为Python提供了mkl-fft，它替代了数字.fft. 您只需键入：

pip install mkl-fft

再次运行程序，不做任何更改。在

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

^{bq}$

Answer 2

是的，通过接口pyfftw使用FFTW可能会比numpy.fft或{}减少计算时间。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的尺寸，在我那台普通的电脑上，我得到了：

^{pr2}$

因此pyfftw明显比numpy.fft和{}快。使用padding更快，但计算的内容不同。在

最后，pyfftw在第一次运行时可能看起来比较慢，因为它根据documentation使用了标志FFTW_MEASURE。当且仅当多个相同尺寸的DFT被连续计算时，这是一件好事。在