基于多进程的矩阵向量同步乘法可行性咨询

需求描述 #

我有一个满足m远小于n的m×n numpy数组A，要在拥有20个共享内存CPU的节点上运行计算。每个CPU需要将A与不同的n×1向量v相乘，矩阵A全程保持不变。

约束条件 #

矩阵A体积过大，无法在每个CPU上单独加载，因此必须存入节点的共享内存中。由于A*v的结果是m×1向量，不需要在每个CPU上存储m×n规模的矩阵，只需要共享内存中的一份A即可。

问题咨询 #

如果给每个CPU分配一个工作进程，能否通过multiprocessing模块让各进程同时计算A×v（每个进程对应不同的v）？我担心各进程同时访问共享内存时，multiprocessing会把矩阵A复制到每个CPU，引发内存不足的问题。

解决方案 #

完全可以通过multiprocessing模块实现多进程共享矩阵A并并行计算，且不会将A复制到每个进程的内存空间，核心是利用共享内存机制避免数据复制：

1. 使用multiprocessing.shared_memory（Python 3.8+） #

这是官方推荐的共享内存方案，能直接让numpy数组挂载到共享内存：

• 主进程中创建共享内存，并将矩阵A写入共享内存；
• 每个子进程通过共享内存名称，将共享内存区域映射为自己进程内的numpy数组（无数据复制，只是映射内存地址）；
• 子进程使用映射后的数组完成与v的乘积计算，返回结果。

示例代码片段：

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

def worker(shm_name, shape, dtype, v, result_queue):
    # 挂载共享内存到当前进程
    existing_shm = shared_memory.SharedMemory(name=shm_name)
    # 映射为numpy数组，无数据复制
    A = np.ndarray(shape, dtype=dtype, buffer=existing_shm.buf)
    # 计算矩阵-向量乘积
    result = A @ v
    result_queue.put(result)
    # 关闭共享内存（主进程负责销毁）
    existing_shm.close()

if __name__ == "__main__":
    # 初始化大矩阵A（示例规模）
    m, n = 100, 1000000
    A = np.random.rand(m, n)
    
    # 创建共享内存，大小匹配A的字节数
    shm = shared_memory.SharedMemory(create=True, size=A.nbytes)
    # 将A写入共享内存
    A_shared = np.ndarray(A.shape, dtype=A.dtype, buffer=shm.buf)
    A_shared[:] = A[:]
    
    # 准备20个不同的向量v
    vectors = [np.random.rand(n) for _ in range(20)]
    result_queue = Queue()
    processes = []
    
    # 启动子进程
    for v in vectors:
        p = Process(target=worker, args=(shm.name, A.shape, A.dtype, v, result_queue))
        processes.append(p)
        p.start()
    
    # 收集结果
    results = [result_queue.get() for _ in processes]
    
    # 等待所有进程结束，销毁共享内存
    for p in processes:
        p.join()
    shm.close()
    shm.unlink()

2. 避免数据复制的关键 #

• 子进程只是映射共享内存地址，而非复制整个矩阵A；
• Unix/Linux下multiprocessing的默认fork模式会通过写时复制（Copy-On-Write）优化，但矩阵A全程只读，不会触发复制；Windows系统无fork机制，显式使用共享内存是更稳妥的方案；
• 确保子进程不对A进行修改，维持其只读状态，彻底避免复制触发。

3. 注意事项 #

• Python版本需3.8及以上才能使用shared_memory模块；
• 共享内存需要主进程负责创建和销毁，子进程仅挂载和关闭；
• 向量v需要传递给子进程，由于v是n×1向量，体积远小于A，复制开销可忽略。

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