一、准备

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

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

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

  2. 下载预训练权重(来自 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
  3. 准备好自己的数据集(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 无需参数
    

七、进阶操作

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