使用mp.Process/Queue并行np.einsum反而比串行慢，如何优化？

以下是针对你的np.einsum并行提速问题的具体优化手段：

1. 用共享内存避免大数组重复复制 #

你当前仅传递切片索引，但第一个完整数组会在每个子进程启动时被完整复制（Python多进程的写时复制机制在数组被访问后会触发实际复制），这是大数组下并行耗时飙升的核心原因。解决方法是用numpy的shared_memory模块让所有进程共享同一份数组内存：

• 主进程中创建共享内存，将第一个数组写入共享区域；
• 子进程通过共享内存名称映射出数组，无需序列化/复制；
• 处理完后释放共享内存。

示例代码片段：

import numpy as np
from multiprocessing import Pool
from multiprocessing.shared_memory import SharedMemory

# 全局变量，子进程中用于映射共享数组
shared_shm = None
arr1 = None

def init_worker(shm_name, arr_shape, arr_dtype):
    global shared_shm, arr1
    # 子进程中映射共享内存
    shared_shm = SharedMemory(name=shm_name)
    arr1 = np.ndarray(arr_shape, dtype=arr_dtype, buffer=shared_shm.buf)

def process_slice(slice_idx):
    # 从第二个数组的切片索引获取子数组
    arr2_slice = arr2[slice_idx]
    # 执行einsum计算
    result = np.einsum('ijk,ikl->ijl', arr1, arr2_slice)
    return result

if __name__ == '__main__':
    # 假设arr1是第一个大三维数组，arr2是第二个待切片的三维数组
    arr1 = np.random.rand(1000, 1000, 1000)
    arr2 = np.random.rand(1000, 1000, 1000)
    # 创建共享内存存储arr1
    shm = SharedMemory(create=True, size=arr1.nbytes)
    arr1_shared = np.ndarray(arr1.shape, dtype=arr1.dtype, buffer=shm.buf)
    arr1_shared[:] = arr1[:]  # 将数据写入共享内存

    # 划分arr2的切片（假设沿第0轴切分为4份，对应4核）
    num_slices = 4
    slices = []
    step = arr2.shape[0] // num_slices
    for i in range(num_slices):
        start = i * step
        end = start + step if i < num_slices-1 else arr2.shape[0]
        slices.append(slice(start, end))

    # 初始化进程池，传入共享内存参数
    with Pool(processes=num_slices, initializer=init_worker,
              initargs=(shm.name, arr1.shape, arr1.dtype)) as pool:
        results = pool.map(process_slice, slices)

    # 拼接结果
    final_result = np.concatenate(results, axis=0)

    # 释放共享内存
    shm.close()
    shm.unlink()

2. 禁用numpy内部多线程 #

numpy的einsum默认会利用OpenBLAS/MKL等线性代数库的多线程加速，这会导致多进程+内部多线程的CPU资源竞争，反而降低效率。需要在进程启动前设置环境变量，强制每个进程仅用1个线程：

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"

将这段代码放在脚本最开头，确保numpy加载前生效。

3. 调整任务粒度 #

避免将数组切分得过于细碎，否则进程调度、结果拼接的开销会抵消并行收益。建议根据CPU核心数（比如4核）将数组划分为对应数量的大块，每个子进程处理完整的一块，减少任务切换次数。

4. 选择合适的进程池方法 #

如果结果的顺序不影响后续处理，可使用imap_unordered替代map，它能在子进程完成任务后立即返回结果，减少整体等待时间；若需要传递多个参数，用starmap更方便。

5. 验证内存占用 #

确保系统有足够的空闲内存承载共享数组+多个子进程的计算临时数据，内存不足会触发磁盘交换，导致速度骤降。可通过htop或任务管理器实时监控内存使用情况。

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