基于pytorch的YOLOv5训练
一、数据集制作
在data
文件夹放入数据集
方法一(推荐)
- images文件夹存放训练图片:xxx.jpg
labels文件夹存放标注文件:xxx.txt - 生成train.txt,test.txt,val.txt
文件中内容为图片的路径,训练代码通过读取三个txt中的图片路径索引对应的文件图片。例如:xxx/xxx/xxx.jpg。 - labels文件夹中存放同名的txt标注文件
在yolo中文件的标注信息包括:【cls、 x_center、 y_center、 width、 height】
方法二
- 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)
]
四、修改训练参数
- 训练模型
终端中输入训练参数,进行训练;或直接运行train.py
python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
yolov5m 48
yolov5l 32
yolov5x 16
- 可视化训练
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流