如何正确使用Python multiprocessing包优化scipy CWT大规模信号计算

核心错误点 #

• 进程池控制逻辑位置错误：pool.close()和pool.join()被写在任务提交的for循环内部，第一次循环提交单个任务后就直接关闭并销毁了进程池，后续79个任务完全无法提交，是程序无CPU占用、无法正常运行的核心原因。
• 异步调用参数格式错误：apply_async的args参数要求传入元组，args=(i)会被解析为单个整数而非元组，导致参数传递异常，子进程无法正常调用cwt函数，正确写法为args=(i,)。
• cwt函数内部变量引用错误：cwt函数定义的入参为f，但内部计算width时错误引用了外部循环变量i，依赖全局作用域变量的写法在多进程场景下会导致子进程无法拿到正确的输入值，单进程场景也存在隐式bug。
• 隐式依赖全局变量：cwt函数依赖全局变量trace、sampling_freq，在Windows系统的spawn进程启动模式下，子进程不会继承主进程的全局变量，会直接触发运行错误。

修正后完整可运行代码 #

import numpy as np
import random
from scipy import signal
import multiprocessing as mp

# 初始化全局参数与测试信号
random.seed(0)
comp = 10
trace = np.zeros(1000)
t = np.linspace(0, 1, 1000)
sampling_freq = 500  # Hz

for i in range(comp):
    frq = random.gauss(40, 10)
    ampl = random.uniform(0, .01)
    fi = random.uniform(-np.pi, np.pi)
    comp_i = ampl * np.sin(2 * np.pi * frq * t + fi)
    trace += comp_i

def width(freq, sampling_freq):
    return sampling_freq * (2.0 * np.pi * freq) ** (-.5)

# 重构cwt函数，显式传入所有依赖参数，避免全局变量依赖
def cwt(f, trace, sampling_freq):
    cwt_res = signal.cwt(trace, signal.ricker, [width(f, sampling_freq)])
    return cwt_res

if __name__ == '__main__':
    results = []
    
    def collect_results(result):
        results.extend(result.tolist())
    
    pool = mp.Pool(processes=mp.cpu_count())
    # 先批量提交所有任务
    for i in range(1, 81):
        pool.apply_async(cwt, args=(i, trace, sampling_freq), callback=collect_results)
    # 所有任务提交完成后再关闭进程池，等待全部任务执行完成
    pool.close()
    pool.join()

    print(f"计算完成，共得到{len(results)}组CWT结果")

可选优化写法 #

使用pool.starmap替代apply_async+回调的写法，代码更简洁不易出错：

if __name__ == '__main__':
    pool = mp.Pool(processes=mp.cpu_count())
    # 构造所有任务的参数元组列表
    task_params = [(i, trace, sampling_freq) for i in range(1, 81)]
    # 批量执行任务，直接获取所有结果
    raw_results = pool.starmap(cwt, task_params)
    pool.close()
    pool.join()
    # 结果展平为和原写法一致的格式
    results = [res for sub_res in raw_results for res in sub_res]

内容的提问来源于stack exchange，提问作者Kuando