SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers论文解读

SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

论文：[2105.15203] SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers (arxiv.org)

代码：NVlabs/SegFormer: Official PyTorch implementation of SegFormer (github.com)

期刊/会议：NeurIPS 2021

摘要

我们提出了SegFormer，一个简单，高效而强大的语义分割框架，它将transformer与轻量级多层感知器(MLP)解码器统一起来。SegFormer有两个吸引人的特点:1)SegFormer包括一个新颖的层次结构transformer编码器，输出多尺度特征。它不需要位置编码，从而避免了位置编码的插值，从而导致测试分辨率与训练分辨率不同时性能下降。2) SegFormer避免了复杂的解码器。所提出的MLP解码器聚合来自不同层的信息，从而结合局部注意力和全局注意力来呈现强大的表示。我们表明，这种简单和轻量级的设计是transformer上有效分割的关键。我们扩大了我们的方法，获得了从SegFormer-B0到SegFormer-B5的一系列模型，达到了比以前的模型更好的性能和效率。例如，SegFormer-B4的参数64M，并在ADE20K上mIoU达到50.3%，比之前的最佳方法参数上缩小了5倍，提高了2.2%。我们的最佳模型SegFormer-B5在Cityscape验证集上达到84.0% mIoU，并在Cityscape-C上表现出出色的零样本 鲁棒性。

1、简介

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

由于图像分类和语义分割之间有很强的联系，许多最先进的语义分割框架都是ImageNet上流行的图像分类架构的变体。因此，骨干网络架构的设计一直是语义分割的一个活跃领域。的确，从早期使用VGGs的方法开始，到具有更深入和更强大的骨干网络的最新方法，骨干网络的发展极大地推动了语义分割性能的推进。除了骨干网络架构之外，另一项工作将语义分割作为结构化预测问题，并专注于设计能够有效捕获上下文信息的模块和操作符。这方面的一个典型例子是空洞卷积，它通过用孔“膨胀”内核来增加感受野。

见证了自然语言处理(NLP)的巨大成功，最近人们对将transformer引入视觉任务的兴趣激增。Dosovitskiy等人提出了用于图像分类的Vison Transformer(ViT)。遵循NLP中的Transformer设计，作者将图像分割为多个线性嵌入的patch，并将它们输入到带有位置嵌入(PE)的标准Transformer中，从而在ImageNet上获得了令人印象深刻的性能。在语义分割中，Zheng等人提出了SETR，以证明在该任务中使用transformer的可行性。

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

本文介绍了SegFormer，一种用于语义分割的前沿Transformer框架，它同时考虑了效率、准确性和鲁棒性。与以前的方法相比，我们的框架重新设计了编码器和解码器。我们的方法的关键新颖之处在于:

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

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

我们在三个公开可用的数据集(ADE20K、Cityscape和COCO-Stuff)上展示了SegFormer在模型大小、运行时间和准确性方面的优势。在Citysapces上，我们的轻量级模型SegFormer-B0在没有TensorRT等加速实现的情况下，在48 FPS的情况下产生了71.9%的mIoU，与ICNet相比，分别在延迟和性能方面提高了60%和4.2%。我们最大的模型，SegFormer-B5，产生84.0% mIoU，这意味着相对1.8%的mIoU改善，同时比SETR快5倍。在ADE20K上，该模型的最高水平为51.8% mIoU，比SETR小4倍。此外，我们的方法对常见的损坏和扰动比现有的方法更健壮，因此适用于安全关键型应用。代码将是公开的。

2、相关工作

语义分割：语义分割可以看作是图像分类从图像级到像素级的扩展。在深度学习时代，FCN是语义分割的基础工作，它是一个完全卷积的网络，以端到端的方式进行像素到像素的分类。之后，研究者从不同方面关注FCN的改善，如:扩大感受野；强化上下文信息；引入边界信息；设计大量的attention模块；使用AutoML技术。这些方法显著提高了语义分割的性能，但代价是引入了许多经验模块，使得得到的框架计算要求高且复杂。最近的方法已经证明了基于transformer的架构在语义分割方面的有效性。然而，这些方法仍然需要大量的计算。

