SegFormer: Simple and Efficient Design for SemanticSegmentation with Transformers(SegFormer)翻译

摘要

我们提出了SegFormer，这是一个简单、高效且强大的语义分割框架，它将Transformer与轻量级多层感知机（MLP）解码器相结合。SegFormer具有两个吸引人的特点：1）SegFormer包含一个新颖的层次结构Transformer编码器，该编码器能够输出多尺度特征。它不需要位置编码，从而避免了当测试分辨率与训练分辨率不同时，位置编码的插值导致性能下降的问题。2）SegFormer避免了复杂的解码器。所提出的MLP解码器从不同层聚合信息，从而结合了局部注意力和全局注意力，以产生强大的表示。我们证明，这种简单且轻量级的设计是实现Transformer高效分割的关键。我们扩展了我们的方法，得到了一系列从SegFormer-B0到SegFormer-B5的模型，与之前的同类模型相比，显著提高了性能和效率。例如，SegFormer-B4在ADE20K上达到了50.3%的mIoU，具有6400万个参数，比之前的最佳方法小5倍，性能好2.2%。我们的最佳模型SegFormer-B5在Cityscapes验证集上达到了84.0%的mIoU，并在Cityscapes-C上显示出优异的零样本鲁棒性。代码将在GitHub上发布：http://github.com/NVlabs/SegFormer

1. 引言

语义分割是计算机视觉中的一项基础任务，支持许多下游应用。它与图像分类相关，因为它产生的是逐像素类别预测，而不是图像级别的预测。这一关系在开创性工作[1]中被指出并系统地进行了研究，作者在该工作中使用了全卷积网络（FCN）进行语义分割任务。自那以后，FCN激发了众多后续研究，并成为密集预测任务的主要设计选择。由于分类和语义分割之间存在很强的联系，因此许多最先进的语义分割框架都是ImageNet上流行分类架构的变体。因此，设计骨干架构在语义分割领域一直是一个活跃的研究方向。事实上，从早期使用VGGs[1, 2]的方法，到最新的具有更深更强大骨干的方法[3]，骨干架构的演进极大地推动了语义分割性能边界的扩展。除了骨干架构外，另一类工作将语义分割表述为一个结构化预测问题，并专注于设计模块和操作，以有效捕获上下文信息。这一领域的代表性例子是空洞卷积[4, 5]，它通过“膨胀”卷积核来增加感受野。

见证了在自然语言处理（NLP）中的巨大成功，最近出现了将Transformer引入视觉任务的浓厚兴趣。Dosovitskiy等人[6]提出了用于图像分类的视觉Transformer（ViT）。遵循NLP中的Transformer设计，作者将图像分割成多个线性嵌入的补丁，并将它们与位置嵌入（PE）一起输入到标准Transformer中，从而在ImageNet上取得了令人印象深刻的性能。在语义分割领域，Zheng等人[7]提出了SETR，以证明在该任务中使用Transformer的可行性。SETR采用ViT作为骨干，并结合了几个卷积神经网络（CNN）解码器来扩大特征分辨率。尽管性能良好，但ViT存在一些局限性：1）ViT输出的是单尺度低分辨率特征，而不是多尺度特征；2）在处理大图像时计算成本较高。为了解决这些局限性，Wang等人[8]提出了金字塔视觉Transformer（PVT），这是ViT的自然扩展，具有金字塔结构，适用于密集预测。PVT在目标检测和语义分割任务上相对于ResNet模型表现出了显著的改进。然而，与其他新兴方法（如Swin Transformer[9]和Twins[10]）一样，这些方法主要考虑了Transformer编码器的设计，而忽视了解码器对进一步改进的贡献。

本文介绍了SegFormer，这是一个用于语义分割的尖端Transformer框架，它综合考虑了效率、准确性和鲁棒性。与以前的方法不同，我们的框架重新设计了编码器和解码器。我们方法的关键创新点包括：

• 一个无需位置编码且具有层次结构的Transformer编码器。
• 一个轻量级的全多层感知机（MLP）解码器设计，它无需复杂且计算量大的模块即可产生强大的表示。

