PyTorch自定义Simple_NN运行报错mat1和mat2形状无法相乘怎么解决

错误原因 #

你遇到的是矩阵乘法维度不匹配问题，核心逻辑如下：

• nn.Linear(in_features, out_features) 的运算规则要求输入张量的最后一维必须等于in_features，你初始化模型时传入input_dim = outputs.shape[0] = 2，因此你的Linear层要求输入最后一维为2。
• 你传入的outputs张量形状为(2,6)，在dim=0维度排序后形状保持(2,6)不变，最后一维是6，和Linear要求的2不匹配，因此触发矩阵乘法报错。

解决方法 #

根据你的实际业务需求，可选择以下两种方案： #

方案1：对dim=0方向的2个元素做加权求和 #

该方案适配你当前初始化的长度为2的权重参数，需要调整forward方法中张量的维度，将需要加权的维度放到最后一位：

def forward(self, x, dim = 0):
    sort = torch.sort(x, dim, descending = True)[0]
    # 转置后形状变为(6,2)，匹配Linear层输入要求
    sort = sort.transpose(0, 1)
    sort = self.layer(sort)
    # 可根据需求调整输出形状，比如转置为(1,6)
    return sort

也可以直接用广播乘法实现，逻辑更直观：

def forward(self, x, dim = 0):
    sort = torch.sort(x, dim, descending = True)[0]
    # 权重和排序后张量广播相乘，在dim=0维度求和得到加权结果
    return (self.layer.weight.T * sort).sum(dim=0)

方案2：对dim=1方向的6个元素做加权求和 #

如果你实际需要对每行的6个元素做加权，修改模型初始化参数即可：

# 传入行维度作为Linear层输入维度
model = Simple_NN(input_dim = outputs.shape[1])
# 调用时指定dim=1对每行排序
model(outputs, dim=1)

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