PyTorch CNN训练时仅输出单一结果的问题排查

问题现象 #

训练255×255 RGB图像的5分类CNN时，出现以下异常：

• 损失函数平滑下降，但模型批量输出单一类别（先全输出0，最终仅输出类别2）
• 类别2占比约50%，但模型完全放弃学习其他类别，输出缺乏多样性

典型训练输出（批量大小10） #

LABELS                                 OUTPUT                                 CORRECT
tensor([2, 0, 2, 2, 2, 0, 2, 2, 2, 4]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([2, 2, 2, 2, 3, 4, 1, 2, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 0 / 10
tensor([2, 2, 2, 0, 2, 4, 3, 1, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 1 / 10
tensor([3, 4, 2, 2, 0, 4, 4, 3, 2, 0]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([1, 2, 2, 4, 2, 0, 1, 0, 0, 0]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 4 / 10
tensor([2, 2, 2, 3, 2, 0, 0, 1, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([1, 1, 0, 1, 2, 2, 1, 1, 0, 1]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([0, 2, 1, 3, 3, 2, 1, 0, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([2, 3, 2, 2, 3, 1, 0, 1, 0, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 2 / 10
tensor([3, 2, 3, 1, 1, 2, 0, 4, 2, 2]) tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) 1 / 10
tensor([2, 1, 0, 3, 1, 2, 2, 1, 2, 0]) tensor([2, 2, 2, 2, 2, 0, 2, 2, 0, 2]) 2 / 10
tensor([3, 0, 2, 1, 3, 1, 2, 4, 2, 2]) tensor([2, 2, 2, 2, 2, 2, 2, 2, 2, 2]) 4 / 10
tensor([2, 2, 1, 2, 1, 1, 1, 4, 3, 2]) tensor([2, 2, 2, 2, 2, 2, 2, 2, 2, 2]) 4 / 10

# Remaining predictions are always [2, 2, 2...]
# Loss function is not shown, but it declines smoothly and looks well behaved

问题根源与修复方案 #

1. 训练循环的致命错误 #

训练循环中未清空梯度：

#optimizer.zero_grad()  # 这行被注释了！

• 梯度会在每次迭代中累积，导致模型参数更新方向彻底混乱，最终只能收敛到最简单的单一类别预测。
• 修复：取消注释并将其放在迭代最开始处：

  for X, y in batches:   
      model.train()  
      optimizer.zero_grad()  # 必须放在参数更新流程最前面
      
      pred = model(X)
      loss = loss_fn(pred, y)
      
      loss.backward()
      optimizer.step()
  

2. 数据加载与类别不平衡问题 #

• 未开启数据洗牌：full_dataloader = DataLoader(dataset, batch_size=10, shuffle=False)
  • 模型会连续学习大量类别2的样本，直接偏向该类别，无法学习其他类别的特征。
  • 修复：设置shuffle=True
• 类别不平衡未处理：类别2占比50%，属于严重不平衡，模型会优先学习占比高的类别以降低损失。
  • 修复：给损失函数添加类别权重，例如：

    class_weights = torch.tensor([1.0, 1.0, 0.5, 1.0, 1.0])  # 降低类别2的权重
    loss_fn = nn.CrossEntropyLoss(weight=class_weights)
    

  • 也可采用过采样/欠采样、Focal Loss等方法。

3. 优化器设置问题 #

• SGD学习率lr=0.001可能过低，加上动量0.9，在梯度累积的情况下，参数更新异常缓慢或方向错误。
• 修复：尝试提高SGD学习率到0.01，或改用Adam优化器：

  optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  

4. 网络结构验证 #

• 输入是RGB图像（3通道），conv1的输入通道数n_layers=3设置正确，需确认输入张量X的形状为[batch_size, 3, 256, 256]。
• 全连接层维度计算正确：经过3次卷积+池化后，特征图尺寸为28×28，对应16*28*28的输入维度无误。

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