PyTorch训练CNN时报AddmmBackward返回无效梯度错误求解

错误根源分析 #

你遇到的Function AddmmBackward returned an invalid gradient错误核心是全连接层输入维度和卷积输出维度不匹配，同时存在一处变量名笔误，和数据预处理环节无关，属于网络结构和训练逻辑问题：

• 维度不匹配：你参考的CIFAR10官方教程输入图像为32×32尺寸，两次卷积池化后输出尺寸为16×5×5，对应全连接层输入维度为400。但你的数据集预处理后是224×224的单通道图像，经过两层卷积池化后输出尺寸为16×53×53，flatten后长度为44944，和你定义的fc1输入维度完全不符，矩阵乘法维度错误导致反向传播梯度异常。
• 变量名笔误：验证阶段你写了model.eval()，但你定义的模型实例变量名为net，此处会触发未定义变量错误。

修复方案 #

1. 修正全连接层输入维度 #

你可以先打印卷积输出的尺寸确认，再修改fc1的输入参数：
在forward函数的flatten操作前新增打印逻辑：

def forward(self, x):
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    print(x.shape) # 正常输出为[batch_size, 16, 53, 53]
    x = torch.flatten(x, 1)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

根据输出结果修改fc1定义：

# 原写法：self.fc1 = nn.Linear(16 * 5 * 5, 120)
# 修改为：
self.fc1 = nn.Linear(16 * 53 * 53, 120)

2. 修正验证逻辑笔误 #

把model.eval()改为net.eval()，同时建议验证阶段关闭梯度计算节省显存，修改为如下逻辑：

valid_loss = 0.0
net.eval()
with torch.no_grad(): # 关闭梯度计算，大幅降低显存占用
    for data, labels in validloader:
        # 原有验证推理逻辑保持不变

同时要在每个epoch的训练循环前新增net.train()，把模型切回训练模式。

可选优化 #

• 验证集预处理不要加随机旋转、随机翻转这类随机增强操作，避免验证结果波动，建议把valid_transforms改成和test_transforms一致。
• 如果觉得全连接层参数量太大，可以在第二层卷积后新增几层卷积+池化操作，缩小特征图尺寸后再进入全连接层，减少参数量降低过拟合风险。

内容的提问来源于stack exchange，提问作者Naga Saideep Lanka