超详细！带你轻松掌握 MMSegmentation 整体构建流程

目录

1.语义分割介绍

2.语义分割的应用

自动驾驶

遥感图像分析

医学图像分析

3.MMSegmentation 算法库框架介绍

3.1 MMSegmentation 目录结构

3.2 MMSegmentation 模型实现

4 总结

---

大家好呀，今天我们将开启新的解读文章，带领大家了解 MMSegmentation 算法库的整体框架。MMSegmentation 是 OpenMMLab 开源的基于 PyTorch 实现的功能强大的语义分割工具箱，2020 年 7月 开源至今， Google citation 已有上百引用 (截止到 2022 年 5 月)，很多有影响力的研究工作，例如 Swin Transformer、ConvNeXt 都有用到 MMSegmentation。

MMSegmentation 的主要特性如下：

• 丰富的语义分割模型： 已支持 11 种主干网络和 34 种算法，例如常用模型 FCN， PSPNet 和 DeepLabV3；Transformer 模型，Swin Transformer、Segmenter 和 SegFormer； Real-Time 实时分割模型， ICNet、BiSeNet 和 STDC 等；以及最近流行的网络 ConvNeXt 和 MAE。
• 大量开箱即用的模型权重：在 16 个常用的语义分割数据集上提供了 590 个训练好的模型。
• 统一的性能评估框架：优化和统一了训练和测试的流程，方便公平比较各个模型在特定任务上的表现。

https://github.com/open-mmlab/mmsegmentation​github.com/open-mmlab/mmsegmentation

1.语义分割介绍

在开始介绍 MMSegmentation 整体架构前，我们先简单了解下分割任务和语义分割任务。分割任务的本质是对图像中的每个像素 pixel 做分类，可以细分为语义分割、实例分割和全景分割，它们之间的不同如下图所示：

