PyTorch中批量张量矩阵乘法：3x3xn与3x1xn运算异常求助

问题根源 #

PyTorch的torch.matmul对高维张量遵循固定运算规则：最后两个维度被当作矩阵的行/列维度，其余维度为批量维度。做矩阵乘法时，要求第一个张量的最后一维等于第二个张量的倒数第二维，同时批量维度需支持广播。

你的张量R(3,3,125)最后两维是(3,125)，T(3,1,125)最后两维是(1,125)，完全不满足矩阵乘法的维度要求，因此报错。而torch.mul是逐元素相乘，只会输出和输入同维度的结果，不符合你的需求。

可行解决方案 #

方案1：调整维度顺序，将批量维度前置 #

把批量维度（125）移到第一维，让最后两维符合矩阵乘法的(3,3)和(3,1)结构：

import torch

T = torch.zeros(3,1,125)
R = torch.zeros(3,3,125)

# 重排维度：将batch维度放到最前面
T_reshaped = T.permute(2, 0, 1)  # 从(3,1,125)变为(125,3,1)
R_reshaped = R.permute(2, 0, 1)  # 从(3,3,125)变为(125,3,3)

out = torch.matmul(R_reshaped, T_reshaped)
# 再将维度转回目标格式
out = out.permute(1, 2, 0)
print(out.size())  # 输出: torch.Size([3, 1, 125])

方案2：转置T的矩阵维度，适配matmul规则 #

调整T的后两维，让R的最后一维（125）匹配T的倒数第二维（125）：

import torch

T = torch.zeros(3,1,125)
R = torch.zeros(3,3,125)

# 转置T的第1和第2维度，得到(3,125,1)
T_reshaped = T.transpose(1, 2)
out = torch.matmul(R, T_reshaped)
# 转置回目标维度
out = out.transpose(1, 2)
print(out.size())  # 输出: torch.Size([3, 1, 125])

方案3：使用批量矩阵乘法torch.bmm #

torch.bmm专门处理3维批量矩阵乘法，要求第一维是批量维度，因此同样需要调整维度顺序：

import torch

T = torch.zeros(3,1,125)
R = torch.zeros(3,3,125)

# 调整维度为(batch, row, col)格式
T_reshaped = T.permute(2, 0, 1)
R_reshaped = R.permute(2, 0, 1)

out = torch.bmm(R_reshaped, T_reshaped).permute(1, 2, 0)
print(out.size())  # 输出: torch.Size([3, 1, 125])

内容的提问来源于stack exchange，提问作者Rémi.T