【图像分割】SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

IRevers

已于 2024-08-06 22:03:12 修改

阅读量777

点赞数 25

文章标签： transformer 深度学习计算机视觉图像处理人工智能

于 2024-08-06 22:02:40 首次发布

本文链接：https://blog.csdn.net/qq_40734883/article/details/140967261

版权

SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

论文链接：http://arxiv.org/abs/2105.15203

代码链接：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上展现出了出色的零样本鲁棒性。

二、创新点

一种新颖的无位置编码和分层Transformer编码器。
一个轻量级的全连接多层感知机（All-MLP）解码器设计，能够提供强大的表征，而无需复杂且计算量大的模块。
SegFormer 在三个公开的语义分割数据集中，在效率、准确性和鲁棒性方面取得了新的SOTA。

三、原理

在这里插入图片描述

如图2所示，SegFormer由两个主要模块组成：（1）分层Transformer编码器，用于生成高分辨率的粗特征和低分辨率的细特征；（2）轻量级的全连接多层感知机（All-MLP）解码器，用于融合这些多层特征以生成最终的语义分割掩码。

给定尺寸为 $\times W \times 3$ 的图像，首先将其划分为 $\times 4$ 大小的块。与使用 $16 \times 16$ 大小块的ViT不同，使用较小的块有利于密集预测任务。然后，将这些块作为输入传递给分层Transformer编码器，以获得原始图像分辨率的 ${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 Hierarchical Transformer Encoder

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

分层特征表示。与只能生成单分辨率特征图的ViT不同，该模块的目标是在给定输入图像的情况下生成类似CNN的多级特征。这些特征提供了高分辨率的粗特征和低分辨率的细粒度特征，通常可以提升语义分割的性能。更具体地说，给定分辨率为 $\times W \times 3$ 的输入图像，执行补丁合并以获得分辨率为 $\frac{H}{2^{i+1}} \times \frac{W}{2^{i+1}} \times C_{i}$ 的分层特征图 $F i$ ，其中 $\in {1, 2, 3, 4}$ ，且 $C_{i+1}$ 大于 $C_{i}$ 。

重叠补丁合并。在ViT中使用的patch合并过程将 $\times N \times 3$ 的补丁统一为 $\times 1 \times C$ 的向量。可以将其扩展为将 $\times 2 \times C_{i}$ 的特征路径统一为 $\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}）$ ，然后对层次结构中的任何其他特征图进行迭代。该过程最初计用于合并非重叠的图像或特征补丁。因此，它无法保持这些补丁周围的局部连续性。而使用重叠补丁合并过程可以保持。为此，定义了 $K$ ， $S$ 和 $P$ ，其中 $K$ 是补丁大小， $S$ 是相邻补丁之间的步幅， $P$ 是填充大小。实验设置 $K = 7$ ， $S = 4$ ， $P = 3$ 和 $K = 3$ ， $S = 2$ ， $P = 1$ ，以执行重叠补丁合并，以产生与非重叠过程相同大小的特征。

高效的自注意力。编码器的主要计算瓶颈是自注意力层。在原始的多头自注意力过程中，每个头 $Q$ , $K$ , $V$ 具有相同的维度 $\times C$ ，其中 $\times W$ 是序列的长度，自注意力估计为：
$Softmax(\frac{QK^{T}}{\sqrt{d_{head}}})V. \tag1$
该过程的计算复杂度为 $O(N^{2})$ ，对于大尺寸图像来说是不可行的。因此，采用了[8]中介绍的序列缩减过程。该过程使用缩减比例 $R$ 来缩短序列的长度，具体如下：
$\hat{K} = Reshape(\frac{N}{R}, C \cdot R)(K) \\ K = Linear(C \cdot R, C)(\hat{K}),\tag2$

