训练YoloX¶

训练代码¶

训练环境¶

docker push wojiazaiyugang/yolox:latest

进入YOLOX文件夹pip install -e .

数据集准备¶

voc格式标注。可以用scripts/datasets/generate_yolox_label.py来将voc格式的数据文件夹生成对应的label文件。然后按照VOC的格式把数据、标注文件和label文件放置到对应的位置，例如

root@pai-worker1:/mnt/nfs-storage/yujiannan/data# tree VOC2023
VOC2023
|-- Annotations # 里面是.xml文件
|-- ImageSets
|   `-- Main # 里面是test.txt和trainval.txt
`-- JPEGImages # 里面是.jpg文件

注意yolox会在dataset_dir目录下新建一个annotations_cache文件夹，里面是数据的缓存，所以如果更新了数据，要删除这个缓存文件夹否则新的数据不会被加载。如果检测的类别个数发生了变化，需要修改对应的类别配置yolox/data/datasets/voc_classes.py

配置文件¶

root@senseport-2080ti:/home/senseport0/Workspace/HiAlgorithm/YOLOX# tree -L 1 exps
exps
├── default
├── example
├── __pycache__
├── yolox_l_basketball_detection.py
├── yolox_l_human_detection.py
└── yolox_s_basketball_detection.py

3 directories, 3 files

以yolox_l_basketball_detection.py为例

root@senseport-2080ti:/home/senseport0/Workspace/HiAlgorithm/YOLOX# cat exps/yolox_l_basketball_detection.py 
import os

from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 1
        self.width = 1
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0] # 本次运行的名称，用于生成日志文件夹的命名

        # Define yourself dataset path
        # 数据目录
        self.data_dir = "/mnt/nfs-storage/yujiannan/data/bas_data/train_data" or "datasets/basketball"
        # 训练标注文件
        self.train_ann = "train_21.8.16_train.json" or "train_21.7.19.json"
        # 验证标注文件
        self.val_ann = "val_21.8.16.json" or "train_21.7.19_val.json"
        # 检测类别数
        self.num_classes = 1

        self.max_epoch = 300
        self.data_num_workers = 4
        self.eval_interval = 1

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

        # ---------- transform config ------------ #
        self.mosaic_prob = 1.0
        self.mixup_prob = 1.0
        self.hsv_prob = 1.0
        self.flip_prob = 0.5

        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.eval_interval = 1
        self.data_num_workers = 4
        self.max_epoch = 300
        self.dataset_dir = "/mnt/nfs-storage/yujiannan/data" # 数据root目录
        self.basketball_detection_dir = "2021" # 数据root目录下的VOC2021文件夹

从头训练¶

python -m yolox.tools.train -f exps/basketball_detection/yolox_l.py -d 0 -b 8 --fp16

yolox_l batch=28时大约占用28G显存 yolox_s batch=80时大约占用26G显存

载入预训练模型¶

在从头训练的命令基础上再添加-c path/to/checkpoint.pth参数

迁移训练¶

配置文件中冻结对应的层，配置如下。具体可以参考YOLOX官方文档

from exps.boxing_gloves.yolox_l import Exp as MyExp
class Exp(MyExp):
    def __init__(self):
        super().__init__()
        self.num_classes = 1
        self.exp_name = "human_transfer"
        self.basketball_detection_dir = "2023"

    def get_model(self):
        from yolox.utils import freeze_module
        model = super().get_model()
        freeze_module(model.backbone.backbone)
        return model

work dir是YOLOX_outputs

模型转onnx¶

python tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth