Python中NumPy计算加速：如何利用多进程避免数据拷贝瓶颈

import time
import numpy as np
from tqdm.auto import tqdm
from tqdm.contrib.concurrent import process_map, thread_map
from multiprocessing import Pool, Manager

def mydataset(size, length):
    """生成指定大小和数量的随机NumPy矩阵数据集"""
    for _ in range(length):
        yield np.random.rand(*size)

def calc(mat):
    """模拟对NumPy矩阵的重度计算"""
    for _ in range(1000):
        _ = np.mean(mat)
        _ = np.std(mat)
    return True # 简化返回值，原问题返回avg, std

def main_initial_test():
    ds = list(mydataset((500, 500), 100)) # 生成100个500x500的矩阵

    print("--- 原始方法性能测试 ---")

    t0 = time.time()
    for mat in tqdm(ds, desc="For Loop"):
        calc(mat)
    print(f'For Loop: {time.time() - t0:.2f}s')

    t0 = time.time()
    list(map(calc, tqdm(ds, desc="Native Map")))
    print(f'Native Map: {time.time() - t0:.2f}s')

    t0 = time.time()
    process_map(calc, ds, desc="Process Map")
    print(f'Process Map: {time.time() - t0:.2f}s')

    t0 = time.time()
    thread_map(calc, ds, desc="Thread Map")
    print(f'Thread Map: {time.time() - t0:.2f}s')

if __name__ == '__main__':
    # main_initial_test()
    pass # 暂时注释，后面会展示优化后的代码

For Loop: 51.88s
Native Map: 52.49s
Process Map: 71.06s
Thread Map: 42.04s

import time
import numpy as np
from multiprocessing import Pool, Manager

def mydataset(size, length):
    """生成指定大小和数量的随机NumPy矩阵数据集"""
    for _ in range(length):
        yield np.random.rand(*size)

def calc_with_shared_data(idx, mat_list_proxy):
    """
    模拟对NumPy矩阵的重度计算，通过索引访问共享数据。
    mat_list_proxy 是 Manager.list 的代理对象。
    """
    mat = mat_list_proxy[idx] # 通过索引获取共享列表中的矩阵
    # 模拟一些重度计算
    for _ in range(1000):
        _ = np.mean(mat)
        _ = np.std(mat)
    return True # 简化返回值
    # return avg, std # 如果需要返回计算结果

def main_optimized():
    ds = list(mydataset((500, 500), 100)) # 生成100个500x500的矩阵

    # 1. 创建Manager实例
    manager = Manager()
    # 2. 将原始数据集转换为Manager管理的共享列表
    # 数据在此处被一次性拷贝到Manager的服务器进程内存中
    shared_mat_list = manager.list(ds)

    # 3. 创建进程池，通常设置为CPU核心数
    # 这里使用4个进程进行演示，可根据实际CPU核心数调整
    with Pool(processes=4) as mypool:
        t0 = time.time()
        # 4. 使用starmap传递多个参数：任务索引和共享列表的代理对象
        # zip(range(len(ds)), [shared_mat_list] * len(ds)) 为每个任务生成 (索引, 共享列表代理) 对
        results = mypool.starmap(calc_with_shared_data, 
                                 zip(range(len(ds)), [shared_mat_list] * len(ds)))
        print(f"Manager Pool Starmap: {time.time() - t0:.2f}s")

    # 注意：Manager在with Pool块结束后会自动清理，
    # 如果不使用with语句，需要手动调用manager.shutdown()

if __name__ == '__main__':
    print("--- 优化后方法性能测试 ---")
    main_optimized()

Manager Pool Starmap: 1.94s