transformer骨干网络：ViT是第一个证明纯Transformer可以在图像分类中实现最先进性能的工作。ViT将每张图像视为一系列标记，然后将它们提供给多个Transformer层进行分类。随后，DeiT进一步探索了数据高效的训练策略和ViT的蒸馏方法。更近期的方法，如T2T ViT， CPVT， TNT， CrossViT和LocalViT，引入了对ViT的定制变化，以进一步提高图像分类性能。

除了分类，PVT是第一个在Transformer中引入金字塔结构的工作，展示了纯Transformer骨干网络在密集预测任务中与CNN骨干网络相比的潜力。之后，采用Swin、CvT、CoaT、LeViT、Twins等方法增强特征的局部连续性，去除固定大小的位置嵌入，提高transformer在密集预测任务中的性能。

transformer在专门任务的表现：DETR是第一个使用transformer构建非极大值抑制(NMS)的端到端对象检测框架的工作。其他相关工作也在各种任务中使用了transformer，如跟踪、超分辨率、ReID、着色、检索和多模态学习。在语义分割方面，SETR采用ViT作为骨干网络抽取特征，取得了不错的性能。然而，这些基于transformer的方法效率非常低，因此很难部署到实时应用中。

3、方法

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

给定一张大小为$H\times W\times 3$的图像，我们首先将其分成大小为$4\times 4$的小块。与使用大小为$16\times 16$的patch的ViT相反，使用较小的patch有利于密集的预测任务。然后，我们使用这些patch作为分层Transformer编码器的输入，以获得原始图像分辨率的 { 1 / 4 , 1 / 8 , 1 / 16 , 1 / 32 } \{ 1/4,1/ 8,1 / 16,1 /32 \} {1/4,1/8,1/16,1/32}的多级特征。然后，我们将这些多级特征传递给All-MLP解码器，以$\frac{H}{4}\times \frac{W}{4}\times N_{cls}$分辨率预测分割掩码，$N_{cls}$是目标的类别数。在剩余的部分，我们将详细的介绍编码器和解码器的设计，并总结我们所提出的方法和SETR的不同。

3.1 层级Transformer编码器

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

层级特征表示：与只能生成单分辨率特征图的ViT不同，该模块的目标是，给定一个输入图像，生成类似CNN的多级特征。这些特征提供了高分辨率的粗特征和低分辨率的细粒度特征，通常可以提高语义分割的性能。更精确地说，给定分辨率为$H\times W\times 3$的输入图像，我们进行patchh merge，得到分辨率为$\frac{H}{2^i+1}\times \frac{W}{2^i+1}\times C_i$的分层特征图$F_i$，其中 i ∈ { 1 , 2 , 3 , 4 } i \in \{1,2,3,4\} i∈{1,2,3,4}，且$C_{i+1}$大于$C_i$。

重叠patch合并。给定一个图像patch，ViT中使用的patch merge过程将$N\times N\times 3$个patch统一为$1\times 1\times C$矢量。这可以很容易地扩展到将$2\times 2\times C_i$特征路径统一为$1\times 1\times C_{i+1}$向量，以获得分层特征图。使用这种方法，我们可以将层次结构特征从$F_1(\frac{H}{4}\times \frac{W}{4}\times C_1)$缩小到$F_2(\frac{H}{8}\times \frac{W}{8}\times C_2)$，然后迭代层次结构中的任何其他特征映射。这个过程最初被设计用来组合不重叠的图像或特征patch。因此，它无法保持这些patch周围的局部连续性。相反，我们使用overlapped patch merge过程。为此，我们定义$K,S,P$，其中$K$是patch大小，$S$是两个相邻patch之间的步长，$P$是填充大小。在我们的实验中，我们设置$K=7,S=4,P=3$和$K=3,S=2,P=1$进行重叠patch合并，得到与非重叠过程相同大小的特征。

