SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

https://github.com/NVlabs/SegFormer
https://arxiv.org/abs/2105.15203

我们介绍了SegFormer，这是一种简单，有效但功能强大的语义分割框架，它将变压器与轻型多层感知器 (MLP) 解码器统一在一起。SegFormer具有两个吸引人的功能:

• 1. SegFormer包括一种新颖的分层结构变压器编码器，该编码器输出多尺度特征。它不需要位置编码，从而避免了位置码的插值。如果使用位置编码，当测试分辨率与训练不同时，这会导致性能下降。
• 2. SegFormer避免了复杂的解码器。拟议的MLP解码器聚合来自不同层的信息，从而将本地注意力和全局注意力结合在一起，以呈现强大的表示形式。我们证明了这种简单轻巧的设计是对变压器进行有效分割的关键。

我们扩大了我们的方法，以获得从SegFormer-B0到SegFormer-B5的一系列模型，达到比以前的同行明显更好的性能和效率。例如，SegFormer-B4在64M参数的ADE20K上实现了50.3% mIoU，比以前的最佳方法小5倍，2.2% 更好。我们的最佳模型SegFormer-B5在Cityscapes验证集上实现了84.0% mIoU，并在Cityscapes-C上显示了出色的零镜头鲁棒性。

1 Introduction

语义分割是计算机视觉中的一项基本任务，并且可以实现许多下游应用程序。它与图像分类有关，因为它产生每像素类别预测而不是图像级别预测。在开创性的工作中指出并系统地研究了这种关系 [1]，其中作者使用完全卷积网络 (FCNs) 进行语义分割任务。从那时起，FCN激发了许多后续工作，并已成为密集预测的主要设计选择。

由于分类和语义分割之间存在很强的关系，因此许多最先进的语义分割框架都是ImageNet上用于图像分类的流行体系结构的变体。因此，设计主干体系结构一直是语义分割的活跃领域。确实，从使用VGGs的早期方法 [1,2] 开始，到具有明显更深，更强大的主干的最新方法 [3]，主干的演变极大地推动了 语义分割。
除了主干体系结构外，另一工作领域还将语义分割作为结构化的预测问题，并着重于设计 modules and operators ，从而可以有效地捕获上下文信息。该区域中的一个代表性示例是膨胀卷积 [4,5]，它通过 “膨胀” 带有孔的内核来增加感受野。

见证了自然语言处理 (NLP) 的巨大成功，最近人们对将变形金刚引入视觉任务的兴趣激增。Dosovitskiy等人 [6] 提出了用于图像分类的视觉变换器 (ViT)。按照NLP中的Transformer设计，作者将图像分成多个线性嵌入的补丁，然后将它们馈送到带有位置嵌入 (PE) 的标准变压器中，从而在ImageNet上取得了令人印象深刻的性能。在语义分割中，Zheng等人 [7] 提出了SETR，以证明在此任务中使用变压器的可行性。

SETR采用ViT作为骨干，并结合了多个CNN解码器以扩大特征分辨率。尽管性能良好，但ViT仍有一些局限性: 1) ViT输出单尺度低分辨率功能，而不是多尺度功能。2) 在大图像上具有较高的计算成本。

为了解决这些局限性，Wang等人 [8]（Pyramid vision transformer: A versatile backbone for dense prediction without convolutions） 提出了金字塔视觉变换器 (PVT)，这是ViT的自然扩展，具有金字塔结构，用于密集预测。PVT在对象检测和语义分割方面比ResNet对应物有了相当大的改进。但是，与其他新兴方法 (例如Swin Transformer [9] （Swin transformer: Hierarchical vision transformer using shifted windows）和Twins [10]（Twins: Revisiting spatial attention design in vision transformer）) 一起，这些方法主要考虑了Transformer编码器的设计，而忽略了解码器对进一步改进的贡献。

