rock3a: 基于自建数据集+yolov5s模型的rknn模型训练部署全流程

上一篇文章其实已经详述了模型训练到部署的整个流程,但是数据集到模型都是用的官方的coco数据集,这里为了记录开发板的模型训练到部署的整个流程,重新开了一篇文章进行记录。

首先准备数据集和rockchip官方推荐的yolov5源代码

这里需要注意的是数据集中的标签文件内容 一行只能有5个数,分别是类别和(x,y,w,h)

但是我平时用的labelimg.exe 生成的xml文件转化为txt标签文本文件后 一行有6个数,分别为类别,(x,y,w,h),以及置信度,这里需要做一下转换。

然后需要将coco/yaml128.yaml配置文件中的部分参数做一下修改 主要是训练类别数和训练类别名称

 一般做目标检测(比如我)只做一类检测,这里需要将类别数改为1,并将类别标签改为实际的标签名。

接下来就是训练模型,在yolov5-master路径下,键入命令

python3 train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 64

该命令的意思是读取coco128.yaml配置文件里面的数据集路径和类别数、类别名,模型架构为yolov5s,也就是small版本的模型,考虑到开发板的性能,所以这里选择的是small版本的模型,然后重新训练模型。训练模型完成后,需要将模型进行转换,转换为onnx模型。

这里需要将model/yolo.py文件某个部分进行修改,才能正确的将pt模型导出为onnx模型。

注:训练时不需要修改该部分代码  导出模型时需要修改该部分代码

椭圆处是需要注释的部分。

在yolov5-master 主目录下,键入命令

python3   export.py --weights yolov5s.pt --img 640 --batch 1 --opset 12 --include onnx

可以将pt模型转化为onnx模型。将转化后的onnx模型拷贝出来。

接下来在已经安装rknn-toolkit的虚拟机上,对onnx模型进行转换。

首先采用Netron.exe软件打开转化好的onnx模型查看其网络结构

并记下其三个输出的标号。

然后将onnx模型拷贝到/home/rock/workspace/rknn-toolkit2-1.4.0/examples/onnx/yolov5 这个目录下。

该目录下的test.py是瑞芯微官方已经写好的demo文件。

我在这里类比test.py写了一个onnx2rknn.py  ,功能与test.py差不多,只是少了连板推理部分。

注意需要在onnx2rknn.py文件中做部分修改

修改完成后,键入命令 python3 onnx2rknn.py文件,正常情况下会生成一个rknn结尾的文件。

接下来将该rknn文件拷贝到开发板上的对应目录下,我这里是

/home/rock/software/rknpu2-main/examples/rknn_yolov5_demo/install/rknn_yolov5_demo_Linux

相对于原始文件,这里我们需要修改如下几个部分:

  1. 需要在源代码中修改类别数

 

  1. 需要在配置文件中修改类别数

配置文件一共有两个,如上图所示。

我这里将配置文件里面的类别数修改为我自己的类别数

然后在开发板重新编译,键入命令

sh build-linux_RK356X.sh

然后进入 install/rknn_yolov5_demo_Linux/ 路径下,去执行可执行文件

这里说一句:

模型训练初期和模型多次迭代后,其运行速度是不一样的。

实测,当模型开始迭代时,运行一次目标检测耗时大约120ms,当模型迭代300个epoch时,运行一次目标检测耗时60ms,这其中的原因不得而知,如果有大佬知道的欢迎讨论

到此为止,基于自建数据集的模型训练到模型转化,模型部署已经搭建完毕。

猜你喜欢

转载自blog.csdn.net/linxizi0622/article/details/128615305