Python NumPy重计算的并行优化：利用数据共享避免性能瓶颈

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

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

# 模拟重计算函数
def calc(mat):
    # 模拟一些重计算，例如对大型矩阵进行多次统计分析
    for ii in range(1000):
        avg = np.mean(mat)
        std = np.std(mat)
    return avg, std

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

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

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

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

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

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

if __name__ == '__main__':
    main_problematic_example()

For Loop: 100%|████████████████████████████████████████████| 100/100 [00:51<00:00,  1.93it/s]
单线程for循环: 51.88s
Native Map: 100%|████████████████████████████████████████████| 100/100 [00:52<00:00,  1.91it/s]
原生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

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

# 优化的计算函数：通过索引访问共享数据
def calc_optimized(idx, mat_list):
    # 模拟一些重计算
    for ii in range(1000):
        avg = np.mean(mat_list[idx])  # 从共享列表中获取矩阵
        std = np.std(mat_list[idx])
    return avg, std

def main_optimized_example():
    ds = list(mydataset((500, 500), 100))

    print("\n--- 优化后方法性能测试 ---")

    # 创建一个进程池，通常设置为CPU核心数
    # 为了演示效果，这里使用4个核心
    num_cores = 4 
    mypool = Pool(num_cores)

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

    t0 = time.time()
    # 使用starmap，因为它允许我们将多个参数传递给工作函数
    # 这里传递的是数据的索引和共享列表本身
    res_optimized = mypool.starmap(calc_optimized, zip(range(len(ds)), [mylist] * len(ds)))
    print(f"Manager共享内存方式: {time.time() - t0:.2f}s")

    mypool.close()
    mypool.join()
    manager.shutdown() # 关闭管理器

if __name__ == "__main__":
    main_optimized_example()

--- 优化后方法性能测试 ---
Manager共享内存方式: 1.94s