YOLOX模型训练

一、环境配置

参考源码：https://github.com/Megvii-BaseDetection/YOLOX.git

二、准备数据

1.准备数据放到YOLOX/datasets下，具体格如下：

　　datasets/VOCdevkit/VOC2007下存放图片及对应得xml文件，分别对应得文件夹为：JPEGImages  Annotations

2.运行下面代码生成对应label文件：

import os
import sys
from sklearn.model_selection import train_test_split

#获取当前路径

def main():
    # 获取图片路径
    image_list = os.listdir(data_dir+"/JPEGImages")
    train_fns, val_fns = train_test_split(image_list, test_size=0.1, shuffle=True, random_state=23)
    label_dir = data_dir+"/ImageSets/Main"
    if not os.path.exists(label_dir):
        os.makedirs(label_dir)
    with open(label_dir+"/trainval.txt","w",encoding='utf-8') as f1 ,open(label_dir+"/test.txt","w",encoding='utf-8') as f2:
        for train_img in train_fns:
            f1.write(train_img.split(".jpg")[0] + "\n")
        for test_img in train_fns:
            f2.write(test_img.split(".jpg")[0] + "\n")


if __name__ == '__main__':
    data_dir = "/home/lwh/YOLOX/datasets/VOCdevkit/V0C2007"
    main()

最终生成 目录ImageSets/Main,有两个对应得文件 test.txt  trainval.txt

三、创建配置文件，方便训练时调整训练参数

　　1.创建config.yaml文件

　　

CLASSES:
- blackboard   # 数据集的标签，本教程两个类别
- screen
# ---------------- model config ---------------- #
num_classes: 2       # 待检测的类别个数
# ---------------- dataloader config ---------------- #
input_size_w: 896
input_size_h: 896
# --------------- transform config ----------------- #
mosaic_prob:  1.0
mixup_prob:   0.4
hsv_prob: 1.0
flip_prob: 0.5
degrees: 10.0
translate: 0.1
mosaic_scale_low: 0.1
mosaic_scale_high: 2
mixup_scale_low: 0.5
mixup_scale_high: 1.5
shear: 2.0
enable_mixup: True

# --------------  training config --------------------- #

warmup_epochs: 5
max_epoch: 50
warmup_lr: 0
no_aug_epochs: 15
min_lr_ratio: 0.05
ema: True
weight_decay: 5e-4
momentum: 0.9
print_interval: 10
eval_interval: 1

# -----------------  testing config ------------------ #
test_conf: 0.01
nmsthre: 0.65

　　2.对应得写入代码 parseYaml.py

import yaml
import os
from easydict import EasyDict as edict


class YamlParser(edict):
    """ This is yaml parser based on EasyDict.
    """

    def __init__(self, cfg_dict=None, config_file=None):
        if cfg_dict is None:
            cfg_dict = {}

        if config_file is not None:
            assert (os.path.isfile(config_file))
            with open(config_file, 'r',encoding="utf-8") as fo:
                cfg_dict.update(yaml.load(fo.read(), Loader=yaml.FullLoader))

        super(YamlParser, self).__init__(cfg_dict)

    def merge_from_file(self, config_file):
        with open(config_file, 'r') as fo:
            self.update(yaml.load(fo.read()))

    def merge_from_dict(self, config_dict):
        self.update(config_dict)


def get_config(config_file=None):
    return YamlParser(config_file=config_file)

3.代码和配置文件分别存放到YOLOX/config目录下

四、修改对应的代码

　　1、D:\python_project\YOLOX\yolox\data\datasets\voc_classes.py

import os
from config import parseYaml
config_dir = os.path.abspath(os.getcwd())

cfg = parseYaml.get_config(config_dir+"/config/config.yaml")
# VOC_CLASSES = ( '__background__', # always index 0
# VOC_CLASSES = (
#   "blackboard",
#   "screen"
# )

VOC_CLASSES = cfg.CLASSES

　　2.D:\python_project\YOLOX\yolox\exp\yolox_base.py

增加配置文件的导入

from config import parseYaml
config_dir = os.path.abspath(os.getcwd())

