MMSegmentation的用法（手把手入门教程）搭配Colab，对自己的数据进行训练

本文所用到的所有资料和代码已上传到Github，欢迎star：GitHub - zlr-zmm/MMSeg

1 介绍

MMSegmentation是一个基于Pytorch的开源语义分割工具箱，他是OpenMMLab项目的一部分 GitHub：GitHub - open-mmlab/mmsegmentation: OpenMMLab Semantic Segmentation Toolbox and Benchmark. MMseg提供的模块化设计极大便利了使用者，特别是用于入门学习，对于各类常用网络模型的即插即用集合，非常适合快速上手。

 这是mmsegmentation文件的预览，接下来，本文将以自己的数据集，利用谷歌云盘和谷歌colab服务器来重现一下MMsegmentation的使用方法

2 数据集准备

将mmsegmentation文件下载到本地来进行操作

本次数据集采用MICCAI2023牙齿分割挑战的数据集

下载来的数据集按照CHASE_DB1数据集的格式分配好，CHASE_DB1数据集格式如下：

CHASE_DB1
│   ├── images
│   │   ├── training
│   │   ├── validation
│   ├── annotations
│   │   ├── training
│   │   ├── validation

然后把数据压缩，命名为CHASEDB1.zip，使用mmseg给的数据集转换工具来转换：

使用时注意路径，需要先cd到该文件的路径，在我电脑上就是：

cd E:\mmsegmentation\mmsegmentation  # 这条指令就是把当前路径放到这个文件夹这里
python tools/dataset_converters/chase_db1.py /path/to/CHASEDB1.zip   # 使用数据转换工具

转换后的数据标签的深度是8位的，看上去是全黑的，具体显示效果如下：

(右边是给定的标签，可以将图像像素值x255，这样就看得到标签了）

这里仅仅展示这一个数据集的转换方法，其他数据集转换方法也类似，读者可以按照自己的需求进行选择。数据转换完成之后，工作就完成了一大半，接下来介绍如何利用Google Colab和Google Drive完成数据训练

3 Google Colab和Google Drive介绍

！！！使用需要科学上网！！！

Google Colab是谷歌提供的一个在线的工作平台，用户可以通过浏览器访问，最最最重要的是他提供免费的GPU算力，可以帮助落魄学生进行简单的模型训练，Colab的用法和Jupyter NoteBook一直，都是基于.ipynb格式的文件。

Colab可以最长运行时间为12小时，使用的时候需要把控这个时间，免费的GPU算力可能无法满足你的项目需求，Colab也支持氪金，但是略贵

如果没有visa支付的国内用户，可以去某宝下单，会更便宜，一个项目可以买一个pro。
对于存储，推荐使用谷歌云盘（Google Drive），谷歌云盘提供免费的5GB内存（很适合白嫖），如果数据集很大，或者运行保存的模型很大，推荐氪金，没有很贵

（MOP是澳币，这里是梯子的问题，所以显示的是澳币），200GB一般够用，如果不够用订阅2T就行了，某宝可以帮助用户订阅，而且如果需要更多的内存，补差价就行了，还是很方便的。

了解了谷歌平台的使用，读者需要在Google Drive中新建一个文件夹，点击左上角的新建按钮，然后新建文件夹，命名为Dataset，这里我已经新建好了这个文件夹，把自己的命名为CHASE_DB1的数据集上传到这里，不需要压缩上传，直接拖进来就可以上传，上传速度与网速有关。

4 MMSegmentation使用（以SegFormer为例）

接下来关注到configs这个文件夹，打开mmsegmentation/configs/segformer文件夹，可以看到里面有很多文件，（mmseg中SegFormer地址：https://github.com/open-mmlab/mmsegmentation/tree/main/configs/segformer）

 这里讲解一下这些文件命名规则：segformer是网络模型，mit-b0是使用的backbone，160k是训练轮次，cityscapes是使用的数据集，512x512是crop size。点开segformer_mit-b0_8xb2-160k_ade20k-512x512.py可以看到有如下内容：

_base_ = [
    '../_base_/models/segformer_mit-b0.py', '../_base_/datasets/ade20k.py',
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
]
crop_size = (512, 512)
data_preprocessor = dict(size=crop_size)
checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segformer/mit_b0_20220624-7e0fe6dd.pth'  # noqa
model = dict(
    data_preprocessor=data_preprocessor,
    backbone=dict(init_cfg=dict(type='Pretrained', checkpoint=checkpoint)),
    decode_head=dict(num_classes=150))

optim_wrapper = dict(
    _delete_=True,
    type='OptimWrapper',
    optimizer=dict(
        type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.01),
    paramwise_cfg=dict(
        custom_keys={
            'pos_block': dict(decay_mult=0.),
            'norm': dict(decay_mult=0.),
            'head': dict(lr_mult=10.)
        }))

param_scheduler = [
    dict(
        type='LinearLR', start_factor=1e-6, by_epoch=False, begin=0, end=1500),
    dict(
        type='PolyLR',
        eta_min=0.0,
        power=1.0,
        begin=1500,
        end=160000,
        by_epoch=False,
    )
]
train_dataloader = dict(batch_size=2, num_workers=2)
val_dataloader = dict(batch_size=1, num_workers=4)
test_dataloader = val_dataloader

其中

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

这里定义了模型、数据集、一些默认参数，运行轮次参数，接下来在config当中新建一个文件夹，给定一个任意的名字，这就是我们自己项目的config文件，并新建一个python文件叫做baseline

将segformer_mit-b0_8xb2-160k_ade20k-512x512.py或者按照你的需求，把整个文件中的内容复制到baseline.py里面

接着来介绍一下_base_里面的内容

 datasets文件存放了各种数据集的数据读取文件，和data_load差不多，点开chase_db1.py，里面有这些内容：

# dataset settings
dataset_type = 'ChaseDB1Dataset'
data_root = 'data/CHASE_DB1'
img_scale = (960, 999)
crop_size = (128, 128)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(
        type='RandomResize',
        scale=img_scale,
        ratio_range=(0.5, 2.0),
        keep_ratio=True),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='PackSegInputs')
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', scale=img_scale, keep_ratio=True),
    # add loading annotation after ``Resize`` because ground truth
    # does not need to do resize data transform
    dict(type='LoadAnnotations'),
    dict(type='PackSegInputs')
]
img_ratios = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
tta_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(
        type='TestTimeAug',
        transforms=[
            [
                dict(type='Resize', scale_factor=r, keep_ratio=True)
                for r in img_ratios
            ],
            [
                dict(type='RandomFlip', prob=0., direction='horizontal'),
                dict(type='RandomFlip', prob=1., direction='horizontal')
            ], [dict(type='LoadAnnotations')], [dict(type='PackSegInputs')]
        ])
]

