论文学习-SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

---

一、SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

Abstract

我们提出了一个SegFormer，一个简单，高效但强大的语义分割框架，它将 Transformers与轻量级多层感知器(MLP)解码器统一起来。前者有两个吸引人的特点：1)前者包括一个新的层次结构的Transformer encoder，它输出多尺度特征。它不需要位置编码，从而避免了位置码的插值，当测试分辨率与训练不同时，会导致性能下降。2)SegFormer避免了复杂的解码器。所提出的MLP解码器聚合了来自不同层的信息，从而结合了局部注意和全局注意，呈现出强大的表示。我们证明，这种简单和轻量级的设计是在Transformers上高效分割的关键。我们将我们的方法进行扩展，获得了从SegFrar-B0到SegFrar-B5的一系列模型，达到了比以前的模型更好的性能和效率。例如，SegFrore-B4在64M参数的ADE20K上达到50.3%的mIoU，比之前的最佳方法小5×，好2.2%。我们最好的模型Segfrar-B5在城市景观验证集上达到了84.0%的mIoU，代码将在：github.com/NVlabs/SegFormer发布。

1 Introduction

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

由于分类和语义分割之间有很强的关系，许多最先进的语义分割框架是在ImageNet上的图像分类化流行架构的变体。因此，设计主干架构仍然是语义分割中的一个活跃领域。事实上，从使用VGGs的早期方法，到具有更深入、更强大的骨干的最新方法，骨干的发展极大地突破了语义分割的性能极限。
除了主干架构外，另一项工作还将语义分割作为一个结构化的预测问题，并专注于设计模块和操作符，从而可以有效地捕获上下文信息。这方面的一个典型例子是dilated convolution，它通过用孔“inflating”核来增加接受场。

亲眼目睹了自然语言处理(NLP)的巨大成功，最近人们对将Transformers 引入视觉任务的兴趣激增。多索维茨基等人提出了vision Transformers用于图像分类。根据NLP中的Transformer设计，作者将一个图像分割成多个 embedded patches，并将其输入一个带有位置嵌入(PE)的标准Transformer，从而在ImageNet上取得了令人印象深刻的性能。在语义分割中，Zheng等人[7]提出了SETR来演示在该任务中使用Transformer的可行性。

SETR采用ViT作为骨干，并合并了几个CNN解码器来扩大特征分辨率。尽管ViT性能良好，但也有一些局限性：1)ViT输出单尺度低分辨率特性，而不是多尺度特性。2)它在大图像上的计算成本较高。为了解决这些限制，Wang等人提出了一种pyramid vision Transformer(PVT)，这是一种具有金字塔结构的ViT的自然扩展，用于密集预测。PVT在目标检测和语义分割方面比ResNet有相当大的改进。然而，与其他新兴的方法如Swin Transformer和Twins一起，这些方法主要考虑了Transformer编码器的设计，而忽略了解码器对进一步改进的贡献。

本文介绍了SegFormer，一个尖端的语义分割变压器框架，共同考虑了效率、准确性和鲁棒性。与以前的方法相比，我们的框架重新设计了编码器和解码器。我们的方法的关键新颖性是：

• 一种新的无位置编码的分层Transformer编码器。
• 一种轻量级的All-MLP解码器设计，产生一个强大的表示，没有复杂的和计算要求的模块。
• 如图1所示，在三个公开的语义分割数据集的效率、准确性和鲁棒性方面设置了一种最先进的新方法。
  

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

我们在三个公开的数据集上展示了SegFormer，它在模型大小、运行时和准确性方面的优势。在城市地图上，我们的轻量级模型SegFormer-B0，没有像TensorRT等加速实现，在48FPS下产生71.9%的mIoU，与ICNet相比，在延迟和性能方面分别相对提高了60%和4.2%。我们最大的模型SegFormer-B5产生84.0%的mIoU，相对提高了1.8%的mIoU，同时比SETR快5倍。在ADE20K上，该模型设置了一个最新的51.8%的mIoU，同时比SETR小4个倍。此外，我们的方法比现有的方法对常见的腐败和扰动更具鲁棒性，因此适用于安全关键型的应用。代码将被公开提供。

