这里记录一下将yolov5用于训练自己的数据集，这里v5的github地址

Environment

本文是直接 pip install -r requirement.txt直接安装成功，如果不容易成功，就一个个安装即可。如果出现报错，多关注一下torch的版本即可。

Data preparation

label的格式:
<object-class> <x_center> <y_center> <width> <height>

举个🌰：
414_735_0.txt

1
2
7 0.503 0.499 0.098 0.451
7 0.465 0.633 0.033 0.053

数据集格式
data
├──coco.yaml
├──coco128.yaml
├──==defect_data==
   ├──images
      ├──train
      ├──val
   ├──labels
      ├──train
      ├──val
├──…

然后制作自己的defect.yaml配置文件(直接复制coco.yaml修改或者新建一个文件将一下代码copy进去)：

1
2
3
4
5
6
train: ./data/defect_data/images/train  # absolute_path_train_img
val: ./data/defect_data/images/val  # absolute_path_val_img

nc: 8  # 类别数量

names: ['quexian1', 'quexian2', ... 'quexian8']

修改models/xxx.yaml，有s, m, l, x四个模型供选择，需要修改类别nc数目；同时最好把anchors也修改为自己数据对应的anchors（需要自己提前聚类），不改也可train。

Pretrained Model

建议爬梯子下载最新的预训练权重，这个repo经常更新，可能和以前的预训练模型不适应。下载好放在weights下，在train.py里修改相应路径。

Training

配置train.py

1
2
3
4
5
6
7
8
9
10
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default='data/defect.yaml', help='data.yaml 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')
    parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

训练好的模型会被保存在./runs/train/exp3/weights/last.pt和best.pt

Inference

配置detect.py

1
2
3
4
5
6
7
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp3/weights/last.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/defect_data/images/val', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

如果需要inference 自己数据集原图大小，需要修改utils/dataset.py中

1
2
3
4
class LoadImages:  # for inference
		# Padded resize
        img = letterbox(img0, new_shape=self.img_size)[0]

修改为：

1
2
3
4
class LoadImages:  # for inference
		# Padded resize
        img = letterbox(img0, new_shape=(math.ceil(img0.shape[0]/32)*32, math.ceil(img0.shape[1]/32)*32))[0]

Visualization of Training

利用tensorboard进行可视化，训练开始会在主目录生成一个runs文件，利用tensorboard打开：
tensorboard --logdir=./runs

Thanks

欢迎评论，每天回复☼