其中， $K$ 是待减少的序列， $Reshape(\frac{N}{R} , C \cdot R)(K)$ 表示将 $K$ 重塑为形状为 $NR \times (C \cdot R)$ 的序列， $Linear(C_{in}, C_{out})(\cdot)$ 表示以 $C_{in}$ 维张量作为输入，生成 $C_{out}$ 维张量作为输出的线性层。因此，新的 $K$ 的维度为 $\frac{N}{R} \times C$ 。因此，自注意力机制的复杂度从 $O(N^{2})$ 降低到 $O(\frac{N^{2}}{R})$ 。在实验中，将 $R$ 设置为 [64, 16, 4, 1]，从stage-1 到stage-4。

Mix-FFN. ViT 使用位置编码（PE）引入位置信息。然而，PE 的分辨率是固定的。因此，当测试分辨率与训练分辨率不同时，位置编码需要进行插值，这通常会导致准确率下降。为了缓解这个问题，CPVT [54] 使用 $\times 3$ 卷积与PE 结合，实现了一种数据驱动的PE。对于语义分割来说，位置编码实际上是不必要的。相反，引入了Mix-FFN，它考虑了零填充对位置信息的泄漏[69]，直接在前馈网络（FFN）中使用 $\times 3$ 卷积。Mix-FFN可以表示为：
$MLP(GELU(Conv_{3×3}(MLP(x_{in})))) + x_{in}, \tag3$
其中， $x_{in}$ 是来自自注意力模块的特征。Mix-FFN 将一个 $\times 3$ 卷积和一个MLP 混合到每个FFN 中。实验中将展示一个 $\times 3$ 卷积足以为 Transformers 提供位置信息。特别地，使用深度可分离卷积来减少参数数量并提高效率。

3.2 Lightweight All-MLP Decoder

SegFormer将一个仅由MLP层组成的轻量级解码器纳入其中，从而避免了其他方法中通常使用的手工制作和计算需求高的组件。实现这样一个简单解码器的关键在于，分层Transformer编码器具有比传统CNN编码器更大的有效感受野（ERF）。

提出的全MLP解码器包括四个主要步骤。首先，来自MiT编码器的多级特征 $F_{i}$ 经过一个MLP层，统一通道维度。然后，在第二步中，特征被上采样到 $1/4$ 大小，并进行拼接。第三步，采用MLP层来融合拼接后的特征 $F$ 。最后，另一个MLP层将融合后的特征用于预测分割掩码 $M$ ，分辨率为 $\frac{H}{4} × \frac{W}{4} × N_{cls}$ ，其中 $N_{cls}$ 是类别数。可以将解码器表示为：

$\hat{F}_{i} = Linear(C_{i}, C)(F_{i}), ∀i hat{F}_{i} = Upsample(\frac{W}{4} × \frac{W}{4})(\hat{F}_{i}), ∀i F = Linear(4C, C)(Concat(\hat{F}_{i})), ∀i M = Linear(C, N_{cls})(F), \tag4$
其中， $M$ 表示预测的掩码，$Linear(C_{in}, C_{out})(\cdot) $表示具有输入和输出向量维度 $C_{in}$ 和 $C_{out}$ 的线性层。

有效感受野分析。对于语义分割任务，保持较大的感受野以包含上下文信息一直是一个核心问题[5, 19, 20]。这里使用有效感受野（Effective Receptive Field，ERF）[70]作为一个工具包来可视化和解释为什么MLP解码器设计在Transformers上如此有效。

图2可视化了DeepLabv3+和SegFormer的四个编码器阶段和解码器头部的ERF。可以得出以下观察结果：

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

CNN中有限的感受野要求人们使用上下文模块，如ASPP [18]，以扩大感受野，但不可避免地变得复杂。我们的解码器设计借鉴了Transformer中的非局部注意力，从而实现了更大的感受野而不复杂。然而，相同的解码器设计在CNN骨干网络上效果不佳，因为整体感受野的上限由第4阶段的有限感受野决定，我们将在表1d中进行验证。