如图1所示，SegFormer在三个公开可用的语义分割数据集上在效率、准确性和鲁棒性方面均达到了最新水平。

首先，所提出的编码器在执行与训练分辨率不同的图像的推理时，避免了位置编码的插值。因此，我们的编码器可以轻松适应任意测试分辨率，而不会影响性能。此外，层次化部分使编码器能够生成高分辨率的精细特征和低分辨率的粗略特征，这与只能生成具有固定分辨率的单个低分辨率特征图的ViT形成对比。

其次，我们提出了一个轻量级的MLP解码器，其核心思想是利用Transformer诱导的特征，其中较低层的注意力往往保持局部性，而最高层的注意力则高度非局部。通过聚合来自不同层的信息，MLP解码器结合了局部注意力和全局注意力。因此，我们获得了一个简单直接的解码器，能够生成强大的表示。

我们在三个公开可用的数据集ADE20K、Cityscapes和COCO-Stuff上展示了SegFormer在模型大小、运行时间和准确性方面的优势。在Cityscapes上，我们的轻量级模型SegFormer-B0在没有使用TensorRT等加速实现的情况下，以48 FPS的速度达到了71.9%的mIoU，与ICNet[11]相比，延迟和性能分别提高了60%和4.2%。我们最大的模型SegFormer-B5达到了84.0%的mIoU，与SETR[7]相比，性能提高了1.8% mIoU，同时速度快了5倍。在ADE20K上，该模型以比SETR小4倍的模型大小达到了51.8%的mIoU，创下了新的最先进水平。此外，我们的方法在面对常见损坏和扰动时比现有方法具有更高的鲁棒性，因此适合用于安全关键型应用。代码将公开提供。

2. 相关工作

语义分割。语义分割可以看作是图像分类从图像级别到像素级别的扩展。在深度学习时代[12-16]，全卷积网络（FCN）[1]是语义分割的基础工作，它以一种端到端的方式进行像素到像素的分类。之后，研究人员从不同方面改进FCN，如：扩大感受野[17-19, 5, 2, 4, 20]；细化上下文信息[21-29]；引入边界信息[30-37]；设计各种注意力模块[38-46]；或使用AutoML技术[47-51]。这些方法显著提高了语义分割的性能，但代价是引入了许多经验性的模块，使得最终的框架计算量大且复杂。最近的方法证明了基于Transformer的架构在语义分割中的有效性[7, 46]，但这些方法仍然计算量大。

Transformer骨干网络。ViT[6]是第一个证明纯Transformer可以在图像分类中实现最先进性能的工作。ViT将每张图像视为一系列标记，然后将它们输入到多个Transformer层中进行分类。随后，DeiT[52]进一步探索了ViT的高效训练策略和蒸馏方法。更近期的方法，如T2T ViT[53]、CPVT[54]、TNT[55]、CrossViT[56]和LocalViT[57]，对ViT进行了定制化的修改，以进一步提高图像分类性能。

除了分类之外，PVT[8]是第一个在Transformer中引入金字塔结构的工作，证明了纯Transformer骨干网络在密集预测任务中与CNN相比的潜力。之后，诸如Swin[9]、CvT[58]、CoaT[59]、LeViT[60]和Twins[10]等方法增强了特征的局部连续性，并移除了固定大小的位置嵌入，以提高Transformer在密集预测任务中的性能。

特定任务的Transformer。DETR[52]是第一个使用Transformer构建端到端目标检测框架的工作，该框架无需非极大值抑制（NMS）。其他工作也将Transformer用于各种任务，如跟踪[61, 62]、超分辨率[63]、ReID[64]、着色[65]、检索[66]和多模态学习[67, 68]。对于语义分割，SETR[7]采用ViT[6]作为骨干网络来提取特征，取得了令人印象深刻的性能。然而，这些基于Transformer的方法效率非常低，因此难以部署在实时应用中。

3.方法

