Yolo v5模型训练那点事

关于Yolo 5的模型训练这点事

作者：胖橙

---

前言

​ 随着人工智能的不断发展，模型训练这门技术也越来越重要，很多人都开启了学习模型训练的篇章，现在我就来介绍以一下模型训练的基础内容。

准备工作

• 数据集
• labelImg
• Yolov5的源代码

关于数据集

区分度

​ 关于数据集，我想说的是区分度越高越好。因为我们每一张照片之间的区分度越大的话，我们的电脑学习效率越高。打个比方：我们学习一个知识，那么如果一直做一个类型的例题的话，换一种新的题型可能就不认识了。同理我们的每一张照片之间的区分度越高，那么我们的电脑就会学习的越好。

训练集(train)和验证集(valid)

​ 在有监督的机器学习中，经常会说到训练集（train)、验证集（validation）和测试集（test），这三个集合的区分可能会让人糊涂，特别是，有些人就会搞不清楚验证集和测试集有什么区别。

​ 如果我们自己已经有了一个大的标注数据集，想要完成一个有监督模型的测试，那么通常使用均匀随机抽样的方式，将数据集划分为训练集、验证集、测试集，这三个集合不能有交集，常见的比例是8:1:1，当然比例是人为的。从这个角度来看，三个集合都是同分布的。

​ 有了模型后，训练集就是用来训练参数的，说准确点，一般是用来梯度下降的。而验证集基本是在每个epoch完成后，用来测试一下当前模型的准确率。因为验证集跟训练集没有交集，因此这个准确率是可靠的。

如何区分开训练集和验证集

​ 准确的来说，我们在分类之前，他们都是普通的照片，只不过我们赋予了他其他的含义而已（也就是说，我说谁是验证集，谁就是）。那么我们按照一个什么比例去划分呢？我们可以按照7:3的比例去进行一个划分。但是我们的图片计数不能太少，一般至少要100张才能出现一个效果。然后我么分别把他们放在两个文件夹之中，文件夹名字固然可以自己来定。这里也可以按照我的来定义。

|——train

​ |——images // 训练集的图片存储

​ |——labels //训练集的txt文件

|——valid

​ |——images //验证集的图片存储

​ |——labels //验证集的txt文件

labelImg

进行批量标注

利用Open Dir按钮可以打开需要被标注的图片的文件夹。
利用Change Save Dir按钮可以打开标注文件存放的文件夹。
利用w快捷键或者点击create RectBox可以开始标注，标注完后需要保存。所有图片标注完后可以得到标注文件。标注文件内容代表被标注图片内所含的物体。

XML文件转换TXT文件

XML文件进行转换的时候，我们可以用写好的脚本来进行转换。

代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir,getcwd
from os.path import join
import glob

classes = ['Aubergine','Chili','Bone','Newspaper','Waste_Butteries','Medical_Cotton_Swabs','Pesticides','Leftover_Food','Radioactive_Material','Expired_Medicines']
#我们要训练的类的名字写好
def convert(size,box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    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_name):
    in_file = open('./valid/labels/'+image_name[:-3]+'xml')
    out_file = open('./valid/labels/'+image_name[:-3]+'txt','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'):
        cls = obj.find('name').text
        if cls not in classes:
            print(cls)
            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))
        bb = convert((w,h),b)
        out_file.write(str(cls_id)+" "+" ".join([str(a) for a in bb]) + '\n')

wd = getcwd()

if __name__ == "__main__":
    for image_path in glob.glob("./valid/images/*.jpg"):
        image_name = image_path.split('\\')[1]
        convert_annotation(image_name)

然后我们把转换好的txt文件放入到我们对应的labels文件夹下面。

yaml的编写

YAML（YAML Ain’t Markup Language）以数据为中心，比json、xml更适合做配置文件。

我们的yaml要放到数据集那一层文件夹中。

|——train

|——valid

|——xxxx.yaml

yaml的书写方式

train: ../Rubbish_Data/train/images
val: ../Rubbish_Data/valid/images

nc: 10
names: ['Aubergine','Chili','Bone','Newspaper','Waste_Butteries','Medical_Cotton_Swabs','Pesticides','Leftover_Food','Radioactive_Material','Expired_Medicines']

开始准备训练

在Pycharm中的train文件加入Parameters文件

可以参考如下。

--data
../Rubbish_Data/Rubbish_data.yaml
--cfg
models/yolov5s.yaml
--batch-size
120

最后开始运行训练代码即可。