[Share your Experience] How to train custom datasets in MMYOLO?

[vansin]

Welcome share your experience about train custom datasets in MMYOLO~

欢迎您分享使用 MMYOLO 训练自定义数据集的经验，将有机会加入 MMYOLO 高级贡献者群，获得技术指导哟~

[Bin-ze]

mmyolo custom dataset 教程

• 目的：
  mmyolo刚刚开源，相信很多小伙伴想快速上手，应用于自己的数据集，此教程从一个实际开发者的角度出发，来分享下如何快速将mmyolo应用与自己的任务。
• 步骤：
  • 环境配置
  • 数据集转换(coco)
  • 配置文件编写
  • 发起训练

下面以训练RTMDet网络为例，开始介绍，let's go!

环境配置

关于环境配置，请参考mmyolo环境配置
这里分享一下我自己的配置环境方式：

1. 首先创建一个shell脚本：

touch env.sh
chmod 777 env.sh

2. 打开env.sh，写入上述环境配置：

    conda create -n open-mmlab python=3.8 pytorch==1.10.1 torchvision==0.11.2 cudatoolkit=11.3 -c pytorch -y
    conda activate open-mmlab
    pip install openmim
    mim install mmengine
    mim install "mmcv>=2.0.0rc1,<2.1.0"
    mim install "mmdet>=3.0.0rc1,<3.1.0"
    git clone https://github.com/open-mmlab/mmyolo.git
    cd mmyolo
    # Install albumentations
    pip install -r requirements/albu.txt
    # Install MMYOLO
    mim install -v -e .
   

3. 运行

bash env.sh

然后等待脚本执行完成吧(这里请注意如果下载速度较慢，请配置国内镜像源)
如果没有报错，到这里环境配置就完成了！

数据集转换

为了更加方便快捷的使用，这里建议将需要训练的数据集转换为coco格式。这里比较自由，可以自行编写脚本进行转换。
如果是voc数据集或者是yolo格式数据(yolov5 repo所要求的数据集格式)，可以参考：
voc2coco.py
yolo2coco.py

下面从我自己的应用角度出发(yolo数据集转coco,再使用mmyolo训练)，开始介绍：

1. 将数据集转换为coco，使用上述脚本yolo2coco.py(注意为了使用该脚本，请在yolo数据集root路径下创建class.txt文件，里面存储自定义数据集类别)

2. 执行完成后将得到与coco 2017完全一致的目录树格式：

    |-datasetname(自己命名的数据集名称)
        |- annotations
            |-instances_val2017.json
            |-instances_train2017.json
        |-train2017
            |-img1.jpg
            ...
        |-val2017
            |-img11.jpg
            ... 
   

3. 创建软链接，将其放置在mmyolo要求的默认位置~/mmyolo/data目录下：

cd ~/mmyolo
mkdir data && cd data
ln -s ~/dataset_dir(转换数据集的位置) ./

此时查看data目录：

ls

将会出现转换数据集文件夹，数据准备完成！

配置文件编写

环境配置以及数据准备完成后，是时候选择合适的网络训练了！
下面我以RTMDet-s为例以及我的自定义14类food dataset为例，开始介绍

1. 打开config文件夹，创建自己的配置文件，这里不建议直接更改mmyolo提供的配置文件，可以复制一份到新的文件夹下，这里我的操作：

