【YOLOv8n部署至RK3588】模型训练→转换rknn→部署全流程
A7bert777 2024-09-20 13:05:02 阅读 70
文件梳理
之前博主发布过YOLOv8转RKNN模型并在开发板上部署的流程,但是遇到了许多问题,发现很难修补,遂决定在官方项目下进行模型转换与部署(之前的转换是在Github个人博主的项目下进行的)
OK,进入正题,模型转换需要以下工具:ultralytics_yolov8、rknn_model_zoo、rknn-toolkit2,现在详细介绍这三者的关系。
以上三个文件均与模型训练无关,而是用于模型转换的,其中,ultralytics_yolov8文件不是必需的,但最好还是下载一份,原因后面会讲;在此,先将以上三个文件的链接放在这:ultralytics_yolov8,rknn_model_zoo ,rknn-toolkit2。
另外说一下这三个文件的版本,全用最新的,rknn_model_zoo用v2.1.0,rknn-toolkit2也用v2.1.0,ultralytics_yolov8只有一个main版本,直接用即可,如下所示:
这三个链接,全是在瑞芯微在Github上的官方仓库,最好科学上网后下载,其中,在Github上搜索rknn-toolkit2会有两个仓库,请以我的链接为准,另一个仓库是瑞芯微之前的官方仓库,并已弃用,现已转为上示链接。
现在讲一下这三个文件的关系,最先用到的是ultralytics_yolov8,在我们训练得到PT模型后,就需要用ultralytics_yolov8实现PT到ONNX模型的转换。其次要用到的是rknn-toolkit2,用其来配置ONNX转RKNN模型的环境。最后,在这个环境下,执行rknn_model_zoo中的py文件,进行模型转换,得到所要的RKNN模型。
最后,再放一张图,捋清RKNPU2(即rknn-toolkit2)与rknn_model_zoo的版本对应关系,避免版本不对导致的不适配问题。
下面开始进入正题,先进行模型训练。
模型训练
YOLOv8的模型训练教程,很多,基础不多叙述,大家可以直接参考其他文章
但有几点关照大家:
1.建议大家在配置环境时,严格按照此文进行配置,尤其是
pip install -e .,这一步尤为重要,不可省略。
2.另外,yolov8的训练方式有很多种,有的是通过命令行进行训练,有的是自己写一个train.py的脚本,在脚本中输入超参数后,再进行训练…,在此,博主建议大家采用脚本训练的方法,并且将脚本放在ultralytics文件夹的同级下,如图所示:
另外,train.py脚本内容如下:
因为在YOLOv8的项目中,存在一个ultralytics文件夹,而我们在配置环境的时候,又有名为ultralytics的环境包,二者存在歧义,导致我们在ultralytics文件夹中进行参数修改后却不生效的情况,(这时如果仍要通过命令行进行训练,必须要在环境配置文件中进行相应修改),比如更改激活函数后,发现得到的PT模型仍为最初的SiLU激活函数,这是因为命令行会优先调用环境配置文件中的参数,所以在ultralytics文件夹中修改激活函数后并不会生效(博主自己踩坑,请大家避免)。
3.最重要的一点:要改激活函数为ReLU,我们如果默认不改,则训练得到的PT模型的激活函数为SiLU,如下所示:
SiLU函数在PT模型中无任何问题,与ReLU函数的训练结果对比非常接近(仅在博主数据集上如此)
改激活函数流程如下:把ultralytics/nn/modules/conv.py中的Conv类进行修改,将default_act = nn.SiLU() # default activation改为
default_act = nn.ReLU() # default activation
如下所示:
大家要是不放心,可以在conv.py中搜索SiLU,然后全部替换为ReLU,如果我们只是引用正常的yolov8.yaml,其实只修改Conv类下即可,因为backbone和neck中并未调用其他类模块。
PT转ONNX
在训练得到PT模型后,我们可以用netron软件打开模型,看一下结构,如图所示:
现在我们的激活函数已经为ReLU了。
在YOLOv8同级目录下放置从Github上下载的ultralytics_yolov8文件,我将其重命名为了ultralytics_yolov8(ptconvert2onnx)。进入文件,将我们的PT模型放置进去,如图所示:
在终端执行命令:yolo export model=yolov8relupt.pt format=rknn
生成如下结果,同时在目录中生成我们想要的ONNX模型
注意!上图的三个红色方框非常重要,YOLOv8在RKNN转换过程中,是要有三个分支的,分别80×80检测头,40×40检测头,20×20检测头。每个检测头下,都有三个输出,分别是box边框输出,即图上的1×64×80×80;还有各类别输出,即1×4×80×80(这里的4就是指我在PT模型训练时的garbage.yaml文件中定义了4个类别),最后是综合置信度输出,即1×1×80×80。(其余检测头同理)
ONNX转RKNN
在进行这一步的时候,如果你是在云服务器上运行,请先确保你租的卡能支持RKNN的转换运行。博主是在自己的虚拟机中进行转换。
先安装转换环境
这里我们先conda create -n rknn210 python=3.8创建环境,创建完成如下所示:
现在需要用到rknn-toolkit2-2.1.0文件。
进入rknn-toolkit2-2.1.0\rknn-toolkit2-2.1.0\rknn-toolkit2\packages文件夹下,看到如下内容:
在终端激活环境,在终端输入pip install -r requirements_cp38-2.1.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
然后再输入pip install rknn_toolkit2-2.1.0+708089d1-cp38-cp38-linux_x86_64.whl
然后,我们的转rknn环境就配置完成了。
现在要进行模型转换,其实大家可以参考rknn_model_zoo-2.1.0\examples\yolov8下的README指导进行转换:
这里我也详细再说一遍转换流程:先进入rknn_model_zoo-2.1.0\examples\yolov8\python文件夹,先打开yolov8.py,进行适配参数修改:
然后修改convert.py,如下所示:
修改完成后,将我们之前得到的onnx模型复制到python文件夹下:
打开终端,激活rknn210环境,输入命令:python convert.py yolov8relupt.onnx rk3588
结果如下:
可以看到,我们的模型转换并量化成功。
这三个红框是和上面图中的三个红框对应的,所以若onnx没有所示的那三个框,则无法转换出适配官方部署文件的rknn模型。
模型部署
保密原因,部署代码不予公开。但大家可以在rknn_model_zoo-2.1.0\examples\yolov8\cpp下的main.cc、postprocess.cc与rknn_model_zoo-2.1.0\examples\yolov8\cpp\rknpu2下的yolov8.cc进行融合,这些都是官方代码,可以直接融合或者根据自己的接口需求进行更改。
部署后检测结果如下所示:
作为对比,我们把没修改激活函数的PT模型转换后的ONNX转换得到的RKNN模型进行部署,结果如下:
可以看到,同一张图片,SiLU激活函数在部署后的目标score量化过程中存在问题(之前博主也遇到类似的问题,主要原因:1.是量化值范围为-128到127,但是127值在量化后变为了0,再传输至sigmoid函数,出来的结果自然是0.5,即图中50.0%的检测结果;2.量化过程中的scale值在SiLU上表现不好,如果修改部署环节的scale值,比如进行放大,则也可以score高于50%,但是放大比例不好确定,属于自娱自乐)。所以在此,建议大家都在训练pt模型时,改为ReLU激活函数。
另外,将模型激活函数转为ReLU后,RKNN模型大小由5.27MB降为3.87M,降低了26.6%。另外检测速度也得到了提升,在50张1080P图像的demo检测中,转为ReLU后在NPU上的处理速度也得到了很大的提升,如图所示。
上述即博主此次更新的YOLOv8部署RK3588,包含PT转ONNX转RKNN的全流程步骤,欢迎交流!
声明
本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。