训练自定义损失函数的SSD-Mobilenet时出现NaN损失值问题

问题现象 #

Epoch 1/10
1/1839 [..............................] - ETA: 709:15:59 - loss: 62.0176Batch 1: Invalid loss, terminating training
2/1839 [..............................] - ETA: 354:31:19 - loss: inf

排查与解决步骤 #

• 数据层面
  • 检查Pascal VOC标注文件：确保XML里的xmin/ymin/xmax/ymax符合逻辑（xmin<xmax、ymin<ymax，且数值不超过对应图片的宽高），排除标注值为0或超出边界的异常样本。
  • 验证预处理流程：确认图片像素值是否正确归一化（比如缩至0-1区间），标注框坐标是否转换为相对图片尺寸的比例值，避免计算时出现除以0的情况。
• 损失计算逻辑
  • 锚框匹配检查：SSD损失依赖锚框与真实框的匹配，若某批次没有锚框匹配到真实框，会导致损失项出现无穷大。添加日志输出每批次匹配成功的锚框数量，定位异常批次。
  • 损失权重调整：确认分类损失和回归损失的权重设置合理，避免某一项权重过大引发数值溢出。比如回归损失用Smooth L1时，要保证输入的回归目标值无异常大的数值。
• 模型训练参数
  • 权重初始化：检查模型权重初始化方式，若初始权重过大导致前向传播输出异常，可换用Xavier或He初始化，缩小初始化方差。
  • 学习率下调：初始学习率过高会导致损失急剧爆炸，建议将初始学习率降至1e-4或更低，逐步尝试调整。
• 硬件与框架
  • 关闭混合精度：RTX 3060支持混合精度，但如果开启后出现数值不稳定，可改用FP32精度训练。
  • 版本适配：确认TensorFlow版本与参考代码兼容，避免API差异引发计算异常。

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