训练流程

相关源文件

• ppocr/data/imaug/__init__.py
• ppocr/data/imaug/label_ops.py
• ppocr/data/imaug/rec_img_aug.py
• ppocr/losses/__init__.py
• ppocr/losses/basic_loss.py
• ppocr/losses/cls_loss.py
• ppocr/losses/combined_loss.py
• ppocr/losses/distillation_loss.py
• ppocr/losses/rec_ctc_loss.py
• ppocr/metrics/distillation_metric.py
• ppocr/modeling/architectures/base_model.py
• ppocr/modeling/architectures/distillation_model.py
• ppocr/modeling/backbones/__init__.py
• ppocr/modeling/backbones/det_mobilenet_v3.py
• ppocr/modeling/backbones/rec_mobilenet_v3.py
• ppocr/modeling/heads/__init__.py
• ppocr/modeling/heads/rec_ctc_head.py
• ppocr/postprocess/__init__.py
• ppocr/postprocess/rec_postprocess.py
• ppocr/utils/save_load.py
• tools/eval.py
• tools/export_model.py
• tools/infer/predict_cls.py
• tools/infer/predict_det.py
• tools/infer/predict_rec.py
• tools/infer/predict_sr.py
• tools/infer/predict_system.py
• tools/infer/utility.py
• tools/infer_det.py
• tools/infer_rec.py
• tools/program.py
• tools/train.py

本文档涵盖了PaddleOCR的训练基础设施，包括主训练循环、数据管道、模型构建和优化系统。训练管道支持多种模型类型（检测、识别、分类、端到端、KIE），具有统一的配置和执行模式。

有关模型架构的信息，请参阅3.1、3.2和3.3。有关数据处理的详细信息，请参阅3.5。有关训练后模型导出，请参阅3.6。

训练系统架构

训练系统由多个相互关联的组件组成，它们协同工作以训练OCR模型。

来源：tools/train.py46-245 tools/program.py184-642

训练入口点

主要的训练入口点是tools/train.py，它负责协调整个训练过程。

main()函数执行以下关键步骤：

1. DataLoader构造：使用build_dataloader(config, "Train", device, logger, seed)构建训练和验证数据加载器。
2. 后处理设置：创建用于评估的后处理类。
3. 模型构建：构建模型架构，并处理识别模型的字符词汇大小。
4. 损失函数：根据算法实例化相应的损失函数。
5. 优化器设置：创建优化器和学习率调度器。
6. 模型加载：加载预训练权重或检查点权重。
7. 训练执行：将实际训练循环委托给program.train()。

来源：tools/train.py46-245 tools/train.py53-72 tools/train.py139-162

核心训练循环

核心训练逻辑位于tools/program.py:train()中，该函数实现了带有集成评估的主训练循环。

训练循环架构

主要训练特性

训练循环支持以下高级特性：

1. 混合精度训练：启用时使用paddle.amp.auto_cast()进行FP16训练。
2. 多种模型类型：支持检测、识别、表格、KIE和公式识别模型。
3. 动态批量处理：根据模型算法要求调整输入处理。
4. 集成评估：训练期间定期评估，并保存最佳模型。
5. 分布式训练：通过paddle.DataParallel支持多GPU。

训练循环以专门的输入处理方式处理不同的模型类型。

来源：tools/program.py184-642 tools/program.py330-372 tools/program.py248-277

数据管道集成

训练管道与数据处理系统紧密集成。

数据管道支持：

• 多种数据集类型：简单的基于文件的数据集、LMDB和专用数据集。
• 可配置的转换：图像增强、缩放和标签编码。
• 动态重采样：某些数据集支持基于epoch的重采样以增加数据多样性。

来源：tools/train.py53-67 tools/program.py302-305

模型架构系统

训练系统采用模块化模型构建方法。

模型构建过程

字符词汇处理

对于识别模型，系统会根据字符词汇自动配置输出通道。

1. 字符计数：从后处理类的字符集中提取。
2. 算法调整：不同算法需要不同的通道数（例如，SAR需要+2用于特殊字符）。
3. 多头支持：处理具有多个预测头的模型（CTC + Attention）。
4. 蒸馏模型：相应地配置教师模型和学生模型。

来源：tools/train.py74-134 tools/train.py139-151

损失和优化框架

训练管道支持一个全面的损失和优化系统。

损失函数架构

优化器配置

优化器系统支持：

• 多种优化器：Adam、AdamW、SGD，支持可配置参数。
• 学习率调度：余弦退火、线性预热、分段衰减。
• 参数组：不同模型组件采用不同的学习率。
• 梯度缩放：用于混合精度训练。

来源：tools/train.py153-162 ppocr/losses/__init__.py66-121

模型持久化和检查点

训练系统提供全面的模型保存和加载功能。

检查点管理

加载机制

系统支持多种加载场景：

1. 恢复检查点：恢复完整的训练状态，包括优化器。
2. 预训练权重：仅加载模型参数以进行微调。
3. 部分加载：加载形状/名称匹配的兼容层。

load_model()函数处理：

• 模型类型检测：KIE模型与标准模型。
• 形状兼容性：自动形状匹配和警告。
• 数据类型转换：根据需要进行FP16到FP32的转换。
• 状态恢复：训练轮次、最佳指标和优化器状态。

来源：ppocr/utils/save_load.py66-169 ppocr/utils/save_load.py214-284 tools/program.py522-552

训练配置系统

训练管道使用分层配置系统。

配置加载过程

配置系统支持：

• YAML配置文件：针对不同实验的结构化配置文件。
• 命令行覆盖：通过-o标志动态修改参数。
• 环境变量：设备和性能调优标志。
• 分层合并：通过点表示法进行嵌套配置更新。

来源：tools/program.py44-116 tools/program.py809-875

评估集成

训练管道在训练期间无缝集成评估。

评估调度

评估系统：

1. 定期评估：每N个训练步执行一次评估。
2. 模型类型处理：根据不同的模型架构调整评估。
3. 指标跟踪：支持各种指标（准确率、HMEAN、BLEU等）。
4. 最佳模型保存：自动保存达到最佳性能的模型。
5. 提前停止：可根据指标配置训练终止。

eval()函数通过相同的模型架构处理验证数据，并为每种模型类型进行适当的后处理。

来源： tools/program.py485-573 tools/program.py645-754

分布式训练支持

训练流程支持跨多个GPU进行分布式训练

分布式训练特性

• 多GPU支持：自动数据并行训练设置
• 同步BatchNorm：跨GPU可选的同步批标准化
• Rank特定操作：仅rank 0执行保存和日志记录操作
• 梯度同步：跨设备的自动梯度同步

来源： tools/train.py48-49 tools/train.py219-223 tools/program.py576-604