2 Related Work

语义分割 语义分割可以看作是图像分类从图像级到像素级的一种扩展。在深度学习时代，FCN是语义分割的基本工作，这是一个完全卷积的网络，以端到端的方式进行像素到像素的分类。之后，研究人员专注于从不同的方面改进FCN，如扩大感受野、细化上下文信息、引入边界信息、设计各种注意力模块、或使用自动感知机技术、这些方法以引入许多经验模块为代价，显著提高了语义分割的性能，使得所得到的框架计算要求高且复杂。最近的方法已经证明了基于Transformer的架构对语义分割的有效性，然而，这些方法仍然需要计算量。
Transformer架构 ViT是第一个证明一个纯Transformer可以在图像分类中达到最先进的性能的工作。ViT将每幅图像视为一系列块，然后将它们输入多个Transformer层以进行分类。随后，DeiT进一步探索了一种数据高效的训练策略和一种精馏方法。最近的一些方法，如T2TViT、CPVT、TNT、CrossViT和LocalViT，都对ViT进行了量身定制的更改，以进一步提高图像分类性能。
除了分类之外，PVT是第一个在变压器中引入金字塔结构的工作，展示了与CNN相比，纯变压器主干在密集预测任务中的潜力。之后，Swin、CvT、CoaT、LeViT和Twins等方法增强了特征的局部连续性，去掉了固定大小的位置嵌入，提高了Transformer在密集预测任务中的性能。
用于特定任务的Transformers DETR是第一个使用Transformer构建一个无非最大抑制(NMS)的终端到端目标检测框架的工作。其他作品也在各种任务中使用了变压器，如跟踪、超分辨率、ReID、着色、检索和多模态学习。在语义分割方面，SETR采用ViT作为骨干来提取特征，取得了令人印象深刻的性能。然而，这些基于Transformer的方法的效率非常低，因此很难在实时应用程序中部署。

3 Method

本节介绍SegFormer，我们的高效、健壮和强大的分割框架，没有手工制作和计算要求很高的模块。如图2所示，SegFormer由两个主要模块组成：（1）是分层Transformer编码器，生成高分辨率粗特征和低分辨率精细特征；（2）是一个轻量级的All-MLP解码器，融合这些多层特性，生成最终的语义分割掩码。
给定一个大小为H×W×3的图像，我们首先将其分成大小为4×4的小块。与使用大小为16×16的块的ViT相反，使用较小的补丁更有利于密集的预测任务。然后，我们使用这些补丁作为分层Transformer编码器的输入，以获得原始图像分辨率的{1/4、1/8、1/16、1/32}的多层次特征。然后我们通过这些多层特性All-MLP解码器预测分割面具$H_4\times W_4\times N_{cls}$，$N_{cls}$指的是分类数。在本节的其余部分中，我们将详细介绍所提出的编码器和解码器的设计，并总结了我们的方法和SETR之间的主要区别。

3.1分级Transformer编码器

