Python如何使用multiprocessing.Pool处理多参数函数的多进程任务

原代码核心问题 #

• 函数名拼写错误：定义的迭代函数名为newtalt，调用时误写为newtonalt
• 多参数传递逻辑错误：Pool.map仅支持单参数目标函数，多参数场景未做适配
• 进程池资源管理错误：在for循环内调用pool.close()，第一次循环结束后进程池就被销毁，后续任务无法提交
• 任务提交效率低下：循环内逐个提交任务，无法充分发挥多进程并行优势

修复后完整代码 #

import matplotlib.pyplot as plt
import numpy as np
from multiprocessing import Pool

def newtalt(fun, x0, err, mit):
    xnew = x0.copy()
    F, dF = fun(x0)
    r = F.copy()
    M = dF.copy()
    sigma = np.linalg.norm(r)
    for k in range(mit):
        if sigma < err:
            break
        d = np.linalg.solve(M, -r)
        xnew = xnew + d
        r, M = fun(xnew)
        sigma = np.linalg.norm(r)
    return k

def fun(x):
    f_r = x[0] ** 3 - 3 * x[0] * (x[1] ** 2) - 1
    f_i = 3 * (x[0]**2) * x[1] - (x[1] ** 3)
    f = np.array([f_r,f_i])
    df_rx = 3 * (x[0] ** 2) - 3 * (x[1] ** 2)
    df_ry = -6 * x[0] * x[1]
    df_ix = 6 * x[0] * x[1]
    df_iy = 3 * (x[0] ** 2) - 3 * (x[1] ** 2)
    df = np.array([[df_rx,df_ry],[df_ix,df_iy]])
    return f, df

if __name__=='__main__':
    err = 1e-5
    mit = 300
    N = 50
    x = np.linspace(-2.5, 2.5, N)
    y = np.linspace(-2.5, 2.5, N)
    A = np.zeros((N, N))
    # 预先生成所有任务的参数元组
    task_list = []
    for i in range(N):
        for j in range(N):
            z = np.array([x[i], y[j]])
            task_list.append((fun, z, err, mit))
    # 上下文管理器自动管理进程池生命周期
    with Pool(processes=10) as pool:
        # starmap自动解包参数元组，适配多入参函数
        iters = pool.starmap(newtalt, task_list)
    # 结果按提交顺序回填到矩阵
    idx = 0
    for i in range(N):
        for j in range(N):
            A[i][j] = iters[idx]
            idx += 1
    plt.imshow(A, cmap='Set1')
    plt.show()

关键修改说明 #

• 批量任务提交：提前将所有坐标对应的参数打包成任务列表，一次性提交给进程池，大幅提升并行效率
• 多参数适配：使用Pool.starmap替代map，可以自动将每个任务的参数元组解包后传入目标函数，无需额外修改原函数定义
• 自动资源管理：使用with语句创建进程池，执行结束后会自动关闭进程池、回收资源，无需手动调用close、join方法，避免时序错误
• 性能优化提示：如果N取值较大，可以使用functools.partial将固定参数fun、err、mit预封装，减少跨进程参数传递的开销

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