MMOCR项目训练与测试全指南

MMOCR项目训练与测试全指南

mmocr OpenMMLab Text Detection, Recognition and Understanding Toolbox 项目地址: https://gitcode.com/gh_mirrors/mm/mmocr

前言

MMOCR作为一款优秀的OCR开源工具库，提供了完整的训练和测试流程支持。本文将详细介绍如何在MMOCR框架下进行模型的训练和测试工作，涵盖从单机单卡到分布式集群的各种使用场景。

单GPU训练与测试

训练流程

MMOCR推荐使用GPU进行模型训练，同时也支持CPU训练模式。训练入口脚本为tools/train.py，其基本使用方式如下：

# 基础训练命令格式
CUDA_VISIBLE_DEVICES= python tools/train.py ${CONFIG_FILE} [PY_ARGS]

# 实际使用示例
# 使用CPU训练DBNet模型
CUDA_VISIBLE_DEVICES=-1 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py

# 使用GPU0训练DBNet模型，指定工作目录并开启混合精度
CUDA_VISIBLE_DEVICES=0 python tools/train.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py --work-dir dbnet/ --amp

关键参数说明

| 参数 | 类型 | 说明 | |------|------|------| | config | str | 必需，配置文件路径 | | --work-dir | str | 指定工作目录，用于保存日志和模型 | | --resume | bool | 是否从最近检查点恢复训练 | | --amp | bool | 是否启用自动混合精度训练 | | --auto-scale-lr | bool | 是否自动调整学习率 | | --cfg-options | str | 覆盖配置文件中的部分设置 |

测试流程

测试入口脚本为tools/test.py，使用方式与训练脚本类似：

# 基础测试命令格式
CUDA_VISIBLE_DEVICES= python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [PY_ARGS]

# 实际使用示例
# 使用CPU测试DBNet模型
CUDA_VISIBLE_DEVICES=-1 python tools/test.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth

# 使用GPU0测试DBNet模型
CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/textdet/dbnet/dbnet_resnet50-dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth

关键参数说明

| 参数 | 类型 | 说明 | |------|------|------| | config | str | 必需，配置文件路径 | | checkpoint | str | 必需，待测试模型路径 | | --show | bool | 是否可视化预测结果 | | --show-dir | str | 可视化结果保存路径 | | --wait-time | float | 可视化间隔时间(秒) |

多GPU分布式训练与测试

对于大型模型，分布式训练可以显著提高效率。MMOCR提供了基于MMDistributedDataParallel实现的分布式脚本。

单机多GPU

# 训练示例：使用4个GPU训练DBNet
./tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 4

# 测试示例：使用4个GPU测试DBNet
./tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 4

单机多任务并行

在拥有多GPU的工作站上，可以同时运行多个任务：

# 在GPU0-3上测试DBNet
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 ./tools/dist_test.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth 4

# 在GPU4-7上训练CRNN
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 ./tools/dist_train.sh configs/textrecog/crnn/crnn_academic_dataset.py 4

多机多GPU

在多机环境下训练DBNet：

# 第一台机器
NNODES=2 NODE_RANK=0 PORT=29500 MASTER_ADDR=10.140.0.169 ./tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 2

# 第二台机器
NNODES=2 NODE_RANK=1 PORT=29501 MASTER_ADDR=10.140.0.169 ./tools/dist_train.sh configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py 2

Slurm集群训练与测试

对于使用Slurm管理的计算集群，MMOCR提供了专用脚本：

# 训练示例：在dev分区上申请1个GPU训练DBNet
GPUS=1 GPUS_PER_NODE=1 CPUS_PER_TASK=5 ./tools/slurm_train.sh dev db_r50 configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py work_dir

# 测试示例：在dev分区上申请1个GPU测试DBNet
GPUS=1 GPUS_PER_NODE=1 CPUS_PER_TASK=5 ./tools/slurm_test.sh dev db_r50 configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth work_dir

高级技巧

从检查点恢复训练

# 从最近检查点恢复训练
python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py --resume

混合精度训练

混合精度训练可以显著提升训练速度：

python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py --amp

自动学习率调整

当使用不同的batch_size时，可以自动调整学习率：

python tools/train.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py --auto-scale-lr

结果可视化

MMOCR支持丰富的可视化功能：

# 实时显示可视化结果
python tools/test.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth --show

# 保存可视化结果
python tools/test.py configs/textdet/dbnet/dbnet_r50dcnv2_fpnc_1200e_icdar2015.py dbnet_r50.pth --show-dir ./vis_results

结语

本文全面介绍了MMOCR框架下的训练和测试方法，从基础的单机单卡操作到复杂的分布式训练场景，涵盖了实际应用中可能遇到的各种情况。掌握这些方法将帮助您更高效地使用MMOCR进行OCR相关的研究和开发工作。

mmocr OpenMMLab Text Detection, Recognition and Understanding Toolbox 项目地址: https://gitcode.com/gh_mirrors/mm/mmocr