《SegFormer：Simple and Efficient Design for Semantic Segmentation with Transformers》论文笔记

Posted 2022-08-24 m_buddy

参考代码：SegFormer

1.概述

介绍：这篇文章提出的分割方法是基于transformer结构构建的，不过这里使用到的transformer是针对分割任务在patch merge、self-attention和FFN进行了改进，使其更加适合分割任务（无需position-encoding，测试图片的尺寸带来的影响更小）对上下文语义信息和局部细节信息的需求并且更加整体模型轻量化，同时得益于transformer强大的上下文感知能力使得对于解码器只使用简单的几层全连接便可达到更好的分割性能。通观全文最大的亮点在分割网络编码器的设计上，若是将对应的解码器再进行进一步调试分割的效果将更好，不过使用如ASPP这样的感受野增强模块效果就不是那么大了。

将文章的方法于其它的一些分割方法进行比较，如下图。文章方法的分割性能更好且参数量更少。

2. 方法设计

2.1 pipeline

文章方法整体上还是一个编解码网络结构，使用新设计的编码器MiT-Bx系列和多层全连接组成，其结构见下图所示：

在上图中最重要的部分便是其中的transformer编码器了，不过这里针对分割任务设定了如下要求：

• 1）能输出多尺度的编码器特征，使得分割网络可以有效使用上下文语义信息和物体细节信息。对应在Swin Transformer中是通过类似于pixle-shuffle反向操作的形式实现多尺度输出的（可以参考这篇文章的解读：图解Swin Transformer）；
• 2）网络设计需要轻量化，本身分割任务就是一种稠密预测任务，轻量化是在实际部署过程中无法回避的问题；
• 3）原本的ViT中存在position-encoding这个在固定输入尺寸的情况下问题不大，但是对于分割这种对于分辨率影响比较大的任务，训练size和infer size不一样会导致position encoding存在问题，因而需要对position-encoding这块进行改进；

2.2 Transformer Block

该模块是文章编码器的基础组成模块，其中包含了：高效self-attention、Mix-FFN和Patch Merge模块，其结构见下图所示：

PS： 这里结构的顺序和文章代码中提供的顺序不太一致，代码中是patch merge在最前面。

patch merge：
在原本的ViT中是单独对每个patch独立进行处理，但是对于分割这类任务来说每个patch单独进行处理且之间无关联是不符合实际情况的，因而需要建立起不同patch之间的联系。在Swin Transformer中是通过Shifted Window Attention来实现不同window上信息的交互的，而在这篇文章中是通过卷积的形式，也就是将每个patch看作是一个卷积窗口，并以这个窗口在输入数据上进行滑动。同时还通过控制其中的stride参数控制输出特征图的大小，从而实现多尺度特征图的输出（卷积带padding）。

class OverlapPatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)

        self.img_size = img_size
        self.patch_size = patch_size
        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
        self.num_patches = self.H * self.W
        # 在输入数据上滑动得到patch并建立不同patch的关联
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
                              padding=(patch_size[0] // 2, patch_size[1] // 2))
        self.norm = nn.LayerNorm(embed_dim)

        self.apply(self._init_weights)
    ...

高效self-attention：
这篇文章中使用的attention是self-attention，对于输入特征$F\in R^{N\ast C},N=H\ast W$其计算量是很高，对此一个办法是通过映射操作减少channel数量。而这篇文章是在$N$维度上进行压缩的，而且对于不同stage上的特征图其压缩的比例是不一样的，也就是特征图大的压缩地更狠一点（这里的压缩操作是通过卷积中stride参数实现的）。

if sr_ratio > 1:
    self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
    self.norm = nn.LayerNorm(dim)

Mix-FFN：
通观文章的算法其中并没有出现position-encoding操作，而是在FFN中使用卷积的形式（带padding），文章对此的解释是卷积足以学习需要的位置信息，其实文章这样使用其观察是来自于这篇文章：

1. How Much Position Information Do Convolutional Neural Networks Encode?
2. NBJL 2020论文导读14：How Much Position Information Do Convolutional Neural Networks Encode ?

在该文章中得出的重要结论为：证明了卷积操作中位置信息来源于Padding。因而这里就将position-encoding操作给省去了，使用带padding参数的卷积进行隐式学习表征，而且这里还是用可分离卷积进一步减少参数和计算量。

class Mlp(nn.Module):  # Mix-FFN
    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)

        self.apply(self._init_weights)
    ...

2.3 分割解码器

对于解码器部分这篇文章使用的几层全连接实现的，其操作可以表达为：

3. 实验结果