利用 YOLOX 训练自己的目标检测数据集

一、准备

首先从 YOLOX 的 GitHub 仓库获取 YOLOX 的代码，设置好 git 之后可以直接 clone 下来

$ git clone https://github.com/Megvii-BaseDetection/YOLOX.git

或者也可以直接下载 zip 文件并解压

下载预训练权重（来自 YOLOX 官方，截至2022/04/16），将其放到 YOLOX 目录下的weights文件夹（自行创建）

Model	size	mAPval 0.5:0.95	mAPtest 0.5:0.95	Speed V100 (ms)	Params (M)	FLOPs (G)	weights
YOLOX-s	640	40.5	40.5	9.8	9.0	26.8	github
YOLOX-m	640	46.9	47.2	12.3	25.3	73.8	github
YOLOX-l	640	49.7	50.1	14.5	54.2	155.6	github
YOLOX-x	640	51.1	51.5	17.3	99.1	281.9	github
YOLOX-Darknet53	640	47.7	48.0	11.1	63.7	185.3	github

准备好自己的数据集（VOC格式），并确认图片和标签除了后缀以外其余都相同

二、创建训练环境

这里推荐使用 conda 环境，具体如何下载和配置可以详见百度，记得将 conda 的 python 和 pip 等添加到环境变量，确保命令行中的 python 命令执行的是 conda 环境中的 Python

添加 yolox 库的依赖

$ cd YOUR_YOLO_DIR
$ pip install -r requirement.txt

三、制作数据集

利用 pycharm 或 VScode 等打开 YOLOX 的工作路径

在./datasets目录下创建如下的 Python 脚本来制作数据集：

# make_dir.py
import os

os.makedirs(r'VOCdevkit\VOC2007\Annotations')  # 标签文件夹
os.makedirs(r'VOCdevkit\VOC2007\ImageSets\Main')  # 训练（验证）/测试数据划分记录文件夹
os.makedirs(r'VOCdevkit\VOC2007\JPEGImages')  # 图片文件夹

然后将所有的xml标签放入Annotations文件夹中，将所有的图片放入JPEGImages文件夹中

在./datasets目录下创建如下的 Python 脚本来划分数据集：

# make_voc_data.py
import random
import os

train_pr = 0.8    # 训练集的比例

xml_names = os.listdir(r'VOCdevkit\VOC2007\Annotations') # 获取所有的标签
nums = len(xml_names)
train_nums = int(train_pr*nums)

list = range(nums)
train_index = random.sample(list,train_nums)

train_val = open(r'VOCdevkit\VOC2007\ImageSets\Main\trainval.txt','w')
test = open(r'VOCdevkit\VOC2007\ImageSets\Main\test.txt','w')

for i in list:
    name = xml_names[i].split(".")[0]+"\n"
    if i in train_index:
        train_val.write(name)
    else:
        test.write(name)

train_val.close()
test.close()

四、相关修改

修改./exps/example/yolox_voc/yolox_voc_s.py：

...

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.num_classes = 7     # 第 14 行的类别数调整为实际的类别数
        self.depth = 0.33
        self.width = 0.50
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.data_num_workres = 0   # 这里笔者在训练的时候没设置没出错，但是有人出错，所以只有一个gpu的话最好设一下
 ...       
# 第29行开始    
        dataset = VOCDetection(
            data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
            # image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
            image_sets=[('2007', 'trainval')],  # 将上一行中的2012删除
            img_size=self.input_size,
            preproc=TrainTransform(
                rgb_means=(0.485, 0.456, 0.406),
                std=(0.229, 0.224, 0.225),
                max_labels=50,
            ),
        )

修改./yolox/data/datasets/voc_classes.py：

# 将其中的 VOC_CLASSES 改成实际的数据集中的类别
VOC_CLASSES = (
    "Cars",
    "Dump truck",
    "Vans",
    "Cement mixer truck",
    "Box van",
    "Buses",
    "Pickup",
)

修改tools/train.py：

# 在代码的最上面加入如下代码
import sys
sys.path.append(r"YOUR_YOLOX_ABSOLUTE_DIR")  # 括号内添加 YOLOX 根目录的绝对路径，该路径下应该还有个yolox的文件夹

五、开始训练

打开终端，确保已经到了 YOLOX 的目录下，若没有则cd YOUR_YOLO_DIR，训练采用命令行的形式，一般采用fine turning的方式训练自己的数据集，命令如下

$ python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 4 -c weights/yolox_s.pth

训练中可选参数及其意义如下：

