如何用Python多进程共享内存处理多维NumPy数组？

嘿，我刚好碰到过一模一样的问题——超大多维NumPy数组用多进程传递时复制成本太高甚至报错，确实不用扁平化也能搞定！给你分享两个实用的方案：

mp.Array本质是一维的，但我们可以通过NumPy的视图转换把它变成多维数组，全程不会复制数据，只是改变索引方式。核心思路是先创建对应总元素数的共享内存，再用np.frombuffer转换成多维数组：

import numpy as np
import multiprocessing as mp

def process_chunk(shared_array, original_shape):
    # 将共享内存转换成目标形状的NumPy数组（无数据复制）
    arr = np.frombuffer(shared_array.get_obj(), dtype=np.float64).reshape(original_shape)
    # 这里写你的并行处理逻辑，比如修改某块数据
    arr[0:10, 0:10, 0:10, 0:10] += 1.0
    print(f"Worker完成部分修改，样本值：{arr[0,0,0,0]}")

if __name__ == "__main__":
    target_shape = (55, 117, 256, 256)
    dtype = np.float64
    # 计算总元素数，匹配共享内存的容量
    total_elements = np.prod(target_shape)
    
    # 创建共享内存数组，类型码要和NumPy dtype对应（'d'对应float64，'f'对应float32）
    shared_array = mp.Array('d', total_elements)
    
    # 初始化共享内存（把原数组数据复制进去，可选）
    original_arr = np.zeros(target_shape, dtype=dtype)
    np.copyto(np.frombuffer(shared_array.get_obj(), dtype=dtype).reshape(target_shape), original_arr)
    
    # 启动多进程处理
    processes = []
    for _ in range(4):  # 按需求设置进程数
        p = mp.Process(target=process_chunk, args=(shared_array, target_shape))
        processes.append(p)
        p.start()
    
    for p in processes:
        p.join()
    
    # 主进程查看最终结果
    final_arr = np.frombuffer(shared_array.get_obj(), dtype=dtype).reshape(target_shape)
    print(f"主进程获取结果，样本值：{final_arr[0,0,0,0]}")

Python 3.8之后引入了专门的共享内存模块，完全为NumPy数组设计，不用手动计算总元素数，直接创建和原数组形状一致的共享视图：

import numpy as np
from multiprocessing import shared_memory, Process

def process_chunk(shm_name, shape, dtype):
    # 连接到已创建的共享内存块
    existing_shm = shared_memory.SharedMemory(name=shm_name)
    # 直接创建对应形状的NumPy数组视图
    arr = np.ndarray(shape, dtype=dtype, buffer=existing_shm.buf)
    # 并行处理逻辑
    arr[10:20, 10:20, 10:20, 10:20] += 2.0
    print(f"Worker完成部分修改，样本值：{arr[10,10,10,10]}")
    # 关闭共享内存（主进程负责最终释放）
    existing_shm.close()

if __name__ == "__main__":
    target_shape = (55, 117, 256, 256)
    dtype = np.float64
    
    # 创建原数组
    original_arr = np.zeros(target_shape, dtype=dtype)
    # 创建共享内存块，大小匹配原数组字节数
    shm = shared_memory.SharedMemory(create=True, size=original_arr.nbytes)
    # 绑定共享内存到NumPy数组视图
    shared_arr = np.ndarray(target_shape, dtype=dtype, buffer=shm.buf)
    # 初始化数据（可选）
    np.copyto(shared_arr, original_arr)
    
    # 启动多进程
    processes = []
    for _ in range(4):
        p = Process(target=process_chunk, args=(shm.name, target_shape, dtype))
        processes.append(p)
        p.start()
    
    for p in processes:
        p.join()
    
    # 查看结果
    print(f"主进程获取结果，样本值：{shared_arr[10,10,10,10]}")
    # 释放共享内存
    shm.close()
    shm.unlink()

• 确保所有进程使用的dtype完全一致，否则会出现数据错乱
• 如果涉及多进程写入同一区域，一定要加锁（比如mp.Lock）避免竞争条件
• 方案1的类型码要和NumPy dtype对应：比如int32对应'i'，uint8对应'B'，可以参考Python的array模块类型码表

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