本文介绍了SegFormer，这是一种用于语义分割的尖端转换器框架，该框架共同考虑了效率，准确性和鲁棒性。与以前的方法相反，我们的框架重新设计了编码器和解码器。我们方法的主要新颖性是:

• 一种新颖的无位置编码分层变压器编码器。 A novel positional-encoding-free and hierarchical Transformer encoder.
• 轻量级的全MLP解码器设计，可产生强大的表示形式，而无需复杂且对计算有要求的模块
• 如图1所示，SegFormer在三个公开可用的语义分割数据集中的效率，准确性和鲁棒性方面设置了新的最新技术。

首先，提出的编码器在对分辨率不同于训练分辨率的图像执行推理时避免插值位置代码。结果，我们的编码器可以轻松地适应任意测试一下分辨率，而不会影响性能。此外，分层部分使编码器能够生成高分辨率的精细特征和低分辨率的粗略特征，这与ViT只能生成具有固定分辨率的单个低分辨率特征图形成对比。其次，我们提出了一种轻量级的MLP解码器，其中的关键思想是利用变压器引起的特征，其中较低层的注意力倾向于保持局部，而最高层的注意力则高度非局部。通过聚合来自不同层的信息，MLP解码器结合了本地和全局关注。结果，我们获得了一个简单明了的解码器，可以呈现强大的表示形式。

我们在三个公开可用的数据集 (ADE20K，Cityscapes和COCO-Stuff) 上展示了SegFormer在模型大小，运行时间和准确性方面的优势。在Citysapces上，我们的轻量级模型SegFormer-B0，没有加速实现，如TensorRT，在48 FPS时产生71.9% mIoU，与ICNet [11] 相比，分别代表了60% 和4.2% 的延迟和性能的相对提高。我们最大的模型SegFormer-B5产生84.0% mIoU，这代表了一个相对1.8% 的mIoU改善，同时比SETR快5倍 [7]。在ad20k上，该模型设置了新的51.8% mIoU，比SETR小4倍。此外，与现有方法相比，我们的方法对常见的损坏和扰动具有明显的鲁棒性，因此适用于对安全至关重要的应用。代码将公开提供。

2 Related Work

Semantic Segmentation。
义分割可以看作是图像分类从图像级别到像素级别的扩展。在深度学习时代 [12-16]，FCN [1] 是语义分割的基础工作，语义分割是一种完全卷积网络，以端到端的方式进行像素到像素的分类。之后，研究人员专注于从不同方面改进FCN，例如: 扩大感受野 [17-19,5，2,4，20]; 完善上下文信息 ；引入边界信息 [30-37]; 设计各种注意模块 [38-46]; 或使用AutoML技术 [47-51]。 这些方法以引入许多经验模块为代价，显着提高了语义分割性能，从而使生成的框架在计算上要求很高且很复杂。最近的方法证明了基于Transformer的体系结构对语义分割的有效性 [7,46]。然而，这些方法仍然需要计算。

Transformer backbones.
ViT [6] 是第一个证明纯变压器可以在图像分类中实现最先进性能的作品。ViT将每个图像视为一系列令牌，然后将其馈送到多个变压器层以进行分类。随后，DeiT [52] 进一步探索了一种数据高效的培训策略和ViT的蒸馏方法。T2T ViT [53]，CPVT [54]，TNT [55]，CrossViT [56] 和LocalViT [57] 等最新方法对ViT进行了量身定制的更改，以进一步提高图像分类性能。

除分类外，PVT [8] 是第一个在变压器中引入金字塔结构的工作，展示了纯变压器主干在密集预测任务中与CNN同行相比的潜力。之后，Swin [9]，CvT [58]，CoaT [59]，LeViT [60] 和Twins [10] 等方法增强了特征的局部连续性，并消除了固定尺寸的位置嵌入，以提高变压器在密集预测任务中的性能。

Transformers for specific tasks

