如何在Python多进程迭代中共享数据？代码问题求助

1. 为什么if __name__ == "__main__"上方代码反复执行 #

在Windows系统中，multiprocessing采用spawn方式创建子进程：每个子进程都会重新导入你的整个脚本文件。此时，if __name__ == "__main__"以外的顶层代码（包括加载数据的部分）会被每个子进程重复执行一遍，导致多次从硬盘加载数据。Linux/macOS下用fork方式不会有这个问题，但跨平台代码必须用if __name__隔离主进程逻辑。

2. 移到if __name__块内后进程挂起的原因 #

如果把加载A、B的代码放到if __name__块里，子进程导入脚本时不会执行这部分代码，导致子进程中的d_shuff_函数无法找到A、B变量。子进程在执行任务时会因缺少必要数据而卡住，表现为进程挂起。

3. 更优雅的任务提交方式 #

不需要通过无意义的索引传参，推荐用Pool的initializer和initargs参数，在子进程启动时一次性把A、B传递给所有子进程并全局初始化，避免重复传递大矩阵，提升效率。

import numpy as np
from numpy.random import default_rng
import scipy.io
from multiprocessing import Pool
from comparisons import D
import time

# 全局变量，用于子进程初始化
global_A = None
global_B = None
ndata = None
nday = None

def init_worker(a, b, n_data, n_day):
    """子进程初始化函数，把主进程加载的数据赋值给全局变量"""
    global global_A, global_B, ndata, nday
    global_A = a
    global_B = b
    ndata = n_data
    nday = n_day

def d_shuff():
    """不需要传参，直接用全局变量"""
    rng = default_rng()
    C = global_B[:, rng.permutation(ndata), :]
    out = np.empty((nday, nday, ndata))
    out[:] = np.nan
    for i in range(ndata):
        out[:, :, i] = D((global_A[:, i, :], C[:, i, :]))
    return out

if __name__ == "__main__":
    nshuff = 1000
    print('Loading data...')
    mat = scipy.io.loadmat(PATH_TO_DATA)
    A = mat["A"]
    B = mat["B"]
    [nbin, ndata, nday] = A.shape

    # 创建进程池时，传入初始化函数和数据
    with Pool(initializer=init_worker, initargs=(A, B, ndata, nday)) as pool:
        t = time.time()
        # 直接生成nshuff个任务，不需要传索引
        out = pool.map(d_shuff, [None]*nshuff)
        print(time.time() - t)
        print('Saving output...')
        scipy.io.savemat(
            PATH_TO_OUTPUT,
            {"out": out}
        )

代码说明 #

• 主进程仅加载一次数据，通过initargs传递给所有子进程，子进程在启动时通过init_worker把数据存到全局变量，避免重复加载和传递。
• d_shuff不再需要接收A、B参数，直接使用全局变量，任务提交时用[None]*nshuff生成任务列表，不需要无意义的索引。
• 确保了子进程各自拥有独立的随机数生成器（default_rng()在每个子进程中会生成不同的种子），洗牌结果不会重复。

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