2024年12月,中国计算机学会在海南博鳌成功举办了第十二届CCF大数据与计算智能大赛(简称2024 CCF BDCI)。本届比赛的算能赛道吸引了1748名选手报名,经过激烈角逐,北京航空航天大学的“常务副SOTA”团队脱颖而出,荣获大赛最佳算法能力奖。
本届CCF BDCI大赛中,算能赛道推出了“基于TPU平台的OCR模型性能优化”的赛题,旨在探索如何通过高效的TPU硬件平台— Duo系列开发板,以及先进的算法优化技术,实现OCR模型在精度、速度和资源消耗等方面的全面优化,进而精确识别目标文本。
针对该赛题,“常务副SOTA”团队选择了YOLOv8关键点模型和PPOCRv2-Rec作为基础模型,因其在速度和精度上的优越表现。团队通过通道剪枝、算子融合和词表剪裁等技术,显著减少了计算量和推理时间。在模型转换中,使用TPU-MLIR工具确保高效运行并保持精度。同时,优化前后处理环节,如使用RVV加速图像处理,进一步提升性能。最终,检测模型和识别模型的推理速度分别提升了近4.9倍和3.5倍,单图推理时间缩短至700ms以内。
以下为“常务副SOTA”团队的研究技术报告,供开发者们参考与学习!
团队简介
我们团队由四位深度学习技术爱好者组成,队名为“常务副SOTA”,希望通过比赛来促进队员的成长,也希望提出的方案能为实际应用添砖加瓦。
摘要
随着深度学习的快速发展,OCR(Optical Character Recognition,光学字符识别)技术已广泛应用于文本识别、自动化办公、信息提取等多个领域。然而,由于OCR任务对计算资源的需求较高,如何在低功耗设备上实现高效的OCR模型成为了一个重要挑战。本文基于TPU平台MilkV-Duo部署,针对OCR模型的性能进行优化,涉及模型选择、优化、转换及前后处理等方面。通过一系列的优化手段,包括通道剪枝、模型量化、算子融合、前处理加速等,成功提升了模型的推理速度和精度,最终实现了在端侧设备上的高效运行。
关键词
OCR,TPU-MLIR,MilkV-Duo 部署,量化
1.比赛分析
本次比赛的目标是将OCR模型优化并部署到TPU平台上,达到高效运行的要求。OCR任务通常包括图像输入、前处理、检测、矫正、分类和识别等步骤。在比赛中,参赛队伍需要对OCR模型进行优化,减少模型的计算量和存储需求,同时提高推理速度和准确率。
比赛提供了一个基线模型,该模型的OCR流程如下:
- 图像输入:读取输入图像并进行预处理。
- 检测前处理:对图像进行一些基本的预处理,例如灰度化、二值化等。
- 检测模型:通过模型输出二值图,进一步进行目标区域的识别。
- 连通域算法:寻找图像中的目标区域并提取最小外接旋转矩形。
- 仿射变换:将图像旋转拉正为矩形。
- 分类前处理:对图像进行分类处理,输出旋转角度。
- 识别前处理:进一步对图像进行旋转,确保识别模型的准确性。
- 识别模型:最后通过识别模型输出识别结果。
比赛的要求是要在TPU平台上实现高效的推理,因此需要优化各个环节,特别是模型的计算量和推理速度。
2.模型选择
为了在TPU平台上实现高效的OCR模型,我们首先从多个候选模型中进行选择。考虑到OCR任务的特点,如需要获取目标的不规则四边形输出,最终我们在PPOCR系列,YOLO目标检测系列,以及YOLO关键点检测系列模型中进行选型。
通过简单的速度实验对比,我们发现PPOCRv4-Det检测模型存在参数量大,模型后处理复杂,且推理速度较慢等明显问题,640*640的输入下端到端需要1080ms才能完成推理,相比之下,YOLOv8检测模型端到端只需要820ms,YOLOv8关键点模型端到端只需要840ms,总的来说使用检测和关键点模型在推理速度上更加有优势。
在选择识别模型时,我们还同时还考虑了精度与速度的平衡,比如PPOCRv4-Rec虽然在精度上有一定提升,但引入Transformer结构后,推理速度较慢,因此我们选择了更早版本,且没有使用Transformer结构的PPOCRv2-Rec作为识别模型。
除此之外,我们考虑到资源占用问题,由于赛题要求在端侧设备上运行,尽量减少OCR任务中的pipeline模型个数,最终我们没有使用方向分类模型,这可以降低资源Milkv-Duo的资源占用。
最终,我们选择使用YOLOv8关键点模型和PPOCRv2-Rec识别模型最为我们团队的OCR任务基座模型,然后在这基础上最更多的优化工作。
3.模型优化
在模型选择之后,我们分别对关键点模型和识别模型进行了多方面的优化,以提高其在TPU平台上的运行效率。具体优化策略如下:
1. 通道剪枝:通过通道剪枝技术,我们减少了YOLOv8n-pose模型的通道为原来的一半,减少了计算量,从而加速了模型的推理速度。我们通过重新训练的方法让通道剪枝能够在不显著降低精度的情况下,减小模型的计算量。
2. 增加特征尺度:在模型的训练过程中,增加了对特征尺度P2(下图黄色部分)的融合,使得模型能够在多个尺度下进行特征提取,从而提高模型对多尺度目标的适应性,进而提升了识别精度。
图1:增加特征尺度
3. 模型算子融合:我们通过算子融合,将多个操作合并为一个操作,如连续的卷积或反卷积和批归一化层、连续的卷积或反卷积和线性变化层等合并为一个算子,还通过将多个复杂琐碎的算子如被拆散的层归一化算子融合为单独的层归一化算子降低了多个算子的计算和访存,通过这一方案,我们成功减少了计算图的复杂性,减小了模型的计算开销。
4. 模型的算子移除:我们通过将一些复杂的后处理算子,如和置信度计算的sigmoid激活以及和提取分类标签的softmax算子从模型中移除,通过和int8输出的反量化融合来减少计算的开销。
5. 词表剪裁:对模型的词表进行了剪裁,将原始的6625个类别减少至3952个,减少了40%的词表大小,缩短了匹配时间。
6. 关键点模型的后处理减枝:我们通过将YOLOv8-pose模型后处理完整裁剪,通过提前使用置信度阈值过滤不必要的计算来提升模型的端到端速度。
通过上述模型层面的优化,我们的方案在保证精度的前提下减少了模型尺寸,减少了模型算子数量,并且避免了大量的冗余计算,更加适合在TPU平台高效运行。
4.模型转换
在模型优化完成后,我们使用TPU-MLIR工具对模型进行转换,以便能够在TPU平台上高效运行。模型转换主要包括以下两个方面:
1. 检测模型转换:将检测模型从onnx转换为mlir格式,最后转换为cvimodel格式。这一过程中,优化了检测模型的输入尺寸,将输入尺寸从640*640调整为320*320,以减少图像前处理的耗时和推理的计算量。我们尝试将检测模型以INT8精度转换,在Milkv-Duo上的精度保持很好,最终我们采用了INT8模型部署,并且将前处理融合到模型道中,以加速端到端的速度。
2. 识别模型转换:将识别模型从onnx转换为mlir格式,最后转换为cvimodel格式。识别模型的转换过程中,由于原始输入尺寸为32x640,转换时遇到了一些问题,最终通过调整为32x320的输入尺寸,解决了转换失败的问题。我们尝试过INT8精度转换,但是在Milkv-Duo上的精度较差,最终我们选择使用BF16精度部署,并且将前处理融合到模型道中,以加速端到端的速度。
5.前后处理优化
在前后处理阶段,我采用了一些优化策略,以进一步提升模型的运行效率。具体优化手段包括:
1. 前处理优化:使用RVV(RISC-V Vector Extension)实现图像resize操作,将原本120ms的图像resize时间减少至90ms,提升了33.33%的速度。
2. 后处理优化:通过使用int8查表实现sigmoid过滤阈值,减少了后处理的计算量,将后处理的时间从20ms减少至18ms,提升了11.11%的速度。
3. anchor生成优化:通过将anchor生成和softmax计算通过int16进行优化,并结合反量化操作,使得后处理速度提升了66.66%(从10ms降至6ms)。
4. OpenCV优化:使用平头哥的OpenCV替代了opencv-mobile,提升了OpenCV算子的执行速度,进而加速了整体模型的运行。
6.实验记录
经过上述优化后,我们对模型进行了全面的实验验证。实验结果如下:
首先是检测模型的实验情况:
模型 | 端到端耗时(ms) | 备注 |
Base | 1080ms | 640x640 |
Base | 650ms | 320x320 |
Pose | 472ms | 320x320 |
Pose | 440ms | +后处理裁剪 |
Pose | 310ms | +int8 融合前处理 |
Pose | 294ms | +算子融合 |
Pose | 260ms | +rvv reize + 库优化 |
Pose | 242ms | +sigmod/softmax优化 |
Pose | 220ms | +阈值 (conf0.3 iou0.7) |
然后是识别模型的实验情况:
模型 | 端到端耗时(ms) | 备注 |
ppocrv2 | 560ms | 32x640 |
ppocrv2 | 380ms | 32x320 |
ppocrv2 | 220ms | +int8 融合前处理 |
ppocrv2 | 209ms | +裁剪softmax |
ppocrv2 | 188ms | +算子融合 |
ppocrv2 | 154ms | +裁剪词表 |
ppocrv2 | 140ms | +rvv reize + 库优化 |
最终,我们优化后的检测模型相比基线版本提升了接近4.9倍,我们优化后的识别模型相比基线版本提升了接近3.5倍。
在单图推理的情况下(包括一次检测和三次识别),优化后的模型能够在700ms以内完成,远低于基线模型的运行时间。这一结果表明,通过上述优化手段,模型在TPU平台上的运行效率得到了显著提升。
致谢
感谢CCF协会、DF平台和SOPHGO公司提供的宝贵机会;感谢主办方免费提供的milkv duo开发板,以及主办方针对开发板板特性和模型部署问题的答疑;感谢平头哥和nihui老师开源OpenCV-THead和OpenCV-Mobile;感谢百度PaddlePaddle团队开源PPOCR工具;感谢Ultralytics对YOLO模型的开源;感谢一起参加此次比赛的友军的分享和讨论。
参考
[1] Ultralytic, YOLOv8,https://github.com/ultralytics/ultralytics
[2] nihui,OpenCV-Mobile,https://github.com/nihui/opencv-mobile
[3] 平头哥团队,OpenCV-Mobile,https://www.xrvm.cn/community/
download?id=4112956065753141248