You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何优化PyFFTW实现的滑动点积使其速度超过SciPy?

如何优化PyFFTW实现的滑动点积使其速度超过SciPy?

我有一个用重叠相加卷积方法实现滑动点积的简单函数:

import numpy as np
from scipy.signal import oaconvolve
import pyfftw
import os

def scipy_sliding_dot(A, B):
    m = A.shape[0]
    n = B.shape[0]
    Ar = np.flipud(A)  # Reverse/flip A
    AB = oaconvolve(Ar, B)

    return AB.real[m - 1 : n]

作为参考,这和下面的朴素实现效果完全一致:

def naive_sliding_dot(A, B):
    m = len(A)
    n = len(B)
    l = n - m + 1
    out = np.empty(l)
    for i in range(l):
        out[i] = np.dot(A, B[i:i+m])
    return out

当我初始化两个*纯实数(非复数)*的随机数组时:

A = np.random.rand(2**6)
B = np.random.rand(2**20)

用%timeit测试scipy_sliding_dot的性能,得到结果:

6.39 ms ± 38.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

之后我尝试用多线程的pyfftw来加速这个操作,实现了如下类:

class pyfftw_sliding_dot(object):
    # 基于相关实现思路
    def __init__(self, A, B, threads=1):
        shape = (np.array(A.shape) + np.array(B.shape))-1
        self.rfft_A_obj = pyfftw.builders.rfft(A, n=shape, threads=threads)
        self.rfft_B_obj = pyfftw.builders.rfft(B, n=shape, threads=threads)
        self.irfft_obj = pyfftw.builders.irfft(self.rfft_A_obj.output_array, n=shape, threads=threads)

    def __call__(self, A, B):
        m = A.shape[0]
        n = B.shape[0]
        Ar = np.flipud(A)  # Reverse/flip A
        rfft_padded_A = self.rfft_A_obj(Ar)
        rfft_padded_B = self.rfft_B_obj(B)

        return self.irfft_obj(np.multiply(rfft_padded_A, rfft_padded_B)).real[m - 1 : n]

然后用下面的代码测试性能:

n_threads = os.cpu_count()
obj = pyfftw_sliding_dot(A, B, n_threads)
%timeit obj(A, B)

得到的结果却是:

33 ms ± 347 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

这意味着多线程的pyfftw比scipy慢了约5倍。我仔细研究了pyfftw的builders相关文档,尝试了各种额外参数(比如planner_effort、overwrite_input等),但pyfftw的性能没有任何改善。

我在pyfftw的使用上哪里出错了?怎样才能让它比scipy更快?

备注:内容来源于stack exchange,提问作者slaw

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.04.14 09:23:09