YOLOv5-6.1从训练到部署（二）：训练自己的数据集

1 数据的标注

一个公司不太可能让算法工程师自己去采集标注数据集，而是交给第三方公司去做，但作为算法工程师，还是有必要了解数据的标注过程。这里假设我们已经有了数据集（即已经采集完成），我们现在要先对它们进行标注。
假如我们想让模型分辨鸟类与无人机，而现在我已经从不同渠道收集了296张图片，其中训练集270张，验证集26张，并且已经将它们分到了不同的文件夹下（更多的时候，我们是先标注，再划分）。接下来是对它们进行标注，这里我们使用标注工具labelimg。

（1）labelimg的安装

首先要安装，在命令行终端输入：

pip install labelimg

安装成功之后，在命令行终端输入labelimg，随后会直接跳出labelimg的界面，如下图所示：

（2）labelimg的基本操作

假设数据的组织形式如下：

在当前目录下打开命令行窗口，然后切换到yolov5虚拟环境（其他安装了labelimg的环境也行），输入labelimg，打开软件界面：

点击右上角的Open Dir

此时默认打开第一张图片，详细界面如下：

按下W，此时鼠标在图像上会显示出参考标线

从目标的左上角到右下角拖动鼠标，会有一个阴影矩形跟着目标移动：

当鼠标拖到目标的右下角时（即已经将目标框住），松开鼠标，会弹出一个小窗口，然后在文本框中输入目标的类别，随后点击OK按钮（或者按回车），则当前目标标注完成。

标注完成后，窗口右边会显示当前图像中已经标注的目标

也可以让已标注的目标在图中隐藏，即在右边将相应的目标勾号取消：

如果发现标错了怎么办？可以在右边右击需要修改或删除的目标，然后在弹出的选项中选择重新编辑类名或者删除边框。

将鼠标移动到图中的角点位置，待其变红色后，还可以通过拖拽对边框的角点位置进行修改：

标完当前图片，按下D键可以进入下一张，在进入下一张之前会弹出提示框问你是否保存，如果点击Yes，那么会弹出窗口设置标签文件的保存位置：

标签文件默认与图片同名，只是后缀名不同，标签文件建议新建一个文件夹专门保存：

如果在每次企图进入下一张时，都跳出窗口让你手动点保存，可以选择自动保存模式，勾选View—Auto Save mode（下次打开labelimg，还需要重新勾选）

进入第二张图片标注，在输类名时，会显示之前出现过的类名，也就是说，所有的类别，在整个数据集中只需要输入一次：

打开刚刚标签文件的保存目录，可以看到xml文件
因为labelimg默认保存的标签文件格式是VOC，因此标签文件的后缀名为xml。如果想以其他格式保存标签文件，可以点击labelimg界面的左边相应位置进行切换，除了VOC，还有YOLO、CreateML另外两种格式，我们这里暂时选择VOC格式标注：

这个软件的基本操作大致就是这样的，作为算法工程师，掌握今天讲的这些操作足够了，如果对labelimg想进一步了解，可以看这篇文章。

2 数据的转换

刚刚说到我们默认保存的标签文件是VOC的数据格式，但我们这里使用的是YOLO框架，所以有必要将其转换成YOLO的格式。所谓的格式，指的是标签文件的保存规范，是不同算法开发者共同遵守的数据协议，只有规定了格式，数据才能在不同算法模型之间切换。在转换之前，我们先来了解一下目标检测常用的数据格式。

首先要声明的是，这里仅是对标签文件的内容进行解析，无论是哪种格式，都对图片和标签的目录结构有所要求，VOC格式和COCO格式数据的目录结构可以看这篇文章，YOLO格式的目录结构本文稍后会介绍。

（1）PascalVOC数据格式

VOC格式的标签文件是XML文件，内容示例（以92.xml为例）如下：

