迈向分割大一统 | OneFormer：一个Transformer统治通用图像分割 原创

摘要

通用图像分割不是一个新概念。过去几十年中统一图像分割的尝试包括场景解析、全景分割，以及最近的新全景架构。然而，这样的全景架构并没有真正统一图像分割，因为它们需要在语义、实例或全景分割上单独训练，以获得最佳性能。理想情况下，真正通用的框架应该只训练一次，并在所有三个图像分割任务中实现SOTA性能。为此，论文提出了OneFormer！！！这是一个通用的图像分割框架，它将分割与一次多任务训练设计相结合。论文首先提出了一种基于任务的联合训练策略，该策略能够在单个多任务训练过程中对每个领域的GT（语义、实例和全景分割）进行训练。其次引入了一个任务token，以在手头的任务上调整模型，使OneFormer的任务是动态的，以支持多任务训练和推理。第三，论文建议在训练过程中使用query-text对比损失来建立更好的任务间和类间区分。值得注意的是，本文的单一OneFormer模型在ADE20k、CityScapes和COCO上的所有三个细分任务中都优于专门的Mask2Former模型，尽管后者使用三倍的资源分别对三个任务中的每一个进行了训练。使用新的ConvNeXt和DiNAT主干，论文观察到了更多的性能改进。我们认为OneFormer是使图像分割更加普遍和流行的重要一步！！！YYDS！

欢迎关注国内首个以自动驾驶技术栈为主线的交流学习社区！
原文出处：微信公众号【自动驾驶之心】

本文的主要贡献如下：

• 论文提出了OneFormer，这是第一个基于transformer的多任务通用图像分割框架，该框架只需要使用单个通用架构、单个模型和单个数据集进行一次训练，就可以在语义、实例和全景分割任务上胜过现有框架，尽管后者需要使用多次资源在每个任务上单独训练；
• OneFormer使用task-conditioned联合训练策略，通过从全景标注中生成所有标签来统一采样不同的GT域（语义、实例或全景），以训练其多任务模型。因此，OneFormer实际上实现了全景分割的原始统一目标[29]；
• 论文通过对三个主要基准的广泛实验来验证OneFormer：ADE20K[15]、Cityscape[14]和COCO[34]。与使用标准Swin-L[38]主干的方法相比，OneFormer在所有三个分割任务上都取得了新SOTA！并使用新的ConvNeXt[39]和DiNAT[21]主干提升了更多性能。

相关工作

图像分割

图像分割是图像处理和计算机视觉中最基本的任务之一。传统的工作通常使用专门的网络架构来处理三种图像分割任务之一（图1a）。

语义分割：语义分割长期以来一直作为一个像素分类问题使用CNN解决[6，7，10，40]。最近的工作[26，27，44，56]展示了基于transformer的方法在语义分割中的成功，继其在语言和视觉中的成功[3，49]。其中，MaskFormer[13]在早期工作[4，16，20]之后，通过使用带有object query的transformer解码器[3]，将语义分割视为mask分类问题。论文还将语义分割公式化为mask分类问题。

实例分割：传统的实例分割方法[2，5，23]也被公式化为mask分类器，其预测二进制mask和每个mask的类标签。论文还将实例分割公式化为mask分类问题。

全景分割：泛视觉分割[29]是为了统一实例和语义分割而提出的。该范围内最早的架构之一是Panoptic FPN[28]，它引入了单独的实例和语义任务分支。随后，基于transformer的架构显著提高了性能[12,13,5,51,59,60]。尽管到目前为止取得了进展，但全景分割模型在性能上仍落后于单个实例和语义分割模型，因此无法充分发挥其统一潜力。受此启发，论文设计了只使用全景标注的OneFormer。

通用图像分割

通用图像分割的概念已经存在了一段时间，首先是图像和场景解析[47，48，58]，然后是全景分割，以统一语义和实例分割[29]。最近，出现了专门为全景分割设计的有前景的架构[12，13，61]，其在语义和实例分割任务上也表现良好。K-Net[61]，一个CNN，使用动态可学习的实例和语义kernel与二分匹配。MaskFormer[13]是一种基于transformer的架构，用作mask分类器。它的灵感来自DETR[3]在transformer范围内对目标检测的重新表述，其中图像被馈送到编码器，解码器基于query输出proposal。Mask2Former[12]改进了MaskFormer，具有可学习的query、解码器中可变形的多尺度注意力[64]、masked cross-attention，并在所有三项任务上取得了新SOTA。不幸的是，它需要在每个任务上单独训练模型以获得最佳性能。因此，在真正统一三个细分任务方面仍然存在差距。据论文所知，OneFormer是第一个使用单一通用模型在所有三个图像分割任务上超越最先进水平的框架。

基于Transformer的架构

