环境:
- tensorRT 6 / tensorRT 7
- Ubuntu 18.04
- CUDA 10.1
一、下载项目
yoloV5 项目:地址
二、安装相应环境
pip install -r requirements.txt
训练关键依赖:
Cython
matplotlib
numpy
opencv-python
pillow
PyYAML
scipy
tensorboard
torch>=1.6.0
torchvision>=0.7.0
tqdm
ps: torch1.6 安装
pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
三、训练数据格式
- datas
-images
- train#存放训练图片
- x.jpg
- val##存放验证图片
-labels
- train#存放训练label
- x.txt
- val#存放验证label
ps: 可以下载 coco128 来进行测试——地址
四、预训练模型下载
Model | APval | APtest | AP50 | SpeedGPU | FPSGPU | | params | FLOPS |
37.0 | 37.0 | 56.2 | 2.4ms | 416 | | 7.5M | 13.2B | |
44.3 | 44.3 | 63.2 | 3.4ms | 294 | | 21.8M | 39.4B | |
47.7 | 47.7 | 66.5 | 4.4ms | 227 | | 47.8M | 88.1B | |
49.2 | 49.2 | 67.7 | 6.9ms | 145 | | 89.0M | 166.4B | |
YOLOv5x + TTA | 50.8 | 50.8 | 68.9 | 25.5ms | 39 | | 89.0M | 354.3B |
45.6 | 45.5 | 65.2 | 4.5ms | 222 | | 63.0M | 118.0B |
五、选择一个模型训练
在项目的 ./models
目录下选择一个模型的配置文件,这里我们选择 yolov5s.ymal
,这是一个最小最快
的模型。关于其他模型之间的比较下面介绍。选择好模型之后,如果你使用的不是coco数据集进行训练,而是自定义的数据集,此时只需要修改*.yaml配置文件
中的 nc: 80
参数和数据的类别列表
下面是 yolo5s.ymal
配置文件的内容:
# parameters
nc: 80 # 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)
]
yolov5s.yaml配置文件
中主要定义了:
- 参数(parameters):类别等
- anchor (这个yoloV5会自动调整,你也可以自己计算,利用darknet计算也可以)
- YOLOv5 backbone
- YOLOv5 head
六、训练
运行下面的命令训练 coco128.ymal
,训练5epochs。可以有两种训练方式,如下参数:
-
--cfg yolov5s.yaml --weights ''
:从头开始训练 -
--cfg yolov5s.yaml --weights yolov5s.pt
:从预训练的模型加载开始训练
YOLOv5 在 coco128 上训练5epochs的命令:
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 --logdir=runs # 默认 http://localhost:6006
# 也可以指定
# tensorboard --logdir=runs --host=127.0.0.1
八、测试
训练结束后,会生成两个预训练的模型:
-
best.pt
:保存的是中间一共比较好模型 -
last.pt
:训练结束后保存的最后模型
python test.py --data ./data/coco128.yaml --weights runs/exp1/weights/best.pt --augment
测试的更多可选参数:
--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,布尔值
九、推理
python detect.py --source 0 # webcam
file.jpg # image
file.mp4 # video
path/ # directory
path/*.jpg # glob
rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp stream
rtmp://192.168.1.105/live/test # rtmp stream
http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8 # http stream
ps:这是基础,后续讲解 TensorRT转化