我们设计了一系列的混合变压器编码器(MiT)，MiT-B0到MiT-B5，具有相同的架构，但大小不同。MiT-B0是我们用于快速推理的轻量级模型，而MiT-B5是我们获得最佳性能的最大模型。我们为MiT的设计部分受ViT的启发，但为语义分割定制和优化。
层次结构的特征表示 与ViT只能生成单分辨率的特征图不同，该模块的目标是，给定一个输入图像，生成类似于cnn的多层次特征。这些特性提供了高分辨率的粗特征和低分辨率的细粒度特征，通常可以提高语义分割的性能。更准确地说，给定一个输入图像分辨率$H\times W\times 3$，我们执行补丁合并获得分层特征地图$F_i$分辨率$H/2^{i+1}\times W/2^{i+1}\times C_i，i\in 1,2,3,4$，$C_{i+1}$大于$C_i$。
重叠的块合并 给定一个图像块，在ViT中使用的块合并过程，将一个$N\times N\times 3$块统一为一个$1\times 1\times C$向量。比如，这可以很容易一个$2\times 2\times C_i$特征路径统一到一个$1\times 1\times C_{i+1}$向量，以获得层次特征映射。使用它，我们可以缩小我们的层次特征由$F_1(H/4\times W/4\times C_1)$到 $F_2(H/8\times W/8\times C_2),$
然后迭代层次结构中的任何其他特征映射。这个过程最初被设计为结合不重叠的图像或特征块。因此，它不能保持这些块周围的局部连续性。相反，我们使用了一个重叠块合并过程。为此，我们定义了K、S和P，其中K是块大小，S是两个相邻patch之间的stride，P是padding大小。在我们的实验中，我们设置K=7、S=4、P=3和K=3、S=2、P=1进行重叠斑块合并，以产生与非重叠过程具有相同大小的特征。
高效的自注意力机制 编码器的主要计算瓶颈是自注意层。在原始的多头自注意过程中，每个头Q、K、V具有相同的维度N×C，其中N=H×W为序列的长度，自注意估计为：
$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{d_{h}ead}})V$$
该过程的计算复杂度为$O(n^2)$，这对于大的图像分辨率是禁止的。相反，我们使用了在[8]中引入的序列缩减过程。此过程使用还原比R来减少序列的长度如下:
$$K=Reshape(\frac{N}{R},C\cdot R)(K)$$
$$K=Linear(C\cdot R,C)(ˆK),$$
其中K是要减少的序列，$Reshape(\frac{N}{R}，C\cdot R)(K)$表示将K形状为$\frac{N}{R}\times (C\cdot R)$，线性$(C_{in}，C_{out})$（·）表示线性层以一个$C_{in}$维张量作为输入，生成一个$C_{out}$维张量作为输出。因此，新的K的维数为$\frac{N}{R}\times c$，因此，自注意机制的复杂性从O($n^2$)降低到O($\frac{n^2}{R}$)。在我们的实验中，我们从第一阶段到第四阶段将R设置为[64,16,4,1]。
Mix-FFN ViT使用位置编码(PE)来引入位置信息。然而，PE的分辨率是固定的。因此，当测试分辨率与训练分辨率不同时，需要插值位置代码，这往往会导致精度下降。为了缓解这个问题，CPVT使用3×3Conv和PE一起实现一个数据驱动的PE。我们认为，位置编码实际上并不是语义分割所必需的。相反，我们引入了Mix-FFN，它考虑了零填充对泄漏位置信息的影响，通过直接在前馈网络(FFN)中使用3×3Conv。Mix-FFN的定义如下：
$$x_{out}=MLP(GELU(Con{v}_{3\times 3}(MLP(x_{in}))))+x_{in}$$
其中，$x_{in}$是来自自我注意模块的特征。Mix-FFN将每个FFN中混合一个3×3的卷积和一个MLP。在我们的实验中，我们将证明一个3×3卷积足以为Transformer提供位置信息。特别是，我们使用深度卷积来减少参数的数量和提高效率。

3.2 轻量级 All-MLP 解码器

SegFormer集成了一个仅由MLP层组成的轻量级解码器，这就避免了在其他方法中通常使用的手工制作和计算要求很高的组件。实现这样一个简单的解码器的关键是，我们的分层Transformer编码器比传统的CNN编码器具有更大的有效接受域(ERF)。
$$F_i=Linear(C_i,C)(F_i),\forall i$$
$$F_i=Upsample(\frac{W}{4}\times \frac{W}{4})(F_i),\forall i$$
$$F=Linear(4C,C)(Concat(F_i)),\forall i$$
$$M=Linear(C,N_{cls})(F),$$

其中M为预测的掩模，线性$(C_{in}，C_{out})（\cdot ）$为分别以$C_{in}$和$C_{out}$为输入和输出向量维数的线性层。

有效的接收域分析 对于语义分割，保持较大的接受域以包含上下文信息一直是一个中心问题。在这里，我们使用有效的接受域(ERF)作为一个工具包来可视化和解释为什么我们的MLP解码器设计在Transformer上如此有效。在图3中，我们可视化了DeepLabv3+和SegFormer的四个编码器阶段和解码器头的erf。我们可以做以下意见：

