如何将训练好的yolov5模型(得到了训练好的pt权重)部署到RK3588开发板上:
先将训练好的模型转换为onxx格式的模型,再转换为rknn模型,最后将得到的rknn模型部署到RK3588开发板上。
1.
得到训练好的pt模型,这里为了方便,直接使用瑞芯微官方仓库的yolov5模型以及训练好的pt权重(使用瑞芯微的yolo项目能更好的适配RK3588芯片)
Yolov5模型下载地址:https://github.com/airockchip/yolov5
下载后对模型进行训练,这里为了方便直接使用官方训练好的pt权重,我们使用的是yolov5s
进入链接后在下方下载权重:

2.模型导出:
将模型导出为onnx格式的模型,首先创建自己的python虚拟环境,再进入刚才下载的yolov5项目,打开yolov5/models/yolo.py文件,将class Detection类的子函数forward函数进行修改:

将图中部分修改为:
python
Def forward(self,x):
Z=[]
for i in range(self.nl):
X[i]=self.m[i](x[i])
return x
注意:在训练时不要修改代码,在训练完成后进行导出时修改这段代码
修完完成后将pt权重放到export.py的同一目录下,打开Anaconda prompt,激活对应python虚拟环境,进入到刚才yolo.py的目录中,运行以下代码:
python export.py --weight yolov5s.pt --img 640 --batch 1 --include onnx
然后会在目录下得到yolov5s.onnx模型文件,这代表成功将模型转换到了onnx格式
3. 安装rknn-toolkit
环境:linux系统
首先用anaconda创建python虚拟环境:
conda create -n rknn python=3.8
新建一个rknn文件夹用来存放rknn-toolkit包:
进入rknn文件夹,使用git命令下载rknn-toolkit包:
Git clone https://github.com/airockchip/rknn-toolkit2.git
下载完毕后进入rknn\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
下载完成后进行rknn_toolkit的下载:
Pip install rknn_toolkit2-2.1.0+708089d1-cp38-cp38-linux_x86_64.whl
下载完成后进行测试rknn_toolkit是否安装成功:
在终端输入python,运行以下代码:
From rknn.api imort RKNN
如果没有报错就代表rknn_toolkit 已经安装成功
4. 将onnx模型转换为rknn格式:
进入rknn\rknn-toolkit2-2.1.0\rknn-toolkit2\exmaples\onnx\yolov5文件夹,将之前生成onnx格式的模型复制到当前目录下
将ONNX_MODEL修改为之前生成onnx格式的模型名,RKNN_MODEL表示生成rknn格式的模型名。
然后修改导出平台的名字:
将test.py的main函数(239行附近)中的 targrt_platform改为rk3588
然后使用命令 python test.py运行,运行成功界面如下:
这说明模型已经被转换成了rknn格式的模型
接下来只需要将rknn格式的模型部署在rk3588平台上
参考链接: