pytorch训练自己的YOLOv5目标检测器（自定义数据集训练）

1.YOLOv5

YOLO系列可以说是单机目标检测框架中的潮流前线了。YOLO系列的原作者虽然放弃了YOLO系列的开发，但是俄罗斯的开发者Alexey接过了YOLO系列的大旗，今年四月份正式推出了YOLOv4，并开源了代码，论文也发了。
YOLOv4是基于darknet平台的，使用官方开源代码需要安装Visual Studio并使用Cmake来编译，inference和train起来都很不方便，so，在github上曾复现YOLOv3的大神ultralytics综合了YOLO系列的各种长处，推出了YOLOv5，作者在coco数据集上速度达到了惊人的140FPS.可以说是非常之快了。

2.环境配置

项目链接为：https://github.com/ultralytics/yolov5.
首先将项目git或者下载到本地，这里我假设平台已经安装了torch1.4以上的版本以及cuda等深度学习基本套件。项目的测试平台为：
操作系统：windows10
IDE：Pycharm
python版本：anaconda Pyhon3.7
pytorch版本：torch 1.4
cuda版本：10.1
显卡:RTX 2070S

2.1 使用pip安装必要模块

项目下载下来后，在terminal中执行如下命令来安装必要环境（建议换成国内的源后进行安装，不然会很慢）：

pip install numpy==1.17
pip install opencv-python
pip install tqdm
pip install pillow
pip install tensorboard
pip install pyyaml
pip install git
pip install pandas
pip install scikit-image
pip install Cython

2.2 pycocotools安装

本项目需要pycocotools模块，但是在windows环境下无法直接通过pip安装pycocotools，安装方法如下：

1. 安装Visual C++ 2015 build tools：点击这里.下载，然后执行即可傻瓜式安装

2. 安装后在Terminal中执行下面命令：

    pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI
   

执行后即可安装完毕

2.3 apex安装

单纯的inference实际上不需要apex模块，如果还要训练自己的数据集，就需要安装这个模块，安装方法如下：

1. 在github上把apex项目下载或者git到本地，链接为https://github.com/NVIDIA/apex.

2. 在terminal中激活pytorch的环境，并且进入到apex的文件夹下，在terminal中执行:

    python setup.py install
   

执行之后即安装完毕

执行后可以执行pip list命令查看当前环境下的所有模块，如果看到环境中有刚才安装的的模块，则环境已经配置完毕！

3.inference

环境配置完毕后可以使用预训练的权重来跑一跑：打开项目根目录下的detect.py,找到下列代码（程序执行入口）：

上面代码中：

1. weights为需要的权重，默认的预训练权重是yolo5s.pt
2. source是我们需要执行推断的文件，可以是图片/文件夹/视频/射像头，当推断图片/文件夹/视频时直接将default改为对应的路径即可，当source为’0’时（注意单引号不可以省略）会调用0号摄像头进行实时inference
3. output是输出路径
4. view-img代表在执行完inference后是否展示结果
5. device可以调整使用显卡或者GPU来执行推断，
   上面的选项均可以修改default值来进行设定，也可以在命令行执行命令中进行指定。

上述代码修改完毕后，在命令行中进入项目根目录并执行python detect.py即可执行推断。示例结果如图（图片来自百度图库，侵删）：

4.数据集准备

4.1 数据集格式转换

本项目要求为每个图片创建txt文件，文件名与图片名对应，后缀为txt，文件内每一行数据为：

类别索引(与yaml文件中的names列表索引对应，从0开始，往下看，4.3有详细介绍) 目标中心坐标 目标ground truth框的宽 目标ground truth框的高

上面所有的数据都要归一化，即中心坐标，宽、高的值都要除以原图的宽、高。
样例格式：

我们自己的数据集一般都是用labelImg等工具对数据进行标定，得到图片的符合VOC格式的xml数据标注。那么就需要对标注进行转换，其实就是单位不同，xml文件是以像素为单位（没有归一化），并且对于目标框的描述也是(左上角坐标，右下角坐标)的格式，不符合本项目要求的数据集标注格式，我们可以手动解析xml完成转换，如果懒的话我也找到了现成的转换代码，如下：

# -*- coding:utf-8 -*-

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

sets = ['train', 'test', 'val']

classes = ['helmet','face','motoperson']  # 填写类别的名字，与后面data/voc.names相同


# 这个函数的主要功能是把xml文件中的标注格式转换成txt需求的格式，并做归一化处理
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    # 计算中心坐标，xml文件中使用(左上角坐标，右下角坐标)标定一个框，需要转换为中心坐标，宽、高的形式
    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)

# 这个函数是读取xml并进行转换
def convert_annotation(image_id):# 输入值为图片的名字，不带后缀
    in_file = open('data/Annotations/%s.xml' % (image_id))# 这里改成自己的xml标注数据集的路径
    out_file = open('data/labels/%s.txt' % (image_id), 'w')# 这里改成自己的需要存放txt转换结果的路径
    print(image_id)
    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 = '0'
        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))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

# 接下来遍历自己的xml文件路径，提取文件名（不带后缀）然后循环调用convert_annotation函数，把文件名输入到此函数中即可

4.2 数据集准备

数据集转换完毕后，我们需要在项目根目录下建立coco128文件夹，文件夹内建立相关子文件夹如下图：

images文件夹下放的是图片文件，labels文件夹下放的是刚才生成的txt标注文件。，train2017代表训练集，val2017代表验证集，这个我们自行划分，但是images和labels文件夹下的文件一定要是一一对应的！下面放张更详细的展开图，以供参考：

4.3 yaml文件更改

进入到data文件夹下，找到coco128.yaml,做如下更改：

上图中train和val是训练集和验证集位置，按照上图进行设置即可。nc代表自定义数据的总类别数，我这里要做头盔识别，一共3类，因此改成3，names是我们的类名列表，注意这里的类名顺序一定要与txt文件中的第一列的索引（从0开始）一一对应！

更改完毕后我们还需要更改模型yaml文件，我们进入到项目的models文件夹：

可以看到里面包含了很多yaml文件，yolov3开头的是YOLOv3的模型，yolov5开头的是YOLOv5的模型，其中s,m,l分别代表小模型，中模型和大模型，训练和推断的速度依次降低，精度以此提高，可以根据自己的需求进行选择(在train.py中进行选择，接下来会讲)，选择我们需要训练的模型的yaml文件，将该文件中的nc也改成我们训练集中的类别数。以yolov5m.yaml为例：

至此，准备工作大功告成！

5.开始训练！

打开项目根目录的train.py，找到运行入口：

1. epochs：代表训练的epochs，更改default值可以改成我们需要的值
2. batch-size：默认是16，可以根据显卡显存进行更改
3. cfg：需要训练的模型文件，对应4.3中的模型yaml，可以根据需要选择训练的模型。
4. data:数据文件夹，按照我们之前的设置，这里不需要更改
5. weights:预训练权重，需要做迁移学习的可以使用其他数据集的权重初始化自己的权重
6. device：使用显卡/cpu来训练，不填的话会根据设备进行自动选择

更改完成后，在pycharm中run这个文件或者在terminal中执行python train.py即可开始训练：

（这个是之前训练的截图。图中提示没装apex模块所以训练详情无法打印，在第一步中有安装Apex的教程，按照教程安装后不会报这个问题）

训练完成后会在weights文件夹下生成best.pt和last.pt分别代表最好的权重和最后一个epoch的权重，我们可以在detect.py中更改对应的权重进行测试：

效果不错。