YOLOX训练代码分析1-COCO与VOC训练

1. YOLOX的网络结构图与代码YOLOv3、YOLOv4、YOLOv5、YOLOx的网络结构图(清晰版)_YMilton的专栏-CSDN博客​​​​​(https://blog.csdn.net/YMilton/article/details/120268832?spm=1001.2014.3001.5502

 (1) 网络结构图

​

(2) yolox代码URL：

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

 2. yolox如何训练coco数据

(1) 第一步(参数理解)：yolox训练数据的入口为YOLOX-main/tools/train.py。main函数中通过get_exp获取训练过程相关的超参数。

其中train.py中的make_parser函数的参数解析如下：

-expn: 训练过程数据保存的位置，默认位置是tools/YOLOX_outputs中的目录，如果指定-expn xxx，则tools/YOLOX_outputs/xxx

-n: 指定训练过程的模型名称，如yolox-s，yolox-m，yolox-l，yolox-x，yolox-tiny等

-b: 训练过程中的batch size大小

-d: 指定训练过程中使用的设备

-f: 设置训练过程中数据加载的python文件，如exps/example/yolox_voc/yolox_voc_s.py

-c: 设置预训练权重文件

-e: 开始训练的周期数

--fp16: 是否允许Apex的混合精度加速，使得训练时间缩短。代码实现如下：

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()

混合精度训练（Mixed Precision）的精髓在于“在内存中用 FP16 做储存和乘法从而加速计算，用 FP32 做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题。

-o: 是否首先占用GPU内存进行训练

opts: 通过命令行的形式修改训练过程的超参数

(2) 第二步(数据加载)：训练过程中如果指定为VOC数据，则通过

-f exps/example/yolox_voc/yolox_voc_s.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 

 其中训练的时候VOC数据需要解压到datasets中，如图

训练过程中如果是指定COCO数据集（COCO数据集不在YOLOX-main/datasets中），需要在example中自己创建yolox_coco_s.py文件(从yolox_voc_s.py中复制内容到yolox_coco_s.py中)，训练的语句

python tools/train.py -f ../exps/example/yolox_coco/yolox_coco_s.py -b 8 -o --fp16

本文中COCO数据集在E:/DataSets/COCO中，其中yolox_coco_s.py文件的配置主要包括网络大小构建(Exp的构造函数)、训练数据集获取(get_data_loader)与验证数据集获取(get_eval_loader)。代码解释如下：

1) 构造函数中的self.depth与self.width控制yolox网络的大小，其中网络大小包括s、m、l、x，不同网络的depth与width不同，s: depth=0.33，width=0.50；m: depth=0.67，width=0.75；l：depth=1.0，width=1.0；x: depth=1.33，width=1.25。其中self.input_size控制网络输入的大小。构造函数如下：

def __init__(self):
    super(Exp, self).__init__()
    self.num_classes = 80
    self.depth = 0.33  #控制网络的深度与宽度
    self.width = 0.50
    self.input_size=(416,416) # 网络输入大小
    self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

2) 加载训练数据集，需要修改VOCDetection为COCODataset，data_dir指定训练数据集的路径，json_file指定训练集标签的json文件，name指定训练的名称，其他的不变。修改如下

# COCO数据的加载
dataset = COCODataset(data_dir='E:/DataSets/COCO', json_file='instances_train2014.json', name='train2014',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))

 3) 加载验证数据集，同样需要修改VOCDetection为COCODataset对象，data_dir指定验证集路径，json_file指定标注json文件。

# 验证数据集加载
valdataset = COCODataset(data_dir='E:/DataSets/COCO', json_file='instances_val2014.json', img_size=self.test_size, preproc=ValTransform(rgb_means=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)))

 最终的网络加载、COCO训练数据集加载、COCO验证数据集加载的整个代码如下：

# encoding: utf-8
import os

import torch
import torch.distributed as dist

from yolox.exp import Exp as MyExp  # Exp:为yolox.exp中的文件yolox_base.py的类Exp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.num_classes = 80
        self.depth = 0.33  #控制网络的深度与宽度
        self.width = 0.50
        self.input_size=(416,416) # 网络输入大小
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

    # 加载COCO训练数据集
    def get_data_loader(self, batch_size, is_distributed, no_aug=False):
        from yolox.data import (COCODataset, TrainTransform, YoloBatchSampler, DataLoader, InfiniteSampler,
                                MosaicDetection)
        # COCO数据的加载
        dataset = COCODataset(data_dir='E:/DataSets/COCO', json_file='instances_train2014.json', name='train2014',
                              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))
        # Mosaic的处理
        dataset = MosaicDetection(dataset, mosaic=not no_aug, img_size=self.input_size,
                                  preproc=TrainTransform(rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_labels=120),
                                  degrees=self.degrees, translate=self.translate,scale=self.scale, shear=self.shear,
                                  perspective=self.perspective, enable_mixup=self.enable_mixup)

        self.dataset = dataset

        if is_distributed:
            batch_size = batch_size // dist.get_world_size()

        sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)

        batch_sampler = YoloBatchSampler(sampler=sampler, batch_size=batch_size, drop_last=False, input_dimension=self.input_size, mosaic=not no_aug)

        dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
        dataloader_kwargs["batch_sampler"] = batch_sampler
        train_loader = DataLoader(self.dataset, **dataloader_kwargs)

        return train_loader

    # 加载COCO验证数据集
    def get_eval_loader(self, batch_size, is_distributed, testdev=False):
        from yolox.data import COCODataset, ValTransform

        valdataset = COCODataset(
            data_dir='E:/DataSets/COCO',
            json_file='instances_val2014.json',
            img_size=self.test_size,
            preproc=ValTransform(rgb_means=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225))
        )

        if is_distributed:
            batch_size = batch_size // dist.get_world_size()
            sampler = torch.utils.data.distributed.DistributedSampler(valdataset, shuffle=False)
        else:
            sampler = torch.utils.data.SequentialSampler(valdataset)

        dataloader_kwargs = {"num_workers": self.data_num_workers,"pin_memory": True,"sampler": sampler}
        dataloader_kwargs["batch_size"] = batch_size
        val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)

        return val_loader

    def get_evaluator(self, batch_size, is_distributed, testdev=False):
        from yolox.evaluators import COCOEvaluator

        val_loader = self.get_eval_loader(batch_size, is_distributed, testdev=testdev)
        evaluator = COCOEvaluator(dataloader=val_loader, img_size=self.test_size, confthre=self.test_conf,
                                  nmsthre=self.nmsthre, num_classes=self.num_classes)
        return evaluator

训练结果输出

训练过程中会在tools/YOLOX_outputs/yolox_coco_s中生成训练过程存储的日志文件。