使用YoloX训练自建数据集

文章目录

1 Yolox 代码环境搭建
2 Yolox 自有数据集训练

随着旷视科技发布 Yolox 的论文和代码后，Yolox 得到了广泛的关注。但由于训练代码和之前的 Yolov3、Yolov4、Yolov5 的代码都不相同。且代码中的训练案例，以 COCO 和 VOC 格式为基准，和平时大家标注的文件，并不是完全相同的格式。而且训练自有数据的讲解流程，很多人不太熟悉。

因此本文以自有标注的人头数据集为案例，一步步和大家一起学习，整体的训练和测试全流程。

1 Yolox 代码环境搭建

在 Yolox 代码训练之前，我们先下载 Yolox 代码，将测试环境搭建起来。

测试的 Demo 跑通了，训练的环境也就没问题了。

1.1 下载 Yolox 代码

Yolox 代码链接：https://github.com/Megvii-BaseDetection/YOLOX

1.2 搭建测试环境

电脑系统为 Ubuntu 18.04 版本。

而 Yolox 测试环境的搭建，其实在代码中的，README.md 中 “Quick Start” 这部分。

首先为了测试环境更加独立，以 conda 为例，新建一个 Yolox 环境。

（1）新建一个 Conda 环境

1	conda create -n Yolox_3.7 python=3.7

（2）进入 Conda 环境

1	source activate Yolox_3.7 # 进入Conda环境中，并到下载好的YOLOX文件夹下。

（3）安装代码依赖的库文件

1	pip3 install -U pip && pip3 install -r requirements.txt

（4）通过 setup.py 安装一些库文件

1	python3 setup.py develop

（5）下载 apex 文件并安装 apex

git clone https://github.com/NVIDIA/apex

cd apex

sudo pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

（6）下载 pycocotools

1
2
3

pip3 install cython

pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

1.3 代码测试：Demo 效果测试

1.3.1 下载 Yolox_s.pth 文件

搭建好环境后，就可以下载官方的预训练模型，对图片进行测试了。

预训练权重的下载链接，在官方代码的说明中。

我们先下载 Yolox_s.pth 文件，尝试测试效果。

https://github.com/Megvii-BaseDetection/YOLOX/blob/main/README.md

下载好 yolox_s.pth.tar 后，放到 YOLOX 代码的根目录下。

1.3.2 Demo 测试

使用代码中自带的图片，进行 Demo 测试。

在 YOLOX 文件夹的终端页面输入：

1	python3 tools/demo.py image -n yolox-s -c yolox_s.pth.tar --path assets/dog.jpg --conf 0.3 --nms 0.5 --tsize 640 --save_result --device [gpu]

YOLOX 的代码中，会新建一个 YOLOX_outputs 文件夹，在其中的 yolox_s/vis_res/，可以看到带有检测效果的图片。

到此 Yolox 的测试环境，以及测试效果都实现了，下面我们再进行 Yolox 的自有数据集训练。

2 Yolox 自有数据集训练

2.1 数据集准备：标注数据

在数据集中，大白采用教室场景下的一个人头数据集，和大家一起尝试整个流程。

① 标注的工具：采用 Labelimg 标注软件

② 标注的图片：3000 张人头图片

③ 标签的类别：head。

④ 下载链接：

PartA of SCUT-HEAD [Google Drive][Baidu Drive]

PartB of SCUT-HEAD [Google Drive] [Baidu Drive]

2.2 数据集准备：训练 & 验证集划分

2.2.1 数据集介绍

当下载好人头数据集后，可以看到 head 的数据集文件夹，数据集只有一个标签：head（人头）。

进入 head 文件夹中，其中包含两个文件夹：

① JPEGImages 文件夹：数据集的图片

② Annotations 文件夹：与图片对应的所有 xml 文件。

我们前面说明，总共有 3000 张图片，即有对应的 3000 个 xml 文件。

2.2.2 模仿 VOC 格式排布

Yolox 的代码中有 VOC、和 COCO 两个数据集加载的格式，这里大白主要演示 VOC 加载的方式。

那么我们首先看一下 VOC 格式的分布：

在 VOC 这些文件夹中，我们主要用到：

① JPEGImages 文件夹：数据集图片

② Annotations 文件夹：与图片对应的 xml 文件

③ ImageSets/Main 文件夹：将数据集分为训练集和验证集，因此产生的 train.txt 和 val.txt。

