如何解决MinneApple项目Faster RCNN训练时‘Loss is nan, stopping training’错误？

问题描述 #

我在运行用于水果计数与检测的MinneApple项目时，多次执行训练命令：

python train_rcnn.py --data_path data --model frcnn --epochs 50 --output-dir data/train/train_results

训练启动后很快出现Loss变为NaN并终止训练的错误，终端输出如下：

Start training
Epoch: [0]  [ 0/20]  eta: 0:04:09  lr: 0.001072  loss: 3.8136 (3.8136)  loss_classifier: 0.9121 (0.9121)  loss_box_reg: 0.6673 (0.6673)  loss_objectness: 2.0488 (2.0488)  loss_rpn_box_reg: 0.1854 (0.1854)  time: 12.4626  data: 4.8800  max mem: 3233
Loss is nan, stopping training
{'loss_classifier': tensor(nan, device='cuda:0', grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(nan, device='cuda:0', grad_fn=<DivBackward0>), 'loss_objectness': tensor(nan, device='cuda:0', grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(nan, device='cuda:0', grad_fn=<DivBackward0>)}

An exception has occurred, use %tb to see the full traceback.

完整回溯信息：

---------------------------------------------------------------------------

SystemExit                                Traceback (most recent call last)

<ipython-input-11-d923a94e9af1> in <module>
      4 start_time = time.time()
      5 for epoch in range(epochs):
----> 6     train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq)
      7     lr_scheduler.step()
      8 

/content/drive/MyDrive/universidad/t3/MinneApple/utility/engine.py in train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq)
     41             print("Loss is {}, stopping training".format(loss_value))
     42             print(loss_dict_reduced)
---> 43             sys.exit(1)
     44 
     45         optimizer.zero_grad()

SystemExit: 1

解决方法 #

针对该问题，可按以下步骤逐一排查修复：

• 降低初始学习率
  当前初始学习率0.001072可能过高，引发梯度爆炸。尝试将学习率降至1e-4或5e-5，可通过修改训练命令的参数（若脚本支持），或直接调整代码中优化器的初始化逻辑。

• 检查数据集标注与完整性
  确认图像与掩码数据是否存在异常：

  • 排查边界框标注是否合规：无负数坐标、宽高为0或超出图像尺寸的情况
  • 验证类别ID是否与模型定义完全匹配，无缺失或无效ID
  • 检查是否存在损坏的图像文件（无法正常读取的图片会导致数据加载异常，进而引发Loss异常）

• 添加梯度裁剪
  在训练循环中加入梯度裁剪，限制梯度的最大范数以避免梯度爆炸。在utility/engine.py的train_one_epoch函数中，反向传播后、优化器更新前添加以下代码：

  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
  

  可根据实际情况调整max_norm的值（如0.5、2.0）。

• 排查数据预处理流程
  检查数据增强与归一化步骤是否存在数值不稳定的操作：

  • 避免出现除以0的情况（如归一化时标准差是否为0）
  • 确认图像归一化的均值、方差是否与数据集匹配，防止数值溢出

• 验证模型初始化与权重
  若使用预训练权重，确认权重文件完整且与模型结构兼容；若从零训练，检查参数初始化方式是否合理，可尝试更换为Xavier或He初始化。

• 关闭混合精度训练（若启用）
  如果脚本中使用了混合精度训练，暂时关闭该功能，避免浮点精度不足导致NaN出现。

内容的提问来源于stack exchange，提问作者Enrique José Carvallo Rioseco