更重要的是，我们的解码器设计本质上利用了Transformer引发的特征，同时产生高度局部和非局部的注意力。通过统一它们，我们的MLP解码器通过添加少量参数产生互补且强大的表征。这是我们设计的另一个关键原因。仅仅使用第4阶段的非局部注意力是不足以产生良好结果的，这将在表1d中进行验证。

3.3 Relationship to SETR

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

只使用ImageNet-1K进行预训练。SETR中的ViT是在更大的ImageNet-22K上进行预训练的。
SegFormer的编码器具有分层结构，比ViT更小，可以捕捉高分辨率的粗糙特征和低分辨率的细节特征。相比之下，SETR的ViT编码器只能生成单一的低分辨率特征图。
在编码器中去除了位置嵌入，而SETR使用固定形状的位置嵌入，当推理时的分辨率与训练时的分辨率不同时，会降低准确性。
MLP解码器比SETR中的解码器更紧凑且计算需求较低。这导致了可忽略的计算开销。相比之下，SETR需要具有多个 $\times 3$ 卷积的复杂解码器。

四、实验

4.1 Experimental Settings

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

实现细节:使用了mmsegmentation代码库，并在一台配备了8块Tesla V100的服务器上进行训练。在Imagenet-1K数据集上对编码器进行了预训练，并随机初始化了解码器。训练过程通过随机调整大小（比例为 $0.5 - 2.0$ ）、随机水平翻转和随机裁剪来进行数据增强，分别将ADE20K、Cityscapes和COCO-Stuff的图像调整为 $512 \times 512$ 、 $1024 \times 1024$ 、 $512 \times 512$ 的大小。根据[9]，在ADE20K上为最大的B5模型设置了裁剪尺寸为 $640 \times 640$ 。在ADE20K上进行了160K次迭代的训练，Cityscapes和COCO-Stuff上进行了80K次迭代的训练。特别地，消融研究进行了40K次迭代的训练。在ADE20K和COCO-Stuff上使用了批量大小为16，而在Cityscapes上使用了批量大小为8。学习率初始值为0.00006，并采用了“poly”学习率调度策略，其衰减因子默认为1.0。为了简化起见，没有采用常用的技巧，如OHEM、辅助损失或类别平衡损失。

在评估过程中，ADE20K和COCO-Stuff将图像的短边缩放到训练裁剪尺寸，并保持纵横比以适应。对于Cityscapes，使用滑动窗口测试方法，裁剪 $1024 \times 1024$ 的窗口进行推理。使用平均交并比（mean Intersection over Union，mIoU）来报告语义分割性能。

4.2 Ablation Studies

在这里插入图片描述

模型大小的影响。首先分析增加编码器大小对性能和模型效率的影响。图1显示了ADE20K数据集的性能与模型效率随编码器大小的变化情况，表1a总结了三个数据集的结果。首先要观察的是解码器与编码器的大小比较。如图所示，对于轻量级模型，解码器只有0.4M 参数。对于MiT-B5 编码器，解码器仅占模型总参数数量的4%。在性能方面，可以观察到，总体上增加编码器的大小可以在所有数据集上获得一致的改进。轻量级模型SegFormer-B0 紧凑高效，同时保持了竞争性能表明我们的方法非常适用于实时应用。另一方面，最大模型SegFormer-B5 在所有三个数据集上取得了最先进的结果，展示了Transformer 编码器的潜力。

MLP解码器通道维度C的影响。现在分析MLP解码器中通道维度 $C$ 的影响，详见第3.2节。表1b展示了性能、FLOPs和参数随该维度变化的情况。可以观察到，设置 $C = 256$ 可以提供非常有竞争力的性能和计算成本。随着 $C$ 的增加，性能也会增加；然而，这会导致更大且不太高效的模型。有趣的是，当通道维度大于768时，性能趋于平稳。基于这些结果，选择 $C = 256$ 作为实时模型SegFormer-B0、B1的通道维度，以及选择 $C = 768$ 作为其他模型的通道维度。

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

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