<annotation>
	<folder>train</folder>
	<filename>92.png</filename>
	<path>C:\data\train\92.png</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>831</width>
		<height>560</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>drone</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>291</xmin>
			<ymin>118</ymin>
			<xmax>777</xmax>
			<ymax>511</ymax>
		</bndbox>
	</object>
</annotation>

这里：
<filename>92.png</filename>是图片名，
<path>C:\data\train\92.png</path>是对应的图片路径，
<size> </size>是图片尺寸，其中<width>831</width>是图片的宽，<height>560</height>是图片的高，<depth>3</depth>图片通道数；
<object> </object>包裹的是单个目标，有多少个目标，就有多少个<object> </object>；在<object> </object>中，<name>drone</name>是目标的名字，<bndbox> </bndbox>包裹的是标签的左上角点和右下角点横纵坐标。
对于XML文件，了解以上信息足够了，如果想知道其他信息，可以查看这篇文章。

而YOLO格式的标签文件是txt文件，接下来我们要对其进行转换。

（2）YOLO数据格式

使用labelimg打开我们之前标注的92.png，然后切换到YOLO数据格式，然后保存，如下操作：

保存之后当前目录下会多出来一个名为92.txt文件：

将其打开

这里有5个数，它们之间用空格隔开，最前面的0是类别编号，因为我们最先标注的是无人机，所以无人机的类别编号就是0，随后的4个数分别表示标注框的中心坐标（x和y）和标注框的宽和高（w和h），他们都被归一化了（即除了图片的宽高）。
YOLO格式的标签文件是txt文件，如果一幅图有多个目标，则txt文件中会有多行。

（3）COCO数据格式

VOC格式和YOLO格式的标签文件，都是一张图片对应一个文件，但COCO不是，COCO格式的数据是把所有图片的标签都写入到一个json文件中（json文件经常用来保存python中的字典）。比如我们这里的数据，如果标签保存为COCO格式，一般会做成两个json文件（train和valid各有一个）。

COCO格式的数据内容组织如下：