有效的self-attention机制：编码器的主要计算瓶颈是自注意层。在原来的多头自注意过程中，每个头$Q,K,V$都具有相同的维数$N\times C$，其中$N=H\times W$为序列长度，则自注意被计算为:
$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{d_{head}}})V$$
该过程的计算复杂度为$O(N^2)$，这对于大图像分辨率来说是禁止的。相反，我们使用一些工作中介绍的序列精简过程。这个过程使用一个缩减比$R$来缩减序列的长度如下:
$$\begin{gathered} \hat{K}=Reshape(\frac{N}{R},C\cdot R)(K), \\ K=Linear(C\cdot R,C)(\hat{K}) \end{gathered}$$
其中$K$为待精简序列，重塑$(\frac{N}{R},C\cdot R)(K)$指将$K$重塑为形状为$\frac{N}{R},(C\cdot R)$的序列，线性$(C_{in},C_{out})(\cdot )$指线性层以一个$C_{in}$维张量作为输入，生成一个$C_{out}$维张量作为输出。因此，新$K$的维数为$\frac{N}{R}\times C$，从而将自注意机制的复杂度从$O(N^2)$降低到$O(N^2/R)$。在我们的实验中，我们将starge-1到stage-4的$R$设为[64,16,4,1]。

Mix-FFN。ViT使用位置编码(PE)来引入位置信息。然而，PE的分辨率是固定的。因此，当测试分辨率与训练分辨率不同时，需要插入位置代码，这往往会导致精度下降。为了解决这一问题，CPVT使用3 × 3 Conv与PE一起实现数据驱动PE。我们认为位置编码对于语义分割实际上是不必要的。相反，我们引入了Mix-FFN，它考虑了零填充去弱化位置信息的影响，直接在前馈网络(FFN)中使用3 × 3 Conv。Mix-FFN可以表述为:
$$x_{out}=MLP(GELU(Con{v}_{3\times 3}(MLP(x_{in}))))+x_{in}$$

其中$x_{in}$是来自self-attention模块的特征。Mix-FFN将一个3 × 3卷积和一个MLP混合到每个FFN中。在我们的实验中，我们将证明3 × 3的卷积足以为transformer提供位置信息。特别是，我们使用深度可分离卷积来减少参数的数量并提高效率。

3.2 轻量的All-MLP解码器

SegFormer集成了一个轻量级的解码器，仅由MLP层组成，这避免了手工制作和在其他方法中通常使用的计算要求高的组件。实现这种简单解码器的关键在于，我们的分层Transformer编码器比传统CNN编码器具有更大的有效感受野(effective reception field, ERF)。

本文提出的All-MLP解码器由四个主要步骤组成。首先，MiT编码器的多级特征$F_i$通过MLP层来统一通道维度。然后，在第二步中，特征被向上采样到1/4，并连接在一起。第三，采用MLP层融合级联特征$F$。最后，另一个MLP层采用融合特征预测$\frac{H}{4}\times \frac{W}{4}\times N_{cls}$分辨率的分割掩模$M$，其中$N_{cls}$为类别数。这让我们可以将解码器表述为:
$$\begin{gathered} {\hat{F}}_i=Linear(C_i,C)(F_i),\forall i \\ {\hat{F}}_i=Upsample(\frac{H}{4}\times \frac{W}{4})({\hat{F}}_i),\forall i \\ F=Linear(4C,C)(Concat({\hat{F}}_i)),\forall i \\ M=Linear(C,N_{cls})(F), \end{gathered}$$
$M$是预测的msk。$Linear(C_{in},C_{out})(\cdot )$指的是线性层，输入和输出维度分别是$C_{in},C_{out}$。

有效的接受野分析。在语义分割中，保持较大的感受野以包含上下文信息一直是一个重要的问题。在这里，我们使用有效感受野(ERF)作为toolkit来可视化和解释为什么我们的MLP解码器设计在transformer上如此有效。在图3中，我们可视化了DeepLabv3+和SegFormer的四个编码器阶段和解码器头的ERF。我们可以观察到:

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