cfg = parseYaml.get_config(config_dir+"/config/config.yaml")

修改Exp类

class Exp(BaseExp):
    def __init__(self):
        super().__init__()

        # ---------------- model config ---------------- #
        self.num_classes = cfg.num_classes
        self.depth = 1.00
        self.width = 1.00
        self.act = 'silu'

        # ---------------- dataloader config ---------------- #
        # set worker to 4 for shorter dataloader init time
        self.data_num_workers = 4
        self.input_size = (cfg.input_size_h, cfg.input_size_w)  # (height, width)
        # Actual multiscale ranges: [640-5*32, 640+5*32].
        # To disable multiscale training, set the
        # self.multiscale_range to 0.
        self.multiscale_range = 5
        # You can uncomment this line to specify a multiscale range
        # self.random_size = (14, 26)
        self.data_dir = None
        self.train_ann = "instances_train2017.json"
        self.val_ann = "instances_val2017.json"

        # --------------- transform config ----------------- #
        self.mosaic_prob = cfg.mosaic_prob
        self.mixup_prob = cfg.mixup_prob
        self.hsv_prob = cfg.hsv_prob
        self.flip_prob = cfg.flip_prob
        self.degrees = cfg.degrees
        self.translate = cfg.translate
        self.mosaic_scale = (cfg.mosaic_scale_low, cfg.mosaic_scale_high)
        self.mixup_scale = (cfg.mixup_scale_low, cfg.mixup_scale_high)
        self.shear = cfg.shear
        self.enable_mixup = cfg.enable_mixup

        # --------------  training config --------------------- #
        self.warmup_epochs = cfg.warmup_epochs
        self.max_epoch = cfg.max_epoch
        self.warmup_lr = cfg.warmup_lr
        self.basic_lr_per_img = 0.01 / 64.0
        self.scheduler = "yoloxwarmcos"
        self.no_aug_epochs = cfg.no_aug_epochs
        self.min_lr_ratio = cfg.min_lr_ratio
        self.ema = cfg.ema

        self.weight_decay = float(cfg.weight_decay)
        self.momentum = cfg.momentum
        self.print_interval = cfg.print_interval
        self.eval_interval = cfg.eval_interval
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # -----------------  testing config ------------------ #
        self.test_size = (cfg.input_size_h, cfg.input_size_w)
        self.test_conf = cfg.test_conf
        self.nmsthre = cfg.nmsthre

　　3.修改D:\python_project\YOLOX\exps\example\yolox_voc\yolox_voc_s.py

　　

增加导入的配置文件

from config import parseYaml
config_dir = os.path.abspath(os.getcwd())

cfg = parseYaml.get_config(config_dir+"/config/config.yaml")

修改EXP初始化函数

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.num_classes = cfg.num_classes
        self.depth = 0.33
        self.width = 0.50
        self.warmup_epochs = cfg.warmup_epochs

        # ---------- transform config ------------ #
        self.mosaic_prob = cfg.mosaic_prob
        self.mixup_prob = cfg.mixup_prob
        self.hsv_prob = cfg.hsv_prob
        self.flip_prob = cfg.flip_prob

        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

修改导入训练和测试数据的路径

with wait_for_the_master(local_rank):
    dataset = VOCDetection(
        data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
        # image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
        image_sets=[('2007', 'trainval')],    #此处修改
        img_size=self.input_size,
        preproc=TrainTransform(
            max_labels=50,
            flip_prob=self.flip_prob,
            hsv_prob=self.hsv_prob),
        cache=cache_img,
    )

def get_eval_loader(self, batch_size, is_distributed, testdev=False, legacy=False):
    from yolox.data import VOCDetection, ValTransform

    valdataset = VOCDetection(
        data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
        image_sets=[('2007', 'test')], #此处默认
        img_size=self.test_size,
        preproc=ValTransform(legacy=legacy),
    )

五、开始训练

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py  -expn TEST  -d 1 -b 8 --fp16 -o -c model/yolox_s.pth

model/yolox_s.pth为下载的预训练模型路径