Detectron2微调大模型时出现训练发散（Inf/NaN）问题求助

核心排查方向及解决措施 #

• 预训练数据与自定义数据集分布差异过大
  PublayNet的通用文档布局和你的自定义文档在排版、元素类型、视觉特征上可能存在显著差异。大模型容量更大，对分布偏移的敏感度远高于小模型，训练初期极易因参数更新幅度过大触发数值异常。
  处理方式：

  • 采用分层微调策略：先冻结backbone（如ResNet-X101的前5层、Dit-B的Transformer底层），仅训练检测头部（RPN、FastRCNN分类/回归头），待loss稳定后，再逐步解冻backbone，用头部学习率1/10的速率微调。
  • 针对性增强数据：添加文档专属的数据增强，比如随机旋转（±15°）、缩放（0.8-1.2倍）、局部模糊、文本遮挡、背景扰动等，缩小与预训练数据的分布差距。

• 大模型梯度爆炸/数值溢出
  大模型参数规模大，梯度更容易累积，即使降低了学习率，部分层的梯度仍可能超出浮点数范围，导致参数更新后出现Inf/NaN。
  处理方式：

  • 启用梯度裁剪：在优化器更新前执行torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)，限制梯度的L2范数，避免梯度爆炸。
  • 更换优化器：将SGD替换为AdamW，AdamW的自适应学习率机制对大模型微调更稳定，初始学习率可设置为1e-6（远低于SGD常用的1e-3）。

• 标注数据存在噪声
  小模型对标注错误的鲁棒性较强，但大模型会拟合标注中的噪声，比如边界框超出图像范围、类别ID错误、重复标注等，直接导致loss计算出现NaN。
  处理方式：

  • 校验标注数据集：遍历所有样本，删除/修正边界框坐标超出[0, img_width]、[0, img_height]的样本；检查类别ID是否在模型定义的类别范围内（避免出现超出num_classes的ID）。
  • 加入标签平滑：在分类损失中设置label_smoothing=0.1，降低模型对错误标注的拟合程度。

• 混合精度训练兼容性问题
  Colab默认可能开启自动混合精度（AMP），但部分大模型（如Cascade RCNN、Dit-B）的某些层对混合精度的兼容性较差，容易触发数值溢出。
  处理方式：

  • 关闭自动混合精度：移除训练代码中torch.cuda.amp.GradScaler()相关逻辑，改用全精度训练。
  • 检查损失计算逻辑：排查FastRCNN的分类损失是否存在除以零的情况，回归损失中的锚框计算是否出现面积为0的异常锚框。

• 显存不足导致隐性异常
  尽管使用A100，大模型的参数规模加上batch数据仍可能导致显存碎片化或不足，表现为数值异常而非直接的OOM错误。
  处理方式：

  • 将batch size降至1，同时启用梯度累积：每4-8步执行一次优化器更新，等效于更大的batch size，同时降低单步显存占用。
  • 训练前清理显存：执行torch.cuda.empty_cache()，减少显存碎片影响。

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