CNN中有限的感受野需要借助ASPP等上下文模块，这些模块扩大了感受野，但不可避免地变得沉重。我们的解码器设计受益于transformer中的非局部关注，并在不复杂的情况下导致更大的感受野。然而，同样的解码器设计在CNN主干上并不能很好地工作，因为在阶段4时，整体接受域的上界是有限的，我们将在后面的表1d中验证这一点。

更重要的是，我们的解码器设计本质上利用了Transformer诱导的特性，同时产生高度局部和非局部注意。通过统一它们，我们的MLP解码器通过添加很少的参数来呈现互补和强大的表示。这是推动我们设计的另一个关键原因。仅从阶段4中获得非局部注意力不足以产生良好的结果，如表1d所示。

3.3 和SETR的关系

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

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

4、实验

4.1 实验设置

数据集：Cityscapes，ADE20K，COCO-Stuff。

实施细节：我们使用mmsegmentation代码库，并在一台8块Tesla V100的服务器上进行训练。我们在Imagenet-1K数据集上预训练编码器，并随机初始化解码器。在训练过程中，我们对ADE20K、Cityscape和COCO-Stuff分别进行了0.5-2.0比例随机调整大小、随机水平翻转和随机裁剪的数据增强，分别为512 × 512、1024×1024、512 × 512。我们在ADE20K上为我们最大的模型B5设置裁剪尺寸为640 × 640。我们使用AdamW优化器训练模型，在ADE20K, Cityscape上进行160K迭代，在COCO-Stuff上进行80K迭代。在消融实验中，我们对模型进行了40K次迭代训练。ADE20K和COCO-Stuff的batch size为16个，Cityscape的batch size为8。学习率被设置为初始值0.00006，然后默认使用因子1.0的“poly”LR schedule。为了简单起见，我们没有采用广泛使用的技巧，如OHEM，辅助损失或类平衡损失。在评估过程中，我们将图像的短边重新缩放为训练裁剪尺寸，并保持ADE20K和COCO-Stuff的纵横比。对于城市景观，我们使用滑动窗口测试进行推断，裁剪1024 × 1024个窗口。我们报告了使用平均交并比(mIoU)的语义分割性能。

4.2 消融实验

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

影响MLP解码器的通道维数$C$。我们现在分析通道维数$C$在MLP解码器中的影响，参见第3.2节。在表1b中，我们展示了性能、flop和参数作为这个维数的函数。我们可以观察到，设置$C=256$提供了非常有竞争力的性能和计算成本。性能随着$C$的增加而增加;然而，这会导致更大且效率更低的模型。有趣的是，当通道数大于768时，这种性能趋于稳定。鉴于这些结果，我们选择$C=256$的实时模型SegFormer-B0, B1和$C=768$的其余部分。

Mix-FFN vs Positional Encoder(PE)。在这个实验中，我们分析了去除transformer编码器中的位置编码以支持使用所提出的Mix-FFN的效果。为此，我们用位置编码(PE)和提出的Mix-FFN训练Transformer编码器，并使用两种不同的图像分辨率对Cityscapes进行推断:768×768使用滑动窗口，1024×2048使用整个图像。

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

有效的接受野评估。在第3.2节中，我们认为与其他CNN模型相比，我们的MLP解码器受益于transformer具有更大的有效感受野。为了量化这种效果，在这个实验中，我们比较了我们的MLP解码器与基于CNN的编码器(如ResNet或ResNeXt)使用时的性能。如表1d所示，我们的MLP解码器与基于CNN的编码器耦合产生的精度明显低于与所提出的transformer编码器耦合。直观地说，由于CNN的感受野比Transformer小(参见3.2节中的分析)，MLP解码器不足以进行全局推理。相比之下，将我们的Transformer编码器与MLP解码器耦合可以获得最佳性能。此外，对于Transformer编码器，需要将低级局部特征和高级非局部特征结合起来，而不是只有高级特征。

