如何用python更快地完成100000次2d-fft？

Question

我有一个形状为（100000，256，256）的3d numpy数组，我想在2d数组的每个堆栈上做FFT，这意味着要进行100000次FFT。在

我已经测试了速度的单一和堆叠的数据与最低代码以下。在

import numpy as np
a = np.random.random((256, 256))
b = np.random.random((10, 256, 256))

%timeit np.fft.fft2(a)

%timeit np.fft.fftn(b, axes=(1, 2,))

得出如下结论：

872 µs ± 19.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

6.46 ms ± 227 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

10万次fft需要一分钟以上的时间。在

有没有更快的方法可以同时进行多个fft或ifft？

更新： 经过一番搜索，我找到了cupy，这似乎有帮助。在

Answer 1

在Python中处理FFT时，CuPy是我的首选包。在处理大尺寸fft时，它有绝对惊人的性能，再加上多次迭代。当然，它依赖于通过调用cuFFT plan函数在内部制定一个1D[2D]计划，但您可能不需要为此担心。当我把你的代码改成这样时，我已经看到了2倍的改进：

import cupy as cp
a = cp.random.random((256, 256))
b = cp.random.random((256, 256))

%timeit cp.fft.fft2(a)

%timeit cp.fft.fftn(b, axes=(0, 1,))

Answer 2

pyfftw包装FFTW库可能比np.fft和scipy.fftpack包装的FFTPACK库更快。 毕竟，FFTW在西方代表最快的傅立叶变换。在

最小的代码是：

import numpy as np
import pyfftw
import multiprocessing
b = np.random.random((100, 256, 256))
bb = pyfftw.empty_aligned((100,256, 256), dtype='float64')
bf= pyfftw.empty_aligned((100,256, 129), dtype='complex128')
fft_object_b = pyfftw.FFTW(bb, bf,axes=(1,2),flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())
bb=b
fft_object_b(bb)

以下是执行np.fft和pyfftw的扩展代码：

^{pr2}$

最后，结果是速度显著加快：pyfftw比np.fft在我的电脑上。，使用2个线程。在

np.fft.fft2, 1 slice 0.00459032058716
np.fft.fftn, 100 slices 0.478203487396
bf[3,42,42] (-38.190256258791734+43.03902512127183j)
number of threads: 2
pyfftw, 1 slice 0.000421094894409
pyfftw, 100 slices 0.0439268112183
bf[3,42,42] (-38.19025625879178+43.03902512127183j)

你的电脑似乎比我的好得多！在