来源: ( Panoptic Segmentation https://arxiv.org/pdf/1801.00868.pdf)

从上图可以看出：

• 语义分割是给图像中的每个像素分配一个类别，得到特定类别的 mask；
• 实例分割是对特定的物体进行分类，与目标检测输出物体的边界框和类别不同，实例分割输出的是特定物体的 mask 和类别；
• 全景分割是语义分割和实例分割的结合，对于可数的对象实例 things 如行人、汽车去做实例分割，对于不可数的语义区域 stuff 如天空、地面做语义分割。

目前，MMSegmentation 支持的分割任务为语义分割 ，MMDetection 中支持了实例分割和全景分割。

2.语义分割的应用

语义分割作为计算机视觉的一项基础任务，在自动驾驶、医学、遥感、视频等各个领域都有着广泛的应用，特定的任务也有常用的数据集和算法模型。虽然具体的任务场景不同，但是它们都属于语义分割任务，因此 MMSegmentation 提供的所有语义分割算法都可以拿来使用。

自动驾驶

图像是自动驾驶中非常重要的数据来源，因为摄像头的成本低于激光雷达，而且相较于点云数据，直观的图像更符合人眼的视觉感受。通过语义分割模型，识别出图像中的特定类别，例如：车道线、车辆和行人，可以辅助自动驾驶系统理解场景，做出决策。目前 MMSegmentation 支持的城市街景数据集 Cityscapes，车道线检测模型 ERFNet，实时语义分割模型 BiSeNet 等都和此相关。

源：Cityscapes 官网示例 https://www.cityscapes-dataset.com/examples/

遥感图像分析

遥感图像主要来源于航空飞行器或卫星的航拍。通过对特定前景进行语义分割，可以获得地面的相关信息如：地表植被的变化情况，停车场车辆或机场飞机的数量变化等。已经在智慧城市和智慧地图中得到了广泛的应用。目前 MMSegmentation 支持了常用的遥感图像分割 RGB 数据集 Potsdam、Vaihingen、和 iSAID。

来源：iSAID官网 https://captain-whu.github.io/iSAID/

医学图像分析

语义分割是医学图像分析中常用的计算机视觉任务，通过分割出相关前景（如器官、肿瘤）可以辅助医生进行诊断，例如评估肿瘤长径的变化来判断放疗或化疗的效果，通过眼底视网膜血管的形态变化来筛查和诊断相关疾病。目前 MMSegmentation 支持了眼底视网膜血管分割数据集 DRIVE、 STARE、CHASE DB1 和 HRF，以及常用于医学图像分割任务的 UNet 算法。

来源：DRIVE 官网 https://drive.grand-challenge.org/

我们希望开源能够让学术界的工作更扎实，让工业界的痛点得到解决。目前 MMSegmentation 在维护的同时，还会不断为一些细分方向如遥感图像、人体解析和医学图像提供更多的支持。欢迎大家加入我们，不管是提建议、需求，还是参与代码贡献，我们都会第一时间响应。

https://github.com/open-mmlab/mmsegmentation​github.com/open-mmlab/mmsegmentation正在上传…重新上传取消

3.MMSegmentation 算法库框架介绍

现在我们就带大家一起了解下 MMSegmentation 的整体架构，进一步降低大家使用和扩展框架的难度，力争将 MMSegmentation 打造为易懂易上手的主流语义分割框架。

本文解读的是 MMSegmentation v0.24.1 的整体架构，如果后续版本有比较大的改动，我们也会适时更新进行新的解读。
预告一下：在下一篇文章里我们会进一步讲解数据集相关的内容，包括语义分割任务常用的数据集和如何处理自己的数据集，方便大家快速上手 MMSegmentation 进行实验。敬请期待哦！

下面我们对每个模块进行详细解读~

3.1 MMSegmentation 目录结构

按照代码目录下的文件夹，MMSegmentation 代码库主要可以包含四个部分：

（1）./tools 包括了调用 MMSegmentation 作为训练和测试入口的 ./tools/train.py 和 ./tools/test.py，预训练模型和数据集准备的转换脚本，以及部署和可视化相关的脚本。

详细介绍可见 Github 里的文档。

（2） ./configs 包括了各个算法的配置文件、存放常用的数据集配置、基础模型以及训练策略的基配置文件 ./configs/_base_。

（3）./mmseg 里面是 MMSegmentation 的算法库，包括核心组件、数据集处理、分割模型代码和面向用户的 API 接口。

（4）./data 指的是存放数据集的路径，在原本的代码库中没有这个文件夹。用户只需指定正确的文件夹路径即可使用数据。

下面是详细的 MMSegmentation 的算法库目录结构：

# MMSegmentation 算法库目录结构的主要部分 
mmsegmentation 
   | 
   |- configs                        # 配置文件 
   |     |- _base_                   ## 基配置文件 
   |     |     |- datasets             ### 数据集相关配置文件 
   |     |     |- models               ### 模型相关配置文件 
   |     |     |- schedules            ### 训练日程如优化器，学习率等相关配置文件 
   |     |     |- default_runtime.py   ### 运行相关的默认的设置 
   |     |- swin                     ## 各个分割模型的配置文件，会引用 _base_ 的配置并做修改  
   |     |- ...                         
   |- data                           # 原始及转换后的数据集文件 
   |- mmseg  
   |     |- core                     ## 核心组件 
   |     |     |- evaluation           ### 评估模型性能代码 
   |     |- datasets                 ## 数据集相关代码 
   |     |     |- pipelines            ### 数据预处理 
   |     |     |- samplers             ### 数据集采样代码 
   |     |     |- ade.py               ### 各个数据集准备需要的代码 
   |     |     |- ... 
   |     |- models                    ## 分割模型具体实现代码 
   |     |     |- backbones             ### 主干网络 
   |     |     |- decode_heads          ### 解码头 
   |     |     |- losses                ### 损失函数 
   |     |     |- necks                 ### 颈 
   |     |     |- segmentors            ### 构建完整分割网络的代码 
   |     |     |- utils                 ### 构建模型时的辅助工具 
   |     |- apis                      ## high level 用户接口，在这里调用 ./mmseg/ 内各个组件 
   |     |     |- train.py              ### 训练接口 
   |     |     |- test.py               ### 测试接口 
   |     |     |- ... 
   |     |- ops                       ## cuda 算子（即将迁移到 mmcv 中） 
   |     |- utils                     ## 辅助工具 
   |- tools 
   |     |- model_converters          ## 各个主干网络预训练模型转 key 脚本 
   |     |- convert_datasets          ## 各个数据集准备转换脚本 
   |     |- train.py                  ## 训练脚本 
   |     |- test.py                   ## 测试脚本 
   |     |- ...                       
   |- ... 

MMSegmentation 的算法库有 3 个关键组件：

1../mmseg/apis/，用于训练和测试的接口

2../mmseg/models/，用于分割网络模型的具体实现

3../mmseg/datasets/，用于数据集处理

本文我们主要介绍算法模型相关的代码，因此涉及内容主要在 ./mmseg/models 里面。

3.2 MMSegmentation 模型实现

Segmentor

MMSegmentation 中将语义分割模型定义为 segmentor， 一般包括 backbone、neck、head、loss 4 个核心组件，每个模块的功能如下：

1. 预处理后的数据输入到 backbone（ 如 ResNet 和 Swin Transformer ）中进行编码并提取特征。
2. 输出的单尺度或者多尺度特征图输入到 neck 模块中进行特征融合或者增强，典型的 neck 是 特征金字塔 (Feature Pyramid Networks， FPN)。
3. 上述多尺度特征最终输入到 head 部分，一般包括 decoder head，auxiliary head 以及 cascade decoder head，用以预测分割结果（它们的区别我们会在下文具体介绍）。
4. 最后一步是计算 pixel 分类的 loss，进行训练。

需要说明的是，上述 4 个组件不是每个算法都需要的，比如很多模型里没有 neck 和 auxiliary head 组件。分割器 segmentor 的具体代码见文件 ./mmseg/models/segmentors/。

MMSegmentation 里面的分割器框架可以分为 “Encoder Decoder” 结构和 “Cascade Encoder Decoder” 结构。现有的大多数模型为 “Encoder Decoder” 结构，即利用 encoder 提取图像特征，再用 decoder 去解码上述特征。 “Cascade Encoder Decoder” 的解码部分不是单独的解码头，而是级联式的 2 个或多个解码头，前一个解码头的输出作为后一个解码头的输入。

关于分割器 segmentor 的训练和测试的基本逻辑，以语义分割经典的 “Encoder Decoder” 结构为例：

class EncoderDecoder(BaseSegmentor): 
   def __init__(...): 
        # 构建 backbone、neck 和 head 
        self.backbone = build_backbone(backbone) 
        if neck is not None: 
            self.neck = build_neck(neck) 
        self._init_decode_head(decode_head) 
        self._init_auxiliary_head(auxiliary_head) 
  def forward_train(...):  
        # 利用 backbone+neck 进行特征提取 
        x = self.extract_feat(img) 
        losses = dict() 
        # decode head 输出预测特征图并计算出 loss 
        loss_decode = self._decode_head_forward_train(x, img_metas, 
                                                      gt_semantic_seg) 
        losses.update(loss_decode) 
        # auxiliary heads 输出预测特征图并计算出 loss 
        if self.with_auxiliary_head: 
            loss_aux = self._auxiliary_head_forward_train( 
                x, img_metas, gt_semantic_seg) 
            losses.update(loss_aux) 
        return losses 
 
  def simple_test(...): 
        # 调用 inference 函数，对输入图片做全图或者滑动窗口的推理，得到 logits 值 
        seg_logit = self.inference(img, img_meta, rescale) 
        # 做 argmax 得到预测的 prediction mask 
        seg_pred = seg_logit.argmax(dim=1) 
 
   def aug_test(...): 
        ... 

EncoderDecoder 里面分别定义了训练和测试的接口，训练时调用 forward_train() 返回一个 dict，包含各种 loss ，测试时则会调用 simple_test() 或者测试时数据增广的 aug_test()，只返回预测的分割结果。

训练时预测结果并计算 loss 的主要逻辑是在 _decode_head_forward_train 中实现：

def _decode_head_forward_train(...): 
    # 调用每个 head 自身的 forward_train 方法, 并计算出 loss 
 losses = dict() 
    loss_decode = self.decode_head.forward_train(x, img_metas, 
                                                 gt_semantic_seg, 
                                                 self.train_cfg) 
 
    losses.update(add_prefix(loss_decode, 'decode')) 
    # 返回 
    return losses 

对于不同的 head，都可以抽象为：seg_logits = self.forward(inputs) ，即：网络前传得到预测的 logtis 值，然后再计算各个 head 的对应 loss：

def forward_train(...): 
 seg_logits = self.forward(inputs) 
    losses = self.losses(seg_logits, gt_semantic_seg) 
    return losses 
 
def losses(self, seg_logit, seg_label): 
 loss = dict() 
    seg_logit = resize( # 将预测得到的 logits 值 resize 成原图大小 
        input=seg_logit, 
        size=seg_label.shape[2:], 
        mode='bilinear', 
        align_corners=self.align_corners) 
    .... 
    for loss_decode in losses_decode: # 分别计算这个 decode head 中的各个 loss 
        if loss_decode.loss_name not in loss: 
            loss[loss_decode.loss_name] = loss_decode( 
                seg_logit, 
                seg_label, 
                weight=seg_weight, 
                ignore_index=self.ignore_index) 
        else: 
            loss[loss_decode.loss_name] += loss_decode( 
                seg_logit, 
                seg_label, 
                weight=seg_weight, 
                ignore_index=self.ignore_index) 
    .... 
    return loss 

接下来，我们详细介绍分割器 segmentor 里4 个核心组件：backbone， neck，head，和 loss。

Backbone

目前 MMSegmengtation 中已经集成了大部分主干网络，具体见文件 ./mmseg/models/backbones/，v0.24.1 已经实现的骨架如下：

通常定义的“主干网络” 是指从上游任务（如 ImageNet ）预训练，然后用于多个下游任务（如目标检测、实例分割、语义分割、姿态估计）中的网络，而在 ./mmseg/models/backbones 里主干网络的定义有所不同，会把一些分割算法的网络结构也作为“主干网络”，如 UNet、 FastSCNN、CGNet、ICNet、BiSeNetV1/V2、ERFNet、STDC。

其中最常用的是 ResNet v1c 系列和 Vision Transformer 系列。如果你需要对骨架进行扩展，可以继承上述网络，然后通过注册器机制注册使用。一个典型用法为 ./configs/_base_/models/segmenter_vit-b16_mask.py 里面的：

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_base_p16_384_20220308-96dfe169.pth'  # noqa 
# model settings 
model = dict( 
    type='EncoderDecoder', 
    pretrained=checkpoint, # 加载的预训练模型，这里为 Google Research提供的由 JAX 训练框架得到的 Vision Transformer 
    backbone=dict( 
        type='VisionTransformer', # 骨架类名，后面的参数都是该类的初始化参数 
        img_size=(512, 512), 
        patch_size=16, 
        in_channels=3, 
        embed_dims=768, 
        num_layers=12, 
        num_heads=12, 
        drop_path_rate=0.1, 
        attn_drop_rate=0.0, 
        drop_rate=0.0, 
        final_norm=True, 
        norm_cfg= dict(type='LN', eps=1e-6, requires_grad=True), 
        with_cls_token=True, 
        interpolate_mode='bicubic', 
    ), 

同 OpenMMLab 其他算法库一样，我们使用了 MMCV 中的模块注册机制，通过修改配置文件的 type ，可以使用在 MMSegmentation 已经实现的 backbone 模型。此外，还可以使用 MMClassification 里面的更多主干网络，如 ShuffleNet、EfficientNet 等，可根据 ./configs/convnext 里面 ConvNeXt 的实现方式，详细的方式可以参考： MMDet居然能用MMCls的Backbone？论配置文件的打开方式。

Neck

neck 可以认为是 backbone 和 head 的连接层，主要负责对 backbone 的特征进行高效融合和增强，能够对输入的单尺度或者多尺度特征进行融合、增强输出等。具体见文件 ./mmseg/models/necks/，v0.24.1 已经实现的 neck 如下：

最常用的应该是 FPN，一个典型用法是 ./configs/_base_/models/pointrend_r50.py 里面：

 neck=dict( 
    type='FPN', 
    in_channels=[256, 512, 1024, 2048], # 骨架多尺度特征图输出通道 
    out_channels=256, # 增强后通道输出 
    num_outs=4), # 输出num_outs个多尺度特征图 

Head

MMSegmentation 的 head 是用来处理 backbone 或 neck 的特征图，对图像里的每个像素 pixel 做分类然后得到分类的结果。具体见文件 ./mmseg/models/decode_heads/，v0.24.1 已经实现的 head 如下：

虽然它们都是用来解码特征图中的信息，但在使用上，可以将它们分为 decoder head，auxiliary head 以及 cascade decoder head：

decoder head 是直接在训练和推理中作为图像预测输出的 head。

auxiliary head 是只在训练过程中输出图像预测用来辅助损失函数计算的 head。

cascade decoder head 是指级联式的2个或多个解码头，前一个解码头的输出作为后一个解码头的输入， OCRNet 和 PointRend 两种算法就使用了 cascade decoder head。

在 MMSegmentation 里每个 head 自己单独计算损失，所以把这个公共的行为抽象成了一个基类： BaseDecodeHead，每个算法的 head 都继承自这个基类，类里面包括了计算 loss 的函数，用于计算 head 输出的 logits 值和 label 的损失。

Loss

MMSegmentation 里的 loss 计算的是每个像素上的 logits 和分割标签之间的差别，使用最多的是 cross entropy loss 和 dice loss，v0.24.1 已经实现的 loss 如下：

除了 ./mmseg/models/losses/ 里的这些 loss 外，计算 loss 时还可以用到一些策略和方法，比如：在线难样本挖掘策略 (OHEM， Online Hard Example Mining) 。

4 总结

本文主要带大家一起解读了 MMSegmentation 中的代码结构以及模型组件，希望大家有所收获。在实际使用中，可以将封装的各个组件自主搭配，或者设计新的某一个组件，实现语义分割网络高效地训练和测试。

下一篇文章我们会介绍目前学术界主流的语义分割数据集在 MMSegmentation中的实现，以及如何用 MMSegmentation 跑自己的数据集，方便大家快速上手使用 MMSegmentation 代码库进行实验。敬请期待哦！

欢迎大家来 MMSegmentation 体验，如果对你有帮助的话，欢迎给我们点个 star~

https://github.com/open-mmlab/mmsegmentation​github.com/open-mmlab/mmsegmentation