使用U版pytorch yolov5代码训练自己的目标检测模型

开源代码地址：https://github.com/ultralytics/yolov5

1、训练数据准备

yolo系列算法的训练数据存储格式不同于voc和coco,yolov5训练数据的构造参考Train Custom Data。

a)、yolov5算法的数据存放结构如下：

其中，train目录中存放训练集数据，test目录中存放验证集数据，train和test目录下的images目录中存放训练集和验证集图片，labels目录中用于存放训练集和验证集的标签label，其中images和labels中的文件的文件名一一对应，示例如下：

train/images/im0.jpg  # image
train/labels/im0.txt  # label

 b)、label中标签的存储方式

label目录中的每个txt文件中每行表示一个目标对象，使用五元组（class_id, x, y, w, h）来表示一个目标对象，其中class_id从0开始，表示目标类别的编号，x,y表示目标区域中心点的坐标，w,h表示目标区域的宽和高。并且，x和w是关于图像的width归一化0到1之间的结果，y和h是关于图像的height归一化到0到1之间的结果。

c)、train.txt中每行存储一个train/images目录下的图片路径，test.txt中每行存储一个test/images目录下的图片路径。

2、修改训练数据配置文件

复制一份data/coco.yaml配置文件，命名为data/custom.yaml，修改配置文件如下：

# 可以配置绝对路径，示例中将train.txt和test.txt放在data目录下
# 训练集
train: ./data/train.txt
# 验证集
val: ./data/test.txt
# 测试集
test: ./data/test.txt

# 类别数量
nc: 2

# 类别名称，类别名称顺序和class_id序号对应
names: [ 'cat', 'dog']

3、重新生成anchor box坐标

在yolo系列算法中，yolov1没有使用anchor,这也导致yolov1算法的召回率和准确率都较低，从yolo2算法开始使用anchor,并且与Faster RCNN算法使用固定尺度的anchor大小不同，yolov2算法使用聚类方法计算anchor box的大小，这样得到的anchor box比使用固定大小的anchor box更符合训练样本中目标的大小分布。yolov5算法中提供了一个计算训练集anchor分布的脚本utils/autoanchor.py,直接调用kmean_anchors计算anchor box：

其中，path参数是第二步中修改后的配置文件，n表示anchor box的数量，也就是聚类中心的数量，img_size表示最终训练模型时使用的图片的长边大小 。

4、修改网络模型配置文件

复制一份models/yolov5s.yaml配置文件，命名为models/yolov5s-custom.yaml，修改配置文件如下：

# parameters
nc: 2  # 类别数量
depth_multiple: 0.33  # model depth multiple，模型深度扩展系数
width_multiple: 0.50  # layer channel multiple，模型宽度（通道数）扩展系数

# anchors, anchor替换为使用kmean_anchors聚类得到的anchor大小
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, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, 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, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, C3, [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, 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)
  ]

5、模型训练train.py

下面看一下train.py的重点参数配置：

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 预训练权重
    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    # 网络模型配置文件models/yolov5s-custom.yaml
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    # 训练数据配置文件data/custom.yaml
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
    # 超参数配置文件，data/hyp.scratch.yaml是从头开始训练的超参数配置
    # data/hyp.finetune.yaml是微调训练模型的超参数配置
    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
    # 训练迭代轮数
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
    # 训练图片大小
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    # 是否使用矩形训练，矩形训练是关于batch size中图像大小不一致时，对图像进行最小噪声填充的方法
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    # 中断恢复训练
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--notest', action='store_true', help='only test final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train', help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    opt = parser.parse_args()

模型训练过程中的输出指标解读：

P：precision，准确率

R：recall，召回率

mAP：mean average precision，平均准确率

[email protected]: IoU为0.5的平均准确率

[email protected]:.95：IoU为0.5， 0.55, 0.6, 0.65, ……，0.95的平均准确率

该版本代码在训练过程中使用metric.py模块中的fitness函数来衡量模型的性能，fitness通过权衡P、R、[email protected]、[email protected]:.95多种指标来得到偏向不同指标的最佳模型，fitness函数如下：

def fitness(x):
    # Model fitness as a weighted combination of metrics
    w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, [email protected], [email protected]:0.95]
    return (x[:, :4] * w).sum(1)

默认情况下， P、R、[email protected]、[email protected]:.95对于模型性能的权重为[0.0, 0.0, 0.1, 0.9]，如果想让得到的最优模型更加关注召回率，可以将R对应的权重设置的大一些，如果更加关注预测结果的准确率，可以将P设置的大一些， P、R、[email protected]、[email protected]:.95权重的和等于1.

6、模型测试集评估test.py

下面看一下test.py的重点参数配置：

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='test.py')
    # 模型地址
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', 
help='model.pt path(s)')
    # 步骤2中配置的data/custom.yaml数据集配置文件
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
    # 同上
    parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
    # 同上
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    # 预测为正确的置信度阈值
    parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
    # IoU的阈值
    parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
    # 使用模型评估训练集、验证集或者测试集
    parser.add_argument('--task', default='val', help='train, val, test, speed or study')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
    # 是否在模型推理阶段使用数据增强，也就是TTA（Test Time Augment）
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
    parser.add_argument('--project', default='runs/test', help='save to project/name')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--half', type=bool, default=False, help='use FP16 half-precision inference')

模型的预测性能和--conf-thresh、--iou-thresh呈负相关， --conf-thresh、--iou-thresh越高，则模型的性能指标就越低。test.py会输出每种目标类别在当前 --conf-thresh、--iou-thresh阈值设置下的P、R、[email protected]、[email protected]:.95具体数值，以及模型在所有目标类别上的平均值。

7、模型预测detect.py

下面看一下detect.py的重点参数配置：

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 模型地址
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    # 待检测的图像目录
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    # 类别预测置信度阈值
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    # 预测bbox的IoU置信度阈值
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    # 每张图片上最多检测的目标数量
    parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    # 是否进行推理阶段的数据增强，同test.py
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--half', type=bool, default=False, help='use FP16 half-precision inference')

detect.py的预测结果保存在run/detect目录下。