如何高效实现Mx3x3线性变换与Mx3x1向量的MATRIX-WISE乘法？

直接利用数值计算库（比如NumPy或PyTorch）的批量矩阵乘法功能即可，完全不需要手动循环M维度，这类操作底层由优化的BLAS/LAPACK或CUDA内核实现，效率比Python循环高几个数量级。

NumPy 实现示例 #

假设你有两个数组：

• transforms: 形状为 (M, 3, 3) 的线性变换集合
• vectors: 形状为 (M, 3, 1) 的向量集合

直接用 @ 运算符或 np.matmul() 完成批量乘法：

import numpy as np

# 生成示例数据
M = 1000
transforms = np.random.rand(M, 3, 3)
vectors = np.random.rand(M, 3, 1)

# 批量计算变换后的向量
result = transforms @ vectors  # 等价于 np.matmul(transforms, vectors)
# result 的形状为 (M, 3, 1)，与预期一致

PyTorch 实现示例 #

如果用PyTorch（支持GPU加速），逻辑完全一致：

import torch

# 生成示例数据
M = 1000
transforms = torch.rand(M, 3, 3)
vectors = torch.rand(M, 3, 1)

# 批量计算变换后的向量
result = transforms @ vectors  # 等价于 torch.matmul(transforms, vectors)
# result 的形状为 (M, 3, 1)

原理说明 #

这类库的批量矩阵乘法会自动识别批量维度（这里的M维度），对每个索引下的3x3矩阵和3x1向量执行矩阵乘法，全程在底层C/CUDA层面并行计算，完全避开了Python解释器的循环开销。

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