有效感受野评估。在第3.2节中，MLP解码器与其他CNN模型相比受益于具有更大有效感受野的Transformer。为了量化这种影响，这个实验比较了MLP解码器与ResNet或ResNeXt等基于CNN的编码器配对时的性能。如表1d所示，将MLP解码器与基于CNN的编码器配对会导致显著较低的准确率，与将其与提出的Transformer编码器配对相比。直观地说，由于CNN的感受野较小（详见第3.2节的分析），MLP解码器无法进行全局推理。相反，将Transformer编码器与MLP解码器配对可以获得最佳性能。此外，对于Transformer编码器，有必要结合低层局部特征和高层非局部特征，而不仅仅是高层特征。

4.3 Comparison to state of the art methods

在ADE20K [72]、Cityscapes [71]和COCO-Stuff [73]数据集上比较将所提方法与现有方法。
在这里插入图片描述

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

如图所示，在ADE20K上，SegFormer-B0仅使用3.8M参数和8.4G FLOPs，mIoU达到37.4%，在参数、FLOPs和延迟方面优于所有其他实时对应方法。例如，与DeeplabV3+ (Mo-bileNetV2)相比，SegFormer-B0的帧率为7.4 FPS，更快且保持了3.4%更好的mIoU。此外，SegFormer-B5优于所有其他方法，包括之前最好的SETR，并建立了一个新的51.8%的最新结果，比SETR高出1.6%的mIoU，同时更加高效。

如表 2 所示，结果在Cityscapes上也成立。SegFormer-B0在输入图像的较短边为1024的情况下，达到了15.2 FPS和76.2%mIoU（即平均交并比），相比DeeplabV3+，mIoU提高了1.3%，速度提升了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倍，体积比SETR小4倍[7]。

在Cityscapes测试集上，遵循常见的设置 [20]，将验证图像合并到训练集中，并使用Imagenet-1K预训练模型进行结果报告，同时使用Mapillary Vistas [76]进行结果报告。

在这里插入图片描述

如表3所示，仅使用Cityscapes fine数据和Imagenet-1K预训练，所提方法达到了82.2%的mIoU，优于包括SETR在内的所有其他方法，后者使用了ImageNet-22K预训练和额外的Cityscapes coarse数据。使用Mapillary预训练，所提方法实现了83.1%的mIoU，实现了最新的最优结果。

在这里插入图片描述

图3 展示了在Cityscapes上的定性结果，SegFormer提供了比SETR更好的细节，并且，比DeeplabV3+具有更平滑的预测。
在这里插入图片描述

COCO-Stuff. 最后，在完整的COCO-Stuff数据集上评估SegFormer。为了比较，由于现有方法没有在该数据集上提供结果，作者重新实现了最具代表性的方法，如DeeplabV3+，OCRNet和SETR。在这种情况下，该数据集上的FLOPs与ADE20K报告的相同。如表4所示，SegFormer-B5仅使用84.7M参数就达到了46.7%的mIoU，比SETR提高了0.9%，且模型尺寸小了4倍。总之，这些结果表明SegFormer在语义分割方面在准确性、计算成本和模型尺寸方面具有优势。

4.4 Robustness to natural corruptions

在这里插入图片描述

模型的鲁棒性对于许多安全关键任务（如自动驾驶）非常重要[77]。这个实验评估了SegFormer对常见破坏和扰动的鲁棒性。为此，按照 [77] 的方法生成了Cityscapes-C数据集，该数据集通过从噪声、模糊、天气和数字类别中生成的16种算法生成的破坏来扩展Cityscapes验证集。将所提方法与DeeplabV3+的变体和其他方法进行了比较，如[77]中所述。这个实验的结果总结在表5中。

所提方法在高斯噪声上相对改进了588%，在雪天气上相对改进了295%。这些结果表明了SegFormer的强大鲁棒性，预计这将有助于鲁棒性对于安全关键应用。

