Numpy如何无循环向量化批量计算Ai.T @ Bi @ Ci矩阵运算

你需要的是批次维度对齐的逐组张量乘法，不需要遍历首维度，有两种高效的向量化实现方式：

• 方案1：使用np.einsum（最直观，性能最优）
  直接按张量维度匹配写求和规则即可，代码如下：

  res = np.einsum('ij,ijk,ik->i', A, B, C)
  

  维度说明：

  • i对应批次维度
  • j对应A的特征维度、B的第一个特征维度
  • k对应B的第二个特征维度、C的特征维度
    最终输出维度为(批次数量,)，和你循环计算的结果完全一致。

• 方案2：使用广播机制配合@运算符
  给A、C扩充维度适配批量矩阵乘法规则即可，代码如下：

  res = (A[:, None, :] @ B @ C[:, :, None]).squeeze(axis=(-2, -1))
  

  维度变换说明：

  • 把形状为(N,2)的A扩充为(N,1,2)，形状为(N,2)的C扩充为(N,2,1)
  • 批量矩阵乘后得到形状为(N,1,1)的结果，调用squeeze去掉多余维度即可得到形状为(N,)的输出。

原有方法问题说明 #

你之前使用tensordot得到的是全组合计算结果，相当于对所有批次的A、B做了交叉计算，仅对角线值有效，时间复杂度为O(N²)，在数十万批次场景下会产生巨量无效计算，完全不适用，不建议采用取对角线的偷懒方案。

结果验证 #

你可以用循环方式计算小批量结果对比，验证一致性：

# 循环实现（仅用于验证，不要用于大批量计算）
res_loop = np.array([a.T @ b @ c for a,b,c in zip(A,B,C)])
print(np.allclose(res, res_loop)) # 输出为True则验证通过

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