{
    "info": info,               # dict
    "licenses": [license],      # list, 内部由多个字典组成
    "images": [image],          # list, 内部由多个字典组成
    "annotations": [annotation],# list, 内部由多个字典组成
    "categories": [category]    # list, 内部由多个字典组成

COCO数据集现在有3种标注类型文件，分别是：object instances（目标实例）、object keypoints（目标上的关键点）、image captions（看图说话）。COCO官方数据集的标签文件如下：

info、images、licenses这三个字段在3种类型的标注文件都是相同的，annotation和category这两种字段在不同类型的JSON文件中是不一样的。作为目标检测，我们只需要知道images、annotations和categories即可。

其中images是一个列表，列表的每个元素均是一个字典，字典里是一张图片的信息，有多少张图片就有多少个字典，字典内容如下：

{      # images是一个list,存放所有图片(dict)信息。image是一个dict,存放单张图片信息 
    "id": int,                  # 图片的ID编号（每张图片ID唯一）
    "width": int,               # 图片宽
    "height": int,              # 图片高
    "file_name": str,           # 图片名字
    "license": int,             # 协议
    "flickr_url": str,          # flickr链接地址
    "coco_url": str,            # 网络连接地址
    "date_captured": datetime,  # 数据集获取日期
}

这里我们只需要知道id、width、height、filename四项即可，其中id是图片的编号，在annotations中也要用到，每张图是唯一的。

接下来是categories，它也是一个列表，列表的每个元素都是字典，有多少个类别就有多少个字典，每个字典内容如下：

{                     # 类别描述
    "id": int,                  # 类别对应的ID（0默认为背景）
    "name": str,                # 子类别名字
    "supercategory": str,       # 主类别名字
}

这里我们一般只需要关注id和name就可以，注意，这里的id是类别的编号，可以认为它就是类别索引，要把它和图片id区别开来。

最后是annotations，它也是列表，其元素是字典，字典里是标注的边框信息，整个数据集有多少个边框，这里就有多少个字典。所谓整个数据集，指的是前面images中所包含的所有图片。每个字典内容如下：

{
    "id": int,                  # 目标对象ID（每个对象ID唯一），每张图片可能有多个目标
    "image_id": int,            # 对应图片ID，与images中的ID对应
    "category_id": int,         # 对应类别ID，与categories中的ID对应
    "segmentation": RLE or [polygon],   # 实例分割，得到的分割对象多边形边界点坐标[x1,y1,x2,y2,....,xn,yn]
    "area": float,              # 分割对象的区域面积（多边形面积）
    "bbox": [xmin,ymin,width,height], # 目标检测，对象定位边框[x,y,w,h]
    "iscrowd": 0 or 1,          # 表示是否是人群，iscrowd=1 用于标记一大群人
}

这里我们只需要知道id、image_id、category_id、bbox即可。

我们要做的是将VOC格式的数据集转成YOLO格式，整个过程不涉及COCO，但实际目标检测工作中，COCO格式的数据集非常常见，所以有必要对此有所了解。如果只是做目标检测，掌握本文中关于COCO数据集的介绍完全够用，但如果是图像分割任务，则还需要做进一步了解，更多关于COCO数据集的解读，可以看这篇文章。

（4）VOC格式数据集转YOLO格式

当所有数据标注完之后，在yolov5-6.1目录下新建一个名为voc2yolo.py的python脚本，内容如下：

import os
import xml.etree.ElementTree as ET


def xml2txt():
    # 创建标签的保存目录
    valid_label_dir = "uav_bird_training/data/labels/valid"
    if not os.path.exists(valid_label_dir):
        os.makedirs(valid_label_dir)

    train_label_dir = "uav_bird_training/data/labels/train"
    if not os.path.exists(train_label_dir):
        os.makedirs(train_label_dir)

    for ann_dir in [r"C:\data\valid", r"C:\data\train"]:
        # 遍历目录
        files = os.listdir(ann_dir)
        for xml_file in files:
            # 筛选出xml文件
            if xml_file.endswith(".xml") is not True:
                continue

            # 解析xml文件
            if os.path.isfile(os.path.join(ann_dir, xml_file)):
                # 拼接xml文件的路径名并解析
                xml_path = os.path.join(ann_dir, xml_file)
                tree = ET.parse(xml_path)
                root = tree.getroot()

                # 根据图像路径获取图像名
                img_name = root.find('filename').text

                # 获取图像文件的文件名和扩展名
                file_name, suffix = os.path.splitext(img_name)

                # 创建txt文件的文件名与路径
                txt_name = file_name + '.txt'
                data_label_text_f = os.path.join(train_label_dir, txt_name)

                # 创建并打开txt文件
                file_write_obj = open(data_label_text_f, 'w')

                # 获取图像宽高
                size = root.find('size')
                w = int(size.find('width').text)
                h = int(size.find('height').text)

                # 将类别和边框信息写入到txt文件中
                for elem in root.iter("object"):
                    # 解析目标
                    name = elem.find('name').text   # 获取标签中内容 # bird,drone
                    print("class name: ", name)
                    if name == "bird":      # 如果是鸟，则类别标签为0
                        clazz_index = 0
                    if name == "drone":     # 如果是无人机，则类别标签为1
                        clazz_index = 1

                    # 获取边框中心点坐标及宽高
                    x1 = float(elem.find("bndbox").find("xmin").text)
                    y1 = float(elem.find("bndbox").find("ymin").text)
                    x2 = float(elem.find("bndbox").find("xmax").text)
                    y2 = float(elem.find("bndbox").find("ymax").text)

                    # 归一化
                    cx = (x1 + (x2 - x1) / 2) / w
                    cy = (y1 + (y2 - y1) / 2) / h
                    sw = (x2 - x1) / w
                    sh = (y2 - y1) / h
                    file_write_obj.write("%d %f %f %f %f\n" % (clazz_index, cx, cy, sw, sh))

                file_write_obj.close()
                print("processed image : ", img_name)


if __name__ == "__main__":
    xml2txt()

这里面类别标签0是为bird、1为drone，和（2）YOLO数据格式介绍的不太一样，之所以如此，是因为我标数据的时候，并不是从第一张图片开始，导致0成为了bird，只能将错就错。这段程序的其他部分仅凭注释就很好理解，并不复杂，这里不再进行解释。

程序运行之后，标签文件将存放到名为uav_bird_training/data/labels的下面（下图为Pycharm2023中的截图）

打开train，就是下面这样

3 训练模型

（1）数据配置

一般成熟的目标检测框架（例如YOLO、SSD）都会规定数据的目录结构，YOLO也规定了自己的目录结构，这其实是YOLO格式数据的一部分。
我们将图片和YOLO格式的标签文件（txt文件）按如下目录结构进行配置：

顺带提一下，如果有不同的数据集，那么只允许图中数据集目录名（uav_bird_training）发生改变，该目录下的文件、文件夹名不允许改变。
这里uav_bird_training/data/images存放的是图片，只需要把图片移动过来就行，uav_bird_training/data/labels存放的是标签，是（3）VOC格式数据集转YOLO格式中程序生成的txt文件保存位置。这里还有个dataset.yaml，里面保存了数据集配置，我们根据已有的yaml文件来仿写的。

在yolov5-6.1/data下有几个yaml文件，里面是常见的公开数据集配置信息：

我们可以打开VOC.yaml查看内容：

这里不需要填写标签文件的目录，此外，如果对VOC数据集熟悉，可以发现train下面的images/train2012、images/train2007、images/val2012、images/val2007等，就是VOC下保存图片的目录，也就是说，path+train[0-3]其实就是训练集图片的相对路径。基于此，我们可以新建自己的data.yaml，内容如下：

# train and val datasets 
train: uav_bird_training/data/images/train/
val: uav_bird_training/data/images/valid/

# number of classes
nc: 2

# class names 顺序要和标签文件（txt文件）中的0、1顺序一致
names: ['bird', 'drone']

（2）训练

我们这里采用迁移学习，即在预训练模型yolov5s.pt的权重上进行微调。在命令行中输入以下命令：

python train.py --img 640 --batch-size 4 --epochs 3 --data uav_bird_training/dataset.yaml --weights yolov5s.pt

这里img参数是模型的输入图片大小，batch-size是一个batch的数据量，data是我们刚刚创建的数据配置文件，weights是预训练权重文件。

如果想从头开始训练模型（即使用随机初始化），则需要指定参数weight为空字符串（如果weights参数省略，则会默认使用yolov5s.pt进行权重初始化，并且train.py中还有一条语句为pretrained = weights.endswith('.pt')，即如果weights参数的字符串以.pt结尾，则会认为使用了预训练权重），同时指定cfg参数（如果weights参数不是以.pt结尾，此时如果省略cfg参数则会报错）。

例如，我想从头开始训练yolov5s模型，可以这样:

python train.py --img 640 --batch-size 4 --epochs 3 --data uav_bird_training/dataset.yaml --weights '' --cfg models/yolov5s.yaml

不过这种方式不推荐，因为我们的数据量太小，很难拟合。
在训练的时候，经常会遇到下面这个问题：

OSError: [WinError 1455] 页面文件太小，无法完成操作。

这是由于磁盘资源太小造成，只需要把磁盘分页改大即可，可以看这篇文章。当然，出现了这个错误，程序也有可能继续运行，程序会从最大分页尝试，如果不行，它会尝试更小的空间，直到试出了一个可以正常运行的空间。

如果命令行窗口出现下面的内容，说明训练成功：

从验证集上的结果可以看到，mAP是比较低的，主要是bird类训练结果太差，主要原因是Epoch太少、batch-size太小导致，毕竟这里只是演示，而且我的本地电脑显存太小。

训练结果的目录结构如下：

这里有很多文件，其中hyp.yaml是模型训练时的超参数，opt.yaml是命令行参数，png结尾的图片是各种性能指标的示意图，jpg图片则是训练、验证数据的示意图（Mosaic增强后），results.csv是训练过程中各个Epoch的各项指标，有点类似于训练日志，events.out.tfevents是tensorboard的事件文件，我们可以使用tensorboard查看，文件夹weights保存的是权重文件：

当然，也可以微调或从头训练yolov5l、yolov5n、yolov5m、yolov5x等其他模型，或者使用yolov5s6，方法与yolov5s一致。（从YOLOv5-6.1开始，支持--img参数为1280，即yolov5s6训练时，输入图像的尺寸可以到1280）

（3）性能评估

上述训练命令虽然也能看到在验证集上的结果，但在实际业务中，如果单纯想看mAP不可能重新跑一遍训练程序。假设一个场景，如果我是甲方公司的工程师，我要对乙方给我的结果进行验收，现在他们给了我一个压缩文件，名为uav_bird_result.zip（假设它是runs/train/exp压缩后得到的），这里面除了模型权重文件之外，还有各种性能指标的示意图以及results.csv，但各种示意图和results.csv文件都可以造假，唯独权重文件不能（给甲方的文件，）。此时，最客观最公正的结果是，使用乙方给到的模型权重在验证/测试集上跑一遍，这个指标多少，它就是多少。

所谓性能评估，指的是评估模型的精度和速度，精度包括准确率、召回率、mAP等，速度则可以用平均每张处理用时来表示。让模型在验证/测试集上评估性能，命令行的代码如下：

python val.py --img 640 --weights runs/train/exp/weights/best.pt --data uav_bird_training/dataset.yaml --batch-size 4

输出如下：

当然，也可以看看在模型在训练集上的结果，只需要指定参数--task为train，命令如下：

python val.py  --task train --img 640 --weights runs/train/exp/weights/best.pt --data uav_bird_training/dataset.yaml --batch-size 4

当然，实际业务当中，甲方不会用训练集或验证集对你提交的模型进行性能评估，而是使用测试集，测试集一般不会给到乙方（防止被加入到训练集里面）。我们这里的案例没有配置测试集，因为数据少，并且都是我们自己使用，不需要提交给别人。假设有测试集，那么只需要把数据配置好（图片和标签文件仿照train和valid组织），然后指定参数--task为test即可。

（4）使用自己训练的模型进行检测

这个和使用自定义相似，只需要指定权重参数和图像参数即可，命令为：

python detect.py --source uav_bird_training/data/images/valid/valdrone-2.jpg --weights runs/train/exp/weights/best.pt

因为我刚刚训练的模型mAP太低，因此检测的效果可能很差。

4 AutoDL训练

不是所有人都有GPU资源，我们总不可能为了学这门课专门去买一台带显卡的电脑吧？这个时候可以在网上获取GPU资源，推荐两个常用的算力平台：一个是谷歌的Colab，这个需要有谷歌账号，这个免费，但最长只能使用12小时，适合几个小时就训练好了的小模型；另一个是AutoDL，像GTX 3090这种显卡一个小时租金只需要一块多，V100只需要两块多，很便宜。

需要注意的事，不管是Colab还是AutoDL，操作系统均为Linux，需要重新下载代码（我们之前演示的是Windows版的代码）。我们这里演示一下如何在AutoDL上训练模型，Colab上的操作也是类似。（最近听说DeepLn也很实惠，V100不到一块钱，也可以尝试，操作和Auto DL完全一样）

（1）租赁显卡并打开终端

首先注册登陆，然后充值，随后选择需要的显卡，这里我们找RTX3090：

随后配置基础环境

配置镜像（虽然Python是3.8版，但问题不是特别大）

点击“立即创建”

创建后，点击JupyterLab

打开终端

在打开的终端中，我们进入到autodl-tmp目录下

（2）下载代码并解压

回到GitHub，找到我们之前下载YOLOv5-6.1代码的页面，然后滑到最下面，找到Source Code (tar.gz)，然后右击——复制链接地址

拿到地址后，就可以在终端中下载代码了，输入以下命令：

wget https://github.com/ultralytics/yolov5/archive/refs/tags/v6.1.tar.gz

接着解压：

tar -xvf v6.1.tar.gz

（3）配置环境与上传数据

进入到项目中

cd yolov5-6.1

安装环境

pip install -r requirements.txt

接下来将本地的数据文件夹（uav_bird_training）压缩后（做成zip文件）上传

上传数据非常的慢，上传后大概是这个结果：

解压数据：

unzip uav_bird_training.zip

（4）测试环境

打开tutorial.pynb

自己创建一个cell，然后在里面输入：

import torch
import utils
!python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images
display.Image(filename='runs/detect/exp/zidane.jpg', width=600)

之所以不运行自带的cell，是因为默认的cell是从克隆开始，需要重新下载很多东西。不知道怎么回事，这里的机器安装环境很快，但从下载东西巨慢，所以即便是字体文件和模型权重文件，下载都需要很长时间。（其实页可以将本地下载权的重也传上去，不过字体文件稍微麻烦一点，因为你不知道该存放到什么位置）

（5）训练

其实如果环境安装顺利，第（4）步可以跳过的，但我直接输入训练命令，总是提示
urllib.error.URLError: <urlopen error [Errno 110] Connection timed out>
所以我想在jupyter notebook中先执行检测程序（detect.py），这样可以看到字体文件和模型文件的下载速度，第（4）步有了结果后，就可以训练了。

因为RTX 3090的显存有24G，因此我们可以把batch-size设置的大一些，同时因为不是在本机上，我们可以设置Epochs为20，命令如下：

python train.py --img 640 --batch-size 16 --epochs 20 --data uav_bird_training/dataset.yaml --weights yolov5s.pt

结果如下：

如果对比刚刚在本地训练的结果，当epochs和batch-size变大之后，mAP有了显著提升，此外，训练速度也快了很多，本地（RTX 3060）3个Epoch是0.025 hours，AutoDL上自己租的设备是20个Epoch0.02 hours。

需要注意的是，batch-size不是越大越好，特别是我们这里数据量较小、且epoch较小，这种情况下如果把batch-size设的很大（假如为64），那么5个batch就能把训练集遍历完，相当于一个Epoch参数仅仅更新了5次，20个Epoch参数只更新了一百次，导致的结果是模型训练不充分，未能有效拟合。

使用刚刚训练的模型进行性能统计和检测，和本地（windows）的命令一摸一样，这里不再赘述。

（6）将训练结果压缩并下载到本地

训练的结果，包括权重文件、各个指标示意图、各个Epoch指标，都在runs/train/exp2中，我们可以用打包压缩，然后下载到本地。

打包命令：

tar -czvf uav_bird_result.tar.gz runs/train/exp2

下载：

然后在本地解压缩即可。

使用完了要记得关机，否则会持续计费，关机按钮在容器实例中：

OK，我们已经介绍完了AutoDL上训练模型的基本操作，其中关于下载、上传文件，速度都十分缓慢，AutoDL推荐使用公网网盘（例如百度网盘、阿里云盘等）进行传输，具体操作可以看AutoDL的帮助文档。

5 总结

终于把训练部分讲完了，本文可谓干货满满，内容主要包括：（1）使用labelimg对自定义数据集进行标注；（2）目标检测常见数据集格式；（3）YOLOv5的数据配置与模型训练；（4）在AutoDL上训练模型。掌握了这四项，那么本文的学习目的就达到。