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

Comparatively slow python numpy 3D Fourier Transformation

对于我的工作,我需要对大图像执行离散傅立叶变换 (DFT)。在当前示例中,我需要 1921 x 512 x 512 图像的 3D FT(以及 512 x 512 图像的 2D FFT)。现在,我正在使用 numpy 包和相关函数 np.fft.fftn()。下面的代码片段通过以下方式示例性地显示了 equal-sized/slightly 较小的 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 秒(数量级)。现在,在工作中,机器要好得多,cluster/grid-architecture 系统和 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 中无法减少计算时间,我必须切换到其他更快的实现。

是的,与 numpy.fftscipy.fftpack 相比,通过接口 pyfftw 使用 FFTW 有可能会减少您的计算时间。这些 DFT 算法实现的性能可以在 this one : some interesting results are reported in 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.fftscipy.fftpack。使用填充甚至更快,但是计算的东西是不同的。

最后,pyfftw 起初 运行 可能看起来较慢,因为它根据 documentation 使用标志 FFTW_MEASURE。当且仅当连续计算许多相同大小的 DFT 时,这是一件好事。

您可以尝试 Intel MKL(数学内核库)中的 FFT,它比 FFTW faster。 Intel 为 Python 提供了 mkl-fft,它取代了 numpy.fft。您需要做的就是输入:

pip install mkl-fft

和 运行 再次使用您的程序,没有任何更改。

此外,numpy 1.17(即将发布)将有新的 FFT 实现:

Replacement of the fftpack-based FFT module by the pocketfft library

Both implementations have the same ancestor (Fortran77 FFTPACK by Paul N. Swarztrauber), but pocketfft contains additional modifications which improve both accuracy and performance in some circumstances. For FFT lengths containing large prime factors, pocketfft uses Bluestein’s algorithm, which maintains O(N log N) run time complexity instead of deteriorating towards O(N*N) for prime lengths. Also, accuracy for real-valued FFTs with near-prime lengths has improved and is on par with complex-valued FFTs.