• DeepLabv3+的ERF，即使在Stage-4，即最深的阶段，也相对较小。
• SegFormer的编码器自然地产生在较低阶段类似卷积的局部注意，同时能够输出高度非局部注意，有效地捕获Stage-4的上下文。
• 如图3中的放大块所示，MLP头部（蓝框）的ERF与第4阶段（红框）不同，除了非局部注意外，还有明显更强的局部注意。

CNN中有限的接受域要求人们求助于背景模块，如ASPP[18]，这些模块扩大了接受域，但不可避免地变得沉重。我们的解码器设计受益于Transformer的非局部注意，并导致一个更大的接受域而不复杂。然而，同样的解码器设计并不能很好地适用于CNN主干，因为在第四阶段，整体接受域是有限的，我们将稍后在表1d中验证这一点。

3.3 与SETR的关系

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

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

4 Experiments

4.1 Experimental Settings

数据集：我们使用了三个公开的数据集：城市景观、ADE20K和COCOStuff。ADE20K是一个场景解析数据集，涵盖了150个细粒度的语义概念，由20210张图像组成。城市景观是一个用于语义分割的驱动数据集，由5000个精细标注的高分辨率图像组成。COCO-Stuff包含172个标签，由164k张图像组成：118k用于训练，5k用于验证，20k用于测试开发，20k用于测试挑战。

实施细节 我们使用了mmsegmentation，并在一个使用8台Tesla V100的服务器上进行了训练。我们在Imagenet-1K数据集上对编码器进行预训练，并随机初始化解码器。在训练过程中，我们分别对ADE20K、城市景观和协比较为2.0、随机水平翻转和随机裁剪为512×512、1024×1024×1024、512×512。在之后，我们将我们最大的型号B5的ADE20K上的作物大小设置为640×640。我们使用AdamW优化器对模型在ADE20K上进行160K迭代，在Coco-Stuff上进行80K迭代。特别的是，对于消融研究，我们对模型进行了40k次迭代的训练。我们使用了ADE20K和COCO-Stuff的批大小为16，而城市景观的批大小为8。学习速率被设置为初始值为0.00006，然后使用默认的因子为1.0的“poly”LR调度。为简单起见，我们没有采用广泛使用的技巧，如OHEM、辅助损失或类平衡损失。在评估过程中，我们将图像的短边重新调整为训练裁剪大小，并保持ADE20K和COCO-Stuff的长宽比。对于城市景观，我们通过裁剪1024×1024窗口，使用滑动窗口测试进行推理。我们使用我来报告语义分割的性能。

4.2消融研究

对模型大小的影响 我们首先分析了增加编码器的大小对性能和模型效率的影响。图1显示了ADE20K的性能与模型效率作为编码器大小的函数，表1a总结了三个数据集的结果。这里首先要观察到的是解码器与编码器相比的大小。如图所示，对于轻量级模型，解码器只有0.4M的参数。对于MiT-B5编码器，解码器只占模型中总参数数的4%。在性能方面，我们可以观察到，总的来说，增加编码器的大小会对所有数据集产生一致的改进。我们的轻量级模型SegFormer-B0是紧凑和高效的，同时保持了具有竞争力的性能，这表明我们的方法对于实时应用非常方便。另一方面，我们最大的SegFormer-B5模型，在所有三个数据集上取得了最先进的结果，显示了我们的Transformer编码的潜力。
MLP解码器信道维度的影响 我们现在分析了在MLP解码器中的通道维度C的影响，见第3.2节。在表1b中，我们显示了性能、flops和参数量作为这个维度的函数。我们可以观察到，设置C=256提供了一个非常有竞争力的性能和计算成本。性能随着C的增加而增加；然而，它会导致更大、效率更低的模型。有趣的是，当通道尺寸大于768时，这种性能趋于稳定。鉴于这些结果，我们选择C=256作为我们的实时模型SegForme-B0，B1和C=768作为其余的。

