python - 如何使用 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(7 次运行的平均值 ± 标准差，每次 1000 次循环)

每次循环 6.46 ms ± 227 µs(7 次运行的平均值 ± 标准偏差，每次 100 个循环)

100000 次 fft 需要一分钟以上。

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

更新: 经过一番搜索，我发现了 cupy ，这似乎有帮助。

Answer 1

<强> pyfftw ，包裹 FFTW库，可能比 FFTPACK 更快由 np.fft 和 scipy.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 的执行进行计时的扩展代码:

import numpy as np
from timeit import default_timer as timer
import multiprocessing
a = np.random.random((256, 256))
b = np.random.random((100, 256, 256))

start = timer()
for i in range(10):
    np.fft.fft2(a)
end = timer()
print"np.fft.fft2, 1 slice", (end - start)/10

start = timer()
for i in range(10):
     bf=np.fft.fftn(b, axes=(1, 2,))
end = timer()
print "np.fft.fftn, 100 slices", (end - start)/10
print "bf[3,42,42]",bf[3,42,42]


import pyfftw

aa = pyfftw.empty_aligned((256, 256), dtype='float64')
af= pyfftw.empty_aligned((256, 129), dtype='complex128')
bb = pyfftw.empty_aligned((100,256, 256), dtype='float64')
bf= pyfftw.empty_aligned((100,256, 129), dtype='complex128')
print 'number of threads:' , multiprocessing.cpu_count()

fft_object_a = pyfftw.FFTW(aa, af,axes=(0,1), flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())

fft_object_b = pyfftw.FFTW(bb, bf,axes=(1,2),flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())


aa=a
bb=b
start = timer()
for i in range(10):
    fft_object_a(aa)
end = timer()
print "pyfftw, 1 slice",(end - start)/10

start = timer()
for i in range(10):
    fft_object_b(bb)
end = timer()
print "pyfftw, 100 slices", (end - start)/10
print "bf[3,42,42]",bf[3,42,42]

最后，结果是显着的加速:事实证明，在我的计算机上，pyfftw 比 np.fft 快 10 倍。，使用 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)

你的电脑看起来比我的好得多!