一、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
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()
代码执行后在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()
运行后会生成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
多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.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进行操作即可