本节介绍了SegFormer，这是一个高效、鲁棒且强大的分割框架，没有手工设计和计算量大的模块。如图2所示，SegFormer由两个主要模块组成：（1）一个层次化的Transformer编码器，用于生成高分辨率的粗略特征和低分辨率的精细特征；（2）一个轻量级的全MLP解码器，用于融合这些多级特征以生成最终的语义分割掩码。

给定一张大小为H×W×3的图像，我们首先将其分割成大小为4×4的补丁。与ViT使用16×16大小的补丁不同，使用较小的补丁更有利于密集预测任务。然后，我们将这些补丁作为层次化Transformer编码器的输入，以获得原始图像分辨率的{1/4, 1/8, 1/16, 1/32}处的多级特征。接着，我们将这些多级特征传递给全MLP解码器，以预测分辨率为H/4×W/4×Ncls的分割掩码，其中Ncls是类别数。在本节的其余部分，我们将详细介绍所提出的编码器和解码器设计，并总结我们的方法与SETR之间的主要差异。

3.1 层次化Transformer编码器

我们设计了一系列Mix Transformer编码器（MiT），从MiT-B0到MiT-B5，它们具有相同的架构但尺寸不同。MiT-B0是我们的轻量级模型，用于快速推理，而MiT-B5是性能最佳的最大模型。MiT的设计部分受ViT启发，但针对语义分割进行了定制和优化。

层次化特征表示：与只能生成单一分辨率特征图的ViT不同，该模块的目标是给定输入图像，生成类似CNN的多级特征。这些特征提供了高分辨率的粗略特征和低分辨率的精细特征，这些特征通常可以提高语义分割的性能。更具体地说，给定分辨率为H×W×3的输入图像，我们执行补丁合并操作以生成具有H/2(i+1) × W/2(i+1) × Ci分辨率的层次化特征图Fi，其中i∈{1, 2, 3, 4}，且Ci+1大于Ci。

重叠补丁合并：在ViT中使用的补丁合并过程将N×N×3的补丁统一成一个1×1×C的向量。这可以很容易地扩展到将2×2×Ci的特征补丁统一成一个1×1×Ci+1的向量，以获得层次化特征图。使用这种方法，我们可以将层次化特征从F1（H/4 × W/4 × C1）缩小到F2（H/8 × W/8 × C2），然后对其他层次中的任何特征图进行迭代。此过程最初被设计为组合非重叠的图像或特征补丁。因此，它无法保留这些补丁周围的局部连续性。相反，我们使用重叠补丁合并过程。为此，我们定义了K、S和P，其中K是补丁大小，S是两个相邻补丁之间的步长，P是填充大小。在我们的实验中，我们设置K=7, S=4, P=3以及K=3, S=2, P=1来执行重叠补丁合并，以生成与非重叠过程相同大小的特征。

高效自注意力：编码器的主要计算瓶颈是自注意力层。在原始的多头自注意力过程中，每个头Q、K、V的维度都是N×C，其中N=H×W是序列的长度，自注意力估计为：

 该过程的计算复杂度为O(N^2)，对于大图像分辨率来说是不可接受的。相反，我们使用了在[8]中引入的序列缩减过程。该过程使用缩减比R来缩减序列的长度，如下所示：

 

 

其中K是需要缩减的序列，Reshape(N/R, C·R)(K)指的是将K重塑为N/R × (C·R)形状的序列，而Linear(C_in, C_out)(·)指的是一个线性层，它以Cin维的张量作为输入，并生成一个Cout维的张量作为输出。因此，新的K的维度为N/R × C。结果，自注意力机制的复杂度从O(N2)降低到O(N2/R)。在我们的实验中，我们从第一阶段到第四阶段将R分别设置为[64, 16, 4, 1]。

Mix-FFN。ViT使用位置编码（PE）来引入位置信息。然而，位置编码的分辨率是固定的。因此，当测试分辨率与训练分辨率不同时，位置代码需要进行插值，这通常会导致精度下降。为了缓解这个问题，CPVT[54]使用3×3的卷积与位置编码一起实现数据驱动的位置编码。我们认为，对于语义分割来说，位置编码实际上并不是必需的。相反，我们引入了Mix-FFN，它考虑了零填充对位置信息泄露的影响[69]，通过在前馈网络（FFN）中直接使用3×3的卷积来实现。Mix-FFN可以表示为：

 

 