自DETR[3]引入以来，基于Transformer编码器-解码器结构[3,31,36,64]的架构在目标检测中被证明是有效的。Mask2Former[12，13]证明了这种架构在使用mask分类公式进行图像分割时的有效性。受到这一成功的启发，论文还将本文的框架制定为基于query的mask分类任务。此外，论文声称，计算任务引导query的query-text对比损失[43，57]可以帮助模型学习任务间差异，并减少模型输出中的类别预测失误。与本文工作同时期的LMSeg[1]使用来自多个数据集分类的文本来计算query-text对比损失，并应对多数据集分割训练挑战。与LMSeg[1]不同，本文的工作侧重于多个任务，并使用训练样本的GT标签中的类来计算query-text对比损失。

方法

OneFormer是一个通用的图像分割框架，它在全景、语义和实例分割方面进行了联合训练，并优于单独训练的模型。如图2所示。OneFormer使用两个输入：示例图像和“the task is {task}”形式的任务输入。在论文的单个联合训练过程中，任务从每个图像的{panoptic、instance、semantic}中统一采样。首先使用主干和像素解码器从输入图像中提取多尺度特征。论文对任务输入进行标记，以获得一个一维任务标记，该标记用于调节object query，从而为每个输入建立任务模型。此外为GT标签中的每个类创建一个表示二进制mask数量的文本列表，并将其映射到文本query表示。注意，文本列表取决于输入图像和{task}。为了监督模型的任务动态预测，论文从全景标注中生成相应的GT。由于GT是任务相关的，论文计算目标和文本query之间的query-text对比损失，以确保object query中存在任务可区分性。object query和多尺度特征被馈送到Transformer解码器以输出最终预测。

Task Conditioned联合训练

现有的用于图像分割的半通用架构[12，13，61]在联合训练所有三个分割任务时面临性能的显著下降（表7）。论文将其未能应对多任务挑战归因于其架构中缺少任务调节。

论文使用任务条件联合训练策略来解决图像分割的多任务训练一次挑战。特别是，论文首先从GT标签的{panoptic、semantic、instance}中统一采样task。通过从全景标注中导出任务特定标签，从而仅使用一组标注，实现了全景标注的统一潜力[29]。

接下来，论文从任务特定的GT标签中提取图像中存在的每个类别的一组二值mask，即，语义任务保证图像中存在每个类别只有一个二值mask，而实例任务表示仅thing类的非重叠二值mask，忽略stuff区域。全景任务表示stuff类的单个不定形mask和thing类的非重叠mask，如图3所示。随后对mask集进行迭代，以创建一个文本列表（$T_{list}$），模板为“a photo with a {CLS}”，其中CLS是对应二值mask的类名。每个样本的二值mask数量随数据集的不同而不同。因此用“a/an {task} photo”条目填充$T_{list}$，以获得长度恒定的$N_{text}$的填充列表（$T_{pad}$），填充条目表示no-object masks。之后使用$T_{pad}$计算query-text对比损失。

论文使用带有模板“the task is {task}”的任务输入（$I_{task}$）来调整任务的架构，该模板被标记化并映射到任务标记（$Q_{task}$）。即使用$Q_{task}$在task上控制OneFormer。

Query表示

在训练期间使用两组query：text queries（文本查询，$Q_{text}$）和object queries（目标查询，Q）。$Q_{text}$是图像中待分割目标的基于文本的表示，而Q是基于图像的表示。

为了获取$Q_{text}$，论文首先对文本条目$T_{pad}$进行标记化，并通过文本编码器[57]传递标记化的表示，该编码器是一个6层transformer[49]。编码的$N_{text}$文本嵌入表示输入图像中的二值mask及其对应类的数量。论文进一步将一组$N_{ctx}$可学习文本上下文嵌入（$Q_{ctx}$）concat到编码文本嵌入，以获得最终的N个text queries（$Q_{text}$），如图4所示。论文使用$Q_{ctx}$的动机是学习样本图像的统一文本上下文[62，63]。注意只在训练期间使用text queries；因此可以在推理过程中删除文本映射器模块，以减小模型大小。

