PyTorch训练时AddmmBackward梯度形状不匹配报错如何解决

错误根因分析 #

这个报错本质是全连接层输入维度不匹配导致的，存在3个核心问题：

• 你的单张输入张量为torch.Size([3, 48, 48])，展平后总元素数为3*48*48=6912，你误将三通道输入当成单通道计算，错写为2304（单通道48*48的尺寸）
• 特征传入全连接层前没有做张量展平操作，主干网络输出的多维张量直接喂给全连接层，就会出现形状不匹配
• 激活函数和损失函数的搭配逻辑错误，多层线性层中间未加激活，且损失函数和输出层逻辑混用

解决方案 #

第一步：修正全连接层结构 #

调整全连接层定义，增加展平层、修正输入维度、补全中间激活、删除冗余输出层：

model.fc = nn.Sequential(
    # 展平层：将(batch, 通道数, 高, 宽)的4维张量转为(batch, 总特征数)的2维张量
    nn.Flatten(),
    # 修正第一层输入维度为三通道展平后的总特征数6912
    nn.Linear(6912, 100),
    nn.ReLU(),
    nn.Linear(100, 50),
    nn.ReLU(),
    nn.Linear(50, 20),
    nn.ReLU(),
    nn.Linear(20, 3)
    # 用CrossEntropyLoss不需要额外加LogSoftmax，该损失函数内置了Softmax+NLLLoss的计算逻辑
)

第二步：预验证维度匹配 #

训练前加一行测试代码，确认模型输出形状符合预期，避免训练中途报错：

test_input = torch.randn(2, 3, 48, 48).to(device)
test_output = model(test_input)
# 正常输出应为torch.Size([2, 3])，和你的分类数匹配
print(test_output.shape)

可选调整：损失函数适配 #

如果你要保留原代码里的nn.LogSoftmax(dim=1)输出层，就把损失函数换回nn.NLLLoss()即可，两类搭配二选一不要混用。

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