如何使用多进程批量对多个NumPy数组执行自定义函数运算

完全可以通过多进程并行完成计算，你的场景属于典型的多独立输入的CPU密集型计算任务，非常适配多进程并行逻辑，具体实现如下：

前置注意事项 #

numpy默认自带多线程运算优化，多进程场景下要先关闭numpy内部线程，避免进程间抢占CPU资源导致性能反而下降。

完整实现代码 #

import numpy as np
import multiprocessing
import os

# 关闭numpy内部多线程，避免和多进程资源冲突
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'

# 你的原计算函数
def s_1d_nb1(x):
    a = np.zeros(len(x))
    for i in range(len(x)):
        a[i] = np.mean((x[i:] - x[:len(x)-i])**2)
    return a

if __name__ == '__main__':
    # 配置参数：n为你要计算的数组数量
    n = 10
    # 统一生成所有输入数组，存为列表，替代你原来单独定义x1~xn的写法
    input_arrays = [np.random.rand(10**6) for _ in range(n)]
    
    # 初始化进程池，默认进程数等于CPU核心数，也可以手动指定processes参数
    with multiprocessing.Pool() as pool:
        # 批量执行计算，返回结果顺序和输入列表顺序完全对应
        output_arrays = pool.map(s_1d_nb1, input_arrays)
    
    # 按索引取对应结果即可，比如new_x1对应第一个输入的输出
    new_x1 = output_arrays[0]
    new_x2 = output_arrays[1]
    # new_xn = output_arrays[n-1]

额外说明 #

• if __name__ == '__main__'是Python多进程的强制要求，尤其是Windows系统下不加会直接报错，全平台都建议保留该判断。
• 如果你的内存不足以一次性存下所有输入输出数组，可以把pool.map替换为pool.imap，迭代式逐个获取结果，降低内存占用。
• 你的原函数里的循环逻辑如果后续要进一步优化，可以用numba的JIT编译加速，结合多进程收益会更高。

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