正确使用 pyfftw 加速 numpy

Question

我正在尝试从 Matlab 飞跃到 numpy，但我迫切需要我的 fft 的速度。现在我知道 pyfftw，但我不知道我是否正确使用它。我的方法类似于

import numpy as np
import pyfftw
import timeit

pyfftw.interfaces.cache.enable()

def wrapper(func, *args):
    def wrapped():
        return func(*args)
    return wrapped

def my_fft(v):
    global a
    global fft_object
    a[:] = v
    return fft_object()

def init_cond(X):
    return my_fft(2.*np.cosh(X)**(-2))

def init_cond_py(X):
    return np.fft.fft(2.*np.cosh(X)**(-2))

K = 2**16
Llx = 10.
KT = 2*K
dx = Llx/np.float64(K)
X = np.arange(-Llx,Llx,dx)

global a
global b
global fft_object
a = pyfftw.n_byte_align_empty(KT, 16, 'complex128')
b = pyfftw.n_byte_align_empty(KT, 16, 'complex128')
fft_object = pyfftw.FFTW(a,b)

wrapped = wrapper(init_cond, X)
print min(timeit.repeat(wrapped,repeat=100,number=1))

wrapped_two = wrapper(init_cond_py, X)
print min(timeit.repeat(wrapped_two,repeat=100,number=1))

我很欣赏通过 pyfftw 调用 scipy 和 numpy fft 的构建器函数和标准接口。不过，这些都表现得非常缓慢。通过首先创建 fft_object 的实例，然后在全局范围内使用它，我能够获得与 numpy 的 fft 调用一样快或略快的速度。

话虽如此，我的工作假设是智慧被隐含地储存起来。真的吗？我需要明确说明吗？如果是这样，最好的方法是什么？

另外，我认为 timeit 是完全不透明的。我是否正确使用它？它是否存储了我所说的重复的智慧？提前感谢您提供的任何帮助。

Answer 1

在交互式（ipython）会话中，我认为以下是您想要做的（ipython 很好地处理了时间）：

In [1]: import numpy as np

In [2]: import pyfftw

In [3]: K = 2**16

In [4]: Llx = 10.

In [5]: KT = 2*K

In [6]: dx = Llx/np.float64(K)

In [7]: X = np.arange(-Llx,Llx,dx)

In [8]: a = pyfftw.n_byte_align_empty(KT, 16, 'complex128')

In [9]: b = pyfftw.n_byte_align_empty(KT, 16, 'complex128')

In [10]: fft_object = pyfftw.FFTW(a,b)

In [11]: a[:] = 2.*np.cosh(X)**(-2)

In [12]: timeit np.fft.fft(a)
100 loops, best of 3: 4.96 ms per loop

In [13]: timeit fft_object(a)
100 loops, best of 3: 1.56 ms per loop

In [14]: np.allclose(fft_object(a), np.fft.fft(a))
Out[14]: True

你读过tutorial吗？你不明白什么？

我建议使用builders interface 来构造FFTW 对象。尝试各种设置，最重要的是线程数。

默认不存储智慧。你需要extract it yourself。

您所有的globals 都是不必要的 - 您要更改的对象是可变的，因此您可以很好地处理它们。 fft_object 总是指向同一个东西，所以这不是全局的没有问题。理想情况下，您只是不希望在ii 上循环。我建议研究如何构建数组，以便您可以在一次调用中完成所有操作

编辑： [编辑编辑：我写了以下段落，只是粗略地看了一下你的代码，显然它是一个递归更新，如果没有一些严重的狡猾，矢量化并不是一种明显的方法。不过，我在底部有一些关于您的实施的 cmets] 我怀疑您的问题是对如何最好地使用像 Python（或实际上是 Matlab）这样的语言进行数值处理的更根本的误解。核心原则是尽可能矢量化。通过这个，我的意思是尽可能少地汇总你的 python 调用。不幸的是，我看不出如何用你的例子来做到这一点（尽管我只考虑了 2 分钟）。如果这仍然失败，请考虑cython - 但请确保您真的想走那条路（即您已经用尽了其他选择）。

关于全局变量：不要那样做。如果要创建具有状态的对象，请使用类（这就是它们的用途）或者在您的情况下使用闭包。全局几乎从来都不是你想要的（我认为在我所有的 python 编写中，我至少有一个模糊的合法用途，那就是在 pyfftw 的缓存代码中）。我建议阅读this nice SO question。 Matlab 是一种蹩脚的语言 - 造成这种情况的众多原因之一是它的垃圾范围设置工具往往会导致不良习惯。

如果你想修改全局引用，你只需要全局。我建议阅读更多关于Python scoping rules 以及python 中的变量really are 的内容。

FFTW 对象携带您需要的所有数组，因此您无需单独传递它们。使用调用接口几乎没有开销（特别是如果您禁用规范化）来设置或返回值 - 如果您处于该优化级别，我强烈怀疑您已经达到了限制（我会警告这一点对于许多非常小的 FFT，这可能完全是正确的，但此时您需要重新考虑您的算法以矢量化对 FFTW 的调用）。如果您发现每次更新数组的开销很大（使用调用接口），这是一个错误，您应该提交它（我会很惊讶）。

底线，不必担心每次调用都更新数组。这几乎肯定不是您的瓶颈，但请确保您了解规范化并根据需要禁用它（与 update_arrays() 和 execute() 方法的原始访问相比，它可能会稍微减慢速度）。

您的代码没有使用缓存。缓存仅在您使用 interfaces 代码时使用，并减少了 Python 在内部创建新 FFTW 对象的开销。由于您自己处理 FFTW 对象，因此没有理由缓存。

builders 代码是获取 FFTW 对象的约束较少的接口。我现在几乎总是使用构建器（从头开始创建 FFTW 对象要方便得多）。您想直接创建 FFTW 对象的情况非常少见，我很想知道它们是什么。

算法实现评论： 我不熟悉您正在实施的算法。但是，我目前有一些关于您如何编写它的方法。 你在每个循环上计算nl_eval(wp)，但据我所知，这与前一个循环中的nl_eval(w) 相同，所以你不需要计算两次（但这需要注意它是当你到处都有全局变量时，很难看到发生了什么，所以我可能会遗漏一些东西。

不要打扰my_fft 或my_ifft 中的副本。只需执行fft_object(u)（在我的机器上为 2.29 毫秒，而前向案例为 1.67 毫秒）。内部数组更新例程使复制变得不必要。此外，正如您所写，您复制了两次：c[:] 表示“复制到数组c”，而您复制到c 的数组是v.copy()，即@ 的副本987654347@（所以一共两份）。

更明智（并且可能有必要）将输出复制到保存数组中（因为这样可以避免在调用 FFTW 对象时破坏中间结果），但请确保您的保存数组正确对齐。我相信您已经注意到这很重要，但复制输出更容易理解。

您可以将所有缩放比例一起移动。 wn 计算中的 3 可以移动到 nl_eval 中的 my_fft 内。您还可以将它与 ifft 中的归一化常数结合起来（并在 pyfftw 中将其关闭）。

查看numexpr 了解基本的数组操作。与普通的 numpy 相比，它可以提供相当多的加速。

无论如何，从这一切中得到你想要的。毫无疑问，我遗漏了什么或说了一些不正确的话，所以请尽可能谦虚地接受它。值得花一点时间研究一下 Python 与 Matlab 相比如何运行（事实上，忽略后者）。