自构建数据集loss出现NaN的问题

[Laki-Cat]

大佬您好，我在使用pytorch移植版ppocrv4 （https://github.com/WenmuZhou/PytorchOCR）进行训练时遇到了一些问题。我使用的是自定义数据集，数据是通过将单字数据合成为文本行的方式构建的。在训练过程中，模型的 loss 出现了 NaN 的情况。而在此之前，我使用 TTF 字体文件合成了大量数据进行训练，并没有出现类似的问题。

下面的图片选自我合成的数据集

我的配置文件是
Global:
device: gpu
epoch_num: 200
log_smooth_window: 20
print_batch_step: 10
output_dir: ./ch_PP-OCRv4_rec_train/student.pth
eval_epoch_step: [0, 1]
cal_metric_during_train: true
pretrained_model:
checkpoints:
use_tensorboard: true
infer_mode: false
infer_img: doc/imgs_words/ch/word_1.jpg
character_dict_path: &character_dict_path torchocr/utils/ppocr_keys_v1.txt
max_text_length: &max_text_length 54
use_space_char: &use_space_char true

Export:
export_dir:
export_shape: [ 1, 3, 48, 1000 ]
dynamic_axes: [ 0, 2, 3 ]

Optimizer:
name: Adam
lr: 0.001
weight_decay: 3.0e-05

LRScheduler:
name: CosineAnnealingLR
warmup_epoch: 5

Architecture:
model_type: rec
algorithm: SVTR_HGNet
Transform:
Backbone:
name: PPHGNet_small
Head:
name: MultiHead
head_list:
- CTCHead:
Neck:
name: svtr
dims: 120
depth: 2
hidden_dims: 120
kernel_size: [1, 3]
use_guide: True
- NRTRHead:
nrtr_dim: 384
max_text_length: *max_text_length

Loss:
name: MultiLoss
loss_config_list:
- CTCLoss:
- NRTRLoss:

PostProcess:
name: CTCLabelDecode
character_dict_path: *character_dict_path
use_space_char: *use_space_char

Metric:
name: RecMetric
main_indicator: acc

Train:
dataset:
name: MultiScaleDataSet
ds_width: false
data_dir: ./PytorchOCR/datasets/val.txt
ext_op_transform_idx: 1
label_file_list:
- /root/lanyun-tmp/PytorchOCR/datasets/datasets.txt
transforms:
- DecodeImage:
img_mode: BGR
channel_first: false
- RecConAug:
prob: 0.5
ext_data_num: 2
image_shape: [48, 1000, 3]
max_text_length: *max_text_length
- RecAug:
- MultiLabelEncode:
gtc_encode: NRTRLabelEncode
- KeepKeys:
keep_keys:
- image
- label_ctc
- label_gtc
- length
- valid_ratio
sampler:
name: MultiScaleSampler
scales: [[1000, 32], [1000, 48], [1000, 64]]
first_bs: &bs 89
fix_bs: false
divided_factor: [8, 16] # w, h
is_training: True
loader:
shuffle: true
batch_size_per_card: *bs
drop_last: true
num_workers: 8
Eval:
dataset:
name: SimpleDataSet
data_dir: /root/lanyun-tmp/PytorchOCR/datasets/
label_file_list:
- ./PytorchOCR/datasets/val.txt
transforms:
- DecodeImage:
img_mode: BGR
channel_first: false
- MultiLabelEncode:
gtc_encode: NRTRLabelEncode
- RecResizeImg:
image_shape: [3, 48, 1400]
- KeepKeys:
keep_keys:
- image
- label_ctc
- label_gtc
- length
- valid_ratio
loader:
shuffle: false
drop_last: false
batch_size_per_card: 100
num_workers: 4

[GreatV]

出现 loss 为 NaN 的问题一般可能由以下几类原因导致。针对你自构建数据集训练 PytorchOCR v4 过程中的 NaN 问题，结合你的配置文件和描述信息，下面是详细的排查建议：

一、数据本身的问题

1. 文本标签异常或不匹配字符集：

   • 请检查数据集中标签内容是否存在非法字符，是否完全在 character_dict_path 指定的字典范围内。
   • character_dict_path 为 torchocr/utils/ppocr_keys_v1.txt，不能有遗漏或超范围字符。
   • 多数场景下 label 中含有字典外字符会导致编码出错，进而模型在前向传播过程中出现除零等异常而导致 loss 为 NaN。

