《SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers》

期刊：NIPS

年份：2021

代码：https://github.com/NVlabs/SegFormer.

摘要

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

Introduction

主要贡献：

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

图1 ADE20K 的性能与模型效率。所有结果都是通过单一模型和单尺度推理报告的。SegFormer 实现了新的最先进的 51.0% mIoU，同时比以前的方法更有效。

Method

SegFormer由两个主要模块组成:

• (1)分层Transformer编码器，生成高分辨率粗特征和低分辨率细特征;
• (2)一个轻量级的All-MLP解码器来融合这些多层次特征，生成最终的语义分割掩码。 

2.1 编码器

设计了一系列混合 Transformer 编码器 (MiT)、MiT-B0 到 MiT-B5，具有相同的架构但大小不同。MiT-B0 是我们用于快速推理的轻量级模型，而 MiT-B5 是最佳性能的最大模型。

2.1.1 Overlap Patch Merging

目的：该模块的目的是将输入图像的补丁合并为更高层次的特征表示，同时保持特征的空间分辨率。这与传统的非重叠补丁合并不同，后者通常用于降低特征的空间维度。

过程：在ViT中，补丁合并过程将一个N x N x 3的补丁统一为一个1 x 1 x C的向量。SegFormer扩展了这个过程，将的特征路径合并为一个的向量，以获得层次化的特征图。

参数：在重叠补丁合并过程中，定义了三个参数：K是补丁大小，S是相邻补丁之间的步长，P是填充大小。这些参数决定了合并过程的具体实现。

重叠的优势：使用重叠补丁合并而不是传统的非重叠方法可以保留局部连续性，这对于语义分割尤其重要，因为它有助于在特征的不同层次上保持空间信息。

实验设置：在实验中，作者设置了特定的K、S、P值来执行重叠补丁合并。例如，他们可能使用\K = 7，S = 4，P = 3，以及K = 3，S = 2，P = 1，以便产生与非重叠过程相同大小的特征。

层次化特征生成：通过重叠补丁合并，SegFormer能够从初始的特征图逐步生成更高层次的特征图，例如，，以及，其中表示第i个层次的特征通道数。

通过这种设计，SegFormer的编码器能够生成具有丰富层次信息的特征表示，为后续的解码和语义分割任务提供了坚实的基础。

2.1.2 Efficient Self-Attention

对传统Transformer中自注意力机制的一个改进，目的是降低计算复杂度并提高模型的效率。

传统自注意力的计算瓶颈：在标准的Transformer模型中，自注意力层的计算复杂度为O(N^2)，其中N是序列长度（在图像中对应于图像的宽度和高度的乘积）。这种计算复杂度对于高分辨率图像来说是不可行的。

序列缩减（Sequence Reduction）：为了解决这个问题，SegFormer采用了序列缩减技术，通过减少序列的长度来降低自注意力的计算复杂度。这是通过一个缩减比率R来实现的。

缩减过程：具体来说，原始的K序列（其中K是自注意力层中的键Key、查询Query和值Value的维度）被重新排列（reshape）为形状为 的张量，然后通过一个线性层（Linear Layer）进一步缩减到C维，其中C是最终的通道数。 

对性能的影响：尽管计算复杂度降低了，但Efficient Self-Attention仍然能够保持自注意力机制的核心优势，即能够捕捉长距离依赖关系，这对于语义分割任务至关重要。

2.1.3 Mix FFN

Mix-FFN（混合前馈网络）用于替代传统Transformer中的Positional Encoding（位置编码）。

背景：在原始的Transformer模型中，位置编码用于向模型提供序列中每个元素（例如图像中的补丁）的位置信息。然而，位置编码的分辨率是固定的，当测试图像的分辨率与训练时不同，就需要对位置编码进行插值，这可能会影响模型的性能。

Mix-FFN的设计：为了解决这个问题，SegFormer引入了Mix-FFN。它通过在前馈网络（Feed-Forward Network, FFN）中使用3x3的卷积层来引入位置信息，从而避免了固定分辨率的位置编码。

有效感受野（ERF）：使用Mix-FFN的Transformer编码器具有较大的有效感受野，这有助于模型捕获更丰富的上下文信息，对于语义分割任务来说非常重要。

性能和效率的平衡：通过使用Mix-FFN，SegFormer能够在不牺牲性能的前提下，提高模型的计算效率，这对于实时或资源受限的应用场景非常有价值。

2.2 轻量级全MLP解码器

轻量级全MLP解码器（Lightweight All-MLP Decoder）用于将编码器生成的多尺度特征转换为最终的语义分割图。

设计动机：传统的解码器通常包含复杂的手工艺设计和计算密集型组件。SegFormer通过使用全MLP解码器避免了这些复杂性，从而实现了一个简单而高效的解码过程。

MLP层的优势：MLP（多层感知器）层因其简单性和强大的学习能力而被选中。它们可以有效地聚合来自不同层的特征，捕获局部和全局的上下文信息。

解码器的工作流程：

• 统一通道维度：首先，来自Mix Transformer编码器的多尺度特征通过MLP层进行处理，以统一通道维度。
• 上采样和特征融合：然后，这些特征被上采样到原始图像的1/4分辨率，并在通道维度上进行拼接。
• 特征融合：接着，使用另一个MLP层来融合拼接后的特征。
• 预测分割掩码：最后，一个MLP层将融合的特征转换为最终的语义分割掩码，其分辨率与输入图像相同。

公式：

有效感受野分析：SegFormer的Transformer编码器具有较大的有效感受野（ERF），这使得MLP解码器能够有效地捕获全局上下文信息，即使它本身结构简单。

与CNN编码器的比较：当MLP解码器与CNN编码器结合使用时，由于CNN的有限感受野，解码器无法有效地进行全局推理，导致性能下降。

性能和效率：轻量级全MLP解码器在保持计算效率的同时，提供了与复杂解码器相媲美的性能，这在实时或资源受限的应用中尤为重要。

Result

与最先进的ADE20K和cityscape方法的比较。SegFormer 在#Params、#Flops、#Speed 和 #Accuracy 方面具有显着优势。请注意，对于 SegFormer-B0，我们将图像的短边缩放到 {1024, 768, 640, 512} 以获得速度和准确性的权衡。

Conclusion

在本文中，我们介绍了 SegFormer，这是一种简单而强大的语义分割方法，它包含一个无位置编码、分层 Transformer 编码器和一个轻量级 AllMLP 解码器。它避免了以往方法中常见的复杂设计，既提高了效率和性能。SegFormer 不仅在常见数据集上取得了新的最先进的结果，而且显示出强大的零样本鲁棒性。我们希望我们的方法可以作为语义分割的可靠基线，并激发进一步的研究。一个限制是，虽然我们的最小 3.7M 参数模型小于已知的 CNN 模型，但尚不清楚它是否可以在只有 100k 内存的边缘设备芯片中很好地工作。我们将其留给未来的工作。