TensorRT (五）yoloV5 训练模型

最新推荐文章于 2025-06-11 10:37:15 发布

原创最新推荐文章于 2025-06-11 10:37:15 发布 · 3.1k 阅读

14 ·

CC 4.0 BY-SA版权

python 同时被 2 个专栏收录

93 篇文章

订阅专栏

TensorRT

6 篇文章

订阅专栏

本文详细介绍YOLOv5目标检测模型的环境搭建、数据准备、模型训练及测试流程。涵盖了YOLOv5项目的下载、环境配置、训练数据格式、预训练模型选择等方面。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

环境：

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