Skip to main content

W&B对您的OCR模型进行训练和调试

在PaddleOCR中使用W&B集成跟踪训练和评估指标以及模型检查点的简要教程
Created on May 31|Last edited on May 31
本报告是作者Manan Goel所写的"Train and Debug Your OCR Models with PaddleOCR and W&B"的翻译



目录






PaddleOCR介绍

PaddleOCR的目标是创建领先的多语言实用OCR工具,帮助用户训练出更好的模型,并通过 PaddlePaddle将其应用于实践。库中的W&B集成可帮助您在训练期间跟踪训练和验证集指标,此外还提供带有适当元数据的检查点。
同时,如果您愿意,我们还提供包含工作代码的Colab,其在OCR流程的文本检测模块中使用了PaddleOCR

在本教程中,我们将在 ICDAR2015数据集(其中包含一个1000张图像训练集和一个使用可穿戴相机获得的500张图像测试集)上训练一个 MobileNetV3主干。



进行设置

安装W&B SDK

首先,我们来安装并登录 W&B 账户:
pip install wandb
wandb login

安装PaddleOCR

接下来我们安装 PaddlePaddle:
pip install paddlepaddle-gpu pyclipper attrdict -qqq
然后是克隆PaddleOCR GitHub存储库,其将被用于安装程序包并获取用于训练预实施模型的训练脚本:
git clone https://github.com/PaddlePaddle/PaddleOCR
cd PaddleOCR
pip install -e .
很棒!既然我们已经完成了W&B 和PaddleOCR的安装工作,接下来就可以去设置数据集并训练文本检测模型了。

如何下载ICDAR2015数据集?

我们将使用这里面的ICDAR2015数据集。为了便于使用,我们已将数据记录为W&B文件:
import wandb
api = wandb.Api()
artifact = api.artifact("manan-goel/icdar2015/icdar2015-dataset:latest")
artifact.download(root="./train_data/icdar2015")

下载已进行预训练的MobileNetV3模型

在本教程中,我们将使用预训练的MobileNetV3模型作为文本检测模型的主干。模型权重将从图像模型的PaddlePaddle库中获取:
wget -P ./pretrain_models/ https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV3_large_x0_5_pretrained.pdparams

训练模型

为了在训练流程中使用W&B自动开始实验跟踪,可将以下代码片段添加到yaml配置文件中,然后将其用作训练脚本的输入:
wandb:
project: CoolOCR
entity: my_team
name: MyOCRModel
所有用作输入的参数都需添加至 wandb.init,此外也可在yaml文件中的wandb标题下进行添加。本教程中用于实验的配置文件可在这里找到。在文件底部添加上述代码行可激活W&B记录器。
如需使用此yaml文��来训练模型,请在PaddleOCR库中使用以下命令:
python tools/train.py -c configs/det/det_mv3_db.yml \
-o Global.pretrained_model=./pretrain_models/MobileNetV3_large_x0_5_pretrained

训练可视化与指标验证

最后,我们将在每10次训练后使用1个评估步骤对模型进行5轮次的评估。下面是部分指标:

训练指标




系统指标

W&B还将自动记录每次运行的 GPU和CPU的使用时间!



验证指标




下载并使用已训练模型

检查点在每个轮次结束时,以及在每个模型的保存步骤中,被记录为W&B文件(含相应元数据与标签)。可使用以下代码进行下载并将其用于进一步的训练和评估目的。
import wandb
artifact = wandb.Api().artifact('manan-goel/text_detection/model-2138qk4h:best', type='model')
artifact_dir = artifact.download()
如需使用已训练的模型来进行文本检测,可使用PaddleOCR库中的推理脚本。
!python3 tools/infer_det.py -c configs/det/det_mv3_db.yml -o Global.infer_img="./doc/imgs_en/" Global.pretrained_model="./artifacts/model-2138qk4h:v9/model_ckpt"
这将对./doc/imgs_en/文件夹中的所有图像进行注释,并将结果存储在输出文件夹中。

补充内容:在您的W&B仪表板中记录已注释图像

在训练结束时对模型性能进行可视化的一种不错方式是记录其输入和输出图像,并将其记录至W&B仪表板中。首先,我们来初始化一个全新的W&B运行,并加载图像路径:
import wandb
import glob

wandb.init(project="text_detection")
wandb.use_artifact('manan-goel/text_detection/model-2138qk4h:best')
table = wandb.Table(columns=["Input Image", "Annotated Image"])

inp_imgs = sorted(glob.glob("./doc/imgs_en/*.jpg"), key=lambda x: x.split("/")[-1])
out_imgs = sorted(glob.glob("./output/det_db/det_results/*.jpg"), key=lambda x: x.split("/")[-1])
然后向wandb表格添加图像并将其记录至W&B中。
for inp in inp_imgs:
for out in out_imgs:
if out.split("/")[-1] != inp.split("/")[-1]:
continue
table.add_data(
wandb.Image(inp),
wandb.Image(out)
)
wandb.log({
"Predictions": table
})
wandb.finish()


Run set
11


结束语

本教程简要介绍了如何将W&B与PaddleOCR结合使用,以支持您在OCR模型开发方面的所有需求。别忘了在Colab上看看本报告的可执行代码版本。

相关工作


Iterate on AI agents and models faster. Try Weights & Biases today.