An image to text model based on transformer used on OCR task.
在这里集成了 NVIDIA FasterTransformer 用于预测加速。同时集成了 FasterTransformer float32 以及 float16 预测。以下是使用 FasterTransformer 的说明。
-
本项目依赖于 PaddlePaddle 2.3.2 GPU版本
-
如果使用NVIDIA FasterTransformer则需安装如下依赖：
CMake >= 3.10
CUDA 10.1 或以上版本 （需要 PaddlePaddle 框架一致）
gcc 版本需要与编译 PaddlePaddle 版本一致，比如使用 gcc8.2
推荐使用 Python3，python依赖：
pip install attrdict pyyaml paddlenlp -
对于图灵架构GPU(CUDA版本为10.1)也可以通过docker镜像来安装paddle运行环境：
拉取cuda运行环境镜像（已经配置好对应的gcc和cmake等常用工具）
docker pull registry.cn-shanghai.aliyuncs.com/janelu9/dl:cuda-10.1-cudnn8-ubuntu18.04启动容器
docker run -d --gpus all --name paddle --shm-size=32g --ulimit memlock=-1 -v /mnt:/mnt -p 8888:8888 -it 146273117745 /bin/bash安装Anaconda
./Anaconda3-*-Linux-x86_64.shsource ~/.bashrc安装paddle及其依赖
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simplepip install paddlepaddle-gpu==2.3.2.post101 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html-
使用了SwinTransformer、CSwinTransformer等作为图像部分的编码器；
-
使用了KLDivLoss和CTCLoss相结合的损失函数（可以仅使用Attention解码，也可以使用CTC前缀束搜索加Attention重排序的方式解码）；
-
使用ERNIE-3.0等基于TansformerEncoder/Decoder的中文预训练模型作为TrOCR文本部分的解码器以适用于中文OCR识别任务；
-
可以改用较小的字典（如
ocr_keys_v1.txt，包含6623个常见字符），脚本在加载预训练Decoder模型时会重构wordembedding以适应新的字典； -
集成了 NVIDIA FasterTransformer 用于预测加速，以解决当模型解码器的维度、束搜索空间、层数，较高、大、深时可能出现的推断效率问题；
-
增加了一个自定义的Beam Search方法，通常情况下更快更准。
image2text.py中包含了用于训练的Image2Text模型和用于快速推断的FasterTransformer, 模型基于paddlepaddle开发. 这时你可以从paddlenlp模型库中加载各种基于中文数据集的预训练模型.
image_aug.py中包含了多种数据增强策略，可以有效增加模型泛化能力。用户也可以自行添加基于pillow或opencv的图像增强策略，注意好img和ndarray间的转换即可。
注：FasterTransformer会在FasterTransformer第一次被调用时自动编译。
- 在
train.py中配置好数据目录、训练集标签、测试集标签和预训练模型位置等参数。 - 训练模型、评估和保存：
单机多卡启动，默认使用当前可见的所有卡
python -m paddle.distributed.launch train.py单机多卡启动，设置当前使用的第0号和第1号卡
python -m paddle.distributed.launch --gpus '0,1' train.py单机多卡启动，设置当前使用第0号和第1号卡
export CUDA_VISIBLE_DEVICES=0,1 && python -m paddle.distributed.launch train.py模型训练过程会自动保存验证集效果最好的参数于文件best_model.pdparams中。
- 在部署推理服务时可以通过
paddle.jit将模型转换为静态图，然后使用paddle inference加载以提升推理效率。
paddle.jit.save(layer=infer,path="inference",input_spec=[paddle.static.InputSpec(shape=[None,3,384,384],dtype='float32')])注：该模型在用于处理不规则排列文字的OCR任务（如印章、公式、多行文本识别）时简单有效。准确率在样本呈正态分布的测试集上（比较客观的抽样样本集上）通常可以达到90%以上。
If you find it useful or use my TrOCR code in your research, please cite it in your publications.
@misc{Image2Text,
author = {Jian Lu},
title = {Image2Text: An image to text model based on transformer used on OCR task},
year = {2022},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://gitee.com/janelu9/TrOCR.git}},
}