SegFormer——无需位置编码，轻量级的transformer架构，并有很强的零样本鲁棒性

题目：SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

作者：Enze Xie ，The University of Hong Kong
code：https://github.com/NVlabs/SegFormer

1.研究背景

Introduction

由于分类和语义分割之间有很强的联系，许多先进的语义分割框架是ImageNet上流行的图像分类架构的变体。因此，设计主干架构一直是语义分割中的一个活跃领域。事实上，从早期的使用VGGs [ 1 , 2]的方法开始，到最新的具有显著更深和更强大的骨干的方法[ 3 ]，骨干的进化极大地推动了语义分割的性能边界。除了主干架构之外，另一类工作将语义分割描述为一个结构化的预测问题，并专注于设计模块和运算符，这可以有效地捕获上下文信息。这方面比较有代表性的例子是空洞卷积[ 4、5]，它通过对带有空洞的内核进行"膨胀"来增加感受野。
近年来，随着自然语言处理( Natural Language Processing，NLP )领域的巨大成功，将Transformer引入视觉任务引起了人们极大的兴趣。Dosovitskiy等人[ 6 ]提出视觉转换器( ViT )用于图像分类。遵循NLP中的Transformer设计，作者将一幅图像分割成多个线性嵌入的块，并将其输入到一个标准的带有位置嵌入( PE )的Transformer中，从而在ImageNet上获得了令人印象深刻的性能。在语义分割方面，Zheng等人[ 7 ]提出SETR来证明在这个任务中使用Transformers的可行性。
SETR采用ViT作为主干，并集成了多个CNN解码器，以扩大特征分辨率。尽管ViT具有良好的性能，但也存在一些局限性：1 ) ViT输出的是单尺度低分辨率特征，而非多尺度特征。2 )在大图像上计算代价高。为了解决这些局限性，Wang等人[ 8 ]提出了金字塔视觉变压器( PVT )，一种具有金字塔结构的ViT的自然扩展，用于密集预测。与ResNet相比，PVT在目标检测和语义分割方面有很大的改进。然而，与其他新兴方法如Swin Transformer [ 9 ]和Twins [ 10 ]等一起，这些方法主要考虑Transformer编码器的设计，忽略了解码器对进一步改进的贡献。

Related work

语义分割可以看作是图像分类从图像级别到像素级别的延伸。在深度学习时代[ 12-16 ]，FCN [ 1 ]是语义分割的基础工作，它是一种全卷积网络，以端到端的方式进行像素到像素的分类。此后，研究者从不同方面对FCN进行改进，如：扩大感受野；精炼上下文信息；引入边界信息[ 30-37 ]；设计各种注意力模块[ 38-46 ]；或者使用AutoML技术。这些方法以引入许多经验模块为代价，显著提高了语义分割性能，使得得到的框架计算量大且复杂。最近的一些方法已经证明了基于Transformer架构的语义分割[ 7、46]的有效性。然而，这些方法仍然在计算上要求很高。
ViT [ 6 ]是第一个证明纯Transformer在图像分类中可以达到最先进性能的工作。ViT将每幅图像视为一个令牌序列，然后将其馈送到多个Transformer层进行分类。随后，Dei T [ 52 ]进一步探索了一种数据高效的ViT训练策略和蒸馏方法。最近的一些方法如T2T ViT [ 53 ]、CPVT [ 54 ]、TNT [ 55 ]、Cross ViT [ 56 ]和Local ViT [ 57 ]等通过对ViT进行定制化的修改来进一步提高图像分类性能。
除了分类，PVT [ 8 ]是第一个在Transformer中引入金字塔结构的工作，展示了在密集预测任务中，与CNN对应方相比，纯粹的Transformer骨干的潜力。之后，Swin [ 9 ]、CvT [ 58 ]、CoaT [ 59 ]、LeViT [ 60 ]和Twins [ 10 ]等方法通过增强特征的局部连续性和去除固定大小的位置嵌入来提高Transformer在稠密预测任务中的性能。
DETR [ 52 ]是第一个使用Transformers构建无非极大值抑制( NMS )的端到端目标检测框架的工作。其他工作也在跟踪[ 61,62]、超分辨率[ 63 ]、ReID [ 64 ]、彩色化[ 65 ]、检索[ 66 ]和多模态学习[ 67,68]等任务中使用Transformer。对于语义分割，SETR [ 7 ]采用ViT [ 6 ]作为骨架提取特征，取得了不错的性能。然而，这些基于Transformer的方法效率很低，难以在实时应用中部署。

2.研究目的

3.研究方法

创新点

•一种新颖的无需位置编码的分层Transformer编码器。
•一种轻量级的All - MLP解码器设计，无需复杂的计算模块，具有强大的表达能力。
•如图1所示，SegFormer在三个公开的语义分割数据集上在效率、准确性和鲁棒性方面都取得了最新的成就。

