PyFFTW 在多维数组上的性能

Question

我有一个 nD 数组，比如维度：(144, 522720)，我需要计算它的 FFT。

PyFFTW 似乎比 numpy 和 scipy 慢，这不是预期的。

我做错了什么吗？

下面是我的代码

import numpy
import scipy      
import pyfftw
import time

n1 = 144
n2 = 522720
loops = 2

pyfftw.config.NUM_THREADS = 4
pyfftw.config.PLANNER_EFFORT = 'FFTW_ESTIMATE'
# pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'

Q_1 = pyfftw.empty_aligned([n1, n2], dtype='float64')
Q_2 = pyfftw.empty_aligned([n1, n2], dtype='complex_')
Q_ref = pyfftw.empty_aligned([n1, n2], dtype='complex_')

# repeat a few times to see if pyfft planner helps
for i in range(0,loops):
    Q_1 = numpy.random.rand(n1,n2)

    s1 = time.time()
    Q_ref = numpy.fft.fft(Q_1, axis=0)
    print('NUMPY - elapsed time: ', time.time() - s1, 's.')

    s1 = time.time()
    Q_2 = scipy.fft.fft(Q_1, axis=0)
    print('SCIPY - elapsed time: ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    Q_2 = pyfftw.interfaces.numpy_fft.fft(Q_1, axis=0)
    print('PYFFTW NUMPY - elapsed time = ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    Q_2 = pyfftw.interfaces.scipy_fftpack.fft(Q_1, axis=0)
    print('PYFFTW SCIPY - elapsed time = ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    fft_object = pyfftw.builders.fft(Q_1, axis=0)
    Q_2 = fft_object()
    print('FFTW PURE Elapsed time = ', time.time() - s1, 's')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

Answer 1

首先，如果您在主循环之前打开缓存，接口将基本按预期工作：

pyfftw.interfaces.cache.enable()
pyfftw.interfaces.cache.set_keepalive_time(30)

有趣的是，尽管 wisdom 应该被存储，但当缓存关闭时 pyfftw 对象的构造仍然相当慢。不管了，这正是缓存的目的。在您的情况下，您需要使缓存 keep-alive 时间很长，因为您的循环很长。

其次，将fft_object的构建时间包括在最终测试中是不公平的比较。如果你把它移到定时器之外，那么调用 fft_object 是一个更好的措施。

第三，有趣的是，即使启用了缓存，对 numpy_fft 的调用也比对 scipy_fft 的调用慢。由于代码路径没有明显差异，我认为这是缓存问题。这是 timeit 试图缓解的问题。这是我建议的更有意义的时序代码：

import numpy
import scipy
import pyfftw
import timeit

n1 = 144
n2 = 522720

pyfftw.config.NUM_THREADS = 4
pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'

Q_1 = pyfftw.empty_aligned([n1, n2], dtype='float64')

pyfftw.interfaces.cache.enable()
pyfftw.interfaces.cache.set_keepalive_time(30)

times = timeit.repeat(lambda: numpy.fft.fft(Q_1, axis=0), repeat=5, number=1)
print('NUMPY fastest time = ', min(times))

times = timeit.repeat(lambda: scipy.fft.fft(Q_1, axis=0), repeat=5, number=1)
print('SCIPY fastest time = ', min(times))

times = timeit.repeat(
    lambda: pyfftw.interfaces.numpy_fft.fft(Q_1, axis=0), repeat=5, number=1)
print('PYFFTW NUMPY fastest time = ', min(times))

times = timeit.repeat(
    lambda: pyfftw.interfaces.scipy_fftpack.fft(Q_1, axis=0), repeat=5, number=1)
print('PYFFTW SCIPY fastest time = ', min(times))

fft_object = pyfftw.builders.fft(Q_1, axis=0)
times = timeit.repeat(lambda: fft_object(Q_1), repeat=5, number=1)
print('FFTW PURE fastest time = ', min(times))

在我的机器上，输出如下：

NUMPY fastest time =  0.6622681759763509
SCIPY fastest time =  0.6572431400418282
PYFFTW NUMPY fastest time =  0.4003451430471614
PYFFTW SCIPY fastest time =  0.40362057799939066
FFTW PURE fastest time =  0.324020683998242

如果您不通过将 Q_1 更改为 complex128:

强制将输入数组复制到复杂数据类型，您可以做得更好

NUMPY fastest time =  0.6483533839927986
SCIPY fastest time =  0.847397351055406
PYFFTW NUMPY fastest time =  0.3237176960101351
PYFFTW SCIPY fastest time =  0.3199474769644439
FFTW PURE fastest time =  0.2546963169006631

那个有趣的 scipy slow-down 是可重复的。

就是说，如果您的输入是真实的，您应该进行真实的转换（对于 pyfftw 的 >50% speed-up）并操纵合成的复杂输出。

这个例子的有趣之处在于（我认为）缓存在结果中的重要性（我认为这就是为什么切换到真正的转换在加快速度方面如此有效）。当您将数组大小更改为 524288（您认为这可能会加快速度，但不会显着降低速度的下一个 2 的幂）时，您也会看到一些戏剧性的事情。在这种情况下，一切都变慢了很多，特别是 scipy。我觉得 scipy 对缓存更敏感，这可以解释将输入更改为 complex128 时速度变慢的原因（不过 522720 对于 FFTing 来说是一个相当不错的数字，所以也许我们应该期待速度变慢） .

最后，如果速度比准确性次要，您始终可以使用 32 位浮点数作为数据类型。如果将其与进行真正的转换相结合，您将比上面给出的初始 numpy 好 10 speed-up 倍：

PYFFTW NUMPY fastest time =  0.09026529802940786
PYFFTW SCIPY fastest time =  0.1701313250232488
FFTW PURE fastest time =  0.06202622700948268

（numpy 和 scipy 变化不大，因为我认为它们在内部使用 64 位浮点数）。

编辑：我忘记了 Scipy 的 fftpack 真实 FFT 有一个奇怪的输出结构，pyfftw 复制时速度有些慢。这在 new FFT module.

中更改为更明智

新的 FFT 接口是 implemented in pyFFTW，应该是首选。不幸的是，重建文档时出现问题，因此文档已经过时很长时间并且没有显示新界面 - 希望现在已修复。

PyFFTW 在多维数组上的性能

PyFFTW perfomance on multidimensional arrays

python

numpy

fft

pyfftw