Python多进程共享大型NumPy数组以降低通信开销的方案问询

核心原因：之前并行方案慢的根源 #

你遇到的并行耗时高于串行的问题，本质是数据拷贝/序列化开销远超过并行计算的收益：

• Pool的map/apply会通过pickle序列化大型数组，每个子进程都会收到一份完整拷贝，序列化、反序列化和内存拷贝的成本极高。
• Process+Array虽然基于共享内存，但如果未正确构建NumPy内存视图，依然会触发数据拷贝；且ctypes数组的dtype匹配容易出错，带来额外转换开销。

针对Python 3.7.2（M1 Mac）的解决方案 #

Python 3.8才将multiprocessing.shared_memory纳入标准库，3.7.2需通过第三方库shared_memory（pip install shared_memory）或手动用mmap实现共享内存。以下是两种可行方案：

方案1：使用第三方shared_memory库（推荐） #

该库实现了与Python 3.8+标准库一致的API，无需修改即可兼容3.7.2：

import numpy as np
from multiprocessing import Process, Queue
from shared_memory import SharedMemory

def solve_task(shm_name_mat, mat_shape, mat_dtype, shm_name_y, y_shape, y_dtype, n, result_queue):
    # 打开共享内存，构建无拷贝的NumPy视图
    shm_mat = SharedMemory(shm_name_mat)
    mat = np.ndarray(mat_shape, dtype=mat_dtype, buffer=shm_mat.buf)
    shm_y = SharedMemory(shm_name_y)
    y = np.ndarray(y_shape, dtype=y_dtype, buffer=shm_y.buf)
    
    # 构造(nI + mat)，仅创建临时副本（共享内存原数据只读，避免进程冲突）
    temp_mat = mat.copy()
    np.fill_diagonal(temp_mat, temp_mat.diagonal() + n)
    
    # 求解方程组
    x = np.linalg.solve(temp_mat, y)
    result_queue.put((n, x))
    
    # 关闭共享内存句柄
    shm_mat.close()
    shm_y.close()

if __name__ == "__main__":
    # 模拟大型输入数组（替换为你的真实数据）
    mat = np.random.rand(1000, 1000)
    y = np.random.rand(1000)
    
    # 创建共享内存并写入数据（仅拷贝一次）
    shm_mat = SharedMemory(create=True, size=mat.nbytes)
    np.ndarray(mat.shape, dtype=mat.dtype, buffer=shm_mat.buf)[:] = mat[:]
    shm_y = SharedMemory(create=True, size=y.nbytes)
    np.ndarray(y.shape, dtype=y.dtype, buffer=shm_y.buf)[:] = y[:]
    
    # 初始化结果队列
    result_queue = Queue()
    processes = []
    
    # 启动子进程处理每个n值
    for n in [0, 1, 2]:
        p = Process(target=solve_task, args=(
            shm_mat.name, mat.shape, mat.dtype.name,
            shm_y.name, y.shape, y.dtype.name,
            n, result_queue
        ))
        processes.append(p)
        p.start()
    
    # 等待进程完成并收集结果
    for p in processes:
        p.join()
    results = {}
    while not result_queue.empty():
        n, x = result_queue.get()
        results[n] = x
    
    # 验证结果（可选）
    for n in results:
        temp_mat = mat.copy()
        np.fill_diagonal(temp_mat, temp_mat.diagonal() + n)
        assert np.allclose(temp_mat @ results[n], y)
    
    # 清理共享内存
    shm_mat.close()
    shm_mat.unlink()
    shm_y.close()
    shm_y.unlink()

方案2：手动用mmap实现共享内存 #

无需第三方库，直接基于操作系统的内存映射机制：

import numpy as np
from multiprocessing import Process, Queue
import mmap
import os

def solve_task(mat_fd, mat_shape, mat_dtype, y_fd, y_shape, y_dtype, n, result_queue):
    # 从文件描述符映射共享内存，构建NumPy视图
    mat_mmap = mmap.mmap(mat_fd, length=0, access=mmap.ACCESS_READ)
    mat = np.ndarray(mat_shape, dtype=mat_dtype, buffer=mat_mmap)
    y_mmap = mmap.mmap(y_fd, length=0, access=mmap.ACCESS_READ)
    y = np.ndarray(y_shape, dtype=y_dtype, buffer=y_mmap)
    
    # 构造矩阵并求解
    temp_mat = mat.copy()
    np.fill_diagonal(temp_mat, temp_mat.diagonal() + n)
    x = np.linalg.solve(temp_mat, y)
    result_queue.put((n, x))
    
    # 关闭映射
    mat_mmap.close()
    y_mmap.close()

if __name__ == "__main__":
    mat = np.random.rand(1000, 1000)
    y = np.random.rand(1000)
    
    # 创建内存文件描述符（macOS/Linux高效共享方式）
    mat_fd = os.memfd_create()
    os.write(mat_fd, mat.tobytes())
    os.lseek(mat_fd, 0, os.SEEK_SET)
    y_fd = os.memfd_create()
    os.write(y_fd, y.tobytes())
    os.lseek(y_fd, 0, os.SEEK_SET)
    
    result_queue = Queue()
    processes = []
    for n in [0,1,2]:
        p = Process(target=solve_task, args=(
            mat_fd, mat.shape, mat.dtype.name,
            y_fd, y.shape, y.dtype.name,
            n, result_queue
        ))
        processes.append(p)
        p.start()
    
    for p in processes:
        p.join()
    
    # 收集结果
    results = {}
    while not result_queue.empty():
        n, x = result_queue.get()
        results[n] = x
    
    # 清理文件描述符
    os.close(mat_fd)
    os.close(y_fd)

SharedMemory的实现原理 #

SharedMemory本质是利用操作系统级别的共享内存机制，核心流程：

1. 内存分配：创建时调用系统API（macOS用shm_open）在内存中分配一块连续区域，返回唯一名称和文件描述符。
2. 内存映射：子进程通过名称打开共享内存，将其映射到自身虚拟地址空间，无需拷贝数据。
3. NumPy视图：通过np.ndarray的buffer参数直接绑定共享内存缓冲区，仅传递数组元数据（形状、dtype），数据本身留在共享内存中。
4. 生命周期管理：创建进程调用unlink()标记内存为待回收，所有进程关闭句柄后，操作系统自动释放内存。

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