其中x_in是来自自注意力模块的特征。Mix-FFN将3×3卷积和全连接层（MLP）结合到每个前馈网络（FFN）中。在我们的实验中，我们将证明3×3卷积足以为Transformer提供位置信息。特别是，我们使用深度可分离卷积来减少参数数量并提高效率。

3.2 轻量级全MLP解码器

SegFormer包含一个仅由MLP层组成的轻量级解码器，从而避免了其他方法中常用的手工设计和计算量大的组件。实现如此简单的解码器的关键是，我们的层次化Transformer编码器具有比传统CNN编码器更大的有效感受野（ERF）。

所提出的全MLP解码器主要由四个主要步骤组成。首先，来自MiT编码器的多级特征Fi通过一个MLP层来统一通道维度。然后，在第二步中，将特征上采样到原始大小的1/4并拼接在一起。第三步，采用一个MLP层来融合拼接后的特征F。最后，另一个MLP层接收融合后的特征来预测分割掩码M，其分辨率为H/4 × W/4 × N_cls，其中N_cls是类别数。这使我们能够将解码器表述为：

 

其中M指的是预测的掩码，而Linear(C_in, C_out)(·)指的是输入和输出向量维度分别为C_in和C_out的线性层。有效感受野分析。

对于语义分割来说，保持较大的感受野以包含上下文信息一直是一个核心问题[5, 19, 20]。在这里，我们使用有效感受野（ERF）[70]作为工具包来可视化和解释为什么我们的MLP解码器设计在Transformer上如此有效。在图3中，我们可视化了DeepLabv3+和SegFormer的四个编码器阶段和解码器头的ERF。我们可以得出以下观察结果：

• DeepLabv3+的ERF即使在第四阶段（最深的阶段）也相对较小。
• SegFormer的编码器自然产生类似于较低阶段卷积的局部注意力，同时能够在第四阶段输出高度非局部的注意力，有效地捕获上下文信息。
• 如图3中的放大补丁所示，MLP头（蓝色框）的ERF与第四阶段（红色框）不同，除了非局部注意力外，还包含了显著更强的局部注意力。
• CNN中有限的感受野要求人们求助于诸如ASPP[18]之类的上下文模块来扩大感受野，但这不可避免地增加了计算量。我们的解码器设计受益于Transformer中的非局部注意力，并在不增加复杂性的情况下获得了更大的感受野。然而，相同的解码器设计在CNN骨干网络上效果不佳，因为整体的感受野受到第四阶段有限感受野的上限限制，我们将在表1d中稍后验证这一点。
• 更重要的是，我们的解码器设计本质上利用了Transformer诱导的特征，该特征同时产生了高度局部和非局部的注意力。通过将它们统一起来，我们的MLP解码器通过添加少量参数生成了互补且强大的表示。这是我们设计的另一个关键原因。仅从第四阶段获取非局部注意力还不足以产生好的结果，我们将在表1d中验证这一点。

3.3 与SETR的关系

与SETR[7]相比，SegFormer包含了多个更高效且强大的设计：

• 我们仅使用ImageNet-1K进行预训练。SETR中的ViT是在更大的ImageNet-22K上进行预训练的。SegFormer的编码器具有层次结构，比ViT更小，并且能够捕获高分辨率的粗略特征和低分辨率的精细特征。相比之下，SETR的ViT编码器只能生成单个低分辨率特征图。
• 我们移除了编码器中的位置嵌入，而SETR使用固定形状的位置嵌入，当推理时的分辨率与训练时不同时，这会降低准确性。
• 我们的MLP解码器比SETR中的解码器更紧凑且计算量更小，导致可忽略不计的计算开销。相比之下，SETR需要包含多个3×3卷积的重解码器。

4. 实验

