open-mmlab/mmocr 环境搭建、推理和训练入门教程【一】

news2024/9/30 15:29:29

文章目录

- 博文基础信息
- Linux 搭建 open-mmlab/mmocr 运行环境
- - - 准备数据集
    - 准备必要的预训练模型
    - 推理
    - 训练
    - 测试
    - 可视化输出
- 📙 预祝各位前途似锦、可摘星辰

博文基础信息

https://mmocr.readthedocs.io/zh_CN/dev-1.x/get_started/quick_run.html
显卡，11G 1080Ti
本博文只是跟着官方文档进行一个推理、训练、测试代码试运行
涉及代码修改的地方，请参考官方文档

Linux 搭建 open-mmlab/mmocr 运行环境

conda create -n open-mmlab python=3.8 pytorch=1.10 cudatoolkit=11.3 torchvision -c pytorch -y
conda activate open-mmlab
pip3 install openmim
git clone https://github.com/open-mmlab/mmocr.git
cd mmocr
mim install -e .

准备数据集

wget https://download.openmmlab.com/mmocr/data/icdar2015/mini_icdar2015.tar.gz
mkdir -p data/
tar xzvf mini_icdar2015.tar.gz -C data/

# 我这里选择 把 mini_icdar2015  重命名为 icdar2015
cd  data
mv mini_icdar2015 icdar2015

准备必要的预训练模型

如果服务器网络不通畅，可以复制链接，自行到浏览器下载之后，再复制到 /home/用户名/.cache/torch/hub/checkpoints 目录下

# 下载链接
https://download.pytorch.org/models/resnet18-f37072fd.pth

0-6

推理

0-7

推理命令如下

# 推理会自动下载 DBNet  和 CRNN 模型
python tools/infer.py demo/demo_text_ocr.jpg --det DBNet --rec CRNN --show --print-result

测试GPU 占用如下

1-0

训练

修改配置文件

vi configs/textdet/dbnet/dbnet_resnet18_fpnc_1200e_icdar2015.py

添加如下代码

# 每 10 个 epoch 储存一次权重，且只保留最后一个权重
default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=10,
        max_keep_ckpts=1,
    ))
# 设置最大 epoch 数为 400，每 10 个 epoch 运行一次验证
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=400, val_interval=10)
# 令学习率为常量，即不进行学习率衰减
param_scheduler = [dict(type='ConstantLR', factor=1.0),]

1-00

单卡训练命令如下

CUDA_VISIBLE_DEVICES=0 python tools/train.py configs/textdet/dbnet/dbnet_resnet18_fpnc_1200e_icdar2015.py

单卡训练GPU 占用如下

1-1
1-2

测试

对刚刚测试生成的模型进行测试

CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/textdet/dbnet/dbnet_resnet18_fpnc_1200e_icdar2015.py work_dirs/dbnet_resnet18_fpnc_1200e_icdar2015/epoch_400.pth

data_root='icdar2015_textdet_data_root', 时输出如下

1-9

data_root='tests/data/det_toy_dataset', 时输出如下

1-3

可视化输出

可视化结果保存在 imgs/ 目录下

python tools/test.py configs/textdet/dbnet/dbnet_resnet18_fpnc_1200e_icdar2015.py work_dirs/dbnet_resnet18_fpnc_1200e_icdar2015/epoch_400.pth --show-dir imgs/