多维数组的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对象的构造仍然相当缓慢。不要紧，这正是缓存的目的。在您的例子中，您需要使缓存保持活动时间相当长，因为您的循环非常长。

其次，将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减速是可重复的。

也就是说，如果您的输入是真实的，那么您应该进行真实的转换(使用pyfftw加速50%以上)，并操作生成的复杂输出。

这个例子的有趣之处在于(我认为)缓存在结果中的重要性(我建议这就是为什么切换到真正的转换在加快速度方面是如此有效)。当您使用将数组大小更改为524288 (2的下一个幂，您认为这可能会加快速度，但不会显著减慢)时，您也会看到一些戏剧性的东西。在这种情况下，一切都慢了不少，特别是scipy。在我看来，scipy对缓存更敏感，这可以解释为什么将输入改为complex128会变慢(522720对于FFTing来说是一个很好的数字，所以也许我们应该预料到会变慢)。

最后，如果速度次于准确性，则可以始终使用32位浮点数作为数据类型。如果你将它与实际的转换结合起来，你会得到比上面给出的初始numpy更好的10倍的加速：

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中已更改为更合理。

新的快速傅立叶变换接口是implemented in pyFFTW，应该优先考虑。不幸的是，正在重建的文档有一个问题，所以文档已经过时很长一段时间了，没有显示新的界面--希望现在已经修复了。