OneFormer：统治通用图像分割的Transformer

温馨提示：
本篇文章已同步至"AI专题精讲" OneFormer：统治通用图像分割的Transformer

摘要

通用图像分割并不是一个新概念。在过去的几十年中，统一图像分割的尝试包括场景解析、全景分割，以及最近的新型全景架构。然而，这些全景架构并未真正实现图像分割的统一，因为它们需要在语义分割、实例分割或全景分割任务上分别训练，才能达到最佳性能。理想情况下，一个真正通用的框架应当只训练一次，就能在三种图像分割任务中都达到SOTA性能。为此，我们提出了OneFormer，这是一个通过多任务一次性训练设计实现统一分割的通用图像分割框架。我们首先提出了一种任务条件化联合训练策略，使得可以在单次多任务训练过程中，利用来自语义分割、实例分割和全景分割的标注数据进行训练。其次，我们引入了一个task token，用于将模型条件化于当前任务，使得我们的模型在训练和推理过程中都具有任务动态性。第三，我们在训练过程中引入了query-text对比损失，用于更好地区分不同任务和类别。值得注意的是，我们的单一OneFormer模型在ADE20K、Cityscapes和COCO数据集上，在所有三种分割任务中均优于专用的Mask2Former模型，尽管后者在每个任务上分别训练，资源使用量是我们的三倍。在使用新的ConvNeXt和DiNAT主干网络时，我们观察到了更进一步的性能提升。我们认为，OneFormer是实现图像分割通用化与易用性的重要一步。为促进进一步研究，我们开源了我们的代码与模型：https://github.com/SHI-Labs/OneFormer。

1. 引言

图像分割是一项将像素分组为多个区域的任务。此类分组可以是基于语义的（例如：道路、天空、建筑物），也可以是基于实例的（具有明确边界的物体）。早期的分割方法 [7,23,40] 分别处理这两类分割任务，使用专门的架构，因此需要分别投入研究资源。为统一语义分割与实例分割，Kirillov 等人 [29] 提出了全景分割方法，用于将像素分组为无明确形状的背景区域（标记为“stuff”）和具有清晰边界的物体区域（标记为“thing”）。然而，这一努力催生了新的全景专用架构 [11]，并未真正统一先前的任务（见图1a）。最近的研究趋势转向通过新型全景架构统一图像分割，例如 K-Net [61]、MaskFormer [13] 和 Mask2Former [12]。这些全景架构可以在三种任务上训练，并在不更改架构的情况下获得良好性能。然而，为了获得最佳性能，它们仍需在每个任务上分别训练（见图1b）。这种分别训练策略不仅需要额外的训练时间，还需为每个任务生成不同的模型权重。因此，这些方法只能被视为“半通用”方法。例如，Mask2Former [12] 在ADE20K [15]数据集上分别为语义、实例和全景分割任务训练160K次，总训练迭代达到480K，并需保存和部署三个模型用于推理。

为真正实现图像分割的统一，我们提出了一个多任务通用图像分割框架（OneFormer），在仅在一个全景数据集上训练一次的情况下，超越现有的state-of-the-art方法，在三种图像分割任务上都取得了优异表现（见图1c）。通过本研究，我们旨在回答以下问题：

(i) 为什么现有的全景架构 [12,13] 无法通过单次训练或单一模型来处理三种任务？我们假设现有方法之所以需要为每个任务分别训练，是因为其架构中缺乏任务引导，导致在联合训练或使用单一模型时难以学习任务间的领域差异。为解决这一问题，我们引入了形式为文本的任务输入token：“the task is {task}”，用于将模型条件化于当前任务，使我们的架构在训练阶段具备任务引导性，在推理阶段具备任务动态性，且仅使用一个模型。我们从 {panoptic, instance, semantic} 中均匀采样 {task}，以及相应的ground truth，在联合训练过程中确保模型在任务上的无偏性。受全景数据 [29] 同时包含语义与实例信息的启发，我们在训练时从全景标注中提取出语义标签与实例标签。因此，我们在训练中仅需使用全景数据。此外，我们的联合训练时间、模型参数量和FLOPs与现有方法相当，总体训练时间和存储需求最多可减少3倍，从而使图像分割更加高效和易于使用。

