MMSegmentation是OpenMMLab开源的语义分割工具箱，支持丰富的模型和数据集，广泛应用于自动驾驶、遥感和医学等领域。其核心组件包括backbone、neck、head和loss，常用模型有FCN、PSPNet、SwinTransformer等。本文介绍了MMSegmentation的目录结构、模型实现和应用场景，并强调了其实验流程的统一性和易用性。
目录
大家好呀，今天我们将开启新的解读文章，带领大家了解 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
扫一扫
1452
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
