2025 CCF BDCI｜“基于TPU平台的OCR模型性能优化”一等奖作品

2025 CCF BDCI｜“基于TPU平台的OCR模型性能优化”一等奖作品

在2024 CCF BDCI大赛的算能赛道中，北京航空航天大学的“常务副SOTA”团队凭借其在OCR模型性能优化方面的出色表现，荣获大赛最佳算法能力奖。该团队通过一系列创新性的优化策略，成功将OCR模型在TPU平台上的推理速度提升了近4.9倍（检测模型）和3.5倍（识别模型），展现了他们在深度学习和计算机视觉领域的深厚实力。

摘要

随着深度学习的快速发展，OCR（Optical Character Recognition，光学字符识别）技术已广泛应用于文本识别、自动化办公、信息提取等多个领域。然而，由于OCR任务对计算资源的需求较高，如何在低功耗设备上实现高效的OCR模型成为了一个重要挑战。本文基于TPU平台MilkV-Duo部署，针对OCR模型的性能进行优化，涉及模型选择、优化、转换及前后处理等方面。通过一系列的优化手段，包括通道剪枝、模型量化、算子融合、前处理加速等，成功提升了模型的推理速度和精度，最终实现了在端侧设备上的高效运行。

1.比赛分析

本次比赛的目标是将OCR模型优化并部署到TPU平台上，达到高效运行的要求。OCR任务通常包括图像输入、前处理、检测、矫正、分类和识别等步骤。在比赛中，参赛队伍需要对OCR模型进行优化，减少模型的计算量和存储需求，同时提高推理速度和准确率。

比赛提供了一个基线模型，该模型的OCR流程如下：

1. 图像输入：读取输入图像并进行预处理。
2. 检测前处理：对图像进行一些基本的预处理，例如灰度化、二值化等。
3. 检测模型：通过模型输出二值图，进一步进行目标区域的识别。
4. 连通域算法：寻找图像中的目标区域并提取最小外接旋转矩形。
5. 仿射变换：将图像旋转拉正为矩形。
6. 分类前处理：对图像进行分类处理，输出旋转角度。
7. 识别前处理：进一步对图像进行旋转，确保识别模型的准确性。
8. 识别模型：最后通过识别模型输出识别结果。

比赛的要求是要在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格式。这一过程中，优化了检测模型的输入尺寸，将输入尺寸从640640调整为320320，以减少图像前处理的耗时和推理的计算量。我们尝试将检测模型以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.实验记录

经过上述优化后，我们对模型进行了全面的实验验证。实验结果如下：

首先是检测模型的实验情况：

然后是识别模型的实验情况：

最终，我们优化后的检测模型相比基线版本提升了接近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