train_dataloader = dict(
    batch_size=4,
    num_workers=4,
    persistent_workers=True,
    sampler=dict(type='InfiniteSampler', shuffle=True),
    dataset=dict(
        type='RepeatDataset',
        times=40000,
        dataset=dict(
            type=dataset_type,
            data_root=data_root,
            data_prefix=dict(
                img_path='images/training',
                seg_map_path='annotations/training'),
            pipeline=train_pipeline)))

val_dataloader = dict(
    batch_size=1,
    num_workers=4,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_prefix=dict(
            img_path='images/validation',
            seg_map_path='annotations/validation'),
        pipeline=test_pipeline))
test_dataloader = val_dataloader

val_evaluator = dict(type='IoUMetric', iou_metrics=['mDice'])
test_evaluator = val_evaluator

这里就定义了数据集的一些如图像大小、数据集路径等系列东西，复制里面的所有内容然后粘贴到我们自己的baseline.py里面

models文件里面存放了各类模型定义文件，找到segformer_mit-b0.py，里面存放的东西如下：
 

# model settings
norm_cfg = dict(type='SyncBN', requires_grad=True)
data_preprocessor = dict(
    type='SegDataPreProcessor',
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    bgr_to_rgb=True,
    pad_val=0,
    seg_pad_val=255)
