PyTorch中批量矩阵-矩阵外积计算需求及形状示例

看起来你需要对两个形状为torch.Size([32, 300, 8])的张量，计算每个batch、每个特征维度下的向量外积，最终得到[32, 300, 300, 8]的结果——说白了就是在(batch, feature)这两个维度固定时，让两个长度为300的向量做外积，对吧？

分享几个高效简洁的实现方式，都是PyTorch原生支持的：

方法1：广播机制（最直观，代码最少） #

通过给张量增加维度触发广播，直接对应位置相乘：

import torch

# 示例输入
X1 = torch.randn(32, 300, 8)
X2 = torch.randn(32, 300, 8)

# 给X1加第2维，X2加第1维，相乘后自动广播
output = X1.unsqueeze(1) * X2.unsqueeze(2)
print(output.shape)  # 输出: torch.Size([32, 300, 300, 8])

原理很简单：

• X1.unsqueeze(1)把形状变成[32, 1, 300, 8]
• X2.unsqueeze(2)把形状变成[32, 300, 1, 8]
• 相乘时PyTorch会自动把这两个张量广播到[32, 300, 300, 8]，每个位置的乘积就是对应(batch, i, j, feature)的外积结果。

方法2：爱因斯坦求和（可读性强，适合复杂维度操作） #

用torch.einsum可以清晰定义维度间的运算关系，不用纠结维度扩展：

output = torch.einsum('b i d, b j d -> b i j d', X1, X2)
print(output.shape)  # 输出: torch.Size([32, 300, 300, 8])

这里的符号解释：

• b=batch维度，i/j=向量的两个维度，d=特征维度
• 公式表示：对每个b和d，计算i和j的成对乘积，最终保留所有维度。

方法3：批量矩阵乘法（适合习惯矩阵运算的场景） #

如果更熟悉矩阵乘法，可以先调整维度，用torch.bmm实现：

# 把batch和feature维度合并，变成[32*8, 300, 1]和[32*8, 1, 300]
X1_reshaped = X1.flatten(0, 2).view(-1, 300, 1)
X2_reshaped = X2.flatten(0, 2).view(-1, 1, 300)

# 批量矩阵乘法得到[32*8, 300, 300]
temp = torch.bmm(X1_reshaped, X2_reshaped)

# 恢复原维度结构
output = temp.view(32, 300, 300, 8)
print(output.shape)  # 输出: torch.Size([32, 300, 300, 8])

这种方法和广播的底层逻辑一致，只是换了矩阵乘法的形式表达，适合需要和现有矩阵运算流程整合的场景。

这三种方法的计算结果完全相同，你可以根据自己的代码风格和可读性需求选择～

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