本文介绍如何使用 NumPy 的 np.matmul 实现批量矩阵-向量乘法，避免显式 Python 循环，在保持语义清晰的同时显著提升计算效率。适用于深度学习、几何变换等需对 M 组 3×3 矩阵分别作用于对应 3 维向量的场景。
本文介绍如何使用 numpy 的 `np.matmul` 实现批量矩阵-向量乘法，避免显式 python 循环，在保持语义清晰的同时显著提升计算效率。适用于深度学习、几何变换等需对 m 组 3×3 矩阵分别作用于对应 3 维向量的场景。
在科学计算和机器学习中，常需对一批形状为 (M, 3, 3) 的变换矩阵（如旋转矩阵、仿射子块）分别作用于对应的一批三维向量 (M, 3, 1)，得到输出 (M, 3, 1)。若用 Python 循环逐个调用 np.dot 或 @，不仅代码冗长，还会因解释器开销严重拖慢性能。
幸运的是，NumPy 的 np.matmul（或等价的 @ 运算符）原生支持广播式批量矩阵乘法：当输入为 (M, 3, 3) 和 (M, 3, 1) 时，它会自动对每个 batch index i 执行 A[i] @ b[i]，并堆叠结果，最终返回 (M, 3, 1) —— 完全契合需求，且底层由高度优化的 BLAS 实现，速度远超手动循环。
✅ 正确用法示例：
import numpy as np
# 构造 M=4 批数据：4 个 3x3 矩阵，4 个 3x1 向量
A = np.random.randn(4, 3, 3) # shape: (4, 3, 3)
b = np.random.randn(4, 3, 1) # shape: (4, 3, 1)
# 批量矩阵-向量乘法（无需循环！）
result = np.matmul(A, b) # shape: (4, 3, 1)
# 等价写法：result = A @ b
print("Input A shape:", A.shape)
print("Input b shape:", b.shape)
print("Output shape :", result.shape)
⚠️ 注意事项：
- b 必须是三维数组，即 (M, 3, 1) 而非 (M, 3)；若原始向量是二维 (M, 3)，请先用 b[:, :, None] 或 b.reshape(-1, 3, 1) 升维；
- np.matmul 对高维输入按最后两轴进行矩阵乘，其余前导维度（此处为 M）自动广播匹配，因此 A 和 b 的第 0 维（batch 维）必须严格相等；
- 避免使用 np.dot：它在多维情形下行为不同（会对倒数第二轴求和），易引发意外结果；
- 若后续需展平结果，可用 result.squeeze(-1) 得到 (M, 3)。
总结：np.matmul(A, b) 是解决 “M 个 3×3 矩阵分别乘 M 个 3 维列向量” 这一典型任务的标准、高效、可读性强的方案。掌握其广播规则，可大幅简化张量操作逻辑，并为 PyTorch/TensorFlow 中的类似操作打下坚实基础。