为了获得Q，首先将object queries（Q’）初始化为task-token($Q_{task}$）的N-1次重复。然后使用来自2层transformer内部的展平1/4尺度特征的指导来更新Q’[3，49]。来自transformer的更新后的Q’（富含图像上下文信息）与$Q_{task}$连接，以获得N个query的任务条件表示。与朴素全零或随机初始化[3]不同，任务引导的query初始化和与$Q_{task}$的连接对于模型学习多个分割任务至关重要。

任务导向对比Queries

由于三个任务之间的固有差异，为所有三个分割任务开发一个单一模型具有挑战性。目标查询Q的含义取决于任务。查询应该只关注事物类（实例分割），还是应该只预测图像中每个类的一个不定形目标（语义分割）或两者的混合（全景分割）？现有的基于查询的体系结构[12，13]没有考虑到这些差异，因此，无法有效地在所有三个任务上训练单个模型。

为此，论文建议使用Q和$Q_{text}$计算query-text对比损失。使用$T_{pad}$获得文本查询表示$Q_{text}$，其中$T_{pad}$是给定图像中要检测的每个mask的文本表示列表，“a/an {task} photo”表示Q[3]中的无目标检测。因此，文本查询与目标查询的目的一致，表示图像中存在的objects/segments[3]。因此可以使用基于GT的文本和目标查询之间的对比损失，成功地学习查询表示中的任务间差异。此外，对查询的对比学习使模型能够关注类间的差异并减少类别错误分类。

其他结构组件

主干和像素解码器：使用ImageNet[30]预训练主干[21，38，39]从输入图像中提取多尺度特征表示。像素解码器通过逐步上采样主干特征来辅助特征建模。受最近多尺度可变形注意力[12,64]的成功启发，论文使用相同的基于多尺度可形变Transformer（MSDeformAttn）的架构作为像素解码器。

Transformer解码器：使用多尺度特征，{1/4,1/8,1/16,1/32}四个尺度特征作为输入。

损失函数

整体损失函数如下所示：

实验

论文在Cityscapes、ADE20K和COCO上展开实验，评价指标 PQ、AP、和mIoU。

主要结果

ADE20K：论文将OneFormer与表1中ADE20K[15]值数据集上现有的最先进的伪通用和专用架构进行了比较。使用标准Swin-L主干，OneFormer只训练一次，在所有三个图像分割任务上都优于Mask2Former[12]单独训练的模型，并且与使用相同主干的其他方法相比。

Cityscapes：论文将OneFormer与表2中Cityscapes[15]val数据集上现有的最先进的伪通用和专用架构进行了比较。使用Swin-L主干，OneFormer在PQ和AP指标上的改进分别为+0.6%和+1.9%，优于Mask2Former。此外，凭借ConvNeXt-L和ConvNeXt-XL主干，OneFormer分别创下68.5%PQ和46.7%AP的新水平。

COCO：论文将OneFormer与表3中COCO[34]val2017数据集上现有的最先进的伪通用和专用架构进行了比较。使用Swin L主干，OneFormer的性能与单独训练的Mask2Former[12]相当，PQ分数提高了+0.1%。由于COCO[34]中全景和实例标注之间的差异，论文用从全景标注导出的实例GT来评估AP得分。论文在附录中提供了更多信息。遵循[12]，论文基于从全景标注导出的语义GT来评估mIoU。

消融实验

进一步通过一系列消融研究分析OneFormer的组件。除非另有说明，论文在Cityscapes[14]数据集中使用Swin-L 训练OneFormer。

Task-Conditioned Architecture：论文验证了task token（$Q_{task}$）的重要性，通过在表4中一次删除一个组件，通过重复task token（任务引导查询初始化）和可学习文本上下文（$Q_{ctx}$）初始化查询(-2:7%)。 此外，使用可学习的文本上下文（$Q_{ctx}$）导致PQ分数提高了+4.5%，证明了其重要性。最后，将查询初始化为task token的重复（任务引导的查询初始化），而不是使用全零初始化[3]，PQ提高+1.4%，AP得分提高+1.1%，表明任务调节对查询初始化的重要性。

Contrastive Query Loss：实验结果如表5所示。如果使用分类损失而不是对比损失，观察到所有指标（-0.8%PQ、-0.9%AP和-0.4%mIoU）显著下降。论文将下降归因于分类丢失无法有效捕获任务间差异。不使用对比损失的情况下，性能分别下降-8.4PQ、-3.2AP和-0.5mIoU。

Input Text Template：实验结果如表6所示。

Task Conditioned Joint Training：实验结果如表7所示。

Task Token Input：实验结果如表8所示。

减少类别误分：论文的query-text对比丢失有助于OneFormer学习任务间的区别，并减少预测中类别错误分类的数量。Mask2Former错误地将“wall”预测为第一行中的“fence”，将“vegetation”预测为“terrain”，将"terrain"预测为“sidewalk”。同时，OneFormer在具有类似类别的区域（蓝色框内）中产生更准确的预测，如图5所示。

其他可视化结果如下：

结论

本文提出了OneFormer，这是一个新的多任务通用图像分割框架，具有transformer和任务引导查询，以统一语义、实例和全景分割，并使用单个通用架构、单个模型和单个数据集上的训练。论文联合训练的单一OneFormer模型在跨主要数据集的所有三个分割任务上优于单独训练的专用Mask2Former模型，这是以前的单一架构最先进的模型。因此，OneFormer可以将训练时间、权重存储和推理需求减少三分之一，使图像分割更容易实现。我们相信OneFormer是使图像分割更加普遍和流行的重要一步，并将通过开源本文的代码和模型来支持这一方向的进一步研究。

参考

[1] OneFormer: One Transformer to Rule Universal Image Segmentation