超详细！带你轻松掌握 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/mmsegmentationgithub.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/mmsegmentationgithub.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 个核心组件，每个模块的功能如下：

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