python - 比较慢的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 的幂。是否可以在不使用零填充到 2048 的情况下实现类似 MATLAB 的时间？

我问是因为我将不得不大量使用 FFT(达到这种差异将产生巨大影响的数量!)并且如果无法减少 python 中的计算时间，我将不得不切换到其他更快的实现。

Answer 1

是的，与 numpy.fft 或 scipy.fftpack 相比，通过接口(interface) 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 时，这是一件好事。