如何优化Python循环内multiprocessing的运行效率？

问题描述 #

我编写了以下代码来解决一个易并行（embarrassingly parallel）问题：在计算beta, alpha, range(deltas)对应的解后，我需要计算结果的均值并保存到字典中。然而这段代码运行速度极慢，我认为当前循环中的并行化实现效率低下，请问如何优化这段代码以提升运行速度？

原代码 #

import numpy as np
from scipy import integrate
import pandas as pd
import itertools
import multiprocessing as mtp
import time
from tqdm import tqdm

start = time.time()
def derivative(X, t, alpha, beta, delta, gamma):
    x, y = X
    dotx = x * (alpha - beta * y)
    doty = y * (-delta + gamma * x)
    return np.array([dotx, doty])


gamma = 1.
x0 = 4.
y0 = 2.

Nt = 1000
tmax = 30.
t = np.linspace(0.,tmax, Nt)
X0 = [x0, y0]

betas = np.arange(0.9, 1.4, 0.1)
alphas = np.arange(0.9, 1.4, 0.1)
deltas = np.arange(1, 1.2, 0.01)
# paramlist = list(itertools.product(alphas, betas, deltas))

columns = {"time": t}

def function(params):
    alpha, beta, delta = params

    try:
        res = integrate.odeint(derivative, X0, t, args = (alpha, beta, delta, gamma), atol=1e-9, rtol=1e-9)
    except SystemError:
        # fallback values
        x = -1
        y = -1
        print(f"I have an error for alpha: {alpha}, beta: {beta}")

    return res
    

if __name__ == '__main__':
    nProcess = 8
    
    for beta in betas:
        for alpha in alphas:
            paramlist = list(itertools.product((alpha,), (beta,), deltas))
            with mtp.Pool(processes= nProcess) as pool:
                    results = list( tqdm( pool.imap(function, paramlist), total=len(paramlist)))
                    
        
            calc = np.mean(results, axis=0)
            columns[str(beta) + '_' + str(alpha) + '_x'] = calc[:, 0]
            columns[str(beta) + '_' + str(alpha) + '_y'] = calc[:, 1]
                
                
    df = pd.DataFrame(columns)
    df.to_pickle("test.pickle")

核心低效点分析 #

1. 进程池重复创建销毁：外层25次循环每次都新建mtp.Pool，进程启动/销毁的额外开销占比极高。
2. 任务拆分过细：每个delta单独作为任务提交，进程间通信的成本被大幅放大。
3. 错误处理不严谨：异常场景下未返回同形状的结果，会导致后续均值计算报错；多进程环境中print会引发输出混乱。

优化方案与修改后代码 #

优化要点 #

• 复用进程池：整个程序仅创建一次进程池，避免重复初始化的开销。
• 批量提交任务：生成完整的参数组合列表一次性提交，减少多次调度的通信成本。
• 修正错误处理：异常时返回与正常结果同形状的 fallback 值，避免后续计算报错；替换多进程环境下的print操作。
• 利用CPU核心数：自动获取CPU核心数设置进程数，最大化资源利用率。
• 可选：更换更快的ODE求解器：用scipy.integrate.solve_ivp替代odeint，提升单任务求解速度。

修改后代码 #

import numpy as np
from scipy import integrate
import pandas as pd
import itertools
import multiprocessing as mtp
import time
from tqdm import tqdm

start = time.time()

def derivative(X, t, alpha, beta, delta, gamma):
    x, y = X
    dotx = x * (alpha - beta * y)
    doty = y * (-delta + gamma * x)
    return np.array([dotx, doty])

gamma = 1.
x0 = 4.
y0 = 2.

Nt = 1000
tmax = 30.
t = np.linspace(0., tmax, Nt)
X0 = [x0, y0]

betas = np.arange(0.9, 1.4, 0.1)
alphas = np.arange(0.9, 1.4, 0.1)
deltas = np.arange(1, 1.2, 0.01)

# 生成所有参数组合，一次性提交任务
full_paramlist = list(itertools.product(alphas, betas, deltas))
columns = {"time": t}

def function(params):
    alpha, beta, delta = params
    try:
        res = integrate.odeint(derivative, X0, t, args=(alpha, beta, delta, gamma), atol=1e-9, rtol=1e-9)
    except SystemError:
        # 返回与正常结果同形状的 fallback 值，避免均值计算出错
        res = np.full((Nt, 2), -1.0)
        # 多进程环境下建议用logging替代print，避免输出混乱
        # import logging
        # logging.warning(f"Error for alpha: {alpha}, beta: {beta}, delta: {delta}")
    return res, alpha, beta, delta

if __name__ == '__main__':
    # 用CPU核心数设置进程数，最大化利用资源
    nProcess = mtp.cpu_count()
    print(f"Using {nProcess} processes")

    # 只创建一次进程池
    with mtp.Pool(processes=nProcess) as pool:
        # 用imap_unordered提升速度，不需要保持任务顺序
        results = list(tqdm(pool.imap_unordered(function, full_paramlist), total=len(full_paramlist)))
    
    # 按alpha-beta分组，计算delta维度的均值
    alpha_beta_results = {}
    for res, alpha, beta, delta in results:
        key = (round(alpha,1), round(beta,1))  # 避免浮点数精度问题
        if key not in alpha_beta_results:
            alpha_beta_results[key] = []
        alpha_beta_results[key].append(res)
    
    # 填充结果字典
    for (alpha, beta), res_list in alpha_beta_results.items():
        calc = np.mean(res_list, axis=0)
        col_x = f"{beta}_{alpha}_x"
        col_y = f"{beta}_{alpha}_y"
        columns[col_x] = calc[:, 0]
        columns[col_y] = calc[:, 1]
    
    df = pd.DataFrame(columns)
    df.to_pickle("test.pickle")

    end = time.time()
    print(f"Total runtime: {end - start:.2f} seconds")

额外优化建议 #

• 放宽求解精度：如果业务允许，将atol和rtol调整为1e-6或1e-7，可大幅降低ODE求解时间。
• 更换求解器：使用scipy.integrate.solve_ivp替代odeint，它在多数场景下性能更优。修改示例：

  # 调整derivative函数参数顺序以适配solve_ivp
  def derivative(t, X, alpha, beta, delta, gamma):
      x, y = X
      dotx = x * (alpha - beta * y)
      doty = y * (-delta + gamma * x)
      return np.array([dotx, doty])
  
  # 在function函数中替换求解代码
  from scipy.integrate import solve_ivp
  sol = solve_ivp(derivative, [0, tmax], X0, args=(alpha, beta, delta, gamma), t_eval=t, atol=1e-9, rtol=1e-9)
  res = sol.y.T
  

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