使用mmsegmentation训练自己的模型

1. 学习前言

mmsegmentation是一个基于ptorch的语义分割代码库，其中复现了众多先进的分割模型。由于代码风格统一，功能齐全，我打算用mmseg来跑对比实验。

本文使用mmseg中的segmenter模型来运行自己的数据集。自定义的数据集包含背景有8类(7类目标+1类背景），RGB图，图像大小不定，按照ade20k的数据样式排列。

全文只涉及跑segmenter的最小单元，仅适合入门。我理解错误的地方，还请大家能多多指正。

PS：我跑出来的效果并不好，这大概就是生搬硬套的代价吧

2. 数据准备

首先要模仿ade20k存放数据的风格来规划自己的数据集。
ade20k中图像是jpg格式，ground truth是png格式。存放结构如下图所示， ade20k和ADEChallengeData2016这些名字可以改成自己数据集的名，比如我的就叫tiger和tiger_roi

ade20k/
    |ADEChallengeData2016/
    |               |annotations/
    |               |       |training/
    |               |       |validation/
    |               |images/
    |               |       |training/
    |               |       |validation/

自定义数据：

tiger/
    |tiger_roi/
    |       |annotations/
    |       |       |training/
    |       |       |validation/
    |       |images/
    |       |       |training/
    |       |       |validation/

3. 环境准备

需要安装mmcv和mmsegmentation
github上官方文档有安装方法，我个人实际操作待更新

4. 数据丢进mmsegmentation的文件夹里

来到mmsegmentation文件中，新建一个空的文件夹，命名为data。然后把刚刚自定义的数据集拖进来。

5. 修改config文件

想让程序跑起来，只需要一行命令

cd mmsegmentation
python tools/train.py 配置文件

所以我们现在来看看这个配置文件是什么

对于segmenter模型来说，他的配置文件在mmsegmentation/configs/segmenter/，就是这几个红框圈中py文件，根据需求使用其一就可以。我认为配置文件其实就是设定参数（这几个区别没有仔细研究）

那我们打开第一个看看（segmenter_vit-b_mask_8x1_512x512_160k_ade20k.py）

_base_ = [
    '../_base_/models/segmenter_vit-b16_mask.py',
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_160k.py'
]
optimizer = dict(lr=0.001, weight_decay=0.0)

img_norm_cfg = dict(
    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=True)
crop_size = (512, 512)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=True),
    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg'])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(2048, 512),
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img'])
        ])
]
data = dict(
    # num_gpus: 8 -> batch_size: 8
    samples_per_gpu=1,
    train=dict(pipeline=train_pipeline),
    val=dict(pipeline=test_pipeline),
    test=dict(pipeline=test_pipeline))

看不懂，没关系！这个里面用的是ade20k数据集，我自己的数据集只是和ade20k的类别和路径不同罢了。所以只要找到哪里是数据路径，哪里是类别数，这个实验就完成了。（bushi

这个文件中的

_base_ = [
    '../_base_/models/segmenter_vit-b16_mask.py',
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_160k.py'
]

代表这个配置文件也是继承自以上四个py文件里，以我多年copy代码的经验，我要改的地方就在…/ base/datasets/ade20k.py。

那我们去…/ base/datasets/ade20k.py看看

dataset settings
dataset_type = 'ADE20KDataset'
data_root = 'data/ade/ADEChallengeData2016'
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
crop_size = (512, 512)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=True),
    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(2048, 512),
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=4,
    train=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/training',
        ann_dir='annotations/training',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline))

其中data_root = ‘data/ade/ADEChallengeData2016’数据路径，但是没找到类别数的定义，没关系，我们先把这个改成data_root = ‘data/tiger/tiger_roi’。（或者可以新建一个配置文件，但为了教程简洁，就不这么干了）

类别数的定义在…/ base/models/segmenter_vit-b16_mask.py
和刚刚的方法类似，找到num_classes=150，改成自定义数据集的类别（ade20k数据集中有150类目标+背景，Ground truth中0代表背景，1-150代表不同是物体。所以我也应该在这里设定目标类别，也就是num_classes=7。但是这样训练下来效果特别差，而且推理过程得到的结果0也并非代表实际的背景，希望大佬教教我这里）

到现在，配置文件就算是改完啦

6. 训练

cd mmsegmentation
python tools/train.py segmenter_vit-b_mask_8x1_512x512_160k_ade20k.py --work-dir [放日志文件和训练结果的文件夹]  --seed 0

7. 推理

官方demo

8. 学习后言

这改配置，多是一件美事。

Original: https://blog.csdn.net/gray___/article/details/124247411
Author: 胖琦今天也在摸鱼
Title: 使用mmsegmentation训练自己的模型