基于pytorch的YOLOv5训练

一、数据集制作

在data文件夹放入数据集

方法一(推荐)

1. images文件夹存放训练图片：xxx.jpg
   labels文件夹存放标注文件：xxx.txt
2. 生成train.txt，test.txt，val.txt
   文件中内容为图片的路径，训练代码通过读取三个txt中的图片路径索引对应的文件图片。例如：xxx/xxx/xxx.jpg。
3. labels文件夹中存放同名的txt标注文件
   在yolo中文件的标注信息包括：【cls、 x_center、 y_center、 width、 height】

方法二

1. images文件夹存放训练图片：xxx.jpg
   labels文件夹存放标注文件：xxx.txt
   
   2.不用生成txt文件，在配置文件中修改即可。

二、修改配置文件

下载yolov5-5.0代码。
修改data/xxx.yaml作为数据集的配置文件，为模型训练提供数据接口。

方法一(推荐)

# train and val datasets (image directory or *.txt file with image paths)
train: xx/xx/train2017.txt  # 上面我们生成的train，根据自己的路径进行更改
val: xx/xx/val2017.txt  # 上面我们生成的test
#test: ../coco/test-dev2017.txt  # 20k images for submission to https://competitions.codalab.org/competitions/20794
 
# number of classes
nc: 2   #训练的类别
 
# class names
names: ['apple','orange']

方法二

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: data/coco128/images/train2017/  # 128 images
val: data/coco128/images/train2017/  # 128 images

# number of classes
nc: 2

# class names
names: ['apple','orange']

三、修改模型文件

修改models/xxx.yaml作为模型文件，搭建网络结构；
需要修改类别数nc为数据集包括的类别数。

# parameters
nc: 2  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# anchors
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, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [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, BottleneckCSP, [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, BottleneckCSP, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

四、修改训练参数

1. 训练模型
   终端中输入训练参数，进行训练；或直接运行train.py

python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
                                         yolov5m                              48
                                         yolov5l                              32
                                         yolov5x                              16

2. 可视化训练

tensorboard --logdir=runs

训练参数可以通过train.py文件修改
如cfg、data、weights、device、img-size(训练和测试图片的大小)等

五、命令参数介绍

 python detect.py --source file.jpg  # image 
                            file.mp4  # video
                            ./dir  # directory
                            0  # webcam
                            rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa  # rtsp stream
                            http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8  # http stream

--cfg yolov5s.yaml --weights ''：从头开始训练
--cfg yolov5s.yaml --weights yolov5s.pt：从预训练的模型加载开始训练

命令的中参数也可以通过train.py文件修改

python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ''

训练的更多可选参数：

--epochs：训练的epoch，默认值300
--batch-size：默认值16
--cfg：模型的配置文件，默认为yolov5s.yaml
--data：数据集的配置文件，默认为data/coco128.yaml
--img-size：训练和测试输入大小，默认为[640, 640]
--rect：rectangular training，布尔值
--resume：是否从最新的last.pt中恢复训练，布尔值
--nosave：仅仅保存最后的checkpoint，布尔值
--notest：仅仅在最后的epoch上测试，布尔值
--evolve：进化超参数（evolve hyperparameters），布尔值
--bucket：gsutil bucket，默认值''
--cache-images：缓存图片可以更快的开始训练，布尔值
--weights：初始化参数路径，默认值''
--name：如果提供，将results.txt重命名为results_name.txt
--device：cuda设备，例如：0或0,1,2,3或cpu，默认''
--adam：使用adam优化器，布尔值
--multi-scale：改变图片尺寸img-size +/0- 50%，布尔值
--single-cls：训练单个类别的数据集，布尔值

六、tensorboard可视化

在命令行窗口输入

tensorboard --logdir=runs

然后把返回的url地址粘贴到浏览器中即可
注意：
如果返回拒绝了我们的请求，可以在tensorboard的后面加上参数–port ip：

tensorboard --logdir=runs --host=192.168.0.134

七、测试

测试命令

python test.py --weights yolov5s.pt --data ./data/coco.yaml --img 640

--weights ：预训练模型路径，默认值weights/yolov5s.pt
--data：数据集的配置文件，默认为data/coco.yaml
--batch-size：默认值32
--img-size：推理大小（pixels），默认640
--conf-thres：目标置信度阈值，默认0.001
--iou-thres：NMS的IOU阈值，默认0.65
--save-json：把结果保存为cocoapi-compatible的json文件
--task：默认val，可选其他值：val, test, study
--device：cuda设备，例如：0或0,1,2,3或cpu，默认''
--half：半精度的FP16推理
--single-cls：将其视为单类别，布尔值
--augment：增强推理，布尔值
--verbose：显示类别的mAP，布尔值

训练结束后，会生成两个预训练的模型：
1.best.pt：保存的是中间一共比较好模型
2.last.pt：训练结束后保存的最后模型

八、推理

使用预训练的模型进行推理测试，YOLOv5支持多种数据源推理测试，如下：

• 图像
• 目录
• 视频
• 网络摄像头
• http流
• rtsp流
  关于推理的更多参数使用如下命令查看：

python detect.py -h
--weights ：预训练模型.pt的路径，默认值为：weights/yolov5s.pt
--source：输入的数据源，可以是：图片、目录、视频、网络摄像头、http和rtsp流，默认值为：interence/images
--output： 输出检测结果的路径，默认值为：inference/output
--img-size ：用于推理图片的大小（pixels），默认值为：640
--conf-thres：对象的置信度阈值（object confidence threshold），默认值为：0.4
--iou-thres ：NMS的IOU阈值（ IOU threshold for NMS），默认值为：0.5
--fourcc：输出视频的编码格式（必须是ffmepeg支持的），例如：H264格式，默认格式为：mp4v
--half： 使用版精度F16推理（half precision FP16 inference），布尔值，默认为true
--device：cuda设备，例如：0或0,1,2,3或cpu，默认''
--view-img ：显示结果，‘布尔值，默认为true’
--save-txt ：把结果保存到*.txt文件中
--classes：过滤类别 CLASSES [CLASSES …]，filter by class
--agnostic-nms：类不可知 NMS
--augment：增强推理（augmented inference）

推理图像
python detect.py --source inference/1_input/1_img/bus.jpg --weights ./weights/yolov5s.pt --output inference/2_output/1_img/
推理目录
python detect.py --source inference/1_input/2_imgs --weights ./weights/yolov5s.pt --output inference/2_output/2_imgs
检测中有些图片置信度比较低的可以通过--conf-thres参数过滤掉
python detect.py --source inference/1_input/2_imgs --weights ./weights/yolov5s.pt --output inference/2_output/2_imgs --conf-thres 0.8
推理视频
python detect.py --source test.mp4 --weights ./weights/yolov5s.pt --output test_result/3_video
# 示例语法
python detect.py --source ./file.jpg  # 图片 
                          ./file.mp4  # 视频
                          ./dir  # 目录
                          0  # 网络摄像头
      'rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa' # rtsp流
      'http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8'  # http流