统一、灵活、全面，语义分割开源库 MMSegmentation 全面升级

MMSegmentation 是 OpenMMLab 开源项目里的语义分割领域的算法工具箱，它实现了许多高质量语义分割算法模型和数据集，也为语义分割任务提供了统一的框架和基准测试。

它的主要特点如下：

统一性，提供了统一的模块抽象和数据接口以及基准测试；

灵活性，继承自 openmmlab 的模块化设计，各模块可以轻松替换构建不同模型；

全面性：支持多种应用场景的分割任务，包括医疗、遥感和城市街景等，和学术界最新的模型。

MMSegmentation 于 2020 年 7 月正式开源，目前支持了 44 个算法模型，提供了 598 个预训练权重，支持了 14 个各种应用场景的数据集，在学术界和工业界都有一定的影响力。

在 GitHub 上，我们已收获 star 数 4.6k+，fork 数 1.7k+ ，更有 127 位的国内外 contributor 参与了项目的建设，非常感谢大家的支持与贡献。

2022 年 9 月，我们发布了 MMSegmentation 1.0，全面升级了架构设计，适配全新的 OpenMMLab 2.0 体系。

我们接下来将从新版本特点，重点模块介绍两个方面，对全新升级的 MMSegmentation 1.0 进行介绍。

新版本特点

接口更加统一

OpenMMLab 2.0 最新发布的深度学习训练框架 MMEngine，提供了功能非常强大的执行器，统一了算法库训练测试的启动和执行逻辑，另外规范了数据接口协议，使得各模块间（包括模型、数据集、数据增强、评测、可视化等）的数据传递更流畅，模块间的组合更灵活。

功能更加丰富

OpenMMLab 2.0 引入 MMEngine 算法库，新增了许多功能模块，像参数调度器、日志处理模块、消息中心，可视化模块等。下游算法库例如 MMSegmentation 的用户在训练时可以更加方便地设置不同训练策略。另外模型训练、测试的各种信息也可以收集并输出到不同可视化后端，例如 wandb，tensorboard ，方便用户监控模型训练或者是查看测试结果。

速度更快

升级后的 MMSegmentation 1.0 优化了模型推理和训练速度，一方面 MMSegmentation 1.0 支持 pytorch 原生的混合精度训练，帮助提高模型训练效率，另一方面模型中新增了预处理模块，优化了数据变换的效率。

MMSegmentation 升级前后速度对比

我们比较升级前后几个经典语义分割模型的推理时间，用了同样的设备 a100，相同的数据集 cityscape 和输出尺寸 1024x2048，相比于老版本，升级后的模型推理的时间都缩短了 20% 以上。

文档更丰富

我们进一步重构了 MMSegmentation 的文档，丰富了文档内容，提供了用户文档以方便大家自学语义分割领域经典算法，也提供了进阶文档，将所有的代码实现逻辑都公开介绍，方便大家参考。

除此之外，我们还会介绍如何在 MMSegmentation 中添加自定义模块，方便大家使用 MMSegmentation 构建自己的研究项目和工业产品。

重点模块介绍

数据结构

MMSegmentation 1.0 引入了 SegDataSample 数据结构，将语义分割中的数据封装起来，用于各个功能模块之间的数据传递，SegDataSample 里面的字段有三个：gt_sem_seg，pred_sem_seg 和 seg_logits 。前两个分别是标签和模型预测对应的分割掩膜（segmeation mask），seg logits 是模型最后一层没有经过归一化的输出。

数据集和数据变化操作

MMSegmentation 一共支持了 14 个数据集，除了常见的学术数据集如 ADE20k、CityScapes 外，还有医疗和光学航空遥感等偏向应用的数据集。

MMSegmentation 1.0 新定义了 BaseSegDataset，规范了语义分割数据集功能和接口，是 MMEngine 中 BaseDataset 的子类。数据集主要的功能是加载数据信息，数据信息有两种，一种是数据集的元信息，包括类别信息和调色板信息，就是渲染时类别对应的颜色；另一种是数据信息，保存了具体数据集中图片路径和对应的标签路径。

