RK3588部署Yolov5自训练模型

结果展示：

目录

结果展示：

硬件环境：

一. 训练Yolov5的模型

1.创建训练环境

2.拉取yolov5工程

3. 训练自己的模型

4.模型推理

二. 导出rknn模型

1. 创建导出环境

2. pt转onnx模型

3. onnx转rknn模型

三. 部署到rk3588

1. 安装rknn_toolkit2_lite

2. 推理测试

---

---

硬件环境：

PC:Ubuntu22.04

RK3588开发板:野火的鲁班猫4

一. 训练Yolov5的模型

本文所使用的yolov5模型版本为yolov5-v7版本，git如下：
https://github.com/ultralytics/yolov5/blob/v7.0/.github/README_cn.md

1.创建训练环境

为了和其他环境隔离开，使用虚拟环境进行训练，方法如下：

mkdir yolov5 & cd yolov5
virtualenv y5 #virtualenv可以通过pip来进行安装
source y5/bin/activate

2.拉取yolov5工程

git clone https://github.com/ultralytics/yolov5

3. 训练自己的模型

cd yolov5
python3 train.py --img 640 --batch 16 --epochs 20 --data data/luosi.yaml --cfg models/yolov5s.yaml

训练好的模型路径如下所示：

训练过程参考https://blog.csdn.net/u010492025/article/details/134337251

4.模型推理

使用上一步训练好的模型进行推理：

python3 detect.py --weights runs/train/exp2/weights/best.pt --source ../640/images/N_601.jpg

推理结果保存路径如下图所示：

具体效果如下，相似度还是很高的：

二. 导出rknn模型

rk3588有自己的npu，要想使用npu进行推理，需要将模型转换为rknn模型，上一步我们训练出来的模型为pt模型，要转换为rknn模型需要借助官网提供的工具rknn-toolkit2,github仓库如下：

GitHub - airockchip/rknn-toolkit2

本文中所使用的版本为1.5.0版本

1. 创建导出环境

先创建虚拟环境：

mkdir rknn-toolkit1.5 & cd rknn-toolkit1.5
virtualenv toolkit1.5
source toolkit1.5/bin/activate

下载rknn-toolkit2-1.5.0代码，并解压到rknn-toolkit1.5

安装依赖环境：

cd rknn-toolkit2-master/
#选择对应python版本的文件进行安装，我所使用的python为3.10，所以悬在cp310
pip install -r doc/requirements_cp310-1.5.0.txt
#安装依赖环境之后，安装对应的安装包
pip install packages/rknn_toolkit2-1.5.0+1fa95b5c-cp310-cp310-linux_x86_64.whl

验证环境：

没有报错说明安装成功。

2. pt转onnx模型

pt模型不能直接转换为rknn模型，需要先转为onnx或者touchscript格式，在本文中转换为onnx模型。

需要注意的是此步骤是在训练环境中进行的，而不是在导出环境。

在转换之前需要修改一下yolov5的源码，之所示要修改源码是因为按照正常的导出为onnx模型只有一个输出，而rknn模型需要三个输出，所示在将pt模型转为onnx时就需要让模型有三个输出。

修改文件models/yolo.py，如下图所示：

接下来导出onnx模型：

python3 export.py --weights runs/train/exp2/weights/best.pt --img 640 --batch 1 --include onnx --opset 12

以上导出命令有2个需要注意的地方：

① --batch 1      batch一定要指定为1，否则在使用转换后的模型进行推理的时候所需的图片数量可能不是1，会导致推理出错

② --opset 12    在后面步骤中使用rknn-toolkit2工具进行转rknn的时候可能会报错

导出的j结果如图所示：

3. onnx转rknn模型

此步骤需要在第1步创建的导出环境下进行。

进入rknn-toolkit2-master/examples/onnx目录下，复制yolov5为yolov5_test

我们使用yolov5_test来导出我们自己的模型。

首先将上一步导出的best.onnx模型复制到yolov5_test目录下;

还需要复制一张测试图片到yolov5_test目录下，用来做验证。

cd yolov5_test
cp ~/yolov5/yolov5/runs/train/exp2/weights/best.onnx ./
cp ~/N_601.jpg ./

修改test.py，中文注释的地方为修改的地方

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

ONNX_MODEL = 'best.onnx'  #修改