【MMDetection3.0】训练自己的数据集

本文记录下MMDetection3.0版本，即截至目前最新的版本，训练自定义数据集的过程。当前MMDetection已经封装的很好了，虽然易于使用，但其API也愈发复杂，对于新手不太友好，这里记录下自己的踩坑经历。

数据部分

由于是自定义的数据集，则需要重新更改网络中的参数，以满足训练所需。在MMDetection中，数据相关的内容，主要存在于源码文件夹中的 mmdet 文件夹，下面的 datasets 和 evalution 文件夹。

需要修改的代码如下：

1 datasets/coco.py

将代码首位置的 METAINFO 改为下面内容：

METAINFO = {
    
    
'classes':("class1","class2",...),
'palette':[(220, 20, 60), (119, 11, 32),‘’‘,]
}

2 evaluation\functional\class_names.py

def coco_classes() -> list:
   return ['class1','class2',...]

【ps:在修改之后，需要重新编译mmdet,在根目录使用 python setup.py install ,否则会报错，大概是那个 ValueError: need at least one array to concatenate的异常】

3 configs/base/datasets/coco_detection.py

这部分代码主要是和数据所在文件夹相关，数据文件结构需要按照下面的形式：
data/coco/
train2017
1.jpg
2.jpg
…
val2017
1.jpg
2.jpg
…
annotations
instances_train2017.json
instances_val2017.json
【注意：json中类别的id最好为 int类型，这个官网教程也有说明，且索引从0开始。】
修改代码

# dataset settings
dataset_type = 'CocoDataset'
data_root = 'data/coco/'

# Example to use different file client
# Method 1: simply set the data root and let the file I/O module
# automatically infer from prefix (not support LMDB and Memcache yet)

# data_root = 's3://openmmlab/datasets/detection/coco/'

# Method 2: Use `backend_args`, `file_client_args` in versions before 3.0.0rc6
# backend_args = dict(
#     backend='petrel',
#     path_mapping=dict({
    
    
#         './data/': 's3://openmmlab/datasets/detection/',
#         'data/': 's3://openmmlab/datasets/detection/'
#     }))
backend_args = None

train_pipeline = [
    dict(type='LoadImageFromFile', backend_args=backend_args),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='Resize', scale=(1333, 800), keep_ratio=True),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PackDetInputs')
]
test_pipeline = [
    dict(type='LoadImageFromFile', backend_args=backend_args),
    dict(type='Resize', scale=(1333, 800), keep_ratio=True),
    # If you don't have a gt annotation, delete the pipeline
    dict(type='LoadAnnotations', with_bbox=True),
    dict(
        type='PackDetInputs',
        meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                   'scale_factor'))
]
train_dataloader = dict(
    batch_size=2,
    num_workers=2,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    batch_sampler=dict(type='AspectRatioBatchSampler'),
    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,
        backend_args=backend_args))
val_dataloader = dict(
    batch_size=1,
    num_workers=2,
    persistent_workers=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,
        backend_args=backend_args))
test_dataloader = val_dataloader

val_evaluator = dict(
    type='CocoMetric',
    ann_file=data_root + 'annotations/instances_val2017.json',
    metric='bbox',
    format_only=False,
    backend_args=backend_args)
test_evaluator = val_evaluator

模型部分

MMDetection中主要通过配置文件实现模型训练，如果想训练自己的数据集，则需要更改一些配置。这里以retinanet_r18_fpn_1x_coco.py 为例。

原代码：

_base_ = [
    '../_base_/models/retinanet_r50_fpn.py',
    '../_base_/datasets/coco_detection.py',
    '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
]

# model
model = dict(
    backbone=dict(
        depth=18,
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet18')),
    neck=dict(in_channels=[64, 128, 256, 512]))
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001))

修改后的代码：
其中 prefix 主要是将预训练模型进行过滤，过滤掉含backbone的映射，不然会发现，加载的预训练模型中所有参数名均加了 backbone.conv1.weights 而您本身加载的模型，其参数名为 conv1.weight，所以通过prefix 定义即可实现预训练模型和加载模型的参数映射。此外，也可以微调下学习率，避免学习率太大导致训练不稳定。

_base_ = [
    '../_base_/models/retinanet_r50_fpn.py',
    '../_base_/datasets/coco_detection.py',
    '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
]

# model
model = dict(
    backbone=dict(
        depth=18,
        #init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet18')
        init_cfg=dict(type='Pretrained', checkpoint="./pre_train/retinanet_r18_fpn_1x_coco_20220407_171055-614fd399.pth",prefix='backbone',)),
    neck=dict(in_channels=[64, 128, 256, 512]))
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001))

训练命令

其官方文档中分别给出了单GPU和多GPU进行训练的命令，其多GPU主要通过 dist_train.sh 实现训练。其训练命令，使用了占位符，不太容易理解，具体如下：

#!/usr/bin/env bash
CONFIG=$1
GPUS=$2
NNODES=${
    
    NNODES:-1}
NODE_RANK=${
    
    NODE_RANK:-0}
PORT=${
    
    PORT:-29500}
MASTER_ADDR=${
    
    MASTER_ADDR:-"127.0.0.1"}

PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
python -m torch.distributed.launch \
    --nnodes=$NNODES \
    --node_rank=$NODE_RANK \
    --master_addr=$MASTER_ADDR \
    --nproc_per_node=$GPUS \
    --master_port=$PORT \
    $(dirname "$0")/train.py \
    $CONFIG \
    --launcher pytorch ${
    
    @:3}

将其命令可以修改为：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node 2 --master_port 12355 ./tools/train.py ./configs/retinanet/retinanet_r18_fpn_1x_coco.py