在国产芯片上实现YOLOv5/v8图像AI识别-【1.3】YOLOv5的介绍及使用（训练、导出）更多内容见视频

本专栏主要是提供一种国产化图像识别的解决方案，专栏中实现了YOLOv5/v8在国产化芯片上的使用部署，并可以实现网页端实时查看。根据自己的具体需求可以直接产品化部署使用。

B站配套视频：https://www.bilibili.com/video/BV1or421T74f
在这里插入图片描述

数据训练

上一篇博客里面我们已经获得了标注好的数据以及图片，接下来我们就要开始训练过程。

数据整备阶段

首先在yolov5的目录下创建一个datasets目录，这一步是个人习惯，我们将要训练的数据都会放在这里。

所有的数据需要按照目录规范进行放置，通常train、val的分配比例为8:2，images和labels里面的内容需要对应。
在这里插入图片描述

yaml文件准备阶段

编写数据说明文件和结构说明文件，找到data目录下创建一个yaml文件，此处以我个人创建的举例。

数据yaml文件

可以直接复制目录中的coco128.yaml进行修改，将其中的download部分删掉。然后更具自己数据的存放路径进行配置。
在这里插入图片描述

train 写训练图片的完整路径，经过多次尝试，写相对路径会有问题。
val 写验证图片的完整路径。系统会自动找到对应的labels目录。
test 可以不用写，对训练结果不会有影响。
nc 写你需要识别的数量。
names 写你需要识别的类别，此处循序一定要注意。

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
# path: ../datasets/cx  # dataset root dir
train: /app/docs/yolov5_v7.0/datasets/cx/images/train  # train images (relative to 'path') 128 images
val: /app/docs/yolov5_v7.0/datasets/cx/images/train  # val images (relative to 'path') 128 images
test:  # test images (optional)# Classes
nc: 2  # number of classes
names: ['good','bad']  # class names

结构yaml文件