首先，所提出的编码器在对分辨率与训练图像不同的图像进行推理时，**避免了插值位置码。**因此，我们的编码器可以很容易地适应任意测试分辨率而不影响性能。此外，**分层部分使编码器能够同时生成高分辨率的精细特征和低分辨率的粗略特征，这与ViT只能生成固定分辨率的单一低分辨率特征图相反。其次，我们提出了一个轻量级的MLP解码器，其核心思想是利用Transformer诱导的特征，其中低层的注意力倾向于保持局部，而最高层的注意力则高度非局部。**通过聚合来自不同层的信息，MLP解码器结合了局部和全局注意力。因此，我们得到了一个简单而直接的解码器，可以提供强大的表示。
**我们在ADE20K、Cityscapes和COCO - Stuff三个公开数据集上展示了SegFormer在模型大小、运行时间和准确率方面的优势。在Citysapces上，我们的轻量级模型SegFormer - B0在没有TensorRT等加速实现的情况下，在48 FPS时获得了71.9 %的m Io U，与ICNet [ 11 ]相比，在延迟和性能上分别相对提高了60 %和4.2 %。我们最大的模型SegFormer - B5产生了84.0 % mIoU，相对于SETR提高了1.8 % mIoU，速度提高了5倍[ 7 ]。**在ADE20K上，该模型设定了51.8 % mIoU的最新状态，同时比SETR小4 ×。此外，与现有方法相比，我们的方法对常见的损坏和扰动具有更强的鲁棒性，因此适用于安全关键的应用。代码将公开提供。
给定一幅大小为H × W × 3的图像，我们首先将其分成大小为4 × 4的图像块。与ViT使用大小为16 × 16的图块相反，使用较小的图块更有利于密集预测任务。然后将这些图像块作为分层Transformer编码器的输入，以获得原始图像分辨率的{ 1 / 4，1 / 8，1 / 16，1 / 32 }处的多级特征。然后，我们将这些多级特征传递给All - MLP解码器，以H/4 × W/4 × Ncls的分辨率预测分割掩码，其中Ncls是类别数。

Hierarchical Transformer Encoder

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

Hierarchical Feature Representation

该模块的目标是，给定输入图像，生成类似CNN的多级特征。

Overlapped Patch Merging 重叠块合并

给定一个图像补丁，ViT中使用的补丁合并过程将一个N × N × 3的补丁统一成一个1 × 1 × C的向量。这可以很容易地扩展到将一个2 × 2 × Ci特征路径统一成一个1 × 1 × Ci+1的向量，以获得分层特征图。使用这种方法，我们可以从F1（H/4 × W/4 × C1）缩小到F2（H/8 × W/8 × C2），然后对层次结构中的任何其他特征图进行迭代。此过程最初设计用于合并非重叠的图像或特征补丁。因此，它无法保持这些补丁周围的局部连续性。相反，我们使用重叠补丁合并过程。为此，我们定义了K、S和P，其中K是补丁大小，S是相邻补丁之间的步长，P是填充大小。在我们的实验中，我们设置K = 7，S = 4，P = 3和K = 3，S = 2，P = 1，以执行重叠补丁合并，以生成与非重叠过程相同大小的特征。

Overlap Patch Embeddings是一种图像编码方法，用于将图像分成多个重叠的图像块，并将每个图像块编码为一个向量表示。这种方法常用于图像分割、目标检测和图像生成等计算机视觉任务中。 在Overlap Patch Embeddings中，图像被划分为多个重叠的图像块，每个图像块通常是一个固定大小的正方形或矩形区域。这些图像块之间通常有一定的重叠部分，以确保不会丢失重要的图像信息。然后，每个图像块通过一个预训练的卷积神经网络（CNN）模型进行特征提取，得到一个高维的向量表示。 通过使用重叠的图像块，Overlap Patch Embeddings能够捕捉到图像中更细节的信息，而不仅仅是整个图像的概览。这对于处理具有复杂纹理、细小结构或需要更精确定位的目标非常有用。 使用Overlap Patch Embeddings时，可以将每个图像块的向量表示连接起来，形成整个图像的编码表示。这样，图像中的每个部分都可以被独立地表示和处理，从而提供更多的局部信息和上下文关系。

Efficient Self-Attention

我们使用文献[ 8 ]中介绍的序列约化过程。该过程使用缩减比例R来缩减序列的长度

Mix-FFN

ViT使用了position embedding来引入定位信息。但是，position embedding的大小是固定的，因此，当测试的时候输入分辨率与训练的时候的分辨率不一致的话，positional code就需要被插值，导致精度的下降。所以CPV在positional embedding上使用了3X3卷积，来实现数据驱动的一个positional encoding。本文提出的是:positional encoding对于语义分割是没有必要的。 作者引入的是Mix-FFN模块，通过在Feed-forward network(FFN)上直接使用3X3卷积，减弱了zero-padding会丢失一些定位信息的影响。

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x, H, W):
        x = self.fc1(x)
        x = self.dwconv(x, H, W)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

Lightweight All-MLP Decoder

Decoder是由MLP组成的。能够使用如此简单的decoder是因为hierarchical Transformer encoder比传统的CNN encoders有较大的有效感知域。
计算步骤：