从 Voc 的文件夹排布，和 head 的文件夹排布，可以看出：还缺少一个 ImageSets/Main 文件夹。

因此在 head 文件夹中，新建一个 ImageSets 文件集，再在其中新建一个 Main 子文件夹。

即得到下图这样的文件夹结构：

|--head
  |--Annotations
  |--ImageSets
     |--Main
  |---JPEGImages

2.2.3 划分训练集和验证集

因为自有标注好的图片数据，都是放在一起的。

而训练过程中，需要划分为训练集和验证集。

因此还需要编写脚本，将数据集分为训练集和验证集，并且生成对应的 train.txt，和 val.txt，放在 Main 文件夹中。

# 对head文件夹进行拆分，分为train.txt和val.txt
import os
import random

image_path = "JPEGImages/"
xmls_path = "Annotations/"
train_val_txt_path = "ImageSets/Main/"
val_percent = 0.1

images_list = os.listdir(images_path)
random.shuffle(images_list)

# 划分训练集和验证集的数量
train_images_count = int((1-val_percent)*len(images_list))
val_images_count = int(val_percent*len(images_list))

# 生成训练集的train.txt文件
train_txt = open(os.path.join(train_val_txt_path,"train.txt"),"w")
train_count = 0
for i in range(train_images_count):
    text = images_list[i].split(".jpg")[0] + "\n"
    train_txt.write(text)
    train_count +=1
    print("train_count:" + str(train_count))
train_txt.close()

# 生成验证集的val.txt文件
val_txt = open(os.path.join(train_val_txt_path,"val.txt"),"w")
val_count = 0
for i in range(val_images_count):
    text = images_list[train_images_count + i].split(".jpg")[0] + "\n"
    val_txt.write(text)
    val_count+=1
    print("val_count:" + str(val_count))
val_txt.close()

下载好代码后，将脚本文件 train_val_data_split.py 放在 JPEGImages 同路径下：

|--Annotations
|--ImageSets
|--JPEGImages
train_val_data_split.py

并进行运行后，在 ImageSets/Main 文件夹下，就会生成对应的 train.txt 和 val.txt。

主要注意的是：代码中，训练集和验证集的比例，为 9：1，大家也可以自行调整。

2.3 训练准备：修改训练配置参数

2.3.1 修改类别标签和数量

① 修改类别标签

因此前面自有的数据集只有一个类别，head。

将 yolox/data/datasets/voc_classes.py 中的标签信息，进行修改。

1
2
3

VOC_CLASSES = (
	"head",
)

注意：类别后面都要加逗号，例如 “head” 后面加了一个逗号 “，”。

② 修改类别数量

（1）修改 exps/example/yolox_voc/yolox_voc_s.py 中的 self.num_classes

因为只有 head 一种，所以 self.num_classes=1。

（2）修改 yolox/exp/yolox_base.py 中的 self.num_classes

将 self.num_classes=80 修改为 1。

2.3.2 修改训练集信息

（1）修改 exps/example/yolox_voc/yolox_voc_s.py 中的 VOCDetection。

因为是自己的数据集，所以修改为：

1 2	data_dir = "/mnt/data/head/" image_sets=[('train')],

data_dir 是前面 2.2 节中 head 的绝对路径，images_sets 修改为 train。

此外，max_labels，表示图片最多的目标数量，这里大白因为使用的是人头，数量较多，所以改为 100。

（2）修改 yolox/data/datasets/voc.py 中，VOCDection 函数中的读取 txt 文件。

因为自有的数据集，没有 year 年代的信息，所以修改为：

将：
self._year = year
rootpath = os.path.join(self.root, "VOC" + year)

修改为：
rootpath = self.root

2.3.3 修改验证集信息

修改 exps/example/yolox_voc/yolox_voc_s.py 中的 get_eval_loader 函数。

因为是自己的验证数据集，所以修改为：

将：
data_dir = os.path.join(get_yolox_datadir(),"VOCdevkit"),
image_sets = [('2007','test')]

修改为：
data_dir = "/mnt/data/head/"
image_sets = [('val')]

data_dir 是前面 2.2 节中 head 的绝对路径，images_sets 修改为 val。

2.3.4 修改不同的网络结构

以 Yolox_s 网络为例，比如在 exps/default/yolox_s.py 中，self.depth= 0.33 ，self.width= 0.5 。和 Yolov5 中的不同网络调用方式一样。

