MaX-DeepLab：用于端到端全景分割的双路径 Transformer

2021 年 5 月 27 日 TensorFlow

发布人：Google Research 学生研究员 Huiyu Wang 和高级研究员 Liang-Chieh Chen

全景分割是一个计算机视觉任务，会将语义分割（为每个像素分配类标签）和实例分割（检测和分割每个对象实例）合并。作为实际应用中的核心任务，全景分割通常使用多个代理 (Surrogate) 子任务（如使用边界框检测方法）粗略估计全景分割目标，来实现预测一组不相重叠的遮罩 (mask) 及其相对应的类别标签（例如对象所属的类别，如“汽车”、“交通指示灯”、“道路”等）。

语义分割
https://ai.googleblog.com/2018/03/semantic-image-segmentation-with.html
实例分割
https://cloud.google.com/blog/products/ai-machine-learning/whats-in-an-image-fast-accurate-image-segmentation-with-cloud-tpus

Cityscapes 数据集的示例图像及其全景分割遮罩

以往的方法利用代理子任务树粗略估计全景分割

在此代理树中，每个子任务都会产生额外的手动设计模块，如锚点设计规则、边界框分配规则、非极大值抑制 (NMS)、thing-stuff（thing 类物体和 stuff 类物体）合并，等等。虽然对于单个子任务和模块，不乏一些出色的解决方案，但当我们将这些子任务整合到一个流水线中进行全景分割时，就会产生不需要的构件，在一些比较棘手的情况下更是如此（例如，两个具有相似边界框的人都触发 NMS，从而导致其中一个遮罩缺失）。

以往提出的 DETR 方法简化边界框检测子任务成端到端操作来解决其中一些问题，事实证明，这种做法的计算效率更高，产生的无关构件也更少。然而，训练过程在很大程度上仍然依赖边界框检测，这显然与基于遮罩的全景分割定义不相符。另一种做法是将边界框从管线中彻底移除，这样做的好处是消除了整个子任务及其相关模块和构件。例如，Axial-DeepLab 预测预定义实例中心的像素偏移量时，如果图像平面中有各种各样的形状，或相邻对象的中心点较为接近，例如下面这张狗狗坐在椅子上的图像，那它所使用的子任务则会遇到对象严重变形的挑战。

Axial-DeepLab
https://ai.googleblog.com/2020/08/axial-deeplab-long-range-modeling-in.html

当狗狗的中心和椅子的中心接近重合时，Axial-DeepLab 就会将它们合并成一个对象

在《MaX-DeepLab：利用遮罩 Transformer 实现端到端全景分割 (MaX-DeepLab: End-to-End Panoptic Segmentation with Mask Transformers)》（将提交给 CVPR 2021 大会）一文中，我们首次为全景分割管线提出了完全的端到端方法——将 Transformer 架构扩展到此计算机视觉任务中，直接预测具有类标签的遮罩。这种方法叫做 MaX-DeepLab，利用 Mask Xformer 对 Axial-DeepLab 进行扩展。该方法采用双路径架构，引入了全局内存路径，允许与任何卷积层直接通信。MaX-DeepLab 最终在极具挑战的 COCO 数据集上以及无边界框的状态下实现了 7.1% 的显著全景质量 (PQ) 增益，首次消除了有边界框方法和无边界框方法之间的差距。在不增加测试时长的情况下，就 PQ 方面而言，MaX-DeepLab 在 COCO 测试开发集上达到了 51.3% 的水平，这已是目前的最高水准。

架构
https://ai.googleblog.com/2017/08/transformer-novel-neural-network.html
COCO
https://cocodataset.org/#home

MaX-DeepLab 是完全端到端的：直接从图像中预测全景分割遮罩

端到端全景分割

受 DETR 启发，我们的模型使用经过 PQ 类目标优化的输出遮罩和类，直接预测一组不相重叠的遮罩及其相应的语义标签。具体来说，我们受到评估指标 PQ 的定义：识别质量（预测的类是否正确）乘以分割质量（预测的遮罩是否正确）的启发，以一模一样的方式在两个具有类标签的遮罩之间定义了一个相似度指标。直接通过一对一匹配，最大化实际遮罩和预测遮罩之间的这种相似度，直接对模型进行训练。这种对全景分割的直接建模使端到端训练和推理成为可能，消除了现有的有边界框方法和无边界框方法必须手工设计先验的弊端。

