首页IT科技yolov5-transformer(yolov5篇—yolov5训练pt模型并转换为rknn模型,部署在RK3588开发板上——从训练到部署全过程)

yolov5-transformer(yolov5篇—yolov5训练pt模型并转换为rknn模型,部署在RK3588开发板上——从训练到部署全过程)

时间2025-08-03 11:24:47分类IT科技浏览21874
导读:一、任务介绍   瑞芯微RK3588是一款搭载了NPU的国产开发板。...

一                    、任务介绍

  瑞芯微RK3588是一款搭载了NPU的国产开发板                    。NPU(neural-network processing units)可以说是为了嵌入式神经网络和边缘计算量身定制的                    ,但若想调用RK3588的NPU单元进行推理加速                                 ,则需要首先将模型转换为**.rknn格式**的模型            ,否则无法使用                                 。

  这次我们的任务是将yolov5训练得到的pt模型               ,一步步转换为rknn模型                                ,并将rknn模型部署在搭载RK3588的StationPC M3主机上                  ,使用NPU推理            。(官网NPU教程:NPU使用 — Firefly Wiki)查阅资料和官网后发现转换和使用过程分以下几步          ,首先总结给大家                               ,后文细说:

1.使用正确版本(v5.0)的yolov5进行训练得到pt模型

2.将pt模型使用yolov5工程中的export.py转换为onnx模型

3.将onnx模型使用rknn-toolkit2中onnx文件夹的test.py转换为rknn模型;

4.在板子上使用rknpu2工具调用rknn模型                        ,实现NPU推理加速               。

  接下来进行详细介绍                                。

二                                 、实验过程

2.1 使用正确版本的yolov5进行训练(平台:x86机器ubuntu22.04系统)

  人菜就要跟着官方教程做     ,请看官方教程:

rknn-toolkit/examples/pytorch/yolov5 at master · rockchip-linux/rknn-toolkit · GitHub

  需要指出的是                              ,虽然我们后面要使用的是rknn-toolkit2工具进行模型转换                             ,但教程却在rknn-toolkit工程中,github二者的父目录如图:

  进入rknn-toolkit工程                         ,浏览至/example/pytorch/yolov5                                  ,在README中赫然写着:

  这就是我们本次要使用的“正确版本的yolov5                       ”      ,怕麻烦的读者请一定按照这个id去yolov5官网找到工程并克隆下来                  。进入yolov5官网GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

git clone 获取指定分支的指定commit版本

git clone https://github.com/ultralytics/yolov5 cd yolov5 git reset --hard c5360f6e7009eb4d05f14d1cc9dae0963e949213

  非常棒                    ,至此我们已经找到了正确版本的yolov5!接下来就是把代码克隆到本地或者云服务器                                 ,开始训练自己的pt模型;此处实际上是yolov5 5.0版本的工程            ,其中未包含预训练模型yolov5s.pt               ,需要我们自行下载                                ,在releasesl链接中找到V5.0下载yolov5s.pt                  ,链接如下:

https://github.com/ultralytics/yolov5/releases

关于yolov5的运行环境配置            、参数修改和训练命令          ,网上教程很多                               ,满足package的要求                        ,可以开始训练收敛即可     ,不作为本文重点          。最终在经过若干轮训练后                              ,我们得到了自己的last.pt模型                             ,本部分告一段落                               。

2.2 best.pt转换为best.onnx(平台:x86机器ubuntu22.04系统)

(1)

  此处依然参考rknn-toolkit工程中的教程进行转换,切记以下步骤!!!

注意:在训练时不要修改yolo.py的这段代码                         ,训练完成后使用export.py进行模型导出转换时一定要进行修改                                  ,不然会导致后面的rknn模型转换失败!

(2) 修改export.py

找到 torch.onnx.export 所在行

 

修改后将last.pt移动至export.py同一文件夹下      ,在命令行调用以下命令:

python export.py --weights last.pt --img 640 --batch 1 --include onnx

我们便可以得到成功转换的模型last.onnx                        。

2.3 best.onnx转换为best.rknn(平台:x86机器ubuntu22.04系统)

