使用RKNN部署CRNN模型踩坑优化历程

最新推荐文章于 2024-08-27 23:31:22 发布

三叔家的猫

最新推荐文章于 2024-08-27 23:31:22 发布

阅读量6.8k

点赞数 9

分类专栏： OCR rknn开发记录文章标签：深度学习 lstm 神经网络

本文链接：https://blog.csdn.net/qq_39056987/article/details/122128191

版权

OCR 同时被 2 个专栏收录

18 篇文章 36 订阅

订阅专栏

rknn开发记录

5 篇文章 6 订阅

订阅专栏

序言

前段时间使用RKNN部署一个文字识别模型，因为文字识别模型用的是目前最普遍使用的CRNN模型，结构也相对简单：卷积+LSTM+全连接，都是比较元老级别的算子，本来已经部署的过程会很顺利，结果发现还是踩了很多坑。写篇文章记录下踩坑优化过程。

一、部署前

接到需求后，一开始想到的是直接使用paddleocr开源的通用文字识别模型，paddleocr提供了三个通用的文字识别模型，都是基于CRNN算法架构，分别是：mobilenetV3版、LCNet版、Resnet34版，前两个是移动端模型，后一个是服务端模型，因为部署到板子上，所以服务端模型肯定用不了了，那只能试试移动端模型了，但是在是的过程中发现对与特定场景精度还是有点不够，虽然说这几个模型在整体的通用文字识别中效果还是不错的，但是因为我的场景要求识别率比较高，所以两个移动端的模型都没法白嫖了，没办法，只能自己训练一个了。

基于text_renderer仓库生成了一批用于训练识别模型的基准数据集，数据量在650w左右，囊括了一百多种中文字体的文字图片数据集，用于模型重新训练；

模型训练基于PytorchOCR（paddleOcr实在是用不习惯）。

一开始用的是mobilenetV3版本的CRNN进行训练，但是发行训练后的模型精度和泛化能力都很差，考虑到ResNet18、ResNet34作为主干模型又没法嵌入到移动端，基于此，选择了repvgg的最小模型作为主干进行训练，最后训练出来的模型在精度和速度上有非常的优秀。唯一的缺点是模型权重大了一些，不过问题不大，可以进行int8量化，模型大小可以缩小到原来的四分之一。

最初我的部署模型网络构成为：repvgg(部署）+ LSTM + Linear + ctc，然后开始了辛酸的部署历程。

二、模型部署

刚开始我查阅了rknn的算子支持文档，发现我的模型里用到的算子都是支持的，我以为部署起来会很顺利，结果。。。

2.1 踩坑一：LSTM算子不支持

LSTM作为老牌算子，各种框架应该都很早适配了才对，并且我查看了RKNN的文档也是写着支持的，但是我通过ONNX转换成RKNN的时候缺出现了问题，报错了，如下所示：
在这里插入图片描述
然后我去rknn的官方论坛搜了一下，发现全是转换报错的问题，并且没有解决方案，这不扯犊子嘛？

最后去他们Q群里问，也没有结果，让我去redmine提问，我哪知道在哪，算了，还是想想其他办法解决吧。

那能不能去掉LSTM层呢？答案是可以的，比较典型的方案就是densenet ocr，去掉了lstm，保留了原来的结构，最后的识别效果也是很不错的，不过缺点是少了字符前后序列信息，有些字符可能会识别错误。

然后对上面的结构进行了修改，将LSTM层去掉，只保留了：repvgg + Linear + ctc，然后重新训练。

训练过程中发现去掉lstm后模型收敛的速度更快了，精度基本上都差不多，最后重新得到了一个模型，心想着模型结构已经非常简单了，总不至于还有问题吧，结果。。。

2.2 踩坑二：全连接层不量化

模型转换很顺利，但是发现量化后的模型推理速度异常的慢，如下：
在这里插入图片描述

然后找原因，找啊找啊找，发现在量化转换后的模型，全连接层耗时非常高，可视化显示了全连接层没有量化（深蓝色表示已量化，浅蓝色表示未量化）？？我透，这尼玛也太坑了吧，因为字符分类数比较多，所以全连接层的参数量比较大，没经过量化的话速度会非常的慢。

在这里插入图片描述
可以查看各层输出的耗时情况如下，可以看到，最后的两层全链接层耗时是非常多的：

所以我在就在想着能不能把这两层优化一下呢。

全链接层既然不量化，那我能不能把全链接换成全卷积呢？答案是可以的，因为使用全卷积网络也是可以作为分类使用，在诸多的目标检测中，基本上全是全卷积网络，所以我决定把后面两层改称全卷积网络的输入输出，即在主干网络后再加入两层1x1的卷积用于分类，所以头部输出就改成了如下结构，最后再把经过1x1卷积后的输出reshape一下成我们需要的输出格式：
在这里插入图片描述
然后重新训练，经过测试，发现全链接和全卷积作为分类层，参数量是不变的，精度几乎也是不变的，关键在与我们使用rknn量化后的速度，废话不多说，量化过程正常，直接上测试结果：

输入是32448的大小图片（另外说明下上面全链接层测试的输入是32224的），速度超快有木有！！推理只需要5.4ms对比之前的一百多毫秒，只是换了两个层，直接起飞了！！另外看下各层耗时输出：
在这里插入图片描述
因为每层都经过量化，所以每层耗时都很均匀，推理速度也基本在5ms左右，精度跟没量化后的只是降低了一点点，所以说这次的优化是非常成功。

三、写在最后

本文主要优化了crnn在rknn上部署的问题，对于去掉LSTM层，个人觉得LSTM虽然具有很强的时序信息，在大多数场合下的识别，如果不是特别复杂的场景，其实并不需要太强的时序信息，结合我的业务场景，所以我这里去掉后并无太大影响，反而训练速度收敛更快，精度更高了，当然泛化能力也许没有带LSTM的好，这个可以自己实验权衡。二是1x1的卷积作为分类层代替全链接，这个我测试的时候发现精度并无太大差异，而且我发现好像ncnn的全链接层int8量化也有些问题，所以有必要的话，改称1x1卷积去做分类层对于部署来说速度提升还是很大的。