MaX-DeepLab 利用 CNN 和遮罩 Transformer 直接预测 N 遮罩和 N 类

双路径 Transformer

与卷积神经网络 (CNN) 上堆叠传统 Transformer 的方案不同，我们提出了一种结合 CNN 与 Transformer 的双路径框架。具体而言，我们通过一个双路径 Transformer 元件，使 CNN 层能够从全局内存中读写数据。这里所说的这个元件采用了 CNN 路径和内存路径之间的所有四种注意力 (Attention) 类型，可以插入到 CNN 中的任意位置，从而允许在任何层与全局内存通信。MaX-DeepLab 还利用堆叠的沙漏式解码器，可将多种尺度特征聚合成高分辨率输出。然后系统会将该输出与全局内存特征相乘，形成遮罩组预测。至于遮罩类别，则使用另一种 Transformer 进行预测。

双路径 Transformer 架构概览

结果

我们在极具挑战性的 COCO 全景分割数据集上，分别使用先进的无边界框方法 (Axial-DeepLab) 和有边界框方法 (DetectoRS) 对 MaX-DeepLab 进行了评估。在不增加测试时长的情况下，就 PQ 方面而言，MaX-DeepLab 在 COCO 测试开发集上达到了 51.3% 的水平，这已是目前的最高水准。

COCO 测试开发集测试对比情况

在无边界框状态下，就 PQ 方面而言，这一结果比 Axial-DeepLab 高出 7.1%，比 DetectoRS 高出 1.7%，第一次消除了有边界框方法和无边界框方法之间的差距。为了与 DETR 进行一致的比较，我们还评估了与 DETR 参数数量和算力均一致的轻量级 MaX-DeepLab 版本。就 PQ 方面而言，此轻量级 MaX-DeepLab 在 val 集上的表现优于 DETR 3.3%，在测试开发集上的表现优于 DETR 3.0%。此外，我们还对端到端表达式、模型伸缩、双路径架构和损失函数进行了广泛的消融研究和分析。此外，MaX-DeepLab 也不像 DETR 一样需要超长训练计划。

如下图所示，MaX-DeepLab 正确地分割了一只坐在椅子上的狗。Axial-DeepLab 依赖于回归对象中心偏移量的代理任务。它之所以失败，是因为狗和椅子的中心太过于接近。作为代理任务，DetectoRS 会将对象的边界框而非遮罩进行分类。由于椅子的边界框置信度较低，所以就被滤除了。

Axial-DeepLab
https://ai.googleblog.com/2020/08/axial-deeplab-long-range-modeling-in.html

关于 MaX-DeepLab 和先进的无边界框及有边界框方法的案例研究

还有一个例子可以证明 MaX-DeepLab 可以在充满挑战的条件下正确分割图像。

MaX-DeepLab 正确分割了相互重叠的斑马。此例对于其他方法也非常具有挑战性，因为这两只斑马的边界框很相似，对象的中心也很接近

结论

我们首次证明了全景分割可以进行端到端训练。MaX-DeepLab 使用遮罩 Transformer 直接预测遮罩和类，消除了手工设计对于经验的依赖，如对象边界框、thing-stuff（thing 类物体和 stuff 类物体）合并等。借助 PQ 式损失函数和双路径 Transformer，MaX-DeepLab 在极具挑战性的 COCO 数据集上取得了最高水准的结果，消除了有边界框方法和无边界框方法之间的差距。

致谢

本研究系与 Yukun Zhu、Hartwig Adam 和 Alan Yuille 合作完成，在此表示感谢。同样感谢 Maxwell Collins、Sergey Ioffe、Jiquan Ngiam、Siyuan Qiao、Chen Wei、Jieneng Chen 以及 Mobile Vision 团队的大力支持和宝贵意见。