目录
一、配置环境
安装好pytorch后在项目目录下执行 pip install -r requirements.txt 即可完成环境配置。pytorch的安装可参考：PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】_哔哩哔哩_bilibili 课程的P1-P3
前置知识：了解最基本的项目环境配置、模型推理、模型训练，如果不了解这些操作可以先学习一下这个视频目标检测 YOLOv5 开源代码项目调试与讲解实战【土堆 x 布尔艺数】_哔哩哔哩_bilibili
二、数据集的准备
以上一篇文章的UA-DETRAC数据集为例，处理好的数据集路径树如下所示：
UA-DETRAC_yolo
├── UA-DETRAC_train
│ ├── images
│ └── labels
└── UA-DETRAC_test
├── images
└── labels
yolov5要求数据集的格式为yolo txt格式，在上一篇文章中已经将UA-DETARC数据集转成了这种格式。其中UA-DETRAC_train是训练集，UA-DETRAC_test是测试集。
一般来说，符合yolov5需要的yolo txt格式的数据集会有images和labels两个文件夹，分别存放图片和标注，标注文件为txt格式，标注文件的内容举例如下图所示。
标注文件中的每一行都表示图片中的一个目标，每行的第一个数表示目标的类别，后面四个数依次表示框住目标的边框的x_center、y_center、w、h，即边框的中心x坐标、中心y坐标、边框宽度、边框高度。
YOLOv5工程下有两个路径是与数据集相关的，分别是data和datasets，data文件夹下存放数据集的yaml配置文件，datasets文件夹下存放数据集。将数据集移动到datasets文件夹下，然后在data文件夹内写好与数据集同名的yaml文件，就可以使用这个数据集进行训练了。
以UA-DETRAC_train为例，在将这个训练集整体移动到datasets文件夹下之后，在data文件夹下创建文件UA-DETRAC_train.yaml（注意一定要与datasets下的数据集同名），其内容如下：
# 待填写
train:
val:
test:
# Classes
nc: 4 # 数据集类别数，要检测的目标有几种就是几个类别
names: ["car", "bus", "van", "others"] # 类别名，注意顺序
在yaml配置文件中需要指明图片和标签的存放地址，即要给train、val、test三个变量赋好值。存放地址可以是目录、txt文件和列表，这里采用txt文件的方式指明存放地址。
这里提供两个脚本，split_train_val.py可按照指定的比例划分出训练集、验证集、测试集，dataset_path_set.py可以按照划分后的结果生成train.txt、val.txt、test.txt分别存不同集图片的存放地址。
split_train_val.py
import os
import random
import argparse
parser = argparse.ArgumentParser()
# 标记文件的地址，根据自己的数据进行修改
parser.add_argument('--label_path', default='../labels', type=str, help='input label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='../ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1 # 训练集和验证集所占比例。
train_percent = 0.8 # 训练集在全部数据中的比例，可自己进行调整
label_file_path = opt.label_path
txt_save_path = opt.txt_path
total_label = os.listdir(label_file_path)
if not os.path.exists(txt_save_path):
os.makedirs(txt_save_path)
num = len(total_label)
list_index = range(num)
trainvalNum = int(num * trainval_percent)
trainNum = int(num * train_percent)
trainval = random.sample(list_index, trainvalNum)
train = random.sample(trainval, trainNum)
file_trainval = open(txt_save_path + '/trainval.txt', 'w')
file_test = open(txt_save_path + '/test.txt', 'w')
file_train = open(txt_save_path + '/train.txt', 'w')
file_val = open(txt_save_path + '/val.txt', 'w')
for i in list_index:
name = total_label[i][:-4] + '\n'
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()dataset_path_set.py
import os
import re
sets = ['train', 'val', 'test']
abs_path = os.getcwd()
print(abs_path)
for image_set in sets:
image_ids = open('../ImageSets/Main/%s.txt' % image_set).read().strip().split()
if not os.path.exists('../dataset_path/'):
os.makedirs('../dataset_path/')
list_file = open('../dataset_path/%s.txt' % image_set, 'w')
for image_id in image_ids:
# 数据集图片位置的绝对路径 按需修改
list_file.write(r'C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\images\%s.jpg' % image_id)
list_file.write('\n')
list_file.close()
依次运行完两个脚本后，就在UA-DETARC_train目录下生成了dataset_path文件夹，文件夹内有需要的train.txt、val.txt、test.txt三个表示图片存放地址的txt文件。最后修改yaml文件。
# 绝对路径按需修改 注意冒号后面要跟一个空格，否则冒号后内容会被识别为字符串而非列表
train: C:/.../.../.../yolov5-6.2/datasets/UA-DETRAC_train/dataset_path/train.txt
val: C:/.../.../.../yolov5-6.2/datasets/UA-DETRAC_train/dataset_path/val.txt
test: C:/.../.../.../yolov5-6.2/datasets/UA-DETRAC_train/dataset_path/test.txt
# Classes
nc: 4 # 数据集类别数，要检测的目标有几种就是几个类别
names: ["car", "bus", "van", "others"] # 类别名，注意顺序
生成的txt文件内容如下图所示，内容为绝对路径。
大功告成，现在可以使用UA-DETRAC_train数据集进行训练了。需要注意的是UA-DETARC_train是训练集，所以在split_train_val.py脚本中测试集的比例被设置成了零，仅在其中按8：2的比例划分出了训练集和验证集，所以后面的test.txt中其实是空的。
如果需要使用测试集UA-DETRAC_test，同样重复一边上面的流程即可，只需要注意将脚本split_train_val.py中控制划分比例的两个变量都设成零，并且注意流程中的绝对路径要按实际情况修改就行了。
三、模型训练与测试
YOLOv5-Lite项目中提供了用于训练、测试、推理的程序train.py、test.py、detect.py。使用这些程序指定参数可以很方便的进行模型的训练、测试、推理。
1.模型训练
train.py中参数各有作用，这里仅介绍比较重要且常用的几个。
--weights 模型预训练权重
--cfg 模型文件，需要指定为存有模型结构的yaml文件
--data 使用哪个数据集训练，需要指定为数据集的yaml配置文件
--epochs 训练轮次
--batch-size 训练过程中一次往gpu中放入多少张图片，这个参数决定了显存占用，一般设置为2的n次幂
--img-size 输入图片的宽高
--resume 是否接续训练，指定为True就是从上次结束的地方继续训练
--device 使用的设备，指定为cpu就是使用cpu进行训练，指定为0/1/2/3就是使用对应的显卡进行训练
--workers dataloader进行数据装载时cpu使用的线程数
在命令行中输入命令进行训练：
python train.py --weights ./weights/v5Lite-e.pt --cfg ./models/v5Lite-e.yaml --data ./data/UA-DETRAC_train.yaml --epochs 40 --batch-size 4 --device 0这里以使用UA-DETRAC_train来训练YOLOv5-Lite的v5Lite-e模型作为例子，预训练权重选作者提供的v5Lite-e.pt权重，使用这个权重进行训练会方便许多；模型结构选择v5Lite-e.yaml；训练轮次按需设置；batchsize根据显卡情况选择，显卡性能好就设置的大一些，通常设置为2的n次幂；device要选对显卡，可以打开任务管理器查看自己的nvidia显卡是GPU几。另外--worker参数按需设置，这个参数会占用内存，并且也并非越大越好太大了会爆内存，可以把这个参数从零开始往上加，当GPU跑满时就不要继续继续增加了，通常设置为2的n次幂。
2.模型测试
--weights 使用哪个权重进行测试
--data 进行测试的数据集
--batch-size 测试过程中一次往gpu中放入多少张图片，这个参数决定了显存占用，一般设置为2的n次幂
--conf-thres 置信度阈值设置
--iou-thres IOU阈值设置
--task 使用train, val, test, speed or study的哪一个进行测试
--device 使用何种设备进行测试，同train.py中的--device参数
在命令行中输入命令进行测试：
python test.py --weights ./weights/best.pt --data ./data/UA-DETRAC_test.yaml --batch-size 4 --task test以使用UA-DETRAC_test对训练出的best.pt权重进行测试作为例子，注意--task参数一定要设置成test。如果你不是使用UA-DETRAC_test进行测试，那么按照实际情况设置。
四、Tips
在使用UA-DETARC_train对模型进行训练，然后使用UA-DETRAC_test对训练出的模型进行测试之后，会发现测试出的模型精度很低。这是因为在UA-DETRAC数据集中，训练集和测试集是不同路口处的相机拍摄得到的结果，训练集和测试集间的差别比较大，单纯使用训练集训练后用测试集进行测试，效果不会太好。
这时可以考虑将UA-DETRAC的训练集和测试集混合到一起后，再重新划分出训练集和测试集，这样可以有效提升训练出模型的精度。
对UA-DETRAC数据集不同处理方式的实验可以参考下面这篇文章：
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
