PyTorch实现两矩阵逐行matmul的简洁高效方法

完全不需要写显式的Python for循环，靠张量运算内置API就能实现等价逻辑，还能绕开Python循环的性能开销，充分跑满GPU并行能力，常见的可用实现有三种：

• 逐元素乘加实现
  执行代码：scores = (a * b).sum(dim=1)
  逻辑说明：两个同形状张量逐元素相乘后，沿着特征维度（也就是形状里的第1维）做累加求和，每一行的累加结果刚好就是该行两个向量的点积，输出形状直接就是(n1,)，和循环算出来的结果完全等价，是最直观、最不容易搞错维度的写法。
• 批量矩阵乘法实现
  执行代码：scores = torch.bmm(a.unsqueeze(1), b.unsqueeze(2)).squeeze()
  逻辑说明：先给两个张量做维度扩充，把原本(n1, n2)的形状调整为torch.bmm要求的批量输入格式：a扩成(n1, 1, n2)、b扩成(n1, n2, 1)，批量矩阵乘法的计算结果是形状为(n1, 1, 1)的张量，把多余的单维度去掉就是目标输出。这个方法调用的是底层高度优化的批量矩阵算子，在大张量、GPU运行场景下的效率通常最高。
• 爱因斯坦求和实现
  执行代码：scores = torch.einsum('ij,ij->i', a, b)
  逻辑说明：通过einsum的下标规则直接绑定计算逻辑：取两个二维张量同一行、同一列位置的元素相乘，最后对特征维度（j维）求和，只保留行维度（i维）的结果。代码不需要手动调整维度，张量形状适配性很强，就算后续加batch维变成3维张量，只需要改下下标规则就能正常运行。

上面三种方法的计算结果和你写的循环输出的数值精度完全一致，当行数n1比较大时，运行速度比Python层的for循环快一个量级以上。

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