如何优化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
相关产品推荐
相关产品推荐

