zl程序教程

您现在的位置是:首页 >  IT要闻

当前栏目

统一、灵活、全面,这个好用的语义分割开源库全面升级啦

2023-02-18 16:49:15 时间

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

MMSegmentation 效果演示 demo

它的主要特点如下:

  • 统一性,提供了统一的模块抽象和数据接口以及基准测试;
  • 灵活性,继承自 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 对模型反传,计算梯度,并更新参数。

测试时的数据流

总结

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

MMSegmentation 新版本地址:

https://github.com/open-mmlab/mmsegmentation/tree/1.x

想了解更多 MMSegmentation 1.0 相关内容,欢迎查看直播回放视频~

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

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

最后,升级后的 MMSegmentation 1.0 的开发分支在 dev-1.x branch,欢迎大家来试用,提出宝贵建议,有任何问题或需求都可以在 issue 里给我们反馈