为了统一不同的网络结构，继续修改 exps/example/yolox_voc/yolox_voc_s.py 中的，self.depth 和 self.width。

再修改 yolox/exp/yolox_base.py 中的，self.depth 和 self.width。

2.3.5 修改其他相关

（1）删除 year 等信息

因为自有数据集中，没有 year 信息，所以需要删除。

即修改 yolox/data/datasets/voc.py 中，_get_voc_results_file_template 函数。

所以将第三行的 year 等删除，如下图所示：

将：
filedir = os.path.join(self.root, "results", "VOC" + self._year, "Main")

修改为：
filedir = os.path.join(self.root, "results")

在训练过程中，在原始的 head 数据集中，会生成一个 results 的文件夹，保存历史信息。

（２）修改验证 epoch 的数量

目前代码中是训练迭代 10 个 epoch，再对验证集做１次验证，但大白想每迭代 1 个 epoch，即做一个验证，及时看到效果。

参数在 yolox/exp/yolox_base.py 的 class Exp 中：

将
self.print_interval = 10
self.eval_interval = 10

修改为：
self.print_interval = 1
self.eval_interval = 1

设置为每迭代一个 epoch，即使用验证集验证一次。

（３）修改验证时的相关信息

主要对读取验证信息的相关代码进行调整，代码在 yolox/data/datasets/voc.py 中_do_python_eval 函数中。

① 因为自有数据集没有 year 信息，所以将其中的 rootpath 和 name：

将：
rootpath = os.path.join(self.root, "VOC" + self._year)

修改为：
rootpath = self.root

② 因为没有 year 信息，所以将其中的 cachedir:

将：
cachedir = os.path.join(self.root, "annotations_cache", "VOC" + self._year,name)

修改为：
cachedir = os.path.join(self.root, "annotations_cache")

在训练过程中，在原始的 head 数据集中，会生成一个 annotations_cache 的文件夹，保存历史信息。

|--Annotations
|--annotations_cache
|--ImageSets
|--JPEGImages
|--results

③ 因为没有 year 信息，所以修改 use_07_metric 的信息。

将：
use_07_metric = True if int(self._year) < 2010 else False

修改为：
use_07_metric = True

2.4 Yolox 训练及常见问题

2.4.1 开始训练

（1）终端训练

将下载好的 yolox_s.pth.tar 放到 YOLOX 文件夹中，打开终端，在终端中输入：

1	python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 64 -c yolox_s.pth.tar

（2）Pycharm 训练

代码运行时，常常需要 Debug 的方式，进行调试执行。

所以可以修改 train.py 的几个配置参数，采用 Debug 或者 Run 的方式进行执行。

主要需要修改以下参数：

① batch-size

根据自己机器的配置，设置 batch-size 的参数，比如大白这里设置的 64。

② devices 参数

如果 GPU 服务器只有 1 张卡，将 devices 的 default 修改为 0。

③ exp_file 参数

将 exp_file 的 default 修改为 yolox_voc_s.py 的路径（如代码版本更新，可重置路径）。

④ ckpt 参数

如果使用预训练权重，将 ckpt 的 default 修改为模型权重的路径。

2.4.2 常见问题

在运行 tools/train.py 时，可能会出现以下问题，如没有可以跳过：

问题 1：apex 路径报错

1
2
3

File "/mnt/code/YOLOX/yolox/core/trainer.py", line9, in <module>
	from apex import amp 
ImportError: cannot import name 'amp' from 'apex'(unknown location)

因为 YOLOX 内的 apex 文件夹，还有一个 apex 文件夹，所以引用路径有点问题。

解决方法：

这时在每个调用 apex 的地方，添加一个 apex. 即可。

将：
from apex import amp

修改为：
from apex.apex import amp

需要注意的是，不少的地方，需要添加 apex.，大概有 10 处左右，修改完之后，错误即可解决。

问题 2：probubuf 报错

1 2	from google.protobuf.internal import enum_type_wrapper ModuleNotFoundError: No module named 'google.protobuf.internal'

解决方法：

（1）pip3 uninsall probobuf

（2）pip3 install google

（3）pip3 install protobuf

2.5 训练效果测试

在上面训练好模型后，我们可以得到一个精度测试最优的网络模型：

① best_ckpt.pth.tar：在 tools/YOLOX_outputs/yolox_voc_s 文件夹中。

② 为了方便测试，再挑选一张人头测试图片，放到 assets 文件夹中。