如何并行化Numpy数组填充的for循环？

好问题！你顾虑多进程直接写入同一变量的问题太对了——毕竟多进程各自有独立的内存空间，直接写不仅会导致数据混乱，甚至子进程的修改根本不会反映到主进程里。下面给你几种实用的并行实现方案：

这种思路最直观安全：让每个进程独立计算自己负责的行，把结果返回给主进程，最后再把所有行拼接成完整数组。完全避开了共享内存的坑。

示例代码：

import numpy as np
from multiprocessing import Pool

# 把你的耗时计算封装成函数，输入行索引i，返回对应行的数组a
def compute_single_row(i):
    # 这里替换成你的实际耗时计算逻辑
    # 比如模拟耗时操作：
    # import time
    # time.sleep(0.001)
    return np.random.rand(10)  # 示例生成10个元素的一维数组

if __name__ == "__main__":
    # 创建进程池，processes可以设为CPU核心数，或者按需调整
    with Pool(processes=4) as pool:
        # 把0到9999的索引分给进程池，得到每行结果的列表
        row_results = pool.map(compute_single_row, range(10000))
    # 将列表转成numpy数组，就是最终的A
    A = np.array(row_results)

如果你觉得multiprocessing的代码有点啰嗦，joblib库封装了更多细节，写起来更清爽：

示例代码：

import numpy as np
from joblib import Parallel, delayed

def compute_single_row(i):
    # 替换成你的耗时计算逻辑
    return np.random.rand(10)

if __name__ == "__main__":
    # n_jobs=-1表示用满所有可用CPU核心
    row_results = Parallel(n_jobs=-1)(delayed(compute_single_row)(i) for i in range(10000))
    A = np.array(row_results)

Parallel负责管理并行任务，delayed用来包装你的计算函数，把每个索引对应的任务分发出去，最后收集结果拼接即可。

如果你的数组特别大，返回所有行再拼接会占用额外内存，这时候可以用共享内存让进程直接写入指定位置。因为你的场景是每行独立写入，不会有多个进程修改同一位置的情况，所以是安全的：

示例代码：

import numpy as np
from multiprocessing import Pool, RawArray

def init_shared_array(shared_memory):
    # 让子进程能访问到共享内存的数组视图
    global shared_A
    shared_A = np.frombuffer(shared_memory, dtype=np.float64).reshape((10000, 10))

def compute_and_write_row(i):
    # 计算得到你的一维数组a
    a = np.random.rand(10)
    # 直接写入共享数组的对应行
    shared_A[i, :] = a

if __name__ == "__main__":
    # 创建共享内存，'d'对应float64类型，大小是10000*10
    shared_mem = RawArray('d', 10000 * 10)
    # 把共享内存转成numpy数组的视图，主进程可以直接操作
    A = np.frombuffer(shared_mem, dtype=np.float64).reshape((10000, 10))
    
    # 初始化进程池，把共享内存传递给子进程
    with Pool(processes=4, initializer=init_shared_array, initargs=(shared_mem,)) as pool:
        pool.map(compute_and_write_row, range(10000))
    
    # 现在A已经被所有进程填充完成，可以直接使用
    print(A.shape)  # 输出 (10000, 10)

• 优先选方案1或方案2，简单易维护，完全不用操心共享内存的问题，适合绝大多数场景。
• 如果数组大到返回结果会占用太多内存，再考虑方案3，注意要保证每个进程只写自己的行，避免数据竞争。

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