DETR [52] 是第一个使用变压器来构建没有非最大抑制 (NMS) 的端到端对象检测框架的工作。其他作品也在跟踪 [61,62] 、超分辨率 [63] 、ReID [64] 、彩色化 [65] 、检索 [66] 和多模态学习 [67,68] 等多种任务中使用了变形金刚。对于语义分割，SETR [7] 采用ViT [6] 作为主干提取特征，取得了令人印象深刻的性能。但是，这些基于变压器的方法效率非常低，因此难以在实时应用中部署。

3 Method

本节介绍SegFormer，我们的高效、健壮和强大的细分框架，而无需手工制作和计算要求高的模块。如图2所示，SegFormer由两个主要模块组成 ( 1) 分层变压器编码器，用于生成高分辨率的粗略特征和低分辨率的精细特征; 以及 (2) 轻量级的全MLP解码器，用于融合这些多级特征以产生最终的语义分割掩码。

给定尺寸为H × w × 3的图像，我们首先将其分为尺寸为4 × 4的斑块。与使用大小为16 × 16的补丁的ViT相反，使用较小的补丁有利于密集的预测任务。

然后，我们将这些补丁用作分层变压器编码器的输入，以获得原始图像分辨率 {1/4，1/8，1/16，1/32} 的多级特征。

然后，我们将这些多级特征传递给全MLP解码器，以H/4 × W/4 × ncls分辨率预测分段掩码 .其中Ncls是类别的数量。在本节的其余部分中，我们详细介绍了建议的编码器和解码器设计，并总结了我们的方法与SETR之间的主要区别。

3.1 Hierarchical Transformer Encoder

我们设计了一系列MiT-B0 MiT-B5的混合变压器编码器 (MiT)，具有相同的架构但不同的尺寸。MiT-B0是用于快速推理的轻量级模型，而MiT-B5是用于最佳性能的最大模型。我们为MiT设计的部分灵感来自ViT，但针对语义分割进行了量身定制和优化。

Hierarchical Feature Representation. 与只能生成单分辨率特征图的ViT不同，该模块的目标是，给定输入图像，生成类似CNN的多级特征。这些功能提供了高分辨率的粗特征和低分辨率的细粒度特征，这些特征通常会提高语义分割的性能。更准确地说，给定分辨率为h × w × 3的输入图像，我们执行补丁合并以获得分辨率为

Overlapped Patch Merging.
1、给定图像补丁，ViT中使用的补丁合并过程将N × N × 3补丁统一为1 ×1 × C向量。这可以很容易地扩展为将2 × 2 × Ci特征路径统一为1 × 1 ×Ci 向量，以获得分层特征图。使用此，我们可以将层次特征从F1 (H/4 × W/4 × C1) 缩小到F2 (H/8 × W/8 × C2)，然后迭代层次中的任何其他特征图。 此过程最初旨在组合不重叠的图像或特征补丁。因此，它无法保留这些补丁周围的局部连续性。

2、相反，我们使用重叠补丁合并过程。 为此，我们定义K，S和P，其中K是补丁大小，S是两个相邻补丁之间的步幅，P是填充大小。在我们的实验中，我们设置K = 7，S = 4，P = 3和K = 3，S = 2，P = 1来执行重叠补丁合并，以产生与非重叠过程具有相同大小的特征。

Efficient Self-Attention. 编码器的主要计算瓶颈是自我注意层。在原始的多头自我注意过程中，每一头Q，K，V具有相同的维数N × c，其中N = H × w为序列的长度，其自我注意估计为:

此过程的计算复杂度为O(N2)，对于大图像分辨率而言，这是禁止的。相反，我们使用 [8]（Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.） 中引入的序列还原过程。此过程使用缩减比R来减少序列的长度，如下所示:

其中K是要缩小的序列，Reshape(N R，C·R)(K) 是指将K重新塑造为N R × (C · R) 形状的，并且线性 (Cin，coout)(·) 是指以Cin维张量为输入，生成coout维张量为输出的线性层。
r以Cin维张量为输入，并生成coout维张量为输出。因此，新的K具有N/R × C的尺寸。结果，自我注意机制的复杂性从O(N2) 降低到O(N2 /R)。在我们的实验中，我们将R从阶段1设置为 [64、16、4、1] 到阶段4

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

