Nips2021(Segmentation) - 高效的语义分割transformer《SegFormer》

原文地址

原文

论文阅读方法

三遍论文法

初识

本文的工作主要是Transformer在语义分割领域的应用，虽然CVPR21已经有SETR（基于ViT）了，但是其本身还是存在不少问题的：① ViT作为backbone只能输出固定分辨率的特征图，这对于密集预测任务显然不够友好；② 由于self-attention操作的存在，transformer的运算量和参数两都非常大，不利于大尺度图像的分割。

SegFor从两个方面解决问题：1）首先是对transformer进行层次化结构设计，得到多层级的特征图；2）构造轻量级的decoder，仅使用MLP进行特征聚合。

除此之外，SegFormer抛弃了位置信息编码，选择采用MixFCN来学习位置信息，这样可以很好地扩充到不同尺度的测试环境下（避免由于尺寸变化，需要对positional-encoding进行插值，从而影响性能）。最后提出的模型在ADE20k上达到了新sota，并且在速度、性能和鲁棒性上都表现很好。

相知

只介绍主要方法和部分实验

网络结构如上图所示，给定一张HxWx3的图像，先分成4x4大小的patch序列（ViT使用16x16，更小的patch对dense prediction越有益），然后送入SegFormer的编码部分得到不同层级的特征，然后再解码部分融合多层次特征得到最终结果。

下面再分别介绍几个主要模块：

Transformer Encoder

为了得到像CNN那样的层次特征，每个Transformer stage之后都会进行patch融合（起到2x2下采样的效果）。

overlapped Patch Merge: 但如果按照ViT中的融合策略直接进行reshape操作的话（$2\times 2\times C_i$直接reshape成$1\times 1\times C_{i+1}$），这样会损失patch之间的空间一致性。所以作者采用overlapping patch merging，对得到的序列特征进行有重叠的融合操作。看了代码才发现，实际上就是采用卷积操作进行局部特征融合同时下采样，有两种设置K=7,S=4,P=3和K=3，S=2,P=1。

网络中所有的分块与融合均采用上述策略，虽然开始说是分成4x4的patch块，但感受野其实不只4x4。通过卷积操作，增大感受野的同时，也能学习到领域patch间的位置关系。

Efficient Self-Attention： 之前也提到Transformer中的self-attention操作计算消耗很大，因此作者进行了改进：将K和V的特征维度进行缩减。文章中提到的做法就是先将K进行Reshape，然后再用Linear做一个特征映射。

这样就会把复杂度从O($N^2$)缩减至O($N^2/R$)

文章中没有提到V，但实际上transformer是要保证Q和K的特征维度一致，K和V的数量一致。在代码实现中，作者仍然是采用了卷积的方法。

Mix-FFN: 之前也提到，SegFormer抛弃了位置编码模块，而是采用Mix-FFN(代替了Transformer Block中的FFN模块)来学习位置信息。主要是采用3x3卷积来学习leak location information，并且是采用的depth-wise卷积来进一步减少参数和计算量。实验证明了其有效性。

All-MLP Decoder

相比于其他的工作，SegFormer的Decoder显得十分清爽。提取得到encoder阶段每个层级的输出，然后经过MLP操作后，再进行融合，主要包含以下四步：

其实就是得到不同分辨率的特征后，先进行Linear转换到C维，再上采样到W/4, H/4。然后进行拼接，再经过一个Linear进行降维，最后用一个Linear学习类别映射，从之前的图也比较好理解。

整体网络结构

在B0的基础上进行缩放，得到B1-B5

讨论与实验

Effective Receptive Field： 作者可视化了SegFormer和DeepLabV3+的有效感受野，如下图所示：

可以看到SegFormer在前期的stage中主要还是捕获局部信息，后面的stage会捕获到更大的non-local信息；MLP head得到的有效感受野有着更强的局部attention。

DeepLabv3+的有效感受野很小，即使在stage-4阶段也是如此。因此其还需要额外的模块来增大感受野，比如ASPP。而SegFormer受益于self-attention操作，有着天然的感受野优势，并且在编码阶段也结合了不同层级的输出（可以视为一种local和non-local的特征融合）。

实验
作者在Cityscapes，ADE20K以及COCO-stuff上进行了消融实验，以及对比了sota方法。

此外，作者还进行了鲁棒性的测试，在Cityscapes-C数据集(包含16种不同的噪声扰动)，报告了结果。

视频结果可以参看B站链接：https://www.bilibili.com/video/BV1MV41147Ko/
这个效果还是比较惊讶的，在一定程度上说明了Transformer建模的鲁棒性非常强

回顾

SegFormer在语义分割上的应用非常成功，这也证明了在这类dense prediction任务上，感受野非常重要。

与SETR（直接应用ViT作为encoder）不同，SegFormer首先构建了层次化的backbone，得到类似resnet的不同分辨率的特征（可用于解码阶段的融合）；其次就是引入了更适合语义分割的模块：比如更细粒度的patch划分，overlapped的patch融合策略，以及抛弃位置编码采用Conv3x3捕获leak位置信息。这不仅提升了性能，也减少了参数量与计算量。

transformer作为backbone本身就能使得模型能够捕获长距离依赖关系，使得最后输出的特征已经融合了non-local信息。所以在编码阶段，作者采用了简单的融合机制就能达到较好的效果。

同时这也说明，transformer在CV的不同下游任务的应用中，网络结构的调整也非常重要，期待未来更多更好的研究出现。

此外，还可以参考作者本人的知乎文章，写得很好：https://zhuanlan.zhihu.com/p/379054782

代码

官方代码已开源：https://github.com/NVlabs/SegFormer/tree/13bb922f26032c257f3d9bdfe67522c2dd36fcb6

如果不熟悉mmsegmetation的话，看起来还是比较费劲的，可以参看这个简洁版复现：https://github.com/camlaedtke/segmentation_pytorch