4.3 对比SOTA

我们现在将我们的结果与ADE20K、Cityscape和COCO-Stuff数据集上的现有方法进行比较。

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

如图所示，在ADE20K上，SegFormer-B0仅使用3.8M参数和8.4G FLOPs就产生37.4%的mIoU，在参数、FLOPs和延迟方面优于所有其他实时对手。例如，与DeeplabV3+ (MobileNetV2)相比，SegFormer-B0是7.4 FPS，速度更快，mIoU保持3.4%。此外，segprevious - b5优于所有其他方法，包括之前最好的SETR，并建立了51.8%的新技术，比SETR高出1.6% mIoU，同时效率显著提高。

如表2所示，我们的结果也适用于Cityscape。SegFormer-B0输出15.2 FPS和76.2% mIoU(输入图像的短边为1024)，与DeeplabV3+相比，mIoU提高了1.3%，速度提高了2倍。此外，当输入图像的短边为512时，SegFormer-B0运行速度为47.6 FPS, mIoU产量为71.9%，比ICNet快17.3 FPS，好4.2%。SegFormer-B5记录了84.0%的最佳IoU，比所有现有方法至少高出1.8% mIoU，运行速度比SETR快5倍，体积小4倍。

在Cityscape测试集中，我们遵循公共的设置，并将验证图像合并到训练集，并使用Imagenet-1K预训练和Mapillary views报告结果。如表3所述，仅使用Cityscape coarse数据和Imagenet-1K预训练，我们的方法获得了82.2%的mIoU，优于所有其他方法，包括使用ImageNet-22K预训练和额外的Cityscape coarse数据的SETR。使用Mapillary预训练，我们获得了83.1% mIoU的最新结果。图4显示了Cityscape的定性结果，其中SegFormer提供了比SETR更好的细节，比DeeplabV3+更平滑的预测。

COCO-Stuff。最后，我们在完整的COCO-Stuff数据集上评估SegFormer。为了进行比较，由于现有方法无法在此数据集上提供结果，我们重新生成了最具代表性的方法，如DeeplabV3+、OCRNet和SETR。在这种情况下，该数据集上的失败与报告的ADE20K相同。如表4所示，SegFormer-B5 mIoU达到46.7%，参数只有84.7M，比SETR提高0.9%，缩小4倍。综上所述，这些结果证明了SegFormer在语义分割精度、计算成本和模型大小方面的优势。

4.4 自然场景下鲁棒性探究

模型鲁棒性对于许多安全关键任务(如自动驾驶)非常重要。在这个实验中，我们评估了SegFormer对常见破坏和扰动的鲁棒性。在最后，我们遵循先前的工作并生成Cityscape-C，它扩展了Cityscape validation数据集，包含16种算法生成的噪声、模糊、天气和数字类别的损坏。我们将我们的方法与DeeplabV3+的变体中报道的其他方法进行了比较。本实验结果汇总如表5所示。

我们的方法明显优于以前的方法，对高斯噪声的相对改善高达588%，对雪天的相对改善高达295%。结果表明SegFormer具有很强的鲁棒性，我们预计它将使鲁棒性非常重要的安全关键应用程序受益。

5、总结

在本文中，我们提出了SegFormer，一种简单、轻量但功能强大的语义分割方法，它包含一个无位置编码的分层Transformer编码器和一个轻量级的All-MLP解码器。该方法避免了以往方法中常见的复杂设计，效率和性能都很高。SegFormer不仅在通用数据集上实现了最新的结果，而且还显示出强大的零样本鲁棒性。我们希望我们的方法可以作为一个坚实的基线语义分割和激励进一步的研究。一个限制是，虽然我们最小的3.7M参数模型比已知的CNN模型要小，但尚不清楚它是否能在只有100k内存的边缘设备芯片上很好地工作。我们把它留给以后的工作。