(ii) 多任务模型在单次联合训练过程中，如何更好地学习任务间与类别间的差异？借鉴transformer框架 [3,12,21,22,27,38,60] 在计算机视觉中的成功，我们将OneFormer构建为基于transformer的方法，并通过使用query token进行引导。为给模型添加任务特定的上下文，我们将query初始化为任务token（由任务输入获得）的重复，并计算query-text对比损失 [43,57]，其中文本由所采样任务的ground-truth标签构建，如图2所示。我们假设，对query施加对比损失有助于引导模型更加敏感于任务差异。此外，它还能在一定程度上减少类别误分类。

我们在三个主流分割数据集上对OneFormer进行了评估：ADE20K [15]、Cityscapes [14] 和 COCO [34]，每个数据集都包含语义、实例和全景三种分割任务。OneFormer使用单个联合训练模型在所有任务上均达到了新的state-of-the-art水平。

总结来说，我们的主要贡献如下：
• 我们提出了OneFormer，这是第一个基于transformer的多任务通用图像分割框架，仅需一次训练，使用单一架构、单一模型与单一数据集，即可在语义、实例和全景分割任务上超越现有框架，而后者需在每个任务上分别训练，资源消耗为我们数倍。

• OneFormer使用了任务条件化联合训练策略，通过从全景标注中推导所有标签，均匀采样不同任务的ground truth，以训练其多任务模型。因此，OneFormer真正实现了全景分割最初设想的统一目标 [29]。

• 我们通过在三个重要基准数据集（ADE20K [15]，Cityscapes [14]，COCO [34]）上的广泛实验验证了OneFormer的有效性。与使用标准Swin-L [38]主干的其他方法相比，OneFormer在所有三种分割任务上均创下新SOTA，并在引入ConvNeXt [39]和DiNAT [21]主干后表现更佳。

2. 相关工作

2.1. 图像分割

图像分割是图像处理和计算机视觉中最基础的任务之一。传统方法通常针对三类图像分割任务之一，采用专门的网络结构进行处理（见 Fig. 1a）。

Semantic Segmentation（语义分割）。语义分割长期以来被作为一个像素分类问题来处理，常用CNN进行建模 [6, 7, 10, 40]。随着transformer在语言与视觉任务中的成功 [3, 49]，更近期的工作 [26, 27, 44, 56] 展示了transformer方法在语义分割中的优越表现。其中，MaskFormer [13] 在早期工作 [4,16,20] 的基础上将语义分割建模为一个mask分类问题，采用带有object queries的transformer解码器 [3]。我们的方法同样将语义分割建模为一个mask分类问题。

Instance Segmentation（实例分割）。传统的实例分割方法 [2, 5, 23] 也被建模为mask分类器，它们为每个mask预测二值掩码和类别标签。我们的方法同样将实例分割建模为一个mask分类问题。

Panoptic Segmentation（全景分割）。Panoptic Segmentation [29] 被提出用于统一实例分割与语义分割。其中最早的一种架构是 Panoptic-FPN [28]，该方法为实例和语义任务分别引入了独立的分支。后续工作借助transformer架构显著提升了性能 [12,13,50,51,59,60]。尽管已有显著进展，全景分割模型在性能上仍落后于分别针对实例分割和语义分割的专门模型，因此并未完全实现其统一的潜力。受此启发，我们设计的 OneFormer 仅通过全景标注进行训练。

2.2. 通用图像分割

Universal Image Segmentation（通用图像分割） 的概念由来已久，最早可追溯至图像与场景解析任务 [47, 48, 58]，随后是为统一语义分割与实例分割而提出的 panoptic segmentation（全景分割）[29]。近年来，一些专为全景分割设计的优秀架构 [12,13,61] 出现，它们在语义分割与实例分割任务中也表现出色。K-Net [61] 是一种CNN架构，采用了动态可学习的实例与语义kernel，并使用二分图匹配。MaskFormer [13] 是基于transformer的架构，将分割问题建模为mask分类问题。其灵感来自 DETR [3] 对目标检测任务在transformer框架下的重构：将图像输入encoder，由decoder基于query生成proposal。Mask2Former [12] 在MaskFormer的基础上进行改进，引入了可学习的query、decoder中的可变形多尺度注意力机制 [64] 和masked cross-attention，并在三个任务上都达到了新的state of the art。不幸的是，该模型需要对每个任务单独训练以达到最佳性能。因此，真正统一三个分割任务的目标尚未完全实现。据我们所知，OneFormer 是第一个在所有三种图像分割任务上都超越state of the art的单一通用模型。

