Python中Multiprocessing Pool替代Matlab parfor的实现疑问

问题描述 #

我需要用Python的Multiprocessing Pool实现类似Matlab parfor的并行计算，将大量数据点分配给多CPU处理。目前代码通过pool.map(process_data, range(ii_start, ii_end + 1))实现循环，但不确定写法是否符合并行计算的要求；同时对索引传递存在疑问：例如当pts=98时，代码需要调用tmean[98]计算并更新ea[98]，但不清楚如何将ii_start到ii_end的索引正确传入process_data函数。

现有代码 #

import numpy as np
from multiprocessing import Pool, cpu_count
import os
import h5py

filepath = '/actual_file_path/rry.mat'
rry_data = {}
f = h5py.File(filepath)
for k, v in f.items():
    rry_data[k] = np.array(v)

C_sorted        = rry_data['C_sorted'].T
ea              = rry_data['ea'].T
idxnew          = rry_data['idxnew'].T
ii              = rry_data['ii'].T        # prev stop point
ny              = rry_data['ny'].T
rhostarprofile  = rry_data['rhostarprofile'].T
rry             = rry_data['rry'].T
tmean           = rry_data['tmean'].T
ymean           = rry_data['ymean'].T
Ystar           = rry_data['Ystar'].T


n = len(ea)

saveInterval = int(1e6)
saveFilename = '/actual_file_path/ea_temp'

ii_start = 98
ii_end = 110.   # setting 110 for testing, ii_end = 172439820
pts = np.linspace(1,n,n)
# Initialize a global list to store results
all_results = []


def process_data(pts):
    
    rho = tmean[pts]  # Python uses 0-based indexing
    rystar = int(idxnew[pts])
    sortedrhoheight = Ystar[rystar]   

    mod = int((pts-1)%ny)
    ry = int(rry[mod])
    rhoheight = Ystar[ry]  

    if sortedrhoheight > rhoheight:
        rhobar = (1 / (rhoheight - sortedrhoheight)) * -1 * np.sum(rhostarprofile[ry:rystar, 0] * C_sorted[ry:rystar, 1])
    elif sortedrhoheight == rhoheight:
        rhobar = 0
    else:
        rhobar = (1 / (rhoheight - sortedrhoheight)) * np.sum(rhostarprofile[rystar:ry, 0] * C_sorted[rystar:ry, 1])

    ea[pts] = (rhoheight - sortedrhoheight) * (rho - rhobar)

# number of cores you have allocated for your slurm task:
number_of_cores = int(os.environ['SLURM_CPUS_PER_TASK'])
print(number_of_cores)
# number_of_cores = cpu_count() # if not on the cluster you should do this instead

if __name__ == '__main__':
    # Parallel loop for ii_start to ii_end
#    with Pool(processes=int(os.environ['SLURM_CPUS_PER_TASK'])) as pool:
    with Pool(number_of_cores) as pool:
        pool.map(process_data, range(ii_start, ii_end + 1))

核心问题分析与修正方案 #

1. 索引传递的正确性 #

你当前使用pool.map(process_data, range(ii_start, ii_end + 1))的写法是正确的：range(ii_start, ii_end +1)会生成从98到110的整数序列，每个整数会作为独立参数传入process_data函数的pts参数。

但需要注意Python与Matlab的索引差异：Matlab使用1-based索引，而numpy数组默认是0-based。如果从h5py读取的数组保留了Matlab的1-based结构（极少情况），可以直接使用；但通常h5py会将Matlab数据转换为0-based的numpy数组，因此需要将传入的pts减1，才能正确对应到numpy数组的索引位置。

2. 并行计算的关键问题 #

原代码存在一个致命缺陷：子进程无法直接修改主进程中的全局数组ea。多进程间内存相互隔离，子进程内对ea的修改不会同步到主进程，最终计算结果无法保存。

正确的做法是让process_data函数返回计算后的索引和对应值，由主进程统一收集结果并更新ea数组。

3. 优化后的完整代码 #

import numpy as np
from multiprocessing import Pool, cpu_count
import os
import h5py

