如何在Numpy中无需循环实现类似torch.bmm的批量矩阵乘法功能

有两种常用的无循环实现方式，效果和你写的循环逻辑完全一致： #

方法1：维度扩展后矩阵乘法（和PyTorch torch.bmm逻辑对齐） #

只需要给二维的b在最后扩展一个单维度，变成三维张量适配批量矩阵乘法规则，运算后压缩多余维度即可：

import numpy as np

a = np.ones((5, 5, 5))
b = np.random.randint(0, 10, (5, 5))

# 扩展b维度为 (5,5,1)，运算后压缩最后一维
c = (a @ b[..., np.newaxis]).squeeze(-1)
# c.shape 输出为 (5, 5)

方法2：使用np.einsum显式指定维度运算 #

写法更简洁直观，直接指定参与运算的维度和求和规则：

c = np.einsum('ijk,ik->ij', a, b)
# c.shape 输出为 (5, 5)

你可以用以下代码验证结果和循环实现完全一致：

# 原有循环实现
c_loop = []
for i in range(5):
    c_loop.append(a[i] @ b[i].T)
c_loop = np.concatenate([d[:, None] for d in c_loop], axis=1).T

# 验证结果相等
print(np.allclose(c, c_loop)) # 输出 True

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