OneFormer: One Transformer to Rule Universal Image Segmentation 一种新的多任务通用图像分割框架

Abstract（研究背景）：

通用图像分割并不是一个新概念。过去几十年中统一图像分割的尝试包括场景句法分析、全景分割，以及最近的新全景架构。然而，这样的全景架构并没有真正统一图像分割，因为它们需要在语义、实例或全景分割上单独训练以达到最佳性能。理想情况下，一个真正通用的框架应该只训练一次，并在所有三个图像分割任务中实现SOTA性能，性能优于Mask2Former、UPerNet等网络。

为此，我们提出了OneFormer：

• 首先提出了一种任务条件联合训练策略，该策略允许在单个多任务训练过程中对每个领域的ground truths（语义、实例和全景分割）进行训练.
• 其次，引入了一个任务token，以手头的任务为条件来约束我们的模型，使我们的模型任务动态，以支持多任务训练和推理。
• 第三，我们提出在训练期间使用（query-text contrastive loss）查询-文本对比损失来建立更好的任务间和类间区别。

1. Introduction

图像分割与新的全景分割结构统一起来，这种panoptic结构可以在不改变架构的情况下在三个任务上进行训练并获得最佳的性能见图1b

为了真正统一图像分割，我们提出了一个多任务通用图像分割框架（Oneformer），它在所有三个图像分割任务上只训练一次，优于现有的图像分割任务（见图1c）。通过这项工作，我们旨在回答以下问题：

1. 基于全景数据捕获语义和实例信息的能力，在训练过程中从相应的全景注释中推导出语义和实例标签；因此在训练过程中只需要全景数据，此外，我们的联合训练时间、模型参数和FLOPs可与现有方法相媲美，将训练时间和存储需求减少到3×，使得分割所消耗的密集型资源更少。
2. 框架制定为一种基于Transformer的方法，它可以通过使用查询令牌来指导。为了向我们的模型中添加特定于任务的上下文，我们将查询初始化为任务标记（从任务输入中获得）的重复，并计算查询文本对比损失。查询上的对比性损失有助于引导模型对任务更加敏感。此外，它还有助于减少类别错误预测。

综上述，本文的贡献：

1. 提出基于transformer的OneFromer，第一个多任务通用图像分割框架只需要训练一次与一个通用架构，一个模型，在一个数据集，超越现有框架在语义、实例和全景分割任务，尽管后者需要单独训练在每个任务使用多次的资源。
2. OneFormer使用任务条件的联合训练策略，通过从全景注释中提取所有标签，统一采样不同的Ground Truth（语义、实例或全景）来训练多任务模型。因此，前者实际上实现了全景分割[29]的原始统一目标。
3. 与使用标准的Swin-L骨干的方法相比,OneFormer在所有三个分割任务上(ADE20K, Cityscapes, COCO)都设置了一个新的最先进的性能，并且通过新的ConvNeXt 和DiNAT 骨干进行了更多的改进。

1. Realated Work

2. Method

如图2是Oneformer的概述。Oneformer使用两个输入：样本图像和任务输入的形式。首先，我们使用主干网和像素解码器从输入图像中提取多尺度特征。我们对任务输入进行标记化，以获得用于条件对象查询的一维任务标记，因此，我们对每个输入的任务建立模型。此外，我们还创建了一个文本列表，表示GT标签中每个类的二进制掩码的数量，并将其映射到文本查询表示形式中(text query representations)。

由于GT是依赖于任务的，我们计算对象和文本查询之间的查询-文本对比损失(contrastive loss)，以确保对象查询中存在任务区别。对象查询和多尺度特征被输入到一个变压器解码器中，以产生最终的预测。

3.1 Task Conditioned Joint Training  任务条件联合训练

使用任务条件的联合训练策略来解决图像分割的多任务训练一次挑战。特别是，我们首先从GT标签的{全景、语义、实例}中统一地采样任务。我们通过从全景注释中推导出特定于任务的标签来实现全景注释[29]的统一潜力，因此，只使用一组注释。

3.2 Query Representations  查询表示

