Yolov5 系列2--- 如何使用Yolov5训练你自己的数据集

上一篇《Yolov5 系列1— Yolo发展史以及Yolov5模型详解》讲了Yolo的发展历史，这一篇的目的是讲述如何使用Yolo v5训练自定的数据集，并会分析一些常见的选项以及背后的故事。

相比第一部分《Yolov5 系列1— Yolo发展史以及Yolov5模型详解》的内容，本篇文章主要聚焦于使用代码进行训练，并对一些需要关注的点进行了说明。目的是让读者能够非常快的上手代码，训练其自己的目标检测任务。

3. 配置Yolo v5的环境+准备数据集+训练`[1]`

这部分主要参照Yolo v5作者Glenn Jocher发的教程, 下面让我们开始。

3.1 准备工作

下载工程和教程数据集(coco128)，安装依赖(需要注意: Python>=3.8 and PyTorch>=1.6.)

git clone https://github.com/ultralytics/yolov5.git
curl -L -o tmp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip && unzip -q tmp.zip && rm tmp.zip  # 下载 coco128 dataset
cd yolov5
pip install -r requirements.txt  # 安装依赖

3.2 准备数据集

这里以coco128为例进行说明，coco128.yaml 是yolov5需要的数据说明文件，其形式如下，定义了（COCO是80分类的数据集）:

① 训练/验证数据的地址
② 数据集的类别
③ 类别的名称

# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /coco128
#     /yolov5


# download command/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/  # 128 images
val: ../coco128/images/train2017/  # 128 images

# 类别数量
nc: 80

# 类别名称
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

对使用自己的数据集时，我们需要按如下形式组织数据, xxx为数据集名称：

| ---- yolov5
| ---- xxx
|       | --- images
|                |  --- train
|                         | --- train_1.jpg
|                         | --- ...
|                         | --- train_n.jpg 
|                |  --- val
|                         | --- val_1.jpg 
|                         | --- ...
|                         | --- val_n.jpg 
|       | --- labels 
|                |  --- train
|                         | --- train_1.txt
|                         | --- ...
|                         | --- train_n.txt
|                |  --- val
|                         | --- val_1.txt
|                         | --- ...
|                         | --- val_n.txt

当你的xxx.yaml准备好后，下面就是创建标签（自己写个脚本就好~）。注意，Yolo v5的标签形式是独特的，对应于xxx\labels\train或xxx\labels\val中的.txt文件，标签文件有如下特点:

① 一行对应一个物体/bbox
② 每行的格式都是类别c bbox中心坐标_x bbox中心坐标y bbox的宽 bbox的高，以空格作为分隔符
③ 坐标都被归一化都0到1之间，如果你的bbox标注是按照像素的绝对位置来的，可以对其进行归一化: 将bbox中心坐标_x和bbox的宽除以图像的宽度，将bbox中心坐标_y和bbox的高除以图像的高度
④ 类别从0开始计算

如果准备好了之后，就会得到一一对应的原图jpg和其标签txt文件, 以coco128为例：

coco128/images/train2017/000000109622.jpg  # image
coco128/images/train2017/000000109622.txt  # label

3.3 选择模型

这里我们以Yolo v5s 2为例进行，因为这是Yolo v5官方提供的目前为止最smallest and fastest的模型，下面是Yolo v5的small, medium, large, super-large在COCO数据集上的表现。可以看出，Yolo v5s是最快的，参数最少的，但也是精度最低的。。。

关于模型的选择，需要结合任务来定，我自己的理解是: 如果你的任务类似COCO，即复杂场景，目标多样，那还是尽量选择大模型；若检测目标单一，那么可以用Yolo v5m甚至Yolo v5s.

在这里插入图片描述
如果你训练自己的数据集，不要忘记把对应的xxx.yaml的类别nc从80改成你需要的类别数量，以及修改类别的名称。

# parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# anchors
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 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9
  ]