打开model目录，找到下面的yaml文件，此处系统已经默认了一些模型文件。通常不需要进行模型魔改的情况下可以基于pt训练，如果需要魔改模型需要自己重新设置一个yaml文件。此处举例看一下。
在这里插入图片描述
下文是一个修改后的s模型文件，主要修改的就是nc，其他内容如果没有学习过模型魔改就不要动。

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license# Parameters
nc: 2  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:- [10,13, 16,30, 33,23]  # P3/8- [30,61, 62,45, 59,119]  # P4/16- [116,90, 156,198, 373,326]  # P5/32# YOLOv5 v6.0 backbone
backbone:# [from, number, module, args][[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2[-1, 1, Conv, [128, 3, 2]],  # 1-P2/4[-1, 3, C3, [128]],[-1, 1, Conv, [256, 3, 2]],  # 3-P3/8[-1, 6, C3, [256]],[-1, 1, Conv, [512, 3, 2]],  # 5-P4/16[-1, 9, C3, [512]],[-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32[-1, 3, C3, [1024]],[-1, 1, SPPF, [1024, 5]],  # 9]# YOLOv5 v6.0 head
head:[[-1, 1, Conv, [512, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 6], 1, Concat, [1]],  # cat backbone P4[-1, 3, C3, [512, False]],  # 13[-1, 1, Conv, [256, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 4], 1, Concat, [1]],  # cat backbone P3[-1, 3, C3, [256, False]],  # 17 (P3/8-small)[-1, 1, Conv, [256, 3, 2]],[[-1, 14], 1, Concat, [1]],  # cat head P4[-1, 3, C3, [512, False]],  # 20 (P4/16-medium)[-1, 1, Conv, [512, 3, 2]],[[-1, 10], 1, Concat, [1]],  # cat head P5[-1, 3, C3, [1024, False]],  # 23 (P5/32-large)[[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)]

启动训练

训练其实很简单，就是要根据自己的实际电脑配置进行训练参数的调整。如果是常规训练参考以下参数就可以了，参数weights和cfg选择一个配置就可以。

python ./train.py --data ./data/coco_bz.yaml --weights ./weights/yolov5s.pt --cfg ./models/yolov5s_bz.yaml --batch-size 48 --epochs 200 --workers 0 --name 754-200 --project yolo5_bz_s

python ./train.py 是官方写好的训练搅拌文件
– data 指向数据整备的yaml文件
–weights 指向已经训练好的模型，这个参数主要两个作用，一个是基于模型结构训练，另一个是训练大型模型异常终端可以接续训练
–cfg 指向模型配置文件
–batch-size 单次训练的的数据量大小，理论是越大越好，要根据实际电脑配置调整，此处注意v5和v8这个参数同一台电脑也不同
–epochs 总过训练多少轮，设置特别大没有太大意义，代码在发现最后几次训练模型没有改进会自动停止
–workers 此项也是根据显卡性能调整，如果不知道怎么设置可以直接设置成0，系统会自动配置
–name 代表本次定义的名字，默认会写exp，重复训练会自动加上后缀
–project 代表本次工程名字，默认会写runs

启动之后会显示以下内容，如果是第一次训练系统会需要下载几个文件可能比较慢，一个是文字格式文件，一个是yolov5模型用来进行数据处理的。

root@935f1467d228:/app/docs/yolov5_v7.0# python ./train.py --data ./data/coco_bz.yaml --cfg ./models/yolov5n_bz.yaml --batch-size 48 --epochs 200 --workers 0 --name 754-200 --project yolo5_bz_ntrain: weights=yolov5s.pt, cfg=./models/yolov5n_bz.yaml, data=./data/coco_bz.yaml, hyp=data/hyps/hyp.scratch-low.yaml, epochs=200, batch_size=48, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=0, project=yolo5_bz_n, name=754-200, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-1, seed=0, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest
remote: Enumerating objects: 390, done.
remote: Counting objects: 100% (273/273), done.
remote: Compressing objects: 100% (96/96), done.
remote: Total 390 (delta 204), reused 217 (delta 177), pack-reused 117
Receiving objects: 100% (390/390), 476.92 KiB | 599.00 KiB/s, done.
Resolving deltas: 100% (261/261), completed with 88 local objects.
From https://github.com/ultralytics/yolov5920c721e..8003649c  master                  -> origin/master* [new branch]        refactor-20240717220233 -> origin/refactor-20240717220233* [new branch]        snyk-fix-19a9bd869ca677b68dcdaf5f4affcd24 -> origin/snyk-fix-19a9bd869ca677b68dcdaf5f4affcd24* [new branch]        snyk-fix-f5bfc0187c0599da5db2839fa7a5f8f5 -> origin/snyk-fix-f5bfc0187c0599da5db2839fa7a5f8f5
github: ⚠️ YOLOv5 is out of date by 345 commits. Use `git pull` or `git clone https://github.com/ultralytics/yolov5` to update.
YOLOv5 🚀 v7.0-0-g915bbf29 Python-3.8.12 torch-1.12.0a0+2c916ef CUDA:0 (NVIDIA GeForce RTX 3060, 12022MiB)hyperparameters: lr0=0.01, lrf=0.01, 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
ClearML: run 'pip install clearml' to automatically track, visualize and remotely train YOLOv5 🚀 in ClearML
Comet: run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet
TensorBoard: Start with 'tensorboard --logdir yolo5_bz_n', view at http://localhost:6006/from  n    params  module                                  arguments                     0                -1  1      1760  models.common.Conv                      [3, 16, 6, 2, 2]              1                -1  1      4672  models.common.Conv                      [16, 32, 3, 2]                2                -1  1      4800  models.common.C3                        [32, 32, 1]                   3                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                4                -1  2     29184  models.common.C3                        [64, 64, 2]                   5                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               6                -1  3    156928  models.common.C3                        [128, 128, 3]                 7                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              8                -1  1    296448  models.common.C3                        [256, 256, 1]                 9                -1  1    164608  models.common.SPPF                      [256, 256, 5]                 10                -1  1     33024  models.common.Conv                      [256, 128, 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     90880  models.common.C3                        [256, 128, 1, False]          14                -1  1      8320  models.common.Conv                      [128, 64, 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     22912  models.common.C3                        [128, 64, 1, False]           18                -1  1     36992  models.common.Conv                      [64, 64, 3, 2]                19          [-1, 14]  1         0  models.common.Concat                    [1]                           20                -1  1     74496  models.common.C3                        [128, 128, 1, False]          21                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]              22          [-1, 10]  1         0  models.common.Concat                    [1]                           23                -1  1    296448  models.common.C3                        [256, 256, 1, False]          24      [17, 20, 23]  1      9471  models.yolo.Detect                      [2, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [64, 128, 256]]
YOLOv5n_bz summary: 214 layers, 1766623 parameters, 1766623 gradients, 4.2 GFLOPsTransferred 57/349 items from yolov5s.pt
AMP: checks passed ✅
optimizer: SGD(lr=0.01) with parameter groups 57 weight(decay=0.0), 60 weight(decay=0.000375), 60 bias
train: Scanning /app/docs/yolov8/ultralytics/datasets/bz_754/labels/train... 754 images, 0 backgrounds, 0 corrupt: 100%|██████████| 754/754 00:00
train: New cache created: /app/docs/yolov8/ultralytics/datasets/bz_754/labels/train.cache
val: Scanning /app/docs/yolov8/ultralytics/datasets/bz_754/labels/val... 200 images, 0 backgrounds, 0 corrupt: 100%|██████████| 200/200 00:00
val: New cache created: /app/docs/yolov8/ultralytics/datasets/bz_754/labels/val.cacheAutoAnchor: 5.75 anchors/target, 1.000 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅
Plotting labels to yolo5_bz_n/754-2002/labels.jpg... 
Image sizes 640 train, 640 val
Using 0 dataloader workers
Logging results to yolo5_bz_n/754-2002
Starting training for 200 epochs...Epoch    GPU_mem   box_loss   obj_loss   cls_loss  Instances       Size0/199      5.74G     0.1083     0.1563    0.02792       1076        640:  19%|█▉        | 3/16 00:08

完成训练

完成训练后我们会获得以下目录，其中weights中就有我们最优的模型和最后一次训练模型，通常使用best，其他参数信息我会在视频讲解中和大家说明。视频链接见博客顶部。
在这里插入图片描述

导出模型

执行以下指令就可以导出我们onnx的模型，这里面

python export.py --weights weights/yolov5s.pt --include  onnx --opset 11

使用Netron就可以看到整体模型的内容，Netron的安装地址 https://github.com/lutzroeder/netron
在这里插入图片描述