Python中并行实现numpy向量点乘与大型稠密矩阵运算的方法咨询

方案1：启用多线程BLAS后端（无代码修改，成本最低） #

• 单核心运行问题本质是默认numpy绑定的BLAS库没有开启多线程支持，numpy的np.matmul、np.dot底层均调用BLAS的通用矩阵乘法接口，只要使用多线程编译的BLAS库（MKL、OpenBLAS均可），设置对应环境变量即可自动并行：
  • 若使用MKL后端：执行前设置环境变量export MKL_NUM_THREADS=节点可用CPU核心数
  • 若使用OpenBLAS后端：执行前设置环境变量export OPENBLAS_NUM_THREADS=节点可用CPU核心数
  • 通用兼容设置：可以同时设置export OMP_NUM_THREADS=节点可用CPU核心数覆盖大部分BLAS库的线程配置
• 可通过import numpy as np; np.__config__.show()查看当前numpy绑定的BLAS库类型

方案2：使用BLAS秩1更新接口降低计算&内存开销 #

你当前执行的操作本质是复数矩阵的秩1更新，完全不需要先计算完整的40000*40000中间矩阵（该矩阵单存储就需要23.8GB内存，是你第二步减法耗时高的核心原因），直接调用BLAS内置的秩1更新函数即可：

• 安装scipy后调用scipy.linalg.blas.zgeru接口，代码示例：

from scipy.linalg.blas import zgeru
# 直接原位更新A，不需要生成中间大矩阵
zgeru(alpha=-k, x=A[:, 0].ravel(), y=A[1, :].ravel(), a=A)

• 该接口同样支持多线程BLAS加速，且省去了中间矩阵的内存申请、拷贝和后续减法遍历的开销，性能提升远高于单纯开多线程矩阵乘法

方案3：GPU加速（集群有GPU资源时优先选择） #

如果集群配有NVIDIA GPU，可使用CuPy库直接把运算卸载到GPU：

• 代码几乎不需要修改，仅把numpy替换为cupy即可：

import cupy as cp
# 先把A转成CuPy数组
A = cp.array(A)
A = A - k * cp.matmul(A[:, 0:1], A[1:2, :])

• 40000*40000的complex128矩阵仅需23.8GB显存，主流数据中心GPU（A10、A100等）均可轻松承载，运算速度是CPU多核心的数倍到数十倍

方案4：分布式CPU并行（多节点集群无GPU时选择） #

如果需要跨多CPU节点并行，可使用Dask数组分片存储大矩阵，Dask会自动调度多节点运算，接口和numpy高度兼容，仅需要少量代码修改即可适配分布式集群。

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