代码收藏家 yolov5训练pt模型并转换为rknn模型，部署在RK3588开发板上——从训练到部署全过程
目录
2.1 使用正确版本的yolov5进行训练(平台：x86机器windows系统)
2.2 best.pt转换为best.onnx（平台：x86机器window系统）
2.3 best.onnx转换为best.rknn（平台：x86机器Linux系统）
2.4 RK3588部署rknn实现NPU加速（平台：aarch板子Linux系统）
一、任务介绍
瑞芯微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机器windows系统)
在这次任务之前，博主本人甚至不知道yolov5有不同的版本（小白落泪），更不知道用不同版本训练得到的pt模型会决定rknn模型转换的成功与否。一开始我使用的是最新版本的yolov5（大概是v7.0）进行训练，最后模型转换自然是不成功。网上也有许多使用新版本yolov5训练，之后修改onnx模型的教程，但都比较麻烦。人菜就要跟着官方教程做，请看官方教程：
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
在右上角搜索这个commit id并进入“In this repository”，如图所示：
进入commits，点击右边的commit id，再次进入后点击Browse files，如下三图所示：
commit id：
Browse files：
对应版本的源码：
非常棒，至此我们已经找到了正确版本的yolov5！接下来就是把代码克隆到本地或者云服务器，开始训练自己的pt模型；此处实际上是yolov5 5.0版本的工程，其中未包含预训练模型yolov5s.pt，需要我们自行下载，在releasesl链接中找到V5.0下载yolov5s.pt，链接如下：
https://github.com/ultralytics/yolov5/releases
关于yolov5的运行环境配置、参数修改和训练命令，网上教程很多，满足package的要求，可以开始训练收敛即可，不作为本文重点。最终在经过若干轮训练后，我们得到了自己的best.pt模型，本部分告一段落。
2.2 best.pt转换为best.onnx（平台：x86机器window系统）
此处依然参考rknn-toolkit工程中的教程进行转换，切记以下步骤！！！
注意：在训练时不要修改yolo.py的这段代码，训练完成后使用export.py进行模型导出转换时一定要进行修改，不然会导致后面的rknn模型转换失败！
修改后将best.pt移动至export.py同一文件夹下，在命令行调用以下命令：
python export.py --weights best.pt --img 640 --batch 1 --include onnx我们便可以得到成功转换的模型best.onnx。
2.3 best.onnx转换为best.rknn（平台：x86机器Linux系统）
2.3.1 环境准备和工具包安装
有了以上的best.onnx模型作为基础，现在可以正式开始rknn模型的转换了。此处的转换工具rknn-toolkit2只能在Linux系统上运行，这里我使用的是PC上的Ubuntu虚拟机。在上文提到的链接中下载rknn-toolkit2，下载后的文件夹如图：
首先安装rknn-toolkit2的环境， 其环境要求在./doc/requirements_cp38-1.4.0.txt中，这里我使用的是miniconda创建的python3.8虚拟环境，和windows下anaconda使用方法类似，方便且占用空间小，推荐大家在管理python环境时使用。创建环境并命名为rknn，使用pip安装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中转换得到的best.onnx复制到该文件夹下，修改该文件夹下的test.py中的内容为自己模型的名字，要修改的地方如下：
在这个test.py的main函数中（在第239行附近）可以了解到这个python文件的作用：
【 配置——加载onnx模型—导出rknn模型——rknn模型推理——后处理——输出结果】
运行test.py，神奇的转换过程就开始了…
python test.py
最终弹窗得到post process result，若结果正确，说明我们的rknn模型转换成功了！这里的test.py构建了一个虚拟的NPU运行环境，模拟在RK3588上真实运行的情况。结果正确说明离最终完成任务只差部署在板子上这一步了。
这时在当前文件夹./example/onnx/yolov5中可以看到生成的best.rknn模型。至此，我们离胜利就不远了。
2.4 RK3588部署rknn实现NPU加速（平台：aarch板子Linux系统）
终于RK3588板子登场的时候了，但如果它还是原生的安卓系统，请按官网教程给它安装一个Linux系统吧。我这里使用TF卡烧录安装了Ubuntu系统。
依然找到官网，下载其中的rknpu2工具，链接如下：https://github.com/rockchip-linux/rknpu2
这一部分借鉴了该博客，十分详细，大家可以将转换后的best.rknn模型复制到板子上，并按照这个步骤进行操作：yolov5训练并生成rknn模型以及3588平台部署_新鑫信心的博客-CSDN博客_rknn yolov5
最终安装的命令如上图最后一行代码所示，第一个参数为生成的推理程序，第二个参数为我们之前训练和转换得到的best.rknn模型，第三个参数为待检测的图片。最终会输出推理结果out.jpg。若读者想要在自己编写的python代码中调用该程序，可以使用subprocess.Popen()命令，我书写的代码段如下，仅供参考：
p = subprocess.Popen(['./rknn_yolov5_demo', './model/RK3588/best.rknn', './defect_det.jpg'])
p.wait()
picDetected = cv2.imread('out.jpg') 我得到的out.jpg如下，一个简单的二分类缺陷检测：
到了这一步，我们已经成功完成了任务，撒花～～～但调用了NPU之后，与CPU相比究竟有没有实现加速呢？这里进行一个对比：
NPU推理速度（69ms左右）：
CPU推理速度（2.16s左右）：
推理速度快了30多倍！NPU还是很给力的。
三、总结
本文比较好地实现了yolov5模型从训练pt模型，到转换为rknn模型，最终部署在RK3588板子上使用NPU加速推理的过程。其它教程指出也可以使用rknn-toolkit-lite在RK3588上完成部署，这部分没有时间验证了，欢迎大家进行探索。
这个确实不错，老哥，