• YOLOX 的模型及训练配置模板文件位于./yolox/exp/yolox_base.py,对于VOC格式的数据可以仿照./exps/example/yolox_voc/yolox_voc_s.py进行修改

1. 更换yolox_m、yolox_l、yolox_x及其他模型

  • 相应的预训练权重如下:

    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
    from yolox.exp import Exp as MyExp
    
    # 如果其余参数不需要修改的话可以只改以下两行参数
    # yolox_m
    class Exp(MyExp):
      def __init__(self):
          super(Exp, self).__init__()
          self.num_classes = 7
          self.depth = 0.67  # 修改这个参数
          self.width = 0.75  # 修改这个参数
    
    # yolox_m
    class Exp(MyExp):
      def __init__(self):
          super(Exp, self).__init__()
          self.num_classes = 7
          self.depth = 1  # 修改这个参数
          self.width = 1  # 修改这个参数
          
    # yolox_x
    class Exp(MyExp):
      def __init__(self):
          super(Exp, self).__init__()
          self.num_classes = 7
          self.depth = 1.33  # 修改这个参数
          self.width = 1.25  # 修改这个参数
    
  • yolox-Darknet53 还需要重写 get_model 函数 详见github

    import os
    import torch.nn as nn
    from yolox.exp import Exp as MyExp
    
    class Exp(MyExp):
        def __init__(self):
            super(Exp, self).__init__()
            self.depth = 1.0
            self.width = 1.0
            self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
    
        def get_model(self, sublinear=False):
            def init_yolo(M):
                for m in M.modules():
                    if isinstance(m, nn.BatchNorm2d):
                        m.eps = 1e-3
                        m.momentum = 0.03
            if "model" not in self.__dict__:
                from yolox.models import YOLOX, YOLOFPN, YOLOXHead
                backbone = YOLOFPN()
                head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act="lrelu")
                self.model = YOLOX(backbone, head)
            self.model.apply(init_yolo)
            self.model.head.initialize_biases(1e-2)
        
            return self.model
    

2. 更改其他训练参数

  • 不建议直接修改./yolox/exp/yolox_base.py,建议仿照./exps/example/yolox_voc/yolox_voc_s.py进行相应修改

    from yolox.exp import Exp as MyExp
    
    class Exp(MyExp):
        def __init__(self):
            super(Exp, self).__init__()
            # ---------------- model config ---------------- #
            self.num_classes = 7  # 数据中的类别
            self.depth = 1.33  # 模型的深度
            self.width = 1.25  # 模型的宽度
            self.act = "silu"  # 模型的激活函数(relu、silu等)
            
            # ---------------- dataloader config ---------------- #
            self.data_num_workers = 4  # 如果训练过程占用内存太多,可以减少该值
            self.input_size = (640, 640)  # (高、宽)
            self.multiscale_range = 5  # 实际多尺度训练中图片范围为[640 - 5 * 32, 640 + 5 * 32],如果需要取消则设定为0
            # self.random_size = (14, 26)  # 取消注释这行来具体设定尺度范围
            self.data_dir = None  # 数据集图片的位置,如果为None则使用“dataset”文件的路径
            self.train_ann = "instances_train2017.json"  # name of annotation file for training
            self.val_ann = "instances_val2017.json"  # name of annotation file for evaluation
            self.test_ann = "instances_test2017.json"  # name of annotation file for testing
            
            # --------------- transform config ----------------- #
            self.mosaic_prob = 1.0  # 应用 mosaic aug 的概率
            self.mixup_prob = 1.0  # 应用 mixup aug 的概率
            self.hsv_prob = 1.0  # 应用 hsv aug 的概率
            self.flip_prob = 0.5  # 应用 flip aug 的概率
            self.degrees = 10.0  # 旋转角度区间 这里是 (-10,10)
            self.translate = 0.1  # 转换区间 这里是 (-0.1, 0.1)
            self.mosaic_scale = (0.1, 2)  # 马赛克尺度
            self.enable_mixup = True  # 是否应用 mixup aug
            self.mixup_scale = (0.5, 1.5)  # mixup aug 尺度
            self.shear = 2.0  # shear angle 区间,这里是 (-2, 2)
    
            # --------------  training config --------------------- #
            
            self.warmup_epochs = 5  # warmup epoch 数
            self.max_epoch = 300  # 最大训练 epoch
            self.warmup_lr = 0  # warmup 期间最小的 learning rate
            self.min_lr_ratio = 0.05  # 最小的 learning rate
            self.basic_lr_per_img = 0.01 / 64.0  # learning rate for one image. During training, lr will multiply batchsize.
            self.scheduler = "yoloxwarmcos"  # LRScheduler 名字
            self.no_aug_epochs = 15  # 最后 n 轮不使用 augmention like mosaic
            self.ema = True  # 在训练中采用 EMA 
    
            self.weight_decay = 5e-4  # 优化器的 weight_decay
            self.momentum = 0.9  # 优化器的 momentum
            self.print_interval = 10  # 迭代中输出日志的频率,如果设定为1则每次都输出log
            self.eval_interval = 10  # 训练过程中评估的 epoch 频率,这里每10个 epoch 会评估一次
            self.save_history_ckpt = True  # 是否保留训练历史的 checkpoint,如果是 False 则只有 latest and best ckpt
            self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]  # 实验名字
    
            # -----------------  testing config ------------------ #
            self.test_size = (640, 640)  # evaluation/test 期间的图片大小
            self.test_conf = 0.01  # evaluation/test 的输出阈值,可能性大于该阈值的预测结果才会被输出
            self.nmsthre = 0.65 # nms 阈值