【yolov5目标检测】使用yolov5训练自己的训练集

数据集准备

首先得准备好数据集，你的数据集至少包含images和labels，严格来说你的images应该包含训练集train、验证集val和测试集test，不过为了简单说明使用步骤，其中test可以不要，val和train可以用同一个，因此我这里只用了一个images

其中images装的是图片数据，labels装的是与图片一一对应同名的yolo格式txt，即类别号，经过归一化的中心x和y坐标以及宽和高

 

下载yolov5

到GitHub上下载整个项目的压缩包zip下来

 下来解压zip，把我们刚刚的数据集也放进去

再下载一个yolov5的预训练模型，我这里选择yolov5n.pt，下下来也放到解压文件夹中

然后用pycharm打开这个文件夹，看看哪里标红还差什么软件包没安装给安装上

配置yaml

先配置一下yolov5预训练模型的yaml，我下载的是yolov5n.pt模型，因此需要配置一下yolov5n.yaml，修改nc的数值为类别数目，我这里的数据集只有乌骨鸡和狮头鹅，因此改成2

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
 
# Parameters
nc: 2  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.25  # 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)
  ]

然后开始配置数据集的yaml，可以复制coco128.yaml的内容进行修改，新建一个xxx.yaml，修改path为数据集路径，train为训练集的相对路径，val为验证机的相对路径，因为没有用上test，因此把它注释掉了，还有names也要修改为数据集的类别名

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)
 
 
# 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: ikunData # dataset root dir
train: images # train images (relative to 'path') 128 images
val: images # val images (relative to 'path') 128 images
#test:  # test images (optional)
 
# Classes
names:
  0: goose
  1: chicken

开始训练

然后准备开始训练，打开train.py，修改它的参数，主要是这三行代码需要修改，修改预训练模型文件的路径，配置文件的路径以及数据集配置文件的路径

    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5n.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5n.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'ikunData.yaml', help='dataset.yaml path')

噢，还有训练的epochs数目，这个次数由你决定

    parser.add_argument('--epochs', type=int, default=30, help='total training epochs')

然后就可以开始运行train.py。

运行完了会生成一个runs文件夹，里面有训练出来的best.pt，和训练过程的记录

然后开始目标检测，准备好运行detect.py的参数，最基本的就是运行的权重文件，就是我们train出来的best.pt，还有要测试的数据路径

--weights runs/train/exp/weights/best.pt --source ikunData/images

把它写到运行配置中去

然后开始运行detect.py，运行完后会在runs里面生成detect文件夹，里面就有检测结果

你可能会发现有多个框框在同一个目标上，这时我们在detect.py上增加一个参数，这里是nms非极大值抑制，我们将IOU的阈值设置为0，再次运行detect.py

--weights runs/train/exp/weights/best.pt --source ikunData/images --iou-thres 0

 这次的效果要好一点