Python multiprocessing共享numpy二维数组读写结果异常问题求助

问题排查 #

你的实现思路是正确的：通过multiprocessing.RawArray开辟共享内存映射numpy数组，各进程操作独立索引无需加锁，不存在根本性错误，运行未得到预期结果是代码存在几处笔误和逻辑疏漏：

• 键名拼写错误：worker函数中读取共享结果数组时，使用的键是worker_array，但初始化函数init_worker存入全局字典的键是result_array，无法正常读取目标数组
• 异步任务未等待完成：你调用了pool.map_async但没有调用wait()等方法等待任务执行完成，with块退出时进程池直接销毁，worker逻辑还未运行就结束了
• 函数参数不匹配：worker定义了i,j两个入参，但你传入map_async的是一维的行索引序列，每个任务只传入一个参数，会触发参数不匹配错误
• 数组类型未显式指定：调用np.frombuffer时没有指定dtype参数，会默认按float64解析内存，和原数组的整数类型不匹配，导致读取/写入的值错误

修正后可运行代码 #

import numpy as np
import multiprocessing as mp

WORKER_DICT = dict()

def shared_array_from_np_array(data_array, init_value=None):
    raw_array = mp.RawArray(np.ctypeslib.as_ctypes_type(data_array.dtype), data_array.size)
    shared_array = np.frombuffer(raw_array, dtype=data_array.dtype).reshape(data_array.shape)
    if init_value is not None:
        np.copyto(shared_array, np.full_like(data_array, init_value))
    else:
        np.copyto(shared_array, data_array)
    return raw_array, shared_array

def init_worker(data_array, result_array, dtype, shape):
    WORKER_DICT['data_array'] = data_array
    WORKER_DICT['result_array'] = result_array
    WORKER_DICT['dtype'] = dtype
    WORKER_DICT['shape'] = shape

def worker(idx):
    # 按原数组dtype解析共享内存
    data = np.frombuffer(WORKER_DICT['data_array'], dtype=WORKER_DICT['dtype']).reshape(WORKER_DICT['shape'])
    result = np.frombuffer(WORKER_DICT['result_array'], dtype=WORKER_DICT['dtype']).reshape(WORKER_DICT['shape'])
    # 这里按行处理，你也可以根据自己的需求修改索引逻辑
    result[idx] = data[idx] * 2
    return

if __name__ == '__main__':
    input_arr = np.array(
        [[1, 1, 2, 2],
         [1, 1, 2, 2],
         [3, 3, 4, 4],
         [3, 3, 4, 4]]
    )
    sh_in_arr, shared_input_array = shared_array_from_np_array(input_arr)
    sh_res_arr, shared_result_array = shared_array_from_np_array(input_arr, init_value=0)
    init_args = (sh_in_arr, sh_res_arr, input_arr.dtype, input_arr.shape)
    with mp.Pool(processes=2, initializer=init_worker, initargs=init_args) as pool:
        # 等待异步任务全部执行完成
        pool.map_async(worker, range(input_arr.shape[0])).wait()
    print('Input:\n', shared_input_array)
    print('Output:\n', shared_result_array)

运行结果 #

Input:
 [[1 1 2 2]
 [1 1 2 2]
 [3 3 4 4]
 [3 3 4 4]]
Output:
 [[2 2 4 4]
 [2 2 4 4]
 [6 6 8 8]
 [6 6 8 8]]

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