如何显著加速 Python 中的高维矩阵乘法运算-Python教程-PHP中文网

本文介绍在处理大量高维张量（如 shape=(990,1,10,3,3) × (990,1,10,3,1)）逐元素矩阵乘法时，用 np.einsum 替代 np.matmul 可提升约 40% 性能，且无需额外依赖或编译开销。

本文介绍在处理大量高维张量（如 shape=(990,1,10,3,3) × (990,1,10,3,1)）逐元素矩阵乘法时，用 `np.einsum` 替代 `np.matmul` 可提升约 40% 性能，且无需额外依赖或编译开销。

在科学计算与仿真任务中，频繁执行小规模但高维度的矩阵乘法（例如每轮对数千个独立的 3×3 × 3×1 张量进行批处理乘法）极易成为性能瓶颈。如示例所示：对形状为 (990, 1, 10, 3, 3) 的旋转矩阵 precession 和 (990, 1, 10, 3, 1) 的向量 vecMblood 执行 80 万次 np.matmul，单核耗时达 133 秒——即使已启用 float32 和多线程，仍有明显优化空间。

关键洞察在于：np.matmul 在处理带冗余维度（如中间的 1 维）的高维张量时，需进行隐式广播与内存布局适配，带来额外开销；而 np.einsum 允许显式指定求和轴与输出结构，可绕过部分内部调度逻辑，更贴近底层 BLAS 的高效调用路径。

以下为推荐的高性能替代方案：

# 原始低效写法（参考基准）
result = np.matmul(precession, vecMblood)  # shape: (990, 1, 10, 3, 1)

# 推荐优化写法（快约 40%）
result = np.einsum("...ij,...j", precession, vecMblood[..., 0])[..., None]

✅ 原理说明：

"...ij,...j" 表示对最后两个维度执行矩阵-向量乘法（i 为输出行索引，j 为求和轴），... 自动匹配前导批量维度（990, 1, 10）；
vecMblood[..., 0] 将 (990,1,10,3,1) 压缩为 (990,1,10,3)，避免 einsum 对单例维度做冗余广播；
[ ..., None ] 在末尾添加新轴，恢复目标 shape (990,1,10,3,1)，确保下游代码兼容。

? 实测性能对比（基于 numpy 1.26+，Intel i7-11800H，float32）：

macOS UI 自动化技能包 —— 用 Python 控制你的 Mac

让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac，无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启：辅助功能 + 屏幕录制权限

下载

np.matmul:      190 µs ± 10 µs per loop  
np.einsum:      108 µs ± 0.8 µs per loop  → 提速 43%

80 万次运算理论可节省约 65 秒（133 → ~68 秒单核），叠加多线程后收益更显著。

⚠️ 注意事项：

einsum 的性能优势在「中小批量 + 中小矩阵」场景最明显；若单次乘法规模极大（如 >1000×1000），matmul 调用高度优化的 OpenBLAS/Intel MKL 可能反超；
务必提前将输入转为 np.float32（precession.astype(np.float32)），可进一步降低内存带宽压力与计算延迟；
避免在循环内重复调用 einsum 字符串解析——上述表达式已预编译，无运行时解析开销；
若后续需扩展至更高阶张量（如 6D），可沿用相同 ... 模式，保持代码可维护性。

综上，对于仿真、物理引擎、批量微分等典型应用场景，np.einsum 是零依赖、即插即用、效果显著的 np.matmul 升级方案。它不改变语义、不引入新依赖，仅通过更精准的张量代数描述，就撬动了 NumPy 底层的性能潜力。