2. 标签为空：

   • 如果有样本对应的文本为空字符串 ""，NRTRLoss 和 CTCLoss 可能无法处理该 Case。
   • 建议在数据加载和编码阶段加入空标签过滤。

3. 图像尺寸或内容异常：

   • 图像必须满足 Transform 所需尺寸且图像内容不可为纯白/黑等异常干扰图。
   • 建议人工检查几个 batch 的图像数据，是否存在长宽为 0 或填充后尺寸超出模型支持的大小。

4. 图像和标签个数不一致：

   • 查看 /root/lanyun-tmp/PytorchOCR/datasets/datasets.txt 中每行路径和标签是否齐全且一一对应。

二、模型/配置问题

1. 多 Head 的 MultiLoss 配置是否合理：

   • 你配置文件中使用的是 MultiHead（CTCHead + NRTRHead），对应的 MultiLoss（CTCLoss + NRTRLoss），请确认两个分支输入满足要求。
   • CTCLoss 对 GT 中不能有重复字符，且预测序列合理长度才不会溢出。
   • 如果预测序列太短、输入图像太小，也可能造成 log(0) 导致 NaN。
   • 尝试单独开启一个 head（如只用 CTCHead）排查是哪一个 head 导致的 loss nan。

2. 输入图像尺寸设置不当：

   • 图像变换后尺寸为 [48, 1000, 3]，而 NRTRHead 对图像宽度分辨率比较敏感，输入图像过长可能造成 Encoder 阶段浮点上下溢，建议先减小最大图像宽度试试，如将 image_shape 设为 [48, 320, 3]。

三、训练超参问题

1. 学习率过高：

   • Adam lr=0.001 对于 NRTR 结构可能过高，建议尝试调小如 1e-4。
   • 尤其如果你迁移了预训练模型，加载的权重和新数据不匹配易收敛不稳定。

2. Batch Size 太大：

   • 多 Head 模型本身占用显存更大，bs=89 可能使部分 batch 在 forward 时显存耗尽造成 NaN，建议适当减小 batch_size_per_card。

3. Mixed precision 或 cudnn 使用异常（针对 GPU）：

   • 如果你使用了 amp、fp16、或者 cudnn benchmark=True 的策略，建议在 debug 期间关闭以排除 Float overflow 问题。

四、调试建议

1. 在 loss 为 NaN 前的几个 batch 打印 loss、logits、label 等，是否已出现数值异常（如 loss 极大、logits 为 inf）
2. 在代码中开启 torch.autograd.detect_anomaly() 进行异常检查。
3. 逐层添加 print 输出，观察是哪个 Head 的输出为 NaN，定位是 CTC 还是 NRTR 路径。

五、建议的排查操作步骤

按以下顺序尝试修复：

• ✅ 核查 datasets.txt 标签是否为空或含非法字符；
• ✅ 在 MultiLabelEncode（gtc_encode: NRTRLabelEncode）中添加对 label 的长度和合法性检查；
• ✅ 将数据增强 RecConAug 临时关掉试试；
• ✅ 降低学习率为 0.0001；
• ✅ 减少 batch_size，从 89 改为 16 或更小；
• ✅ 只用其中一个 head（如先只用 CTCLoss 和 CTCHead）排查是哪一分支造成的 NaN；
• ✅ 打开 debug 模式或 设置 torch.autograd.detect_anomaly() 帮助定位。

六、补充建议

如果怀疑是 NRTRHead 导致的，也可以考虑临时移除该 Head，仅使用 CTC 通路训练模型，确认是否是 NRTRLoss 计算导致 NaN。

如果需要测试字符集是否匹配，可参考 ppocr/utils/char_dict_filter.py 工具自动筛选非字典内的字符。

如果仍有问题可开启 GitHub Discussions 或 Issue 详细贴出数据样例和加载逻辑。

祝你解决问题！如有进一步报错细节欢迎继续补充。

---

Response generated by 🤖 feifei-bot | chatgpt-4o-latest

⚠️ This is an AI-generated response. Please verify any critical information.