欢迎关注『youcans动手学模型』系列
本专栏内容和资源同步到 GitHub/youcans
【YOLO5 项目实战】（1）YOLO5 环境配置与测试
【YOLO5 项目实战】（2）使用自己的数据集训练目标检测模型
【YOLO5 项目实战】（3）PCB 缺陷检测

上节介绍了 YOLOv5 的下载、配置和推理，本节介绍用本地数据集训练 YOLOv5 模型，建立特定任务的模型。

创建自定义模型来检测对象是一个迭代过程，包括收集和组织图像、标注对象、训练模型、模型部署，使用部署的模型进行推理。

使用自己的数据集训练专用的 YOLOv5 模型，可以使用网络上的开源数据集，也可以使用自己收集的图片进行标注制作数据集。

1. 准备训练数据集

1.1 开源数据集

网络上有很多开源的目标检测数据集。

kaggle：一个庞大的社区发布的模型、数据和代码库
Roboflow：构建和部署计算机视觉模型所需的一切
极市：海量数据集
数据科学银行：数据发布的开放可信全球公共产品

1.2 YOLOv5 训练数据集

YOLOv5 默认使用 COCO2017 数据集进行训练，将数据集保存在与 YOLOv5 同一级的文件目录路径中。结构如下。

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco  ← downloads here (20.1 GB)
 
# 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: ../mydatas/  # dataset root dir
train: mydatas/train.txt  # train images (relative to 'path') 118287 images
val: mydatas/val.txt  # val images (relative to 'path') 5000 images
test: mydatas/test.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

COCO128 是一个小型数据集，由COCO train2017中的前128张图像组成。这128张图像用于训练和验证，以验证我们的训练管道是否能够过拟合。

data/coco128.yaml是数据集配置文件，如下所示，它定义了：
（1）数据集根目录路径和训练/val/test图像目录的相对路径（或具有图像路径的*.txt文件），
（2）类名字典：

# 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/coco128 # dataset root dir
train: images/train2017 # train images (relative to 'path') 128 images
val: images/train2017 # val images (relative to 'path') 128 images
test: # test images (optional)

# Classes (80 COCO classes)
names:
    0: person
    1: bicycle
    2: car
    # ...
    77: teddy bear
    78: hair drier
    79: toothbrush

1.3 从 Roboflow 下载公开数据集

Roboflow Annotate是一个简单的基于网络的工具，用于管理和标记用户的图像集，并以YOLOv5的注释格式导出。无论是否使用 Roboflow 标记图像，都可以使用它将数据集转换为YOLO格式，创建YOLOv5 YAML配置文件，并将其托管以导入训练脚本。

（1）打开 Roboflow 网站，从Roboflow 公开数据集中选择一个数据集，下载到本地。

例如，作者选择 口罩数据集 Mask Wearing Dataset。该数据集是一个物体检测数据集，包括 149张 戴各种口罩和不戴口罩的人脸图像。可以使用这个数据集来训练一个模型，用于检测给定照片中的人是否戴口罩。

（2）Roboflow 提供了多种下载格式，主要是数据集组织和标注格式的区别。注意要以YOLOv5 Pytorch格式导出，如下图所示。

（3）下载后将数据集解压缩，保存在与 YOLOv5 同一级的文件目录路径中。结构如下所示。

- Python
  - PythonProjects
    - DataSetMask
      - test
        - images
        - labels
      - train
        - images
        - labels
      - valid
        - images
        - labels
    - VOLOv5

Mask Wearing Dataset 数据集设有 test，train，valid 三个文件夹，分别用作测试、训练和检验，每个文件夹下设有 images，labels 两个文件夹，分别保存图像文件和标注文件。

一个典型的标注文件的内容如下。共有 3行，每行表示一个检测目标。每行有 5个参数，第1列是类别标签（Yes/No），后 4个参数是 BoundingBox 的坐标位置。

0 0.31875 0.5976190476190476 0.05583333333333333 0.13174603174603175
0 0.31916666666666665 0.2912698412698413 0.12833333333333333 0.3253968253968254
0 0.5270833333333333 0.16746031746031745 0.12416666666666666 0.2841269841269841

如果自己收集和整理数据集，也要按照以上格式来组织数据集和标注文件。

2. 修改数据集配置文件 data.yaml

YOLOv5 模型训练时，要调用数据集配置文件 YAML 文件（ .yaml）。

因此，需要在数据集所在的文件夹（例如…/DataSetMask/）创建一个 data.yaml 的文件来存放我们已经整理好的数据，内容如下：

train: ../DataSetMask/train/images
val: ../DataSetMask/valid/images
test: ../DataSetMask/valid/images


nc: 2
names: ['mask', 'no-mask']

其中，train 表示训练集图像文件夹的路径，val 表示验证集图像文件夹的路径，test 表示测试集图像文件夹的路径。nc:2 表示类别数为 2（分为带口罩、不带口罩 2类），names 表示类别名（类别名的个数为 nc）。

