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

一、准备

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 无需参数
  

七、进阶操作

• 训练除yolo_s以外的模型及训练参数修改

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