# YOLOv5 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, BottleneckCSP, [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, BottleneckCSP, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, BottleneckCSP, [512, False]],  # 20  (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, BottleneckCSP, [1024, False]],  # 23  (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

3.4 模型训练

训练根据计算资源(显卡)的情况，主要分3种: 单机单卡, 单机多卡，多机多卡。

需要说明的是，Yolo v5的作者采用的不是熟悉的DP模式torch.nn.DataParallel，而是torch.nn.parallel.DistributedDataParallel, 即DDP模式。按照PyTorch官方的说法[3]: 即使在单机多卡的环境，也推荐DDP模式，DDP模式会显著快于DP模型，关于DDP的例子可见官网[4].

在这里插入图片描述

单机单卡训练:

python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ''

单机多卡训练:

python train.py --img 640 --batch 16 --epochs 5  --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights '' --device 0,1 # 指定第1,2块卡进行训练

多机多卡训练:

python -m torch.distributed.launch --nproc_per_node 2 train.py --weights yolov5s.pt --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --epochs 3 --img 320 --device 0,1 # DDP

所有的训练结果都存在yolov5/runs/文件夹下，按照exp0, exp1的顺序升序排列。

如果成功训练起来了，那么在yolov5/runs/exp0里面会有如下的结构:
在这里插入图片描述

weights/对应模型文件存放的位置(pytorch的pt模型存放地址，Yolo v5会保存在验证集上表现最好的模型best.pt和训练到最后一个epoch的last.pt这2个模型)
labels.png对应训练数据的类别分布 (这里我放的是我自己的数据集…) ，可以看到类别不均衡的现象比较明显，这也是目标检测任务里常见的情况，后面会解释Yolo v5是如何处理这种情况的。
result.png对应训练的各种精度，损失函数指标
横坐标是epoch，纵坐标是各种指标的精度或者loss值。蓝色表示是从头训练的模型的指标；而橙色则是在pre-train model上进行finetune的结果。
opt.yaml对应训练的参数
比如
① 训练的模式: start from scratch(从头开始) 还是 finetune (有预训练模型)。
② 训练的设置: epoch, batch size, 图片分辨率, 是否开启autoanchor(在系列文章1中有介绍), 优化器是Adam还是SGD, 日志文件夹等内容, 以供debug使用。

下表是知乎上一位同学的详细回答版, 找不到原出处了…
在这里插入图片描述

hyp.yaml对应超参数
包括学习率，loss衰减，anchor和颜色空间转换的相关参数，随机翻转的幅度等。

好了，到这步，属于你的Yolo v5s模型就成功训练了，你可以使用
python detect.py --source ./inference/images/ --weights runs/exp0/weights/yolov5s.pt --conf 0.4 --device cpu --img-size 320
进行测试啦~

4. 常见问题

4.1 mAP的含义

在看目标检测的论文时，经常看到用于评价目标检测的一个度量标准mean average precision (mAP)，请问mAP是什么含义，又是如何计算的？

4.1.1 mAP相关内容说明

mAP: mean Average Precision, 即各类别AP的平均值
AP: Precision-Recall曲线下面积，
PR曲线: Precision-Recall曲线
Precision: $\frac {TP} {(TP + FP)}$ 在你认为的正样本中，有多大比例是真的正样本。
Recall: $\frac {TP} {(TP + FN)}$ 在真的样本中，有多少被我们找到了。

以all detections表示所有的预测框的数量，all gts表示所有的GT的数量，那么Precision和Recall可以被重写为如下形式[5]:

TP: IoU> threshold的检测框数量（同一Ground Truth只计算一次）
FP: IoU<= threshold的检测框，或者是检测到同一个GT的多余检测框的数量
FN: 没有检测到的GT的数量

4.1.2 知乎陈子豪同学的回答`[4]`

mAP解释的通俗版本。
在这里插入图片描述

4.1.3 mAP的具体计算逻辑`[5]`

在这里插入图片描述
在我的试验中, conf取0.4.

