使用 NumPy 的 np.matmul 可对批量矩阵(M×3×3)与对应向量(M×3×1)执行高效逐片矩阵乘法,避免显式 Python 循环,显著提升计算性能。
使用 numpy 的 `np.matmul` 可对批量矩阵(m×3×3)与对应向量(m×3×1)执行高效逐片矩阵乘法,避免显式 python 循环,显著提升计算性能。
在深度学习、几何变换或批量线性系统求解等场景中,常需对一批形状为 M×3×3 的变换矩阵,分别作用于 M×3×1 的向量组(例如 M 个三维点经 M 个旋转/仿射矩阵变换)。若用 Python 循环逐个调用 np.dot 或 @,不仅代码冗长,更因解释器开销导致性能严重下降。
幸运的是,NumPy 的 np.matmul(等价于 @ 运算符)原生支持广播式批量矩阵乘法:当输入为 (M, 3, 3) 和 (M, 3, 1) 时,它会自动对每个 batch index i 执行 A[i] @ b[i],输出形状为 (M, 3, 1),完全符合需求。
✅ 正确用法示例:
import numpy as np # 构造批量数据:M=4 个 3×3 矩阵,和 4 个 3×1 向量 M = 4 A = np.random.randn(M, 3, 3) # shape: (4, 3, 3) b = np.random.randn(M, 3, 1) # shape: (4, 3, 1) # 高效逐片乘法(无需循环) result = np.matmul(A, b) # shape: (4, 3, 1) print(result.shape) # 输出: (4, 3, 1) # 等价写法(推荐,更简洁) result = A @ b
⚠️ 注意事项:
- 输入向量必须是三维:(M, 3, 1) 而非 (M, 3)。若原始数据是二维(如 (M, 3)),请先用 b[:, :, None] 或 b.reshape(-1, 3, 1) 升维;
- np.dot 不具备相同广播行为(对高维数组按最后两轴做积并求和),易引发维度错误,应优先选用 np.matmul 或 @;
- 若需兼容 PyTorch/TensorFlow,对应操作分别为 torch.bmm 和 tf.linalg.matmul(均要求输入为三维张量)。
总结:借助 np.matmul 的隐式批处理能力,可一行代码完成 M 次矩阵-向量乘法,兼具简洁性与高性能,是科学计算中矩阵逐片运算的标准实践。