注意，nc 是由数据集的标注内容决定的，不能自行修改。

3. 修改 YOLOv5 的训练程序 train.py

3.1 下载 YOLOv5 预训练模型

推荐从 YOLOv5 release下载 YOLOv5 预训练模型。

本文选择 YOLOv5s，参数约 7.2M。下载完成后，将下载的预训练模型文件 yolov5s.pt 放在 YOLOv5 项目路径下，例如"C:\Python|PythonProjects\YOLOv5\yolov5s.pt"。

3.2 配置 YOLOv5 模型训练的参数

使用自己的数据集训练 YOLOv5 模型，运行 train.py 需要注意以下参数：

• weight，先选用官方的 yolov5s.pt 权重

• cfg，选择 model 目录下的 yolov5s 模型配置文件

• data，选择自己编写的数据集配置文件 dataset.yaml

• epoch，指整个数据集被训练次数，根据计算机性能和需要确定，新手练习时可以设为 2

• batch_size，每次读入的样本数量，根据计算机性能和需要确定，新手练习时可以设为 2

3.3 命令行运行 YOLOv5 训练程序

通过命令行就可以运行 YOLOv5 训练程序 train.py ：

python train.py --weights yolov5s.pt --cfg models/yolov5s.yaml --data …/DataSetMask/data.yaml --epoch 2 --batch-size 4

3.4 修改训练程序 train.py

对于新手，通过命令行直接输入模型训练参数比较麻烦，也可以直接对 YOLOv5 训练程序进行修改，设置模型训练的参数。

在 train.py 中对 weight，cfg，data 参数进行修改，另存为 train1.py：

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
    parser.add_argument("--cfg", type=str, default=ROOT /"models/yolov5s.yaml", help="model.yaml path")
    parser.add_argument("--data", type=str, default=ROOT / "data/mydata.yaml", help="dataset.yaml path")
    ...

进一步的，也可以在程序中修改默认的 训练参数，例如：

    parser.add_argument("--epochs", type=int, default=5, help="total training epochs")
    parser.add_argument("--batch-size", type=int, default=2, help="total batch size for all GPUs, -1 for autobatch")

4. 训练 YOLOv5 模型

4.1 使用 train1.py 训练模型

使用修改的模型训练程序 train1.py，对本地数据集 DataSetMask 进行训练。结果如下。

训练过程中可能出现一些报错，处理方法详见下节“报错的处理”。

4.2 使用命令行训练模型

使用命令行运行 YOLOv5 训练程序 train.py ，运行结果如下：

(torch) PS C:\Python\PythonProjects\YOLOv5> conda activate torch
(torch) PS C:\Python\PythonProjects\YOLOv5> python train1.py --weights yolov5s.pt  --cfg models/yolov5s.yaml  --data ../DataSetMask/data.yaml --epoch 10 --batch-size 4 --img 640   --device cpu