Q1: yolov5的conf的概念?
A1: 表示预测的bbox与实际的某个ground truth (GT_a)的IoU, 若IoU > conf (默认为0.4), 则认为该bbox为TP, 若
GT_a有多个预测的bbox, 则认为IoU最大的且 > conf的那个bbox为TP, 其他的为FP.

从上图的图1可以看出(绿色bbox为GT, 红色为预测的bbox)，bboxA是FP, bboxB为TP, bboxC为FP.

Q2: FN是啥？
A2: FN按照我的理解，只会出现在样本有noise的情况，即标注有错误的情况，有漏标等情况。

通过下表，我们可以绘制出 P-R 曲线（因为 AP 就是 P-R 曲线下面的面积），但是在此之前我们需要计算出 P-R 曲线上各个点的坐标，根据置信度从大到小排序所有的预测框，然后就可以计算 Precision 和 Recall 的值。
在这里插入图片描述
见下表。（需要记住一个叫累加的概念，就是下图的 ACC TP 和 ACC FP）
以Image5的bbox R为例, Recall = 1/ 15(15是ground truth的数量) = 0.0666, Precision = 1/ 1 = 1.
以此类推…

这下, 目标检测任务的里mAP计算就非常清楚了. 下面我们看看模型的推理速度.

4.2 Yolov5的模型速度profiling

一个常见的需求是将模型放到CPU上去跑,这下TensorRT就不能用. 而我们想要分析算子占据的时间, 以便进行模型的裁剪, 用蒸馏或者剪枝的方式压缩模型的flops, 以达到速度提升(精度不明显下降时).

因为Yolo v5所耗得时间大部分是网络推理部分而非NMS. 速度慢的原因主要是由于backbone要提取的图像太大了, 这里利用PyTorch 1.6.0 提供的profiling工具进测试[6].

其方式也很简单, 即在[8]的代码的第71行加入如下2行代码

import argparse
import os
import platform
import shutil
import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import select_device, load_classifier, time_synchronized


def detect(save_img=False):
    ...
    for path, img, im0s, vid_cap in dataset:
        ...
        with torch.autograd.profiler.profile() as prof:
	        # Inference
	        t1 = time_synchronized()
	        pred = model(img, augment=opt.augment)[0]
		print(prof.key_averages().table(sort_by="self_cpu_time_total"))
		...

以我的Yolov5s测试来看, 主要的时间耗费在卷积上, 所以, 如果想达到更快的速度, 需要考虑删减BottleneckCSP层以及卷积层的channel.
在这里插入图片描述

4.3 pytorch模型导出onnx

在之前的依赖版本基础上, 安装ONNX>=1.7. 以训练好的模型为例, 下面是转换的脚本代码.

cd yolov5
export PYTHONPATH="$PWD"  # add path
python models/export.py --weights runs/exp0/weights/best.pt --img 640 --batch 1  # export

如果成功执行, 那么会在runs/exp0/weights目录下多了一个best.onnx模型文件, 这个模型通常会比PyTorch原版的pt格式模型大一点.

用户可以用Netron来查看对应的ONNX模型:

在这里插入图片描述
需要注意, 默认情况下, 不会将Detect层纳入到onnx模型中, 我们需要在models\export.py中, 修改如下代码[7]:

    # Update model
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
        if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish):
            m.act = Hardswish()  # assign activation
        # if isinstance(m, models.yolo.Detect):
        #     m.forward = m.forward_export  # assign forward (optional)
    # 默认是True.
    model.model[-1].export = False # set Detect() layer export=True
    y = model(img)  # dry run

参考内容

[1] Train Custom Data Tutorial #12
[2] Yolo v5s yaml
[3] PyTorch 1.6.0 torch.nn.DataParallel
[4] 目标检测中的mAP是什么含义？–陈子豪
[5] rafaelpadilla/Object-Detection-Metrics
[6] torch.autograd.profiler.profile
[7] yolo v5 models/export.py
[8] yolo v5 detect.py