其中xin是来自自我注意模块的特征。Mix-FFN将3 × 3卷积和MLP混合到每个FFN中。在我们的实验中，我们将证明3 × 3卷积足以为变压器提供位置信息。特别是，我们使用深度卷积来减少参数数量并提高效率

3.2 Lightweight All-MLP Decoder

SegFormer集成了仅由MLP层组成的轻量级解码器，从而避免了其他方法中通常使用的手工制作且对计算要求很高的组件。启用如此简单的解码器的关键是，我们的分层变压器编码器比传统的CNN编码器具有更大的有效接收场 effective receptive field ((ERF)。

所提出的全MLP解码器包括四个主要步骤。首先，来自MiT编码器的多级功能Fi通过MLP层以统一通道尺寸。然后，在第二步中，将特征上采样到1/4并串联在一起。第三，采用MLP层来融合串联特征F。最后，另一个MLP层采用融合特征来预测具有H 4 × W 4 × Ncls分辨率的分割掩模M，其中Ncls是类别数。这让我们将解码器公式化为:

其中M指的是预测的掩码，线性 (Cin，cot)(·) 指的是分别以Cin和cot作为输入和输出向量维度的线性层。

Effective Receptive Field Analysis. 对于语义分割，保持较大的接受字段以包含上下文信息一直是一个中心问题 [5,19，20]。在这里，我们使用有效感受野 (ERF) [70] 作为工具包来可视化和解释为什么我们的MLP解码器设计在变压器上如此有效。在图3中，我们可视化了DeepLabv3和SegFormer的四个编码器级和解码器头的erf。我们可以做出以下观察:

• DeepLabv3的ERF即使在最深处的阶段-4也相对较小。
• ·SegFormer的编码器自然产生类似于较低阶段卷积的局部注意力，同时能够输出高度非本地的注意力，从而在阶段4有效地捕获上下文。
• 如图3中的放大补丁所示，MLP头的ERF (蓝色框) 与第4阶段 (红色框) 不同，除了非本地关注外，本地关注也明显增强。

CNN中的有限接受场要求人们诉诸于诸如ASPP [18] 之类的上下文模块，这些模块扩大了接受场，但不可避免地变得沉重。我们的解码器设计得益于变压器中的非局部关注，并导致更大的感受野而不复杂。然而，相同的解码器设计在CNN主干上效果不佳，因为在第4阶段，整体接受场由有限场界定，我们将在表1d中稍后验证这一点，

更重要的是，我们的解码器设计实质上利用了变压器感应功能，该功能同时产生高度本地和非本地关注。通过统一它们，我们的MLP解码器通过添加少量参数来呈现互补和强大的表示形式。这是推动我们设计的另一个关键原因。仅从第4阶段获得非本地关注不足以产生良好的结果，这将在表1d中得到验证。

3.3 Relationship to SETR

SegFormer包含与SETR相比更高效、更强大的多个设计 [7]:

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

4 Experiments

4.1 Experimental Settings

数据集: 我们使用了三个公开可用的数据集: Cityscapes [71]，ADE20K [72] 和COCOStuff [73]。Ad20k是一个场景解析数据集，涵盖了150由20210图像组成的细粒度语义概念。Cityscapes是用于语义分割的驱动数据集，由具有19个类别的5000个精细注释的高分辨率图像组成。COCO-Stuff涵盖172个标签，由164k图像组成: 训练118k，验证5k，测试一下-dev 20k和测试一下挑战20k。

实现细节: 我们使用了mmsegmentation1代码库，并在具有8 Tesla v100的服务器上进行了训练。我们在Imagenet-1K数据集上预先训练编码器，并随机初始化解码器。在训练过程中，我们通过随机调整比例0.5 2.0大小、随机水平翻转和随机裁剪分别为ad20k、Cityscapes和COCO-Stuff的512 × 512、1024 × 1024、512 × 512来应用数据增强。在 [9] 之后，我们为最大的型号b5在ADE20K上将裁剪尺寸设置为640 × 640。我们使用AdamW优化器训练了模型，在ad20k、Cityscapes上进行160K迭代，在COCO-Stuff上进行80K迭代。例外地，对于消融研究，我们训练了40k迭代的模型。我们对ad20k和COCO-Stuff使用了16的批量大小，对Cityscapes使用了8的批量大小。学习率被设置为0.00006的初始值，然后在默认情况下使用具有因子1.0的 “poly” LR计划。为简单起见，我们没有采用广泛使用的技巧，例如OHEM，辅助损失或类别平衡损失。在评估过程中，我们将图像的短边重新缩放为训练裁剪尺寸，并保持ad20k和COCO-Stuff的宽高比。对于城市景观，我们使用滑动窗口测试一下通过裁剪1024 × 1024窗口来推断。我们报告了使用并集的平均交集 (mIoU) 的语义分割性能。

4.2 Ablation Studies

Influence of the size of model. 我们首先分析增加编码器尺寸对性能和模型效率的影响。图1显示了ADE20K的性能与模型效率与编码器大小的关系，表1a总结了三个数据集的结果。这里首先要观察的是解码器与编码器相比的大小。如图所示，对于轻量级模型，解码器只有0.4M个参数。对于MiT-B5编码器，解码器仅占用模型中参数总数的4%。就性能而言，我们可以观察到，总体而言，增加编码器的大小会在所有数据集上产生一致的改进。我们的轻量级模型，SegFormer-B0，是紧凑和高效的，同时保持竞争力的性能，表明我们的方法是非常方便的实时应用。另一方面，我们的SegFormer-B5是最大的模型，在所有三个数据集上都获得了最先进的结果，这显示了我们的Transformer编码器的潜力。

Influence ofC, the MLP decoder channel dimension. 我们现在分析MLP解码器中信道维度C的影响，参见第3.2节。在表1b中，我们显示了性能，触发器和参数作为此维度的函数。我们可以观察到设置C = 256提供了非常有竞争力的性能和计算成本。性能随着C的增加而增加; 然而，它导致更大、效率更低的模型。有趣的是，该性能平台对于比768更宽的通道尺寸。鉴于这些结果，我们为实时模型SegFormer-B0选择C = 256，其余选择B1和C = 768。

Mix-FFN vs. Positional Encoder (PE) 在此实验中，我们分析了在变压器编码器中删除位置编码以支持使用建议的Mix-FFN的效果。为此，我们使用位置编码 (PE) 和建议的混合FFN训练变压器编码器，并在具有两种不同图像分辨率的城市景观上执行推理: 使用滑动窗口的768 × 768和使用整个图像的1024 × 2048。

表1c显示了该实验的结果。如图所示，对于给定的分辨率，我们使用Mix-FFN的方法显然优于使用位置编码的方法。此外，我们的方法对测试一下分辨率的差异不太敏感: 当使用具有较低分辨率的位置编码时，精度下降了3.3%。相反，当我们使用建议的Mix-FFN时，性能下降减少到只有0.7%。从这些结果中，我们可以得出结论，使用建议的Mix-FFN可以比使用位置编码的编码器产生更好，更强大的编码器。

Effective receptive field evaluation. 在第3.2节中，我们认为，与其他CNN模型相比，我们的MLP解码器受益于具有更大有效感受野的变压器。为了量化这种效果，在本实验中，我们比较了MLP解码器与基于CNN的编码器 (例如ResNet或ResNeXt) 一起使用时的性能。如表1d所示，将我们的MLP解码器与基于CNN的编码器耦合，与将其与所提出的变压器编码器耦合相比，精度明显降低。直观地，由于CNN具有比变压器小的感受野 (请参见第3.2节中的分析)，因此MLP解码器不足以进行全局推理。相比之下，将我们的变压器编码器与MLP解码器耦合可带来最佳性能。而且，对于变压器编码器，需要结合低级局部特征和高级非局部特征，而不仅仅是高级特征。

4.3 Comparison to state of the art methods

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

ADE20K和Cityscapes: 表2总结了我们的结果，包括ADE20K和Cityscapes的参数，FLOPS，延迟和准确性。在表格的顶部，我们报告了实时方法，其中包括最先进的方法以及使用MiT-B0轻量级编码器的结果。在底部，我们将重点放在性能上，并使用更强的编码器报告我们的方法和相关工作的结果。

如图所示，在ad20k上，SegFormer-B0仅使用3.8M参数和8.4G FLOPs即可产生37.4% mIoU，在参数，flops和延迟方面优于所有其他实时同行。例如，与DeeplabV3 (MobileNetV2) 相比，SegFormer-B0为7.4 FPS，速度更快，并保持3.4% 更好的mIoU。此外，SegFormer-B5优于所有其他方法，包括以前的最佳SETR，并建立了一种新的最先进的51.8%，这1.6% 比SETR更好，同时效率更高

如表2所示，我们的结果也适用于城市景观。SegFormer-B0产生15.2 FPS和76.2% mIoU (输入图像的较短侧为1024)，与DeeplabV3 + 相比，这代表1.3% 的mIoU改进和2倍加速。此外，随着输入图像的较短侧被512，SegFormer-B0以47.6 FPS运行，并产生71.9% mIoU，这比ICNet快17.3 FPS，4.2% 更好。SegFormer-B5存档了84.0% 年最好的IoU，比所有现有方法至少1.8% mIoU，并且比SETR运行速度快5倍，小4倍 [7]。

在Cityscapes测试一下集上，我们遵循通用设置 [20]，并将验证图像合并到火车集，并使用Imagenet-1K预训练以及Mapillary远景报告结果 [76]。如表3所述，仅使用Cityscapes精细数据和Imagenet-1K预训练，我们的方法实现了82.2% mIoU优于所有其他方法，包括使用ImageNet-22K预训练和其他Cityscapes粗略数据的SETR。使用Mapillary预训练，我们设置了83.1% mIoU的最新结果。图4显示了Cityscapes上的定性结果，其中SegFormer提供了比SETR更好的细节，并且比DeeplabV3更平滑的预测。

COCO-Stuff。最后，我们在完整的COCO-Stuff数据集上评估SegFormer为了进行比较，由于现有方法未在此数据集上提供结果，因此我们重现了最具代表性的方法，例如DeeplabV3，OCRNet和SETR。在这种情况下，此数据集上的触发器与ad20k报告的触发器相同。如表4所示，SegFormer-B5达到46.7% mIoU，只有84.7m参数，比SETR好0.9%，小4倍。总而言之，这些结果证明了SegFormer在语义分割方面在准确性，计算成本和模型大小方面的优势。

4.4 Robustness to natural corruptions

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

我们的方法明显优于以前的方法，在高斯噪声上产生了高达588% 的相对改进，在雪天气上产生了高达295% 的相对改进。结果表明，SegFormer具有强大的鲁棒性，我们设想将使鲁棒性非常重要的安全关键应用程序受益。

5 Conclusion

在本文中，我们介绍了SegFormer，这是一种简单，干净但功能强大的语义分割方法，其中包含无位置编码的分层变压器编码器和轻量级AllMLP解码器。它避免了以前方法中常见的复杂设计，从而提高了效率和性能。SegFormer不仅在常见数据集上实现了最新的结果，而且还显示出强大的零射鲁棒性。我们希望我们的方法可以作为语义分割的坚实基线，并激发进一步的研究。一个限制是，尽管我们的最小3.7M参数模型比已知的CNN模型小，但尚不清楚它是否可以在仅具有100k内存的边缘设备的芯片中很好地工作。我们把它留给以后的工作。