Mix-FFN vs. Positional Encoder (PE) 在本实验中，我们分析了去除变压器编码器中的位置编码的效果，从而有利于使用所提出的Mix-FFN。为此，我们使用位置编码(PE)和提出的Mix-FFN训练变压器编码器，并对具有两种不同图像分辨率的城市景观进行推断：使用滑动窗口的768×768和使用整个图像的1024×2048。
本实验的结果为表1c。如图所示，对于给定的分辨率，我们使用Mix-FFN的方法明显优于使用位置编码的方法。此外，我们的方法对测试分辨率的差异不那么敏感：当使用较低分辨率的位置编码时，准确率下降了3.3%。相比之下，当我们使用提议的Mix-FFN时，性能下降降低到只有0.7%。从这些结果中，我们可以得出结论，使用所提出的Mix-FFN产生了比那些使用位置编码的更好和更鲁棒的编码器。
Effective receptive field evaluation 在第3.2节中，我们认为，与其他CNN模型相比，我们的MLP解码器受益于Transformer具有更大的有效接受域。为了量化这种效应，在这个实验中，我们比较了我们的MLP-解码器与基于cnn的编码器，如ResNet或ResNeXt一起使用时的性能。
使用基于cnn的编码器的MLP-解码器与与所提出的Transformer编码器耦合相比，产生的精度明显较低。直观地说，由于CNN的接受域比变压器要小（见第3.2节的分析），MLP-解码器不足以进行全局推理。相比之下，耦合我们的Transformer编码器和MLP解码器可以得到最好的性能。此外，对于Transformer编码器，需要将低级局部特征和高级非局部特征相结合，而不是只结合高级特征。

4.3与最先进的方法的比较

我们现在将我们的结果与ADE20K、城市景观和COCO-Stuff数据集上的现有方法进行了比较。
ADE20K and Cityscapes:表2总结了我们的结果，包括ADE20K和城市景观的参parameters, FLOPS, latency,and accuracy。在表的顶部部分，我们报告了实时方法，其中我们包括了最先进的方法和我们使用MiT-B0轻量级编码器的结果。在底部，我们关注性能，并报告了使用更强的编码器的我们的方法和相关工作的结果

如图所示，在ADE20K上，SegFormer-B0仅使用3.8M参数和8.4GFLOPs产生37.4%的mIoU，在参数、失败和延迟方面优于所有其他实时对应。例如，与DeeplabV3+(MobileNetV2)相比，SegFormer-B0为7.4FPS，速度更快，保持mIoU好3.4%。此外，SegFormer-B5优于所有其他方法，包括之前的最佳SETR，并建立了一个新的最先进的51.8%，即比SETR好1.6%的mIoU，同时明显更有效。

如表2所示，我们的结果也适用于城市景观。SegFrorar-B0产生15.2FPS和76.2%的mIoU（输入图像较短的一侧为1024），与DeeplabV3+相比，提高了1.3%的mIoU，加速了2×。此外，由于输入图像的短边为512，Segfror-B0在47.6FPS下运行，产生71.9%的mIoU，比ICNet快17.3FPS，快4.2%。SegFrort-B5的最佳IoU为84.0%，至少比所有现有方法高出1.8%，运行速度比SETR[7]快5×，比SETR[7]小4×.

在城市景观测试集上，我们遵循通用设置[20]，并将验证图像合并到训练集中，并使用Imagenet-1K预训练和辅助视觉[76]报告结果。如表3所示，仅使用城市景观精细数据和Imagenet-1K预训练，我们的方法达到了82.2%的mIoU，优于所有其他方法，包括SETR，它使用了ImageNet-22K预训练和额外的城市景观粗数据。使用ImageNet-22K，我们建立了一个新的83.1%mIoU的最先进的结果。图4显示了城市景观的定性结果，其中seg前者提供了比SETR更好的细节，比DeeplabV3+更平滑的预测。

$K_i$: the patch size of the overlapping patch embedding in Stage i;
$S_i$: the stride of the overlapping patch embedding in Stage i;
$P_i$: the padding size of the overlapping patch embedding in Stage i;
$C_i$: the channel number of the output of Stage i;
$L_i$: the number of encoder layers in Stage i;
$R_i$: the reduction ratio of the Efficient Self-Attention in Stage i;
$N_i$: the head number of the Efficient Self-Attention in Stage i;
$E_i$: the expansion ratio of the feed-forward layer in Stage i;

---