优化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 cpu_count

# 生成模拟数据集
def mydataset(size, length):
    for ii in range(length):
        yield np.random.rand(*size)

# 模拟重度计算函数
def calc(mat):
    # 模拟一些耗时的NumPy计算
    for ii in range(1000):
        avg = np.mean(mat)
        std = np.std(mat)
    return avg, std

def main():
    # 生成100个500x500的随机NumPy矩阵
    ds = list(mydataset((500, 500), 100))

    print(f"当前系统CPU核心数: {cpu_count()}")

    # 1. 传统for循环
    t0 = time.time()
    res1 = []
    for mat in tqdm(ds, desc="For Loop"):
        res1.append(calc(mat))
    print(f'for loop: {time.time() - t0:.2f}s')

    # 2. 原生map函数
    t0 = time.time()
    res2 = list(map(calc, tqdm(ds, desc="Native Map")))
    print(f'native map: {time.time() - t0:.2f}s')

    # 3. process_map (多进程)
    t0 = time.time()
    # 默认使用所有可用CPU核心
    res3 = process_map(calc, ds, desc="Process Map")
    print(f'process map: {time.time() - t0:.2f}s')

    # 4. thread_map (多线程)
    t0 = time.time()
    # 默认使用所有可用CPU核心
    res4 = thread_map(calc, ds, desc="Thread Map")
    print(f'thread map: {time.time() - t0:.2f}s')

if __name__ == '__main__':
    main()

当前系统CPU核心数: 28
For Loop: 100%|████████████████████████████████████████████| 100/100 [00:51<00:00,  1.93it/s]
for loop: 51.88s
Native Map: 100%|████████████████████████████████████████████| 100/100 [00:52<00:00,  1.91it/s]
native map: 52.49s
Process Map: 100%|████████████████████████████████████████████| 100/100 [01:10<00:00,  1.41it/s]
process map: 71.06s
Thread Map: 100%|████████████████████████████████████████████| 100/100 [00:41<00:00,  2.39it/s]
thread map: 42.04s

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

# 生成模拟数据集（与之前相同）
def mydataset(size, length):
    for ii in range(length):
        yield np.random.rand(*size)

# 修改后的计算函数：现在接收索引和共享列表
def calc_shared(idx, mat_list):
    # 模拟一些耗时的NumPy计算
    # 通过索引访问共享列表中的矩阵
    mat = mat_list[idx]
    for ii in range(1000):
        avg = np.mean(mat)
        std = np.std(mat)
    return avg, std

def main_optimized():
    ds = list(mydataset((500, 500), 100))
    print(f"当前系统CPU核心数: {cpu_count()}")

    # 创建一个进程池，通常设置为CPU核心数
    # 为了演示效果，这里使用4个核心进行测试
    num_processes = 4 # 可以根据实际CPU核心数调整
    mypool = Pool(num_processes)

    # 使用Manager创建共享列表
    manager = Manager()
    # 将原始数据放入Manager管理的列表中，数据仅拷贝一次
    mylist = manager.list(ds)

    t0 = time.time()
    # 使用starmap传递多个参数：任务的索引和共享列表
    # zip(range(len(ds)), mylist) 会生成 (0, mylist[0]), (1, mylist[1]), ... 的迭代器
    res_optimized = mypool.starmap(calc_shared, zip(range(len(ds)), [mylist]*len(ds)))
    print(f"map with Manager (using {num_processes} cores): {time.time() - t0:.2f}s")

    mypool.close()
    mypool.join()
    manager.shutdown() # 关闭Manager进程

if __name__ == "__main__":
    main_optimized()

当前系统CPU核心数: 28
map with Manager (using 4 cores): 1.94s