4.1 实验设置

数据集：我们使用了三个公开可用的数据集：Cityscapes[71]、ADE20K[72]和COCOStuff[73]。ADE20K是一个包含150个细粒度语义概念的场景解析数据集，共有20210张图像。Cityscapes是一个用于语义分割的驾驶数据集，包含5000张带有19个类别的高分辨率精细标注图像。COCO-Stuff包含172个标签，共有164k张图像：118k张用于训练，5k张用于验证，20k张用于测试开发集，20k张用于测试挑战集。

实现细节：我们使用mmsegmentation1代码库，并在配备8个Tesla V100的服务器上进行训练。我们在Imagenet-1K数据集上对编码器进行预训练，并随机初始化解码器。在训练过程中，我们通过以0.5-2.0的比例随机调整大小、随机水平翻转和随机裁剪到512×512、1024×1024、512×512（分别对应于ADE20K、Cityscapes和COCO-Stuff）来进行数据增强。遵循[9]，我们为我们的最大模型B5在ADE20K上将裁剪大小设置为640×640。我们使用AdamW优化器在ADE20K、Cityscapes上训练模型160K次迭代，在COCO-Stuff上训练80K次迭代。特别地，对于消融研究，我们对模型进行了40K次迭代的训练。我们在ADE20K和COCO-Stuff上使用批量大小为16，在Cityscapes上使用批量大小为8。学习率最初设置为0.00006，然后使用默认的因子为1.0的“poly”LR调度。为了简化，我们没有采用广泛使用的技巧，如OHEM、辅助损失或类别平衡损失。在评估过程中，我们将ADE20K和COCO-Stuff图像的短边缩放到训练裁剪大小，并保持宽高比。对于Cityscapes，我们使用滑动窗口测试通过裁剪1024×1024窗口进行推理。我们使用平均交并比（mIoU）来报告语义分割性能。

4.2 消融研究

模型大小的影响。我们首先分析增加编码器大小对性能和模型效率的影响。图1显示了ADE20K上性能与模型效率随编码器大小的变化，表1a总结了三个数据集的结果。

 

首先要注意的是解码器与编码器的大小比较。如表所示，对于轻量级模型，解码器只有0.4M参数。对于MiT-B5编码器，解码器仅占模型总参数数的4%。在性能方面，我们可以观察到，总体上，增加编码器的大小在所有数据集上都能带来一致的性能提升。我们的轻量级模型SegFormer-B0紧凑且高效，同时保持了具有竞争力的性能，表明我们的方法非常适合实时应用。另一方面，我们最大的模型SegFormer-B5在所有三个数据集上都达到了最先进的结果，展示了我们的Transformer编码器的潜力。

MLP解码器通道维度C的影响。我们现在分析MLP解码器中通道维度C的影响（见第3.2节）。在表1b中，我们展示了性能、浮点运算次数和参数随此维度的变化。我们可以观察到，设置C=256提供了非常有竞争力的性能和计算成本。随着C的增加，性能也会提高；然而，这会导致更大且效率更低的模型。有趣的是，当通道维度超过768时，性能趋于平稳。根据这些结果，我们为我们的实时模型SegFormer-B0、B1选择C=256，为其他模型选择C=768。

 

 

Mix-FFN与位置编码器（PE）的比较：在这个实验中，我们分析了在Transformer编码器中移除位置编码而使用提出的Mix-FFN的效果。为此，我们训练了带有位置编码（PE）的Transformer编码器和提出的Mix-FFN，并在Cityscapes数据集上以两种不同的图像分辨率进行推理：使用滑动窗口的768×768分辨率和使用整图的1024×2048分辨率。

表1c显示了该实验的结果。如表所示，对于给定的分辨率，我们使用Mix-FFN的方法明显优于使用位置编码的方法。此外，我们的方法对测试分辨率的差异不太敏感：当使用较低分辨率的位置编码时，精度下降了3.3%。相比之下，当我们使用提出的Mix-FFN时，性能下降减少到仅0.7%。从这些结果中，我们可以得出结论，使用提出的Mix-FFN可以产生更好且更鲁棒的编码器，而不是使用位置编码。