cd ~/mmyolo/config
mkdir food(以自己的数据集名称创建了将要存放配置的文件夹) && cp -r rtmdet/* food (将需要改写的配置复制一份到我的配置文件夹)
cd food
mv rtmdet_s_syncbn_8xb32-300e_coco.py rtmdet_s_syncbn_8xb32-300e_food.py (这里看个人习惯吧，这是我的命名习惯，更改文件名的时候需要查看内部配置的继承关系，防止找不到需要继承的配置)

2. 打开rtmdet_s_syncbn_8xb32-300e_food.py，配置文件如下：

    _base_ = './rtmdet_l_syncbn_8xb32-300e_coco.py'
    checkpoint = 'https://download.openmmlab.com/mmdetection/v3.0/rtmdet/cspnext_rsb_pretrain/cspnext-s_imagenet_600e.pth'  # noqa
   
    deepen_factor = 0.33
    widen_factor = 0.5
    img_scale = _base_.img_scale
   
    model = dict(
        backbone=dict(
            deepen_factor=deepen_factor,
            widen_factor=widen_factor,
            init_cfg=dict(
                type='Pretrained', prefix='backbone.', checkpoint=checkpoint)),
        neck=dict(
            deepen_factor=deepen_factor,
            widen_factor=widen_factor,
        ),
        bbox_head=dict(head_module=dict(widen_factor=widen_factor)))
   
    train_pipeline = [
        dict(type='LoadImageFromFile', file_client_args=_base_.file_client_args),
        dict(type='LoadAnnotations', with_bbox=True),
        dict(
            type='Mosaic',
            img_scale=img_scale,
            use_cached=True,
            max_cached_images=40,
            pad_val=114.0),
        dict(
            type='mmdet.RandomResize',
            scale=(img_scale[0] * 2, img_scale[1] * 2),
            ratio_range=(0.5, 2.0),  # note
            resize_type='mmdet.Resize',
            keep_ratio=True),
        dict(type='mmdet.RandomCrop', crop_size=img_scale),
        dict(type='mmdet.YOLOXHSVRandomAug'),
        dict(type='mmdet.RandomFlip', prob=0.5),
        dict(type='mmdet.Pad', size=img_scale, pad_val=dict(img=(114, 114, 114))),
        dict(
            type='YOLOXMixUp',
            img_scale=img_scale,
            use_cached=True,
            ratio_range=(1.0, 1.0),
            max_cached_images=20,
            pad_val=(114, 114, 114)),
        dict(type='mmdet.PackDetInputs')
    ]
   
    train_pipeline_stage2 = [
        dict(type='LoadImageFromFile', file_client_args=_base_.file_client_args),
        dict(type='LoadAnnotations', with_bbox=True),
        dict(
            type='mmdet.RandomResize',
            scale=img_scale,
            ratio_range=(0.5, 2.0),  # note
            resize_type='mmdet.Resize',
            keep_ratio=True),
        dict(type='mmdet.RandomCrop', crop_size=img_scale),
        dict(type='mmdet.YOLOXHSVRandomAug'),
        dict(type='mmdet.RandomFlip', prob=0.5),
        dict(type='mmdet.Pad', size=img_scale, pad_val=dict(img=(114, 114, 114))),
        dict(type='mmdet.PackDetInputs')
    ]
   
    train_dataloader = dict(dataset=dict(pipeline=train_pipeline))
   
    custom_hooks = [
        dict(
            type='EMAHook',
            ema_type='ExpMomentumEMA',
            momentum=0.0002,
            update_buffers=True,
            priority=49),
        dict(
            type='mmdet.PipelineSwitchHook',
            switch_epoch=40,
            switch_pipeline=train_pipeline_stage2)
    ]
   

发现_base_ = './rtmdet_l_syncbn_8xb32-300e_coco.py'，说明其继承了rtmdet_l的配置
继续观察，发现这里并没有关与dataloder(自定义数据集路径需要放在这里面)的配置，可以判断这些配置在它父亲的配置文件下

3. 打开./rtmdet_l_syncbn_8xb32-300e_coco.py，配置如下(这里只展示为了适应新数据集需要更改的部分，且需要更改的地方已经在下面做了标注）：

_base_ = '../_base_/default_runtime.py'

data_root = 'data/coco/'     #**更改为自定义数据集路径**
dataset_type = 'YOLOv5CocoDataset'

img_scale = (640, 640)  # height, width
deepen_factor = 1.0
widen_factor = 1.0
max_epochs = 300            # **设置训练epoch**
stage2_num_epochs = 20      # **这里难度倒炼丹大师吧**
interval = 10               # **多少epoch评估一次呢**

train_batch_size_per_gpu = 32 # **卡多可以使劲造**
train_num_workers = 10
val_batch_size_per_gpu = 5
val_num_workers = 10
# persistent_workers must be False if num_workers is 0.
persistent_workers = True
strides = [8, 16, 32]
base_lr = 0.004
.....

bbox_head=dict(
type='RTMDetHead',
head_module=dict(
    type='RTMDetSepBNHeadModule',
    num_classes=80,        # **自定义数据集多少类？**
    in_channels=256,
    stacked_convs=2,
    feat_channels=256,
    norm_cfg=dict(type='BN'),
    act_cfg=dict(type='SiLU', inplace=True),
    share_conv=True,
    pred_kernel_size=1,
    featmap_strides=strides),

......

train_dataloader = dict(
batch_size=train_batch_size_per_gpu,
num_workers=train_num_workers,
persistent_workers=persistent_workers,
pin_memory=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=dict(
    type=dataset_type,
    data_root=data_root,
    ann_file='annotations/instances_train2017.json', # **标注文件在哪？**
    data_prefix=dict(img='train2017/'),
    filter_cfg=dict(filter_empty_gt=True, min_size=32),
    pipeline=train_pipeline))

val_dataloader = dict(
    batch_size=val_batch_size_per_gpu,
    num_workers=val_num_workers,
    persistent_workers=persistent_workers,
    pin_memory=True,
    drop_last=False,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='annotations/instances_val2017.json', # **标注文件在哪？**
        data_prefix=dict(img='val2017/'),
        test_mode=True,
        pipeline=test_pipeline))

4. 为了适应自定义数据集的类别，需要修改继承的coco数据集的元数据，mmcv提供了如下方法：

    metainfo = {
    'CLASSES': ('1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14',),
    'PALETTE':[
            (220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),
            (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),
            (100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30)
    ]
    }
   

即在配置文件中添加自定义数据集的类别CLASS,PALETTE(我理解的是画框时候的色彩)，构建成字典，并命名为metainfo，这样这个配置将在build dataset的时候覆盖原始的coco元数据。
注：CLASSES是数据集类别的名称，如'bus,person'等，我的数据集标注使用字符串的'1'代表类别1，所以别误解

5. 好了，加上这个，我们的自定义数据集配置就算完成了，最终得到的配置为：
   base = '../base/default_runtime.py'

    data_root = 'data/0301/food/'    # **我的数据集位置**
    dataset_type = 'YOLOv5CocoDataset'
    metainfo = {
    'CLASSES': ('1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14',),
    'PALETTE':[
            (220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),
            (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),
            (100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30)
    ]
    }    # **我的类别设置**
   
   
    img_scale = (640, 640)  # height, width
    deepen_factor = 1.0
    widen_factor = 1.0
    max_epochs = 50   # **没时间跑那麽多了**
    stage2_num_epochs = 10 # **这不稍微调整下？**
    interval = 10
   
    train_batch_size_per_gpu = 16 # **穷人设置16**
    train_num_workers = 10
    val_batch_size_per_gpu = 5
    val_num_workers = 10
    # persistent_workers must be False if num_workers is 0.
    persistent_workers = True
    strides = [8, 16, 32]
    base_lr = 0.004
   
    model = dict(
        type='YOLODetector',
        data_preprocessor=dict(
            type='mmdet.DetDataPreprocessor',
            mean=[103.53, 116.28, 123.675],
            std=[57.375, 57.12, 58.395],
            bgr_to_rgb=False),
        backbone=dict(
            type='CSPNeXt',
            arch='P5',
            expand_ratio=0.5,
            deepen_factor=deepen_factor,
            widen_factor=widen_factor,
            channel_attention=True,
            norm_cfg=dict(type='BN'),
            act_cfg=dict(type='SiLU', inplace=True)),
        neck=dict(
            type='CSPNeXtPAFPN',
            deepen_factor=deepen_factor,
            widen_factor=widen_factor,
            in_channels=[256, 512, 1024],
            out_channels=256,
            num_csp_blocks=3,
            expand_ratio=0.5,
            norm_cfg=dict(type='BN'),
            act_cfg=dict(type='SiLU', inplace=True)),
        bbox_head=dict(
            type='RTMDetHead',
            head_module=dict(
                type='RTMDetSepBNHeadModule',
                num_classes=14,                           # **14类** 
                in_channels=256,
                stacked_convs=2,
                feat_channels=256,
                norm_cfg=dict(type='BN'),
                act_cfg=dict(type='SiLU', inplace=True),
                share_conv=True,
                pred_kernel_size=1,
                featmap_strides=strides),
            prior_generator=dict(
                type='mmdet.MlvlPointGenerator', offset=0, strides=strides),
            bbox_coder=dict(type='mmdet.DistancePointBBoxCoder'),
            loss_cls=dict(
                type='mmdet.QualityFocalLoss',
                use_sigmoid=True,
                beta=2.0,
                loss_weight=1.0),
            loss_bbox=dict(type='mmdet.GIoULoss', loss_weight=2.0)),
        train_cfg=dict(
            assigner=dict(
                type='mmdet.DynamicSoftLabelAssigner',
                topk=13,
                iou_calculator=dict(type='mmdet.BboxOverlaps2D')),
            allowed_border=-1,
            pos_weight=-1,
            debug=False),
        test_cfg=dict(
            nms_pre=1000,
            min_bbox_size=0,
            score_thr=0.05,
            nms=dict(type='nms', iou_threshold=0.6),
            max_per_img=100),
    )
   
    train_pipeline = [
        dict(type='LoadImageFromFile', file_client_args=_base_.file_client_args),
        dict(type='LoadAnnotations', with_bbox=True),
        dict(
            type='Mosaic',
            img_scale=img_scale,
            use_cached=True,
            max_cached_images=40,
            pad_val=114.0),
        dict(
            type='mmdet.RandomResize',
            scale=(img_scale[0] * 2, img_scale[1] * 2),
            ratio_range=(0.1, 2.0),
            resize_type='mmdet.Resize',
            keep_ratio=True),
        dict(type='mmdet.RandomCrop', crop_size=img_scale),
        dict(type='mmdet.YOLOXHSVRandomAug'),
        dict(type='mmdet.RandomFlip', prob=0.5),
        dict(type='mmdet.Pad', size=img_scale, pad_val=dict(img=(114, 114, 114))),
        dict(
            type='YOLOXMixUp',
            img_scale=img_scale,
            use_cached=True,
            ratio_range=(1.0, 1.0),
            max_cached_images=20,
            pad_val=(114, 114, 114)),
        dict(type='mmdet.PackDetInputs')
    ]
   
    train_pipeline_stage2 = [
        dict(type='LoadImageFromFile', file_client_args=_base_.file_client_args),
        dict(type='LoadAnnotations', with_bbox=True),
        dict(
            type='mmdet.RandomResize',
            scale=img_scale,
            ratio_range=(0.1, 2.0),
            resize_type='mmdet.Resize',
            keep_ratio=True),
        dict(type='mmdet.RandomCrop', crop_size=img_scale),
        dict(type='mmdet.YOLOXHSVRandomAug'),
        dict(type='mmdet.RandomFlip', prob=0.5),
        dict(type='mmdet.Pad', size=img_scale, pad_val=dict(img=(114, 114, 114))),
        dict(type='mmdet.PackDetInputs')
    ]
   
    test_pipeline = [
        dict(type='LoadImageFromFile', file_client_args=_base_.file_client_args),
        dict(type='mmdet.Resize', scale=img_scale, keep_ratio=True),
        dict(type='mmdet.Pad', size=img_scale, pad_val=dict(img=(114, 114, 114))),
        dict(
            type='mmdet.PackDetInputs',
            meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                    'scale_factor'))
    ]
   
    train_dataloader = dict(
        batch_size=train_batch_size_per_gpu,
        num_workers=train_num_workers,
        persistent_workers=persistent_workers,
        pin_memory=True,
        sampler=dict(type='DefaultSampler', shuffle=True),
        dataset=dict(
            type=dataset_type,
            data_root=data_root,
            metainfo=metainfo,                                # **很重要！很重要！很重要！重要的事情说三遍！**
            ann_file='annotations/instances_train2017.json',  # **自己的标注文件叫什么就写什么**
            data_prefix=dict(img='train2017/'),
            filter_cfg=dict(filter_empty_gt=True, min_size=32),
            pipeline=train_pipeline))
   
    val_dataloader = dict(
        batch_size=val_batch_size_per_gpu,
        num_workers=val_num_workers,
        persistent_workers=persistent_workers,
        pin_memory=True,
        drop_last=False,
        sampler=dict(type='DefaultSampler', shuffle=False),
        dataset=dict(
            type=dataset_type,
            data_root=data_root,
            metainfo=metainfo,                              # **很重要！很重要！很重要！重要的事情说三遍！**
            ann_file='annotations/instances_val2017.json',  # **自己的标注文件叫什么就写什么**
            data_prefix=dict(img='val2017/'),
            test_mode=True,
            pipeline=test_pipeline))
   
    test_dataloader = val_dataloader
   
    # Reduce evaluation time
    val_evaluator = dict(
        type='mmdet.CocoMetric',
        proposal_nums=(100, 1, 10),
        ann_file=data_root + 'annotations/instances_val2017.json',
        metric='bbox')
    test_evaluator = val_evaluator
   
    # optimizer
    optim_wrapper = dict(
        type='OptimWrapper',
        optimizer=dict(type='AdamW', lr=base_lr, weight_decay=0.05),
        paramwise_cfg=dict(
            norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))
   
    # learning rate
    param_scheduler = [
        dict(
            type='LinearLR',
            start_factor=1.0e-5,
            by_epoch=False,
            begin=0,
            end=1000),
        dict(
            # use cosine lr from 150 to 300 epoch
            type='CosineAnnealingLR',
            eta_min=base_lr * 0.05,
            begin=max_epochs // 2,
            end=max_epochs,
            T_max=max_epochs // 2,
            by_epoch=True,
            convert_to_iter_based=True),
    ]
   
    # hooks
    default_hooks = dict(
        checkpoint=dict(
            type='CheckpointHook',
            interval=interval,
            max_keep_ckpts=3  # only keep latest 3 checkpoints
        ))
   
    custom_hooks = [
        dict(
            type='EMAHook',
            ema_type='ExpMomentumEMA',
            momentum=0.0002,
            update_buffers=True,
            priority=49),
        dict(
            type='mmdet.PipelineSwitchHook',
            switch_epoch=max_epochs - stage2_num_epochs,
            switch_pipeline=train_pipeline_stage2)
    ]
   
    train_cfg = dict(
        type='EpochBasedTrainLoop',
        max_epochs=max_epochs,
        val_interval=interval,
        dynamic_intervals=[(max_epochs - stage2_num_epochs, 1)])
   
    val_cfg = dict(type='ValLoop')
    test_cfg = dict(type='TestLoop')
   

接下来就是发起训练了！

发起训练

有了上述准备，只需要运行训练脚本就可以发起训练
单卡：

python tools/train.py {config}

多卡：

./tools/dist_train.sh {config} 4

结束语

以上是我自己将mmyolo应用于custom dataset的过程，希望能帮助到需要的人。但可能包含比较多的不必要的步骤，或者包含一些个人的使用习惯，如有不足，请批评指示！