通过下面源码可知，第一步的Linear是FC层，Upsample是bilinear插值法。第三步的特征融合和第四步预测都是卷积操作。

from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule


class MLP(nn.Module):
    """
    Linear Embedding
    """
    def __init__(self, input_dim=2048, embed_dim=768):
        super().__init__()
        self.proj = nn.Linear(input_dim, embed_dim)

    def forward(self, x):
        x = x.flatten(2).transpose(1, 2)
        x = self.proj(x)
        return x

class SegFormerHead(BaseDecodeHead):
    """
    SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers
    """
    def __init__(self, feature_strides, **kwargs):
        super(SegFormerHead, self).__init__(input_transform='multiple_select', **kwargs)
        assert len(feature_strides) == len(self.in_channels)
        assert min(feature_strides) == feature_strides[0]
        self.feature_strides = feature_strides

        c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels

        decoder_params = kwargs['decoder_params']
        embedding_dim = decoder_params['embed_dim']

        self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim)
        self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim)
        self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim)
        self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim)

        self.linear_fuse = ConvModule(
            in_channels=embedding_dim*4,
            out_channels=embedding_dim,
            kernel_size=1,
            norm_cfg=dict(type='SyncBN', requires_grad=True)
        )

        self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1)

    def forward(self, inputs):
        x = self._transform_inputs(inputs)  # len=4, 1/4,1/8,1/16,1/32
        c1, c2, c3, c4 = x

        ############## MLP decoder on C1-C4 ###########
        n, _, h, w = c4.shape

        _c4 = self.linear_c4(c4).permute(0,2,1).reshape(n, -1, c4.shape[2], c4.shape[3])
        #  F.interpolate(input, size, scale_factor, mode, align_corners)
        _c4 = resize(_c4, size=c1.size()[2:],mode='bilinear',align_corners=False)

        _c3 = self.linear_c3(c3).permute(0,2,1).reshape(n, -1, c3.shape[2], c3.shape[3])
        _c3 = resize(_c3, size=c1.size()[2:],mode='bilinear',align_corners=False)

        _c2 = self.linear_c2(c2).permute(0,2,1).reshape(n, -1, c2.shape[2], c2.shape[3])
        _c2 = resize(_c2, size=c1.size()[2:],mode='bilinear',align_corners=False)

        _c1 = self.linear_c1(c1).permute(0,2,1).reshape(n, -1, c1.shape[2], c1.shape[3])

        _c = self.linear_fuse(torch.cat([_c4, _c3, _c2, _c1], dim=1))

        x = self.dropout(_c)
        x = self.linear_pred(x)

        return x

Effective Receptive Field Analysis

对于语义分割，保持大的感受野以包含上下文信息一直是[ 5、19、20]的中心问题。在这里，我们使用有效感受野( ERF ) [ 70 ]作为工具包来可视化和解释为什么我们的MLP解码器设计在Transformers上如此有效。在图3中，我们可视化了DeepLabv3 +和SegFormer的四个编码器阶段和解码器头部的ERF。我们可以做出以下观察：
·即使在最深的阶段Stage - 4，Deep Labv3 +的ERF也相对较小；
· SegFormer的编码器在较低的阶段自然地产生类似卷积的局部注意力，而在阶段Stage - 4能够输出有效捕获上下文的高度非局部注意力；
·如图3中的放大图所示，MLP头部(蓝色方框)的ERF与阶段Stage - 4 (红色方框)不同，除了非局部注意力外，局部注意力显著更强。

CNN中有限的感受野要求我们使用上下文模块（如ASPP [18]）来扩大感受野，但这不可避免地增加了复杂性。我们的解码器设计**借鉴了Transformer的非局部注意力，并实现了更大的感受野，而不会变得复杂。然而，同样的解码器设计在CNN骨干网络上效果不佳，因为总体感受野的上限由Stage-4的有限感受野决定，我们将在表1d中验证这一点。
更重要的是，我们的解码器设计基本上利用了Transformer产生的既具有高度局部注意力又具有非局部注意力的特征。通过统一它们，我们的MLP解码器通过增加少量参数产生了互补且强大的表示。**这是我们设计的另一个关键原因。仅仅使用Stage-4的非局部注意力是不足以产生良好结果的，这将在表1d中得到验证。

编码器结构细节

4.实验

5.结论

在本文中，**我们提出了一种简单、干净且功能强大的语义分割方法SegFormer，它包含一个位置编码无关的、分层的Transformer编码器和一个轻量级的AllMLP解码器。**该方法避免了以往方法中常见的复杂设计，具有较高的效率和性能。SegFormer不仅在公共数据集上取得了最新的研究成果，而且表现出了很强的零样本鲁棒性。我们希望我们的方法可以作为语义分割的坚实基础，并推动进一步的研究。

6.摘抄

7.思考

无需位置编码，轻量级的transformer架构，并有很强的零样本鲁棒性——SegFormer

8.Q&A