Python多进程并行矩阵运算后，如何将结果存入列表？

问题背景 #

原本的串行代码可以生成N对随机矩阵并将乘积存入列表，现在尝试用multiprocessing并行化，但无法将每个进程生成的矩阵保存到结果列表中。

解决方案 #

需要调整函数的返回逻辑，让进程池能收集每个任务的计算结果，具体修改如下：

1. 修改defs.py中的函数

import numpy as np
    
def g(n):
    A = np.random.rand(n, n)
    B = np.random.rand(n, n)
    # 返回计算后的矩阵乘积，让进程池能捕获结果
    return A * B


def run_complex_operations(operation, input_list, pool):
    # 返回pool.map的结果，该方法会自动按输入顺序收集所有任务的返回值
    return pool.map(operation, input_list)

2. 修改主程序代码

import defs   
import numpy as np
import multiprocessing as mp

n = 100
N = 10
# 构造每个任务的输入参数：N个n，对应生成N次矩阵
input_list = [n] * N

if __name__ == '__main__':
    processes_pool = mp.Pool(N)
    # 接收并行任务返回的结果列表，out就是包含10个矩阵的列表
    out = defs.run_complex_operations(defs.g, input_list, processes_pool)
    # 正确关闭进程池，等待所有任务完成
    processes_pool.close()
    processes_pool.join()
    
    # 验证结果
    print(f"结果列表包含 {len(out)} 个矩阵")
    print(f"单个矩阵的形状：{out[0].shape}")

关键说明 #

• g函数必须返回计算结果，否则pool.map无法捕获进程的输出。
• pool.map会自动将每个任务的返回值按输入顺序整理成列表，直接接收该列表即可得到最终的out。
• 进程池使用完毕后调用close()和join()是良好实践，确保所有子进程正常结束。

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