2.3.1 环境准备和工具包安装

  有了以上的last.onnx模型作为基础                    ,现在可以正式开始rknn模型的转换了     。此处的转换工具rknn-toolkit2只能在Linux系统上运行                                 ,这里我使用的是本机Ubuntu22.04                              。在上文提到的链接中下载rknn-toolkit2

           ,下载后的文件夹如图:

  首先安装rknn-toolkit2的环境               , 其环境要求在./doc/requirements_cp38-1.4.0.txt中                                ,这里我使用的是anaconda创建的python3.8虚拟环境                  ,创建环境并命名为rknn_3588          ,使用pip安装requirements_cp38-1.4.0.txt中的包

pip install -r requirements_cp38-1.4.0.txt

环境满足需求后便可以安装rknn-toolkit2工具包了                               ,浏览至**./packages**中                        ,输入以下命令:

pip install rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl

提示安装完成后我们可以检查是否安装成功     ,在终端中运行python                              ,输入:

from rknn.api import RKNN

若不报错说明我们的工具包已经安装成功                             ,之后便可进行rknn模型的转换了~

2.3.2 onnx转换为rknn

  在rknn-toolkit2工程文件夹中浏览至 ./examples/onnx/yolov5,将我们在2.2中转换得到的last.onnx复制到该文件夹下                         ,修改该文件夹下的test.py中的内容为自己模型的名字                                  ,要修改的地方如下:

在这个test.py的main函数中(在第230行附近)可以了解到这个python文件的作用:

【 配置——加载onnx模型—导出rknn模型——rknn模型推理——后处理——输出结果】

最终弹窗得到post process result      ,若结果正确                    ,说明我们的rknn模型转换成功了!这里的test.py构建了一个虚拟的NPU运行环境                                 ,模拟在RK3588上真实运行的情况                             。结果正确说明离最终完成任务只差部署在板子上这一步了。

这时在当前文件夹**./example/onnx/yolov5中可以看到生成的last.rknn**模型                         。至此            ,我们离胜利就不远了                                  。

2.4 RK3588部署rknn实现NPU加速(平台:aarch板子Linux系统)

  终于RK3588板子登场的时候了               ,但如果它还是原生的安卓系统                                ,请按官网教程给它安装一个Linux系统吧      。我这里使用TF卡烧录安装了Ubuntu系统                    。

  依然找到官网                  ,下载其中的rknpu2工具          ,链接如下:https://github.com/rockchip-linux/rknpu2

3. 3588平台部署

在3588的主目录上获取官方demo

git clone https://github.com/rockchip-linux/rknpu2.git

进入yolov5目录

cd /home/ptay/rknpu2-master/examples/rknn_yolov5_demo

修改include文件中的头文件postprocess.h

#define OBJ_CLASS_NUM 3 #这里的数字修改为数据集的类的个数

修改model目录下的coco_80_labels_list.txt文件                               , 改为自己的类并保存

xxx xxx xxx

将转换后的rknn文件放在model/RK3588目录下

编译                        ,运行shell

bash ./build-linux_RK3588.sh

成功后生成install目录

cd install/rknn_yolov5_demo_linux

在model目录下放入需要推理的图片

运行

./rknn_yolov5_demo ./model/RK3588/last.rknn ./model/0625_xx_005.jpg

注:后面的图片用的全路径     ,因为当前的相对路径识别不到(主要是看图片所在的位置)

在rknn_yolov5_demo_linux获取到结果

参考:

https://blog.csdn.net/m0_57315535/article/details/128250096

https://blog.csdn.net/m0_51714298/article/details/125916417
声明:本站所有文章                              ,如无特殊说明或标注                             ,均为本站原创发布                                 。任何个人或组织,在未征得本站同意时                         ,禁止复制               、盗用                                、采集                  、发布本站内容到任何网站          、书籍等各类媒体平台            。如若本站内容侵犯了原著者的合法权益                                  ,可联系我们进行处理               。

创心域SEO版权声明:以上内容作者已申请原创保护,未经允许不得转载,侵权必究!授权事宜、对本内容有异议或投诉,敬请联系网站管理员,我们将尽快回复您,谢谢合作!

展开全文READ MORE
YOLOv5教程(YOLOv5深度剖析) 百度关键词工具下载(百度关键词采集软件哪个好用?推荐你使用XXX软件)