数据集模块包含了数据变换模块，支持了许多数据变换，在训练或测试时，可以将一系列数据变换组合成为一个列表，称为数据流水线，传给 dataset 的 pipeline 参数，流水线中每一个模块的输出是下一个模块的输入。

下图展示了一个典型语义分割模型训练时的数据变换流水线。每对样本，经过一个数据变换操作，输出的字典中会加入新的字段（标记为绿色）或更新现有的字段（标记为橙色）。

模型

MMSegmentation 中将一个语义分割算法模型称为 segmentor，并且继承 OpenMMLab 一贯的模块化设计，将 segmentor 分为 6 个模块，分别是：
- data_preprocessor，负责将从数据变换流水线输出的数据搬运到指定设备上，并对其做归一化、padding、组 batch 等操作。这样做的好处是，在数据变化的阶段，数据格式为 uint8，数据搬运到 gpu 上后再转为 FP32 做归一化，减少 cpu 上的计算压力。
- Backbone ，从输入的图片中提取特征图，常见的模型有 ResNet，Swin transformer 等。
- Neck ，连接 backbone 和 decode_head，将从 backbone 输出的特征图进一步处理然后输入到解码头，常见的网络有 Feature Pyramid Network FPN。
- decode_head，负责从输入的特征图预测最后的分割结果。
- auxiliary_head（可选），负责从输入的特征图预测分割结果，但结果只是在训练过程中参与损失计算，不参与推理，在推理时只从解码头中预测输出结果。
- loss，负责神经网络输出结果和真值的损失计算，用于反传时模型梯度计算。

segmentor 的模型结构根据是否由多个 decode_head 集联，分为 encoder_decoder 和 cascade_encoder_decoder 两种。二者的区别是：cascade_encoder_decoder 里有多个 decode_head，并且从第二个解码头开始，每个解码头的输入是上一个解码头输出，作用是对输出结果进一步的 refine。

encoder_decoder 数据流

cascade_encoder_decoder 数据流

数据流

全新升级的 OpenMMLab 2.0 定义了数据流协议，分为训练和测试两种。训练的时候，dataloader 搬运经过 data transforms 处理的数据，传给模型 train_step 方法，模型里会先调数据预处理模块，再传给模型的 forward 函数，前传并计算损失。这个 loss dict 会经过 parse_losses 模块解析，得到一个 loss scalar，然后在 opitimizer warpper 里的 update_params 对模型反传，计算梯度，并更新参数。

训练时的数据流

测试时，数据会传给模型的 test_step 方法，同样是先经过预处理，predict 输出 datasample。这里的datasample 就是输入网络的 datasample，只不过新增了 pred_sem_seg 和 seg_logits 两个字段用以保存网络的预测结果。将这个修改后的 data sample 和 inputs 送到评测器计算评测指标，或者送到可视化器中进行处理。

测试时的数据流

总结

我们对 MMSegmentation 1.0 进行一个简单的总结，作为 OpenMMLab 2.0 的语义分割算法库，目前已在开源社区被广泛使用。基于全新的架构和生态，此次新升级的 MMSegmentation 统一了算法模块间接口参数，提升了训练、推理效率，并新增了便于使用的研究工具。

MMSegmentation 之后的发展，除了进一步紧跟学术热点，实时支持学术界最新的 sota 模型之外，我们也希望它能更高效、更强大，更灵活，应用场景更全面，更易上手，希望能助力语义分割新算法的诞生，和科技产业的落地。

这次升级动作较大，我们考虑到了大家项目迁移成本，所以有比较长的双版本迭代维护计划，如下图所示：

最后，升级后的 MMSegmentation 1.0 的开发分支在 dev-1.x branch，欢迎大家来试用，提出宝贵建议，有任何问题或需求都可以在 issue 里给我们反馈，如果觉得好用，也欢迎大家给我们点个 star 哦~