使用YoloX进行自定义训练

昨天因为要准备比赛，所以搞了一阵yolox的custom_train，效果很ok，遂分享出来。

下载并配置YoloX

https://github.com/Megvii-BaseDetection/YOLOX

怎么clone就不需要我多言了。

pip安装完环境，我们需要clone一下NVIDIA的apex，最后执行

python3 setup.py develop

基本的结构介绍

介绍一下我们自定义训练主要用到的几个文件：

模型主要部分：./yolox

各种工具：./tools

其中我们需要使用train.py来训练，demo.py来预测。

命令分别为：

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 8 -b 64 --fp16 -o -c /path/to/yolox_s.pth [--cache]

python tools/demo.py image -f exps/default/yolox_s.py -c /path/to/your/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]

或者视频：

python tools/demo.py video -n yolox-s -c /path/to/your/yolox_s.pth --path /path/to/your/video --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]

注意：之后要把./yolox/exp/yolox_base.py里面的num_classes和dataset改掉。

注意：如果之后模型报不能找到”yolox“这个module的话，需要在前面添加

import sys

sys.path.append(r'你的/根/目录')

使用VOC类型的数据集训练yolo_s模型需要用到的Exp：./exps/example/yolox_voc/yolox_voc_s.py

这个文件主要定义了我们训练的Exp，我们自定义训练需要修改的主要包括分类数（num_classes）、数据集(dataset)。

数据集：./datasets

我们需要在里面放置VOC或者COCO格式的数据集。

我是用的是自制VOC结构的数据集，结构如下：

 .
└── VOCdevkit
    └── VOC2007
        ├── Annotations
        │   ├── 003002_0.xml
        │   ├── 003002_1.xml
        │   ├── 003008_1.xml
        │   └── .......
        ├── ImageSets
        │   └── Main
        │       ├── test.txt
        │       ├── train.txt
        │       ├── trainval.txt
        │       └── val.txt
        └── JPEGImages
            ├── 003002_0.jpg
            ├── 003002_1.jpg
            ├── 003008_1.jpg
            └──.......

而我初始的数据集只有jpg和xml，所以需要转换。

转换数据集：

先把原始图片目录改成这样，创建用于生成数据集的xml2voc2007.py：

. 
└── data
    ├── 003002_0.jpg
    ├── 003002_0.xml
    ├── 003002_1.jpg
    ├── 003002_1.xml
    ├── 003008_1.jpg
    ├── 003008_1.xml
    └── .......
└── xml2voc2007.py

附上用xml和jpg生成VOC格式数据集的代码：

xml2voc2007.py >folded

#命令行执行：  python xml2voc2007.py --input_dir data --output_dir VOCdevkit

import argparse
import glob
import os
import random
import os.path as osp
import sys
import shutil

percent_train = 0.9

# 主程序执行

def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", default="data", help="input annotated directory")
    parser.add_argument("--output_dir", default="VOCdevkit", help="output dataset directory")
    args = parser.parse_args()

    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    print("| Creating dataset dir:", osp.join(args.output_dir, "VOC2007"))
    
    # 创建保存的文件夹
    if not os.path.exists(osp.join(args.output_dir, "VOC2007", "Annotations")):
        os.makedirs(osp.join(args.output_dir, "VOC2007", "Annotations"))
    if not os.path.exists(osp.join(args.output_dir, "VOC2007", "ImageSets")):
        os.makedirs(osp.join(args.output_dir, "VOC2007", "ImageSets"))
    if not os.path.exists(osp.join(args.output_dir, "VOC2007", "ImageSets", "Main")):
        os.makedirs(osp.join(args.output_dir, "VOC2007", "ImageSets", "Main"))
    if not os.path.exists(osp.join(args.output_dir, "VOC2007", "JPEGImages")):
        os.makedirs(osp.join(args.output_dir, "VOC2007", "JPEGImages"))
    
    # 获取目录下所有的.jpg文件列表
    total_img = glob.glob(osp.join(args.input_dir, "*.jpg"))
    print('| Image number: ', len(total_img))
    
    # 获取目录下所有的joson文件列表
    total_xml = glob.glob(osp.join(args.input_dir, "*.xml"))
    print('| Xml number: ', len(total_xml))
    
    num_total = len(total_xml)
    data_list = range(num_total)
    
    num_tr = int(num_total * percent_train)
    num_train = random.sample(data_list, num_tr)
    
    print('| Train number: ', num_tr)
    print('| Val number: ', num_total - num_tr)
    
    file_train = open(
        osp.join(args.output_dir, "VOC2007", "ImageSets", "Main", "train.txt"), 'w')
    file_val = open(
        osp.join(args.output_dir, "VOC2007", "ImageSets", "Main", "val.txt"), 'w')
    
    for i in data_list:
        name = total_xml[i][:-4] + '\n'
        if i in num_train:
            file_train.write(name[5:])
        else:
            file_val.write(name[5:])
    
    file_train.close()
    file_val.close()
    
    if os.path.exists(args.input_dir):
        # root 所指的是当前正在遍历的这个文件夹的本身的地址
        # dirs 是一个 list，内容是该文件夹中所有的目录的名字(不包括子目录)
        # files 同样是 list, 内容是该文件夹中所有的文件(不包括子目录)
        for root, dirs, files in os.walk(args.input_dir):
            for file in files:
                src_file = osp.join(root, file)
                if src_file.endswith(".jpg"):
                    shutil.copy(src_file, osp.join(args.output_dir, "VOC2007", "JPEGImages"))
                else:
                    shutil.copy(src_file, osp.join(args.output_dir, "VOC2007", "Annotations"))
    
    print('| Done!')

if __name__ == "__main__":
    print("—" * 50)
    main()
    print("—" * 50)

开始训练自己的数据集

以VOC格式的数据集举例：

先在./yolox/data/datasets/voc_classes.py里，把VOC_CLASSES改成你的数据集类别

接着把./exps/example/yolox_voc/yolox_voc_s.py里面的num_classes改成你的类别数

好了，然后我们新建个文件夹，去仓库里把预训练权重下载下来，准备工作就完事了。

训练命令：

python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 8 --fp16 -o -c models/yolox_s.pth`

其中：

-d是num_workers，

-b是batch_size，

-c是cache，就是之前下载的权重。

好的，这样我们就可以开始训练了。

可视化训练过程

yolox自带有tensorboard输出，我们可以在安完tensorboard以后，进入./YOLOX_outputs/yolox_voc_s/tensorboard里面，输入

tensorboard --logdir ./ --bind_all

来可视化训练过程。

评估效果

参照上面demo的命令。