PyTorch中2D与3D张量的matmul运算问题求助

可行方案 #

方案1：使用torch.einsum（推荐） #

einsum能直观定义维度间的运算关系，完全不需要reshape操作：

import torch

a = torch.randn(125, 128)    # Shape: (125, 128)
b = torch.randn(128, 8, 64)  # Shape: (128, 8, 64)

result = torch.einsum('ij,jkl->ikl', a, b)
print(result.shape)  # 输出: torch.Size([125, 8, 64])

这里的'ij,jkl->ikl'含义：

• i对应a的第一维（125），j对应a的第二维（128）
• j对应b的第一维（128），k对应b的第二维（8），l对应b的第三维（64）
• 对重叠的j维度求和，最终得到i-k-l的目标形状。

方案2：广播+维度求和 #

通过给a增加一个维度实现广播，再做元素乘法后求和：

result = (a.unsqueeze(-1) * b).sum(dim=1)
print(result.shape)  # 输出: torch.Size([125, 8, 64])

步骤说明：

1. a.unsqueeze(-1)把a的形状变为(125, 128, 1)，和b的(128, 8, 64)广播后，运算形状扩展为(125, 128, 8, 64)
2. 元素乘法后，在第1维（原128维度）求和，得到目标形状。

原方法报错原因 #

torch.matmul的运算规则限制：当第一个张量是2D（(m, n)），第二个是3D（(n, p, q)）时，matmul会将第二个张量的倒数第二维p视为需要和第一个张量的最后一维n匹配，但这里b的倒数第二维是8，和a的128不匹配，所以报错。

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