-expn 或 --experiment-name # 实验的名字 接受str类的参数 默认为None
-n 或 --name # 模型的名字 接受str类的参数 默认为None
--dist-backend # 分布式的后端 接受str类的参数 默认为nccl
--dist-url # 分布式训练的url地址 接受str类的参数 默认为None
-b 或 --batch-size # 训练批量大小 接受int类的参数 默认为64 实际建议设在16一下 具体取决于显存大小
-d 或 --devices # 训练的设备（GPU）数量 接受int类的参数 默认为None
-f 或 --exp_file # 训练的模型声明文件 接受str类的参数 默认为None
--resume # 是否从上一个checkpoint继续训练，一般中断后继续训练时适用，直接输入--resume不需要跟参数
-c 或 --ckpt # 上次训练的结果，继续训练和fine turning时填写check point路径 默认None 接受str类的参数
-e 或 --start_epoch # 指定开始的 epoch 和 --resume 等参数搭配使用 接受int类的参数 默认为None
--num_machines # 分布式训练的机器数 接受int类的参数 默认为1
--machine_rank # 分布式训练中该机器的权重 接受int类的参数 默认为0
--fp16 # 在训练时采用混合精度 默认False 只要输入了--fp16就是True 无需参数
--cache # 是都将图片缓存在RAM，RAM空间够的话推荐开启，能够加速训练 默认False 只要输入了--cache就是True 无需参数
-o 或 --occupy # 在训练开始时是否占用所有需要的GPU内存 默认False 只要输入了就是True 无需参数
-l 或 --logger # 记录训练的logger 接受str类的参数 默认为 tensorboard 可选 wandb

如果出现OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.（笔者未碰到）

则在tools/train.py的开头中加入如下的代码：

import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

如果在训练了10个epoch后测试时找不到文件（笔者未碰到），则可以修改./yolox/evaluators/voc_eval.py

...
def parse_rec(filename):
    """ Parse a PASCAL VOC xml file """
    # tree = ET.parse(filename)
    tree = ET.parse(os.path.join(r"datasets/VOCdevkit/VOC2007/Annotations",filename))  # 第17行改成这样就可以
    objects = []
    for obj in tree.findall("object"):
        obj_struct = {}
        obj_struct["name"] = obj.find("name").text
        obj_struct["pose"] = obj.find("pose").text
        obj_struct["truncated"] = int(obj.find("truncated").text)
        obj_struct["difficult"] = int(obj.find("difficult").text)
        bbox = obj.find("bndbox")
        obj_struct["bbox"] = [
            int(bbox.find("xmin").text),
            int(bbox.find("ymin").text),
            int(bbox.find("xmax").text),
            int(bbox.find("ymax").text),
        ]
        objects.append(obj_struct)

    return objects
...

重新训练时记得删除./datasets/VOCdevkit中处理数据集VOC2007外新生成的文件

六、数据预测

修改tools/demo.py：

# 在代码的最上面加入如下代码
import sys
sys.path.append(r"YOUR_YOLOX_ABSOLUTE_DIR")  # 括号内添加 YOLOX 根目录的绝对路径，该路径下应该还有个yolox的文件夹

假设需要预测的图片在./assets中

打开终端，确保已经到了 YOLOX 的目录下，若没有则cd YOUR_YOLO_DIR，训练采用命令行的形式，一般的命令如下

$ python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s.py -c weights/best_ckpt.pth --device gpu --save_result --path assets/

推理中可选参数及其意义如下：

# 直接在tools/demo.py后输入文件类型：image video 或者 webcam
-expn 或 --experiment-name # 实验的名字 接受str类的参数 默认为None
-n 或 --name # 模型的名字 接受str类的参数 默认为None
--path # 需要预测的文件夹 默认为./assets/dog.jpg
--camid # webcam demo camera id 接受int类的参数 默认为0
--save_result # 是否保存结果 直接输入 不需要跟参数
-f 或 --exp_file # 训练的模型声明文件 接受str类的参数 默认为None
-c 或 --ckpt # 上次训练的结果
--device # 推理的设备 cpu gpu 二选一
--conf # 输出框的最低置信度 接受float类的参数 默认为0.3
--nms # nms阈值 接受float类的参数 默认为0.3
--tsize # 测试图片大小 接受int类的参数 默认为None
--fp16 # 在训练时采用混合精度 默认False 只要输入了--fp16就是True 无需参数
--legacy # 配饰旧版本 默认False 只要输入了就是True 无需参数
--fuse # Fuse conv and bn for testing 默认False 只要输入了就是True 无需参数
--trt # 在测试时是否使用TensorRT模型 只要输入了就是True 无需参数

七、进阶 *** 作 训练除yolo_s以外的模型及训练参数修改

参考文献：https://www.bilibili.com/video/BV1mP4y1L7kQ