YOLOv5训练自己的数据集

1、克隆代码

YOLOv5

2、配置环境

Python>=3.6.0
PyTorch>=1.7

• 新建虚拟环境

  conda create -n yolov5 python=3.8
  

• 安装Pytorch，去 Pytorch 官网自己安装
• 然后进入代码文件夹

  pip install -r requirements.txt
  

3、将数据集转换至可训练格式

按照这里的代码 数据集转换，就可以将 VOC 格式的数据集转换成 YOLOv5 可以训练的格式

将生成的可训练的数据集复制到和 YOLOv5 工程同一根目录下，即：
其中，YOLOv5 文件夹是工程代码，datasets 文件夹是数据集
数据集的文件目录是：
其中，images 文件夹中存储的图片，labels 文件夹中存储的标签

datasets
       ├─ images
       │    ├─ test # 下面放测试集图片
       │    ├─ train # 下面放训练集图片
       │    └─ val # 下面放验证集图片
       └─ labels
              ├─ test # 下面放测试集标签
              ├─ train # 下面放训练集标签
              ├─ val # 下面放验证集标签

4、修改配置文件

• 配置数据集文件
  在 YOLOv5/data 文件夹中添加一个 new_data.yaml，文件内容如下：

  train: D:/ZADD/code/object detection/datasets/images/train   # train images (relative to 'path') 118287 images
  val: D:/ZADD/code/object detection/datasets/images/val  # val images (relative to 'path') 5000 images
  test: D:/ZADD/code/object detection/datasets/images/test  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
  
  # Classes
  nc: 6  # number of classes
  # class names
  names: ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"]  
  

• 配置模型文件
  在 YOLOv5/models 文件夹中添加一个 new_train.yaml，文件内容如下：
  其实就是复制已有的 yaml 文件，修改目标类别 即可。

  # Parameters
  nc: 6  # 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)
    ]
  
  

5、开始训练

运行 train.py 文件，将刚才新建的 yaml 文件添加到文件中，开始训练即可。