yolov5训练与tensorrt模型转换_yolov5转tensorrt

一、YOLOv5模型训练

参考：https://github.com/ultralytics/yolov5

1、所使用的系统环境如下：

• Ubuntu 18.04
• GTX 3090
• cuda 11.1
• cudnn 8.0.6
• python3.8
• pytorch1.8.0
• torchvision0.9.0

2、准备训练yolov5所需环境

$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -r requirements.txt
1
2
3

3、准备自己的数据集（VOC格式）

（1）创建vehicle_dataset文件夹（名字可以自定义），目录结构如下，将之前labelImg标注好的xml文件和图片放到对应目录下

vehicle_dataset
–labels #存放标注转换格式后的txt文件（后面再进行生成）
–images # 存放图片
–Annotations # 存放图片对应的xml文件
–ImageSets/Main #之后会在Main文件夹内自动生成train.txt，val.txt，test.txt和trainval.txt四个文件，存放训练集、验证集、测试集图片的名字（无后缀.jpg）

（2）ImageSets文件夹下面有个Main子文件夹，其下面存放训练集、验证集、测试集的划分，通过脚本生成，可以创建一个split_train_val.py文件（放置于Annotations同一目录下），代码内容如下：

# coding:utf-8

import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

代码执行后在Main中将会生成如下文件：

（3）把数据集xml格式转换成yolo要求的txt格式。即将每个xml标注提取bbox信息为txt格式（这种数据集格式成为yolo_txt格式），每个图像对应一个txt文件，文件每一行为一个目标的信息，包括类别 xmin xmax ymin ymax。格式如下：

创建voc_label.py文件，将训练集、验证集、测试集生成label标签（训练中要用到），同时将数据集路径导入txt文件中，代码内容如下：（代码内路径改成你的路径）

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["car"]   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h

def convert_annotation(image_id):
    in_file = open('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        # difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()
for image_set in sets:
    if not os.path.exists('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/labels/'):
        os.makedirs('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/labels/')
    image_ids = open('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('/media/ds/新加卷/lxy/yolov5-train/vehicle_dataset/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/vehicle_dataset/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

运行后会生成labels文件夹和三个包含数据集的txt文件，其中labels中为不同图像的标注文件，train.txt等txt文件为划分后图像所在位置的绝对路径，如train.txt就含有所有训练集图像的绝对路径。
运行voc_label.py时报错“ZeroDivisionError: float division by zero”的原因是：标注文件中存在width为0或者height为0的数据，检查修改后可解决。

4、修改配置文件

（1）数据集的配置
yolov5/data文件夹下复制coco.yaml，重命名为car.yaml，修改内容如下：train：接上一步生成的train.txt的位置，val同理，nc为类别数量，names为类别名称。（冒号后面需要加空格，否则会被认为是字符串而不是字典而报错）
（2）编辑模型的配置文件
修改yolov5m6.yaml(位于/home/ds/yolov5/models/hub/下)中的nc值为类别数量
yolov5会自动kmeans算出anchors（只要训练时不加–noautoanchor）

5、选择需要的模型

在yolov5目录下的model文件夹下是模型的配置文件，这边提供s、m、l、x版本，逐渐增大（随着架构的增大，训练时间也是逐渐增大），假设采用yolov5s.yaml，只用修改一个参数，把nc改成自己的类别数。

官网下载预训练模型，然后保存到/yolov5/weights目录下，这里使用的是yolov5m6.pt

6、训练过程可视化

（1）wandb。在训练前，可以pip install wandb可视化训练过程。在训练时会有引导：可以创建新的账号，也可以使用我的账号：9e1c7b51546399407777a98bc0d069f1dfb0ce39

但是可能会出现W&B网页打不开的情况。
（2）使用TensorBoard。
训练开始后，yolov5文件夹下会生成runs/train，里面有训练过程记录，使用tensorboard即可查看。
命令为：tensorboard --logdir=runs/train

7、训练命令

python train.py --img 1280 --batch 64 --epochs 300 --weights weights/yolov5m6.pt  --data data/car.yaml --cfg models/yolov5m6.yaml
1

多GPU训练

python -m torch.distributed.launch --nproc_per_node 2 train.py --img 1280 --batch 64 --epochs 1000 --data all_data16/all_car.yaml --weights weights/yolov5m6.pt --cfg all_data16/yolov5m6.yaml --device 0,1,2,3
1

二、出现的错误

1.AssertionError: train: No labels in /home/pythonProject/yolov5-master/madata/train.cache. Can not train without labels.
解决：https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data，数据格式有问题，严格按照这个文档准备数据。
2.

解决：这两个错误都是因为batch设置太大了，改小一点

二、pytorch模型转tensorrt的engine

参考：https://github.com/wang-xinyu/tensorrtx

https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

yolov5生成的模型为pytorch模型，如果要使用tensorrt进行推理加速，则需要转换成engine。具体转换过程参考上面链接中yolov5的README.md进行操作即可