model = dict(
    type='EncoderDecoder',
    data_preprocessor=data_preprocessor,
    pretrained=None,
    backbone=dict(
        type='MixVisionTransformer',
        in_channels=3,
        embed_dims=32,
        num_stages=4,
        num_layers=[2, 2, 2, 2],
        num_heads=[1, 2, 5, 8],
        patch_sizes=[7, 3, 3, 3],
        sr_ratios=[8, 4, 2, 1],
        out_indices=(0, 1, 2, 3),
        mlp_ratio=4,
        qkv_bias=True,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.1),
    decode_head=dict(
        type='SegformerHead',
        in_channels=[32, 64, 160, 256],
        in_index=[0, 1, 2, 3],
        channels=256,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

这里就是一些模型的参数等等，同样复制这个文件粘贴到baseline.py，剩下的'../base/default_runtime.py', '../base/schedules/schedule_160k.py'这两个文件以此类推进行复制粘贴。这里可以发现官方介绍没有CHASEDB1数据集格式，但是我们采用了，读者要根据自己的实际需求去模块化穿插就行了。

经过上述处理，可以得到一个完整的baseline.py，这样把这些内容复制粘贴到一个文件，可以避免对其他文件进行改动，接下来按照实际对baseline.py进行修改

修改1：

 num_classes需要按照自己的实际情况进行修改，我这里是分割牙齿，因此num_classes是1+1=2（需要算上背景），调整完数据集格式，需要到mmsegmentation/mmseg/datasets/chase_db1.py里面去调整自己的标签名字和数据集格式

 修改2：

然后需要更改数据存放的位置，'/content/drive/MyDrive/Dataset/CHASE_DB1/'就是我们数据集在谷歌云盘存储的位置

截至这里，初步的一些操作都已经完成，接下来将整个mmsegmentation文件夹上传到谷歌云盘，也可以将configs文件里面除了我们自己的文件夹(我的叫做teeth_seg)，其他都删掉上传
上传完成之后谷歌云盘如上

5 在Colab中使用数据

1、选择高性能运行环境

 点击右上角的向下箭头，出现菜单栏，接着选择查看资源
 

 点击更改运行时类型
 

选择T4 GPU，这样就可以使(bai)用(piao)谷歌的云算力

 2、链接Colab和谷歌云盘

from google.colab import drive
drive.mount('/content/drive')

 点击这个类似播放按钮运行这个单元，运行完成之后会提示“Mounted at /content/drive“

使用下述指令，将当前目录跳转到mmsegmentation中：

!cd /content/drive/MyDrive/mmsegmentation

安装需要的库：

!pip install openmim
!mim install mmengine
!mim install mmcv
!mim install mmsegmentation

这里需要替换到mmseg中datasets里面chase_db1.py这个文件，因为我们刚才对这个文件夹进行了修改，指令如下：

%cd /usr/local/lib/python3.10/dist-packages/mmseg/datasets
!rm chase_db1.py
!cp /content/drive/MyDrive/mmsegmentation/mmseg/datasets/chase_db1.py chase_db1.py

/usr/local/lib/python3.10/ 这个路径是mmseg下载保存的位置，可能会不一样，可以打印路径查看到底放在lib里面那个版本的python中

这里是打印datasets文件夹内文件的示例，可以按照实际需求进行更改

3、运行baseline.py

首先确保当前路径在mmsegmentation中，以防万一，可以先cd到这个文件夹

!cd /content/drive/MyDrive/mmsegmentation

然后是计算参数量，这一步可以测试baseline.py代码的逻辑是否存在问题

!python /content/drive/MyDrive/mmsegmentation/tools/analysis_tools/get_flops.py /content/drive/MyDrive/mmsegmentation/configs/teeth_seg/baseline.py

如果成功打印出参数量和计算量，就可以开始训练模型

!python tools/train.py configs/teeth_seg/baseline.py --work-dir work_dir/baseline

测试代码：

!python tools/test.py configs/teeth_seg/baseline.py work_dir/baseline/iter_40000.pth  --out test_result/baseline/iter_40000.pkl --show-dir test_result/baseline

这里为了快，只训练了40000轮，按照上述流程应该训练的是160000轮。测试结果（左图label，右图是预测结果）：

这里只是介绍快速入手，更多的训练和测试详细使用可以参阅官方给的文档：教程4：使用现有模型进行训练和测试 — MMSegmentation 1.1.1 文档

这里介绍的是在colab上使用的情况，在本地使用流程相似，在本地新建一个notebook用来执行指令，然后更改文件路径

mmsegmentation是一个很便于使用的工具箱，读者在了解如何进行基础训练之后可以更多关注mmseg这个文件中对各种模型模块化的设计
 

读者可以根据自己的需要，在mmseg中进行修改，需要注意的是：修改完上传到云盘后，记得按照刚才的指令把文件和user中的文件替换一下。

这里以添加loss为例子，读者在新建python文件设计了自己的loss function之后，需要在init.py中声明一下，然后才可以在baseline.py中使用。

以上就是本博客的所有内容，仅限快速入门使用这个工具箱，后续如果有深入使用技巧，会持续更新，有问题欢迎留言或私信提问