【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~3.YOLOv5训练自己的模型（亲测有效）
保姆级教程，看这一篇就够用了。
在翻阅了网络上很多资料后，发现很多版本的信息比匹配。
花了一周的时间配置环境，以及环境验证，然后写了这篇长文。
有过程，有代码，有经验，欢迎大家批评指正。
一、环境准备
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~1.环境准备（亲测有效）
二、环境验证
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~2.环境验证（亲测有效）
三、yolov5 训练自己的模型
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~3.YOLOv5训练自己的模型（亲测有效）
四、导出模型
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~4.导出模型（亲测有效）
五、部署模型
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~5.模型部署（亲测有效）
三、yolov5 训练自己的模型
（一）虚拟环境再次验证
我的电脑上有多个虚拟环境，一个是 pytorch ( 常用的深度学习框架的虚拟环境 ), 一个是 rknn( 就是之前安装 toolkit2 时新建的虚拟环境)。
pytorch虚拟环境信息
(base) jxx@desk:~$ source activate pytorch
(pytorch) jxx@desk:~$ python
Python 3.10.11 (main, Apr 20 2023, 19:02:41) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torchvision
>>> print(torch.__version__)
2.0.0+cu117
>>> print(torchvision.__version__)
0.15.1+cu117
rknn虚拟环境信息
(base) jxx@desk:~$ source activate rknn
(rknn) jxx@desk:~$ python
Python 3.8.16 (default, Mar 2 2023, 03:21:46)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torchvision
>>> print(torch.__version__)
1.10.1+cu102
>>> print(torchvision.__version__)
0.11.2+cu102
>>>
（二）仓库源码
直接下载 rknn_dev yolov5 github，并切换到 rknn_dev 分支
git clone https://github.com/wangqiqi/yolov5
git checkout rknn_dev # 切换到 rknn_dev 分支
预训练模型下载地址
yolov5s
（三）训练
1. 标注数据
可以用 cvat 或者 labelme 或者其他工具对检测数据进行标注。
在 yolov5 同级目录新建一个 datasets 文件夹，将标注好的数据打包到项目名称文件夹，然后放到 datasets 目录下，后续其他项目也可以同样进行。
2. 修改参数
（1）修改 ./data/proj.yaml
在 yolov5 文件下找到 data 目录，新建一个 proj.yaml 文件，类似于 drp.yaml。
修改 proj.yaml 中的 nc 和 names;
nc: 类别数
names: 类别名称列表
（2）修改 config_onekey.sh
根据项目要求，修改 config_onekey.sh 文件中的相关参数；
例如:
__PRJNAME, __EPOCHSZ, __BATCHSZ, __IMSZ 等等。
3. 开始一键训练
前提条件: 已经创建了虚拟环境 pytorch
前提条件: 已经创建了虚拟环境 pytorch
前提条件: 已经创建了虚拟环境 pytorch
下载预训练模型到 yolov5 同级目录 pre-trained 文件夹下，预训练模型下载地址
预训练模型下载地址
yolov5s
在训练之前确定 ./models/yolo.py 文件中的 class Detect(nn.Module): 类 line:55 里 forward 代码是否解注释了。
如果训练， forward 应为
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
## 此处如果是训练, 请解开以下注释
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
y = x[i].sigmoid()
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
return x if self.training else (torch.cat(z, 1), x)
开始一键训练
前提条件: 已经创建了虚拟环境 pytorch
前提条件: 已经创建了虚拟环境 pytorch
前提条件: 已经创建了虚拟环境 pytorch
bash onekey_bash/train.sh
终端显示如下:
(base) jxxx@desk:~/workspace/algo/yolov5_rknn$ bash onekey_bash/train.sh
train: weights=../pre-trained/yolov5s.pt, cfg=, data=data/drp.yaml, hyp=data/hyps/hyp.scratch.yaml, epochs=100, batch_size=32, imgsz=320, rect=True, resume=False, nosave=False, noval=False, noautoanchor=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, adam=False, sync_bn=False, workers=8, project=runs/train, entity=None, name=drp, exist_ok=True, quad=False, linear_lr=False, label_smoothing=0.0, upload_dataset=False, bbox_interval=-1, save_period=-1, artifact_alias=latest, local_rank=-1, freeze=0, patience=100
github: ⚠️ YOLOv5 is out of date by 1203 commits. Use `git pull` or `git clone [email protected]:wangqiqi/yolov5` to update.
YOLOv5 c5360f6 torch 2.0.0+cu117 CUDA:0 (NVIDIA GeForce RTX 4090, 24209.125MB)
hyperparameters: lr0=0.01, lrf=0.2, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=0.05, cls=0.5, cls_pw=1.0, obj=1.0, obj_pw=1.0, iou_t=0.2, anchor_t=4.0, fl_gamma=0.0, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0
Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 runs (RECOMMENDED)
TensorBoard: Start with 'tensorboard --logdir runs/train', view at http://localhost:6006/
Overriding model.yaml nc=80 with nc=1
from n params module arguments
0 -1 1 3520 models.common.Focus [3, 32, 3]
1 -1 1 18560 models.common.Conv [32, 64, 3, 2]
2 -1 1 18816 models.common.C3 [64, 64, 1]
3 -1 1 73984 models.common.Conv [64, 128, 3, 2]
4 -1 3 156928 models.common.C3 [128, 128, 3]
5 -1 1 295424 models.common.Conv [128, 256, 3, 2]
6 -1 3 625152 models.common.C3 [256, 256, 3]
7 -1 1 1180672 models.common.Conv [256, 512, 3, 2]
8 -1 1 656896 models.common.SPP [512, 512, [5, 9, 13]]
9 -1 1 1182720 models.common.C3 [512, 512, 1, False]
10 -1 1 131584 models.common.Conv [512, 256, 1, 1]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
12 [-1, 6] 1 0 models.common.Concat [1]
13 -1 1 361984 models.common.C3 [512, 256, 1, False]
14 -1 1 33024 models.common.Conv [256, 128, 1, 1]
15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
16 [-1, 4] 1 0 models.common.Concat [1]
17 -1 1 90880 models.common.C3 [256, 128, 1, False]
18 -1 1 147712 models.common.Conv [128, 128, 3, 2]
19 [-1, 14] 1 0 models.common.Concat [1]
20 -1 1 296448 models.common.C3 [256, 256, 1, False]
21 -1 1 590336 models.common.Conv [256, 256, 3, 2]
22 [-1, 10] 1 0 models.common.Concat [1]
23 -1 1 1182720 models.common.C3 [512, 512, 1, False]
24 [17, 20, 23] 1 16182 models.yolo.Detect [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
....
....
....
Epoch gpu_mem box obj cls labels img_size
98/99 1.76G 0.02157 0.007322 0 3 320: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 59/59 [00:04<00:00, 13.07it/s]
Class Images Labels P R [email protected] [email protected]:.95: 100%|███████████████████████████████████████████████████████████████████████████████| 30/30 [00:03<00:00, 7.54it/s]
all 1860 3277 0.979 0.962 0.991 0.727
Epoch gpu_mem box obj cls labels img_size
99/99 1.76G 0.0217 0.007213 0 3 320: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 59/59 [00:04<00:00, 12.90it/s]
Class Images Labels P R [email protected] [email protected]:.95: 100%|███████████████████████████████████████████████████████████████████████████████| 30/30 [00:04<00:00, 6.66it/s]
all 1860 3277 0.979 0.967 0.992 0.737
100 epochs completed in 0.244 hours.
Optimizer stripped from runs/train/drp/weights/last.pt, 14.3MB
Optimizer stripped from runs/train/drp/weights/best.pt, 14.3MB
Results saved to runs/train/drp
参考（以下为编撰本文的引用参考）
- yolov5训练pt模型并转换为rknn模型，部署在RK3588开发板上——从训练到部署全过程
- RK3588模型推理总结
- NPU使用
- 一步一步教大家使用RK3588的NPU之在电脑端交叉编译yolov5运行在ARM平台上-哔哩哔哩
- 转换yolov5s.onnx为rknn模型RK3588主板实时推理和电脑端模拟推理-哔哩哔哩
- 摆脱电脑，直接在RK3588平台上推理，帮用户部署 RKNN 模加速 AI 应用的落地-哔哩哔哩】
- 仓库下载 参考 rknn-toolkit 教程, 需要切换到 commit ID
c5360f6e7009eb4d05f14d1cc9dae0963e949213, 可以直接在 官网github下载， 然后终端执行git checkout c5360f6e7。