环境:

  • 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

yolov5 device 参数怎么用GPU跑 不用CPU yolov5使用gpu训练_P4

ps: 可以下载 coco128 来进行测试——地址

 

四、预训练模型下载

Model

APval

APtest

AP50

SpeedGPU

FPSGPU

 

params

FLOPS

YOLOv5s

37.0

37.0

56.2

2.4ms

416

 

7.5M

13.2B

YOLOv5m

44.3

44.3

63.2

3.4ms

294

 

21.8M

39.4B

YOLOv5l

47.7

47.7

66.5

4.4ms

227

 

47.8M

88.1B

YOLOv5x

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

YOLOv3-SPP

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转化