yolov5训练与tensorrt模型转换

Ayu___

已于 2022-09-21 09:50:15 修改

阅读量1.6k

点赞数 1

文章标签：深度学习人工智能

于 2022-06-20 16:35:30 首次发布

本文链接：https://blog.csdn.net/Ayu___/article/details/125375562

版权

一、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] +