train1: weights=yolov5s.pt, cfg=models/yolov5s.yaml, data=../DataSetMask/data.yaml, hyp=data\hyps\hyp.scratch-low.yaml, epochs=10, batch_size=4, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, evolve_population=data\hyps, resume_evolve=None, bucket=, cache=None, image_weights=False, device=cpu, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs\train, name=exp, 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, ndjson_console=False, ndjson_file=False
github: skipping check (not a git repository), for updates see https://github.com/ultralytics/yolov5
YOLOv5  2024-7-29 Python-3.8.19 torch-2.3.1+cu121 CPU

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
Comet: run 'pip install comet_ml' to automatically track and visualize YOLOv5  runs in Comet
TensorBoard: Start with 'tensorboard --logdir runs\train', view at http://localhost:6006/
Overriding model.yaml nc=80 with nc=2

                 from  n    params  module                                  arguments
  0                -1  1      3520  models.common.Conv                      [3, 32, 6, 2, 2]              
  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  2    115712  models.common.C3                        [128, 128, 2]                 
  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   1182720  models.common.C3                        [512, 512, 1]                 
  9                -1  1    656896  models.common.SPPF                      [512, 512, 5]
 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     18879  models.yolo.Detect                      [2, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
YOLOv5s summary: 214 layers, 7025023 parameters, 7025023 gradients, 16.0 GFLOPs

Transferred 342/349 items from yolov5s.pt
optimizer: SGD(lr=0.01) with parameter groups 57 weight(decay=0.0), 60 weight(decay=0.0005), 60 bias
train: Scanning C:\Python\PythonProjects\DataSetMask\train\labels.cache... 105 images, 0 backgrounds, 0 corrupt: 100%|██████████| 105/105 [00:00<?, ?it/s]
val: Scanning C:\Python\PythonProjects\DataSetMask\valid\labels.cache... 29 images, 0 backgrounds, 0 corrupt: 100%|██████████| 29/29 [00:00<?, ?it/s]


AutoAnchor: 5.89 anchors/target, 0.999 Best Possible Recall (BPR). Current anchors are a good fit to dataset
Plotting labels to runs\train\exp2\labels.jpg... 
Image sizes 640 train, 640 val
Using 4 dataloader workers
Logging results to runs\train\exp2
Starting training for 10 epochs...

      Epoch    GPU_mem   box_loss   obj_loss   cls_loss  Instances       Size
        0/9         0G     0.1147    0.06687    0.02794          0        640: 100%|██████████| 27/27 [01:40<00:00,  3.73s/it]
                 Class     Images  Instances          P          R      mAP50   mAP50-95:   0%|          | 0/4 [00:00<?, ?it/s]WARNING  NMS time limit 0.900s exceeded
                 Class     Images  Instances          P          R      mAP50   mAP50-95:  25%|██▌       | 1/4 [00:02<00:07,  2.57s/it]WARNING  NMS time limit 0.900s exceeded
                 Class     Images  Instances          P          R      mAP50   mAP50-95:  50%|█████     | 2/4 [00:05<00:05,  2.78s/it]WARNING  NMS time limit 0.900s exceeded
                 Class     Images  Instances          P          R      mAP50   mAP50-95:  75%|███████▌  | 3/4 [00:09<00:03,  3.27s/it]WARNING  NMS time limit 0.750s exceeded
                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 [00:11<00:00,  2.82s/it]
                   all         29        162    0.00479      0.149     0.0033   0.000824

...

10 epochs completed in 0.301 hours.
Optimizer stripped from runs\train\exp2\weights\last.pt, 14.4MB
Optimizer stripped from runs\train\exp2\weights\best.pt, 14.4MB

Validating runs\train\exp2\weights\best.pt...
Fusing layers... 
YOLOv5s summary: 157 layers, 7015519 parameters, 0 gradients, 15.8 GFLOPs
                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 [00:08<00:00,  2.07s/it]
                   all         29        162      0.744      0.292      0.315      0.146
                  mask         29        142      0.489      0.585      0.497      0.227
               no-mask         29         20          1          0      0.133     0.0649
Results saved to runs\train\exp2

4.3 训练结果

训练好的模型自动保存在 YOLOv5 目录下的 runs/train/exp/weights/ 文件夹中：

训练日志自动保存在 YOLOv5 目录下的 runs/train/exp/ 文件夹中：

5. 使用训练模型进行推理

detect.py 程序默认读取 data\images 路径的图片，结果默认保存到 runs/detect 文件夹中。

直接在 PyCharm 运行 detect.py 程序，或者从控制台运行 detect.py。这里我们刚才使用训练的模型 best.pt，对图片 Mask02.jpg 进行目标检测和识别。

python detect.py --weights runs/train/exp/weights/best.pt --source data/images/Mask02.jpg

运行结果如下。

(torch) PS C:\Python\PythonProjects\YOLOv5> python detect.py --weights runs/train/exp/weights/best.pt --source data/images/Mask02.jpg
detect: weights=['runs/train/exp/weights/best.pt'], source=data/images/Mask02.jpg, data=data\coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_csv=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs\detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5  2024-7-29 Python-3.8.19 torch-2.3.1+cu121 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)

Fusing layers...
YOLOv5s summary: 157 layers, 7015519 parameters, 0 gradients, 15.8 GFLOPs
image 1/1 C:\Python\PythonProjects\YOLOv5\data\images\Mask02.jpg: 480x640 3 masks, 1 no-mask, 126.7ms
Speed: 3.3ms pre-process, 126.7ms inference, 25.2ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs\detect\exp4

运行结果保存在 runs\detect\exp4 路径，结果如下图所示。

6. 错误的处理

报错1：Git 执行发生错误

ImportError: Bad git executable.

ImportError: Bad git executable.
The git executable must be specified in one of the following ways:
- be included in your $PATH
- be set via $GIT_PYTHON_GIT_EXECUTABLE
- explicitly set via git.refresh(<full-path-to-git-executable>)

解决方案：按照报错信息找到git下的cmd文件，添加下方代码

import os
os.environ["GIT_PYTHON_REFRESH"] = "quiet"

报错2：没有正常生成标签

AssertionError: train: All labels empty in E:\pythonProject\y2_data\train.cache, can not start training. See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

解决方案：大小写问题，类别名如果以大写开头则无法生成标签，labels中的txt文件都是空的，要改称小写字母开头。例如：

classes = ["banana", "snake fruit", "dragon fruit", "pineapple"]

【本节完】

版权声明：
欢迎关注『youcans动手学模型』系列
转发请注明原文链接：
YOLO5 项目实战】（2）使用自己的数据集训练目标检测模型
Copyright 2024 youcans, XUPT
Crated：2024-08-08