多进程运行函数并累加结果至共享Numpy数组的实现问题

我定义了测试函数test_func(n)，它会把np.ones生成的3x3x3数组和数值n相乘：

def test_func(n):
    array = n*np.ones((3,3,3))
    return array

我希望并行运行这个函数，得到所有结果数组的总和，并且保持3x3x3的形状。原本的实现方式如下，虽然可行但比较繁琐：

# 多进程代码开始
from multiprocessing import Pool
if __name__ == '__main__':
    pool = Pool()
    grid = np.arange(2,4,1) # 生成数组 [2,3]
    final_array  = pool.map(test_func, grid) # 存储多进程运行的结果
    final_array  = sum(final_array, axis = 0) # 求和方式有点绕
    pool.close()
    pool.join()

print(final_array) 
# 输出：这是我期望得到的结果
array([[[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]],

       [[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]],

       [[5., 5., 5.],
        [5., 5., 5.],
        [5., 5., 5.]]])

我尝试用共享内存的方式优雅实现，但写的代码有问题，需要排查错误：

import ctypes
import numpy as np

L, N, M = 3,3,3 # 数组形状
mp_arr = mp.Array(ctypes.c_double, L * N * M)
final_array = np.frombuffer(mp_arr.get_obj())
final_array = final_array.reshape((L, N, M))

def test_func(n):
    final_array = np.frombuffer(mp_arr.get_obj())
    final_array = TC_p_value.reshape((L, N, M))
    final_array += n*np.ones((3,3,3))

def init(shared_arr_): # 我甚至不知道这个函数是干嘛的
    global mp_arr
    mp_arr = shared_arr_

# 多进程代码开始
from multiprocessing import Pool
if __name__ == '__main__':
    pool = Pool(initializer=init, initargs=(mp_arr,))
    grid = np.arange(2,4,1)
    pool.map_async(test_func, grid)
    pool.close()
    pool.join()

问题代码的错误点 #

• 未导入multiprocessing别名：代码中使用mp.Array但未声明import multiprocessing as mp，会触发NameError。
• 变量名拼写错误：test_func中的TC_p_value是笔误，应为final_array。
• 共享内存未初始化：mp.Array创建的数组默认是未初始化的垃圾值，需要先清零。
• 异步调用未等待完成：map_async是非阻塞方法，直接执行close()+join()可能导致进程未完成就退出，需调用get()等待结果。
• 缺少写入锁：多进程同时修改共享内存会引发竞态问题，需要锁保证操作原子性。
• 未输出结果：代码末尾未打印final_array，无法验证结果。

修正后的共享内存实现代码 #

import ctypes
import numpy as np
from multiprocessing import Pool, Array, Lock

# 定义数组形状
L, N, M = 3, 3, 3
# 创建带锁的共享内存数组，初始化为0
mp_arr = Array(ctypes.c_double, L * N * M, lock=True)
# 主进程生成numpy数组视图并清零
final_array = np.frombuffer(mp_arr.get_obj()).reshape((L, N, M))
final_array[:] = 0

def init(shared_arr_, lock_):
    # 子进程初始化全局共享变量
    global mp_arr, lock
    mp_arr = shared_arr_
    lock = lock_

def test_func(n):
    # 子进程获取共享内存的numpy视图
    arr = np.frombuffer(mp_arr.get_obj()).reshape((L, N, M))
    # 加锁避免多进程写入冲突
    with lock:
        arr += n * np.ones((L, N, M))

if __name__ == '__main__':
    lock = Lock()
    # 使用上下文管理器管理进程池，自动处理close和join
    with Pool(initializer=init, initargs=(mp_arr, lock)) as pool:
        grid = np.arange(2, 4, 1)
        pool.map(test_func, grid)
    
    # 输出最终结果
    print(final_array)

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