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

桃扇骨 2022-10-10 06:05 298阅读 0赞

开源代码地址：[https://github.com/ultralytics/yolov5][https_github.com_ultralytics_yolov5]

**1、训练数据准备**

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

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

![20210627101818386.png][]

其中，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之间的结果。

![Image Labels][]

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：

![20210627111740145.png][]

其中，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，平均准确率

mAP@.5: IoU为0.5的平均准确率

mAP@.5:.95：IoU为0.5， 0.55, 0.6, 0.65, ……，0.95的平均准确率

该版本代码在训练过程中使用metric.py模块中的fitness函数来衡量模型的性能，fitness通过权衡P、R、mAP@.5、mAP@.5:.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, mAP@0.5, mAP@0.5:0.95]
        return (x[:, :4] * w).sum(1)

默认情况下， P、R、mAP@.5、mAP@.5:.95对于模型性能的权重为\[0.0, 0.0, 0.1, 0.9\]，如果想让得到的最优模型更加关注召回率，可以将R对应的权重设置的大一些，如果更加关注预测结果的准确率，可以将P设置的大一些， P、R、mAP@.5、mAP@.5:.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、mAP@.5、mAP@.5:.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目录下。

[https_github.com_ultralytics_yolov5]: https://github.com/ultralytics/yolov5
[Train Custom Data]: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
[20210627101818386.png]: /images/20221005/f0c307f64e5749a69272c8c13a17a5a5.png
[Image Labels]: /images/20221005/52d237a5179f4ec99a7ebdc53dfa3377.png
[20210627111740145.png]: /images/20221005/1dbce1a8fa804188a671ddc5f5e23798.png