在训练期间，我们在体系结构中使用了两组查询：文本查询（Qtext）和对象查询(Q)。Qtext是对图像中的片段的基于文本的表示，而Q是基于图像的表示。

为了获得文本查询Qtext，我们首先将文本条目Tpad标记化，并将标记化的表示传递给文本编码器text-encoder，这是一个6层Transformer。编码的Ntext文本嵌入表示输入图像中二进制掩码及其对应类的数量。我们进一步将一组Nctx可学习的文本上下文嵌入（Qctx）连接到编码的文本嵌入中，以获得最终的N个文本查询（Qtext），如图4所示。我们使用Qctx的动机是为一个示例图像学习一个统一的文本上下文。我们只在训练期间使用文本查询；因此，我们可以在推理过程中删除文本映射器模块，以减少模型的参数。

为了获得对象查询Q，我们首先将对象查询（Q`）初始化为task−token（Qtask）N-1次重复。然后，我们在2层变压器[3,49]内flatten的1/4尺度特性的指导下更新Q`。更新的Q`变压器（丰富的图像上下文信息）与Qtask获得任务条件表示N queries,问与普通的全零或随机初始化[3]，任务引导初始化查询和连接Qtask是至关重要的模型学习多个分割任务。

3.3 Task Guided Contrastive Queries  任务指导的对比查询

我们建议使用Q和Qtext来计算查询文本对比损失。我们使用Tpad来获得文本查询表示，Qtext，其中Tpad是在给定图像中要检测到的每个掩码的文本表示列表，用“a/an{任务}照片”表示无目标检测在Q中。因此，文本查询与对象查询的目的对齐，表示在图像中呈现[3]的对象/段。因此，我们可以使用ground truth派生的文本和对象查询之间的对比损失，成功地学习查询表示中的任务间区别。此外，对查询的对比学习使我们能够关注类间的差异，并减少类别的错误分类。

3.4 Other Architecture Components   其他体系结构组件

Backbone and PixelDecoder:使用广泛使用的ImageNet[30]预训练的主干[21,38,39]从输入图像中提取多尺度特征表示。我们的像素解码器通过逐步对主干特征进行上采样来帮助特征建模。由于最近多尺度可变形注意[12,64]的成功，我们使用了相同的多尺度可变形变压器（Multi-Scale Deformable Transformer基于MSDeformAttn）的架构为像素解码器。

Transformer Decoder：使用多尺度策略[12]来利用我们的transformer decoder内部更高分辨率的maps。具体来说，我们提供对象查询(Q)和来自像素解码器（Fi）的多尺度输出，即i∈{1/4、1/8、1/16、1/32}作为输入。我们使用分辨率为1/8、1/16和1/32的特征，或使用掩蔽交叉注意（CA）操作[12]更新Q，然后是自注意self-attention（SA），最后是前馈网络（FFN）。我们在transformer解码器内执行这一组交替操作。

来自transformer解码器的最终查询输出被映射到一个K+1维空间，用于类预测，其中K表示类的数量,另一个额外的+1，用于无对象预测。为了获得最终的掩模，我们借助Q和F1/4之间的einsum操作，以原始图像的1/4分辨率解码像素特征（F1/4）。在推理过程中，我们采用与[12]相同的后处理技术，获得最终的全景、语义和实例分割预测。

3.5 Losses  

4. Experiment

数据集Datasets：CityScapes,ADE20K,COCO

评级指标：对于所有三个图像分割任务，我们报告了PQ、AP和mIoU分数。当任务是全景的时，报告PQ分数；当任务是实例和语义的时，报告AP和mIoU分数。

 

1. 5.Conclusion

在这项工作中，我们提出了一种新的前者，即多任务的通用图像分割框架，通过transformer和任务引导查询，将语义、实例和全景分割统一为单一的通用架构，单一模型，并在单一数据集上进行训练。我们联合训练的单一OneFormer模型优于单独训练的专业Mask2former模型，这是之前的单一架构状态。因此，OneFormer可以将训练时间、重量存储和推理托管需求减少到三分之一，使图像分割更容易实现。我们相信OneFrormer是使图像分割更加通用和可访问的重要一步，并将通过开源我们的代码和模型来支持这个方向的进一步研究