# 全局变量，用于子进程读取共享数据
C_sorted = None
ea = None
idxnew = None
ny = None
rhostarprofile = None
rry = None
tmean = None
Ystar = None

def init_worker(data):
    """初始化子进程，传递共享数据集合"""
    global C_sorted, ea, idxnew, ny, rhostarprofile, rry, tmean, Ystar
    C_sorted, ea, idxnew, ny, rhostarprofile, rry, tmean, Ystar = data

def process_data(pts):
    # 将Matlab风格的1-based索引转换为Python的0-based索引
    idx = pts - 1
    
    rho = tmean[idx]
    # 转换idxnew的索引为0-based
    rystar = int(idxnew[idx]) - 1
    sortedrhoheight = Ystar[rystar]   

    mod = int((pts - 1) % ny)
    # 转换rry的索引为0-based
    ry = int(rry[mod]) - 1
    rhoheight = Ystar[ry]  

    # 根据高度关系计算rhobar，处理切片方向问题
    if sortedrhoheight > rhoheight:
        # 确保切片范围正确，无论ry和rystar的大小关系
        start, end = min(ry, rystar), max(ry, rystar)
        rhobar = (1 / (rhoheight - sortedrhoheight)) * -1 * np.sum(rhostarprofile[start:end, 0] * C_sorted[start:end, 1])
    elif sortedrhoheight == rhoheight:
        rhobar = 0
    else:
        start, end = min(ry, rystar), max(ry, rystar)
        rhobar = (1 / (rhoheight - sortedrhoheight)) * np.sum(rhostarprofile[start:end, 0] * C_sorted[start:end, 1])

    ea_val = (rhoheight - sortedrhoheight) * (rho - rhobar)
    # 返回索引和计算结果，供主进程更新
    return (idx, ea_val)

if __name__ == '__main__':
    filepath = '/actual_file_path/rry.mat'
    rry_data = {}
    # 使用with语句确保文件正确关闭
    with h5py.File(filepath, 'r') as f:
        for k, v in f.items():
            rry_data[k] = np.array(v)

    C_sorted        = rry_data['C_sorted'].T
    ea              = rry_data['ea'].T
    idxnew          = rry_data['idxnew'].T
    # 将numpy标量转换为Python整数，避免模运算异常
    ny              = rry_data['ny'].T.item()
    rhostarprofile  = rry_data['rhostarprofile'].T
    rry             = rry_data['rry'].T
    tmean           = rry_data['tmean'].T
    Ystar           = rry_data['Ystar'].T

    n = len(ea)
    saveInterval = int(1e6)
    saveFilename = '/actual_file_path/ea_temp'

    ii_start = 98
    # 将ii_end从浮点数改为整数，避免range函数报错
    ii_end = 110  

    # 获取可用CPU数量，优先使用SLURM分配的核心数
    try:
        number_of_cores = int(os.environ['SLURM_CPUS_PER_TASK'])
    except KeyError:
        number_of_cores = cpu_count()
    print(f"使用 {number_of_cores} 个核心进行并行计算")

    # 准备传递给子进程的共享数据
    worker_data = (C_sorted, ea, idxnew, ny, rhostarprofile, rry, tmean, Ystar)
    # 创建进程池并初始化子进程数据
    with Pool(processes=number_of_cores, initializer=init_worker, initargs=(worker_data,)) as pool:
        # 执行并行计算，收集所有结果
        results = pool.map(process_data, range(ii_start, ii_end + 1))
    
    # 主进程统一更新ea数组
    for idx, val in results:
        ea[idx] = val
    
    # 保存计算后的结果
    np.save(saveFilename, ea)

代码优化说明 #

• 索引转换：统一将Matlab的1-based索引转换为Python的0-based索引，避免数组越界或取值错误。
• 进程间数据传递：使用initializer和initargs一次性传递所有共享数据，减少参数传递的开销，提升并行效率。
• 结果收集与更新：子进程返回计算结果，主进程统一更新ea数组，解决多进程内存隔离导致的结果无法保存问题。
• 鲁棒性优化：处理ii_end的类型问题、ny的标量转换问题，以及切片方向不确定的情况，避免运行时异常。

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