2.3. 基于Transformer的架构

基于transformer的encoder-decoder结构的架构 [3,31,36,64] 自从DETR [3] 被提出以来，就已被证实在目标检测任务中非常有效。Mask2Former [12,13] 进一步展示了此类架构在图像分割任务中的效果，通过将其建模为mask分类问题。受此启发，我们也将本框架建模为基于query的mask分类任务。此外，我们主张，在任务引导的query上计算query-text对比损失 [43, 57] 可以帮助模型学习任务间的差异，并减少模型输出中的类别误判。与我们工作同时进行的 LMSeg [1] 使用从多个数据集的类别体系中提取的文本，来计算query-text对比损失，并解决多数据集训练下的分割问题。而与 LMSeg [1] 不同的是，我们的工作专注于多任务学习，并使用训练样本中的ground-truth标签类别来计算query-text对比损失。

3. 方法

在本节中，我们介绍 OneFormer——一个通用图像分割框架，可在全景分割、语义分割和实例分割任务上联合训练，并超越分别训练的模型。我们在图2中展示了 OneFormer 的概览。OneFormer 有两个输入：样本图像以及任务输入，格式为 “the task is {task}”。在我们的单一联合训练过程中，每张图像的任务从 {panoptic, instance, semantic} 中均匀采样。首先，我们使用一个backbone和pixel decoder从输入图像中提取多尺度特征。我们将任务输入进行tokenize，得到一个一维的task token，用于对object query及模型进行任务调控。此外，我们根据GT标签中存在的类别，构建一个文本列表，表示每个类别对应的二值mask数量，并将其映射为text query表示。需要注意的是，该文本列表依赖于输入图像和任务 {task}。为了监督模型在不同任务下的动态预测，我们从panoptic标注中提取相应的ground-truth。由于ground-truth与任务相关，我们在object query与text query之间计算query-text对比损失，以确保object query中包含任务区分性。object query与多尺度特征共同输入到transformer decoder中，生成最终预测结果。我们将在后续小节中提供更多细节。

3.1. 任务调控的联合训练

现有的半通用图像分割架构 [12, 13, 61] 在联合训练三种分割任务时表现显著下降（见表7）。我们认为，它们未能解决多任务挑战的原因在于缺乏任务调控机制。

我们通过任务调控的联合训练策略来解决图像分割中的“一次训练，多任务”问题。具体而言，我们首先从 {panoptic, semantic, instance} 中均匀采样任务，并基于panoptic标注 [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}$进行padding，得到长度为 $N_{text}$ 的固定长度文本列表（$T_{pad}$），其中padding条目表示无物体的mask。我们在第3.3节中将使用$T_{pad}$计算query-text对比损失。

我们使用任务输入（$I_{task}$），格式为 “the task is {task}”，对架构进行任务调控。该输入经过tokenize并映射为task-token（$Q_{task}$），我们使用$Q_{task}$对OneFormer进行任务调控（详见第3.2节）。

3.2. Query 表示

在训练过程中，我们在架构中使用两类query：text query（$Q_{text}$）和object query（Q）。$Q_{text}$ 是图像中各个分割区域的文本表示，而Q是基于图像的表示。

为了获得$Q_{text}$，我们首先对Tpad中的文本条目进行tokenize，并将其tokenized表示输入一个6层transformer的text encoder [57, 49]。编码后的Ntext个文本embedding表示输入图像中所有二值mask及其类别。我们进一步将一组$N_{ctx}$个可学习的text context embedding（$Q_{ctx}$）拼接到文本embedding后，得到最终的N个text query（$Q_{text}$），如图4所示。我们使用$Q_{ctx}$的动机在于为每张图像学习一个统一的文本上下文表示 [62, 63]。由于text query仅在训练时使用，我们可以在推理时移除text mapper模块，从而减少模型体积。

温馨提示：
阅读全文请访问"AI深语解构" OneFormer：统治通用图像分割的Transformer