【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~3.YOLOv5训练自己的模型（亲测有效）

黄金旺铺

已于 2024-04-27 21:32:29 修改

阅读量1.4k

点赞数

分类专栏：边缘设备文章标签： YOLO

于 2023-08-16 15:31:14 首次发布

by: JW

本文链接：https://blog.csdn.net/zhoujinwang/article/details/132320778

版权

边缘设备专栏收录该内容

15 篇文章 1 订阅

订阅专栏

保姆级教程，看这一篇就够用了。
在翻阅了网络上很多资料后，发现很多版本的信息不匹配。
花了一周的时间配置环境，以及环境验证，然后写了这篇长文。
有过程，有代码，有经验，欢迎大家批评指正。

一、环境准备
【边缘设备】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 git@github.com: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     mAP@.5 mAP@.5:.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     mAP@.5 mAP@.5:.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。

黄金旺铺

关注

0
点赞
踩
12

收藏

觉得还不错? 一键收藏
打赏
0
评论
【边缘设备】yolov5训练与rknn模型导出并在RK3588部署~3.YOLOv5训练自己的模型（亲测有效）

在翻阅了网络上很多资料后，发现很多版本的信息比匹配。花了一周的时间配置环境，以及环境验证，然后写了这篇长文。有过程，有代码，有经验，欢迎大家批评指正。
复制链接

扫一扫