有效感受野评估：在第3.2节中，我们论证了我们的MLP解码器受益于Transformer具有比其他CNN模型更大的有效感受野。为了量化这种效果，在本实验中，我们将我们的MLP解码器与基于CNN的编码器（如ResNet或ResNeXt）进行了比较。如表1d所示，将我们的MLP解码器与基于CNN的编码器结合使用时，准确率显著降低，而与提出的Transformer编码器结合使用时则表现出最佳性能。直觉上，由于CNN的感受野比Transformer小（见第3.2节的分析），MLP解码器不足以进行全局推理。相反，将我们的Transformer编码器与MLP解码器结合使用可以达到最佳性能。此外，对于Transformer编码器，有必要结合低级局部特征和高级非局部特征，而不仅仅是高级特征。

与最先进方法的比较：我们现在将我们的结果与ADE20K[72]、Cityscapes[71]和COCO-Stuff[73]数据集上的现有方法进行比较。

ADE20K和Cityscapes：表2总结了我们的结果，包括ADE20K和Cityscapes的参数、浮点运算次数（FLOPS）、延迟和准确率。在表格的上半部分，我们报告了实时方法，包括最新的方法和使用MiT-B0轻量级编码器的我们的结果。在表格的下半部分，我们重点关注性能，并报告了我们的方法和使用更强编码器的相关工作的结果。

如表所示，在ADE20K上，SegFormer-B0仅使用3.8M参数和8.4G FLOPS就达到了37.4%的mIoU，在参数、FLOPS和延迟方面均优于其他所有实时方法。例如，与DeeplabV3+（MobileNetV2）相比，SegFormer-B0快7.4 FPS，且mIoU高出3.4%。此外，SegFormer-B5在所有其他方法中表现最佳，包括之前的最佳方法SETR，并以51.8%的mIoU建立了新的最先进水平，这比SETR高出1.6%的mIoU，同时效率显著提高。表2还显示了我们的结果在Cityscapes上也成立。

COCO-Stuff：最后，我们在完整的COCO-Stuff数据集上评估了SegFormer。由于现有方法没有在该数据集上提供结果，因此我们复现了最具代表性的方法，如DeeplabV3+、OCRNet和SETR。在这种情况下，该数据集上的FLOPS与ADE20K上报告的相同。如表4所示，SegFormer-B5仅以84.7M参数达到了46.7%的mIoU，这比SETR高出0.9%且小4倍。

对自然损坏的鲁棒性：模型的鲁棒性对于许多安全关键任务（如自动驾驶）非常重要。在本实验中，我们评估了SegFormer对常见损坏和扰动的鲁棒性。为此，我们遵循[77]并生成了Cityscapes-C，该数据集扩展了Cityscapes验证集，包含16种算法生成的损坏类型，涵盖噪声、模糊、天气和数字类别。我们将我们的方法与DeeplabV3+的不同变体和其他方法进行了比较，这些方法在[77]中有报道。该实验的结果总结在表5中。

我们的方法显著优于以前的方法，在高斯噪声上实现了高达588%的相对改进，在雪天条件下实现了高达295%的相对改进。结果表明，SegFormer具有很强的鲁棒性，我们预计这将使安全关键型应用受益，其中鲁棒性至关重要。

5. 结论

在本文中，我们提出了SegFormer，这是一种简单、干净且强大的语义分割方法，包含一个无位置编码的层次化Transformer编码器和轻量级的All-MLP解码器。它避免了以前方法中常见的复杂设计，从而实现了高效能和性能。SegFormer不仅在常见数据集上达到了新的最先进水平，而且还表现出了强大的零样本鲁棒性。我们希望我们的方法能够成为语义分割的一个坚实基线，并激发进一步的研究。一个限制是，尽管我们的最小3.7M参数模型比已知的CNN模型小，但尚不清楚它是否能在只有100k内存的边缘设备芯片上正常工作。我们将其留作未来的工作。