五、总结

本文提出了SegFormer，一种简单、干净且强大的语义分割方法，它包含一个无位置编码的分层Transformer编码器和一个轻量级的全连接多层感知机（All-MLP）解码器。它避免了以前方法中常见的复杂设计，从而实现了高效率和性能的双重提升。SegFormer不仅在常见数据集上取得了新的最先进结果，而且展现出了强大的零样本鲁棒性。一个限制是，尽管最小的370万参数模型比已知的CNN模型要小，但不清楚它是否能在只有100k内存的边缘设备芯片上良好运行。

A Details of MiT Series

本节列出了Mix Transformer (MiT)编码器的一些重要超参数。通过改变这些参数，可以轻松地将编码器从B0扩展到B5。

总结起来，MiT的超参数如下所示：

$K_{i}$ : 阶段 $i$ 中重叠patch嵌入的patch尺寸；
$S_{i}$ : 阶段 $i$ 中重叠patch嵌入的步长；
$P_{i}$ : 在第 $i$ 阶段中，重叠的patch嵌入的填充大小；
$C_{i}$ : 第 $i$ 阶段输出的通道数；
$L_{i}$ : 第 $i$ 阶段的编码器层数；
$R_{i}$ : Efficient Self-Attention 在第 $i$ 阶段的减少比率;
$N_{i}$ : Efficient Self-Attention 在第 $i$ 阶段的头数；
$E_{i}$ : 阶段 $i$ 中前馈层[78] 的扩张比率；

在这里插入图片描述

表6展示了MiT系列的详细信息。为了方便高效的讨论，MiT编码器分配了B0到B5的代码名称，其中B0是为实时设计的最小模型，而B5是为高性能设计的最大模型。

B More Qualitative Results on Mask Predictions

在这里插入图片描述

图5，与SETR和DeepLabV3+进行了Cityscapes、ADE20K和COCO-Stuff的更多定性结果对比。

与SETR相比，SegFormer在物体边界附近预测出具有显著更细节的掩码，因为Transformer编码器可以捕捉比SETR更高分辨率的特征，从而保留更多详细的纹理信息。与DeepLabV3+相比，SegFormer通过Transformer编码器具有更大的有效感受野，从而减少了远距离错误。

C More Visualization on Effective Receptive Field

在这里插入图片描述

图6选择了一些代表性图像和DeepLabV3+和SegFormer的有效感受野（ERF）。除了更大的ERF外，SegFormer的ERF对图像的上下文更敏感。可以看到SegFormer的ERF学习到了道路、汽车和建筑物的模式，而DeepLabV3+的ERF显示出相对固定的模式。结果还表明，Transformer编码器具有比ConvNets更强的特征提取能力。

D More Comparison of DeeplabV3+ and SegFormer on Cityscapes-C

本节详细展示了与SegFormer和DeepLabV3+相比的零样本鲁棒性。根据[77]，对4种“噪声”和其余12种污染和扰动进行了3个严重程度的测试。

如图7所示，随着严重程度的增加，DeepLabV3+的性能明显下降。相比之下，SegFormer的性能相对稳定。此外，SegFormer在所有污染/扰动和所有严重程度上都比DeepLabV3+具有明显优势，展示了出色的零样本鲁棒性。
在这里插入图片描述

IRevers

关注

25
点赞
踩
24

收藏

觉得还不错? 一键收藏
打赏
0
评论
【图像分割】SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

文中提出了SegFormer，一个简单、高效且强大的语义分割框架，将Transformer与轻量级多层感知器（MLP）解码器相结合。SegFormer具有两个特点：&emsp;1）SegFormer包含一个新颖的**分层结构**Transformer编码器，输出**多尺度特征**。它**不需要位置编码**，因此避免了位置代码的插值，从而在测试分辨率与训练不同时导致性能下降。
复制链接

扫一扫