[Transformer]SETR:Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Tran


2021cvpr


From 复旦，牛津，萨里大学，腾讯优图实验室，Facebook

Abstract

近期的语义分割框架采用的都是基于编码器-解码器结构的全卷积神经网络（FCN)，encoder会逐渐降低空间分辨率，在更大的感受野范围内学习更抽象的语义特征。因为上下文建模对于分割至关重要，因此最新的工作就聚焦于如何增大感受野，比如通过空洞卷积或插入注意力模块等。但是基于编解码的FCN体系结构没有发生改变。



本文致力于采取一种新的视角，将语义分割任务视为seq2seq的序列预测任务。具体而言，本文采用了一个纯Transformer模型，没有任何卷积操作没有任何分辨率的压缩，将图像编码为一系列patch。


借助每一层transformer layer提取的全局信息，encoder可以结合较为简单的decoder就可以组成强大的分割模型，称之为SEgmentation TRansformer(SETR).本文在一系列实验（ADE20K，Pascqal,Cityscapes）上进行了实验均取得了优异的性能，截止提交时在ADE20K上排名第一。

Section I Introduction

目前分割网络的主流框架是FCN，FCN包含编码器和解码器部分，编码器负责学习图像特征，解码器负责对学习到的特征进行像素级别的预测。这一过程中特征表述是重中之重，就像大多数为了图像理解设计的CNN一样，encoder多由堆叠的卷积组成，学习到的特征分辨率也逐渐减少,encoder因此可以逐渐学习更加抽象的语义特征。


CNN有两大优势：平移不变性和局部性。前者遵循成像的基本性质，增强了模型对没见过图像的泛化能力；后者通过跨空间的参数共享可以降低模型的复杂程度。


但是CNN也有自身局限，即受限于有限的感受野无法有效的学习长程依赖关系，而这对语义分割、理解图像场景至关重要。
为了克服这一局限性，也进行了一系列改进，比如使用更大尺寸的卷积核、使用空洞卷积或特征金字塔等；另一思路是在FCN中加入注意力模块，借助注意力模块来聚合图特征图中所有像素点的依赖关系。


上述两类优化方法都没有改变FCN原始的编解码网络架构。也因此没有改变问题的本质：即卷积使得encoder逐渐降低输入分辨率，decoder再对压缩后的特征图逐步恢复到原始分辨率。


本文旨在从另一视角对待语义分割任务，通过将堆叠的卷积层组成的encoder替换为纯Transforemr结构，本文称之为SETR模型。Transformer encoder将输入图像看做patch序列，通过全局注意力对输入的序列进行学习。具体来说就是将图像分解为一系列固定大小的patch，随后将每一个patch展平并进行线性嵌入，将embedding后的结果作为transformer的输入，transformer encoder编码之后的结果送入decoder用于恢复到原始图像分辨率。关键点在于这一过程没有任何下采样操作来降低图像分辨率，同时transformer每一层都会对全局上下文进行建模。从而为语义分割问题提供了一个全新的视角。



这一纯Transformer框架灵感源自于其在NLP领域取得的成就。近期像ViT也证明了Transformer可以有效进行图像分类。也证明了图像特征提取不一定通过逐渐降低分辨率，从局部学习到全局。




但将Transformer从图像分类扩展到语义分割意义也十分重要，本文的SETR不仅为模型设计提供了一个全新的视角，而且在许多基准测试上均达到了SOTA。




本文的贡献总结如下：




（1）从序列预测角度重新看到语义分割任务，提出一种替代FCN编解码网络结构的替代方案；




（2）探索了将图像序列化后的结果作为特征表述；




（3）广泛验证了SA的特征表示，设计了三种不同复杂程度的decoder结构。




大量实验表明，与没有使用注意力机制的FCN相比，SETR可以学习更丰富的特征表述，在ADE20K,Pascal,Cityspace等数据集上均达到了SOTA。更是在ADE20K上排名第一。

Section II Related Work

语义分割

通过将分类网络中的分类层移除获得全卷积神经网络可用于语义分割，但是这样分割的结果比较粗糙；随后又有提出基于CRF/MRF等进行优化。为了更好的利用语义和位置信息，又提出了多尺度特征融合的编码-解码结构的不同变体。

目前的研究主要聚焦于FCN感受野有限的问题，比如借助空洞卷积的DeepLab，Dilation等；还有基于ASPP的PSPNet等；近期还有基于注意力机制的PSANet,CCNet,DGMN等。


Transformer



Transformer和self-attention机制在NLP领域引发了革命性的变革，也有诸多探索将Transformer用于图像领域。比如Non-local将attention加入到CNN基准网络中，AANet则是在训练时混合了卷积核SA操作；LRNet则是通过局部SA的计算降低注意力计算的复杂度；Axial-Attention则是将全局的注意力分解为两个方向的轴向注意力，也是为了降低计算量。DETR及其变体则是将transformer加载检测头之后进行目标检测。


STTR和LSTR则是分别用Transformer来做视差估计和车道线检测；ViT则是首个将Transformer用于图像分类的工作，也为在语义分割模型中使用Transformer提供了思路。

Section III Method

Part 1 FCN-Based Semantic segmentation

首先回顾一下FCN基本内容。FCN编码器部分由一系列卷积层堆叠而成，对于HxWx3的输入图像，feature map的分辨率会逐层下降，但当前层的感受野则会更大，但是感受野的扩大知识线性增长，因此需要堆叠较深层次才能获得较大感受野，才能建模长程依赖关系。但也有研究证明当网络达到一定深度，再堆叠带来的增益就会迅速减少，因此感受野受限是FCN的一大问题。

近期一系列改进方法提出可以将FCN与注意力机制结合，但是受限于计算复杂度，一般在较高层次也就是较小分辨率时才进行注意力的计算；但这样就缺失了底层细节信息，学习到的特征表述并不是最佳的。

为了克服这一问题，本文提出一种纯Transformer基于SA的编码器-SETR。



Part 2 SETR


Image to sequence

SETR遵循原始Transformer的输入输出格式，因此需要将2D图像转化为1D序列。参考Fig 1(a).


Fig 1展示了SETR的示意图。首先会将图像切分成固定大小的patch，然后对patch进行embedding，将嵌入后的patch序列作为Transformer encoder的输入。Decoder则通过渐进式上采样（PUP）将encoder的输出恢复至原始二维图像分辨率；最后通过多尺度特征聚合（MLA）获得最终的分割图。

序列化的一个简易方法就是将图像压缩成一个3HW的一维向量，但这样会使得序列长度太长，影响计算复杂福，因此将每一个像素作为transformer的输入是不可取的。

因此本文采取切patch的方法，这样将输入序列长度限制在HW/256,也就是将原图切分成16个patch，随后将patch经过linear embedding投影到1D空间 这样就将2D图像转换为1D序列；此外为了编码空间信息，还进行了位置嵌入。

Transformer

1D序列会送入Transformer encoder进行处理，每一层transformer layer都会学习全局的依赖关系，这样就解决了CNN感受野有限的问题。
Transformer encoder包含Le层MSA+MLP结构。注意力计算公式为：

细节在于LN是在MSA和MLP之前进行的。


Decoder



为了验证Transformer encoder提取特征的有效性，本文设计了三种decoder来产生最终的分割图。



（1）Naive upsampling-SETR Naive




原始上采样即使用双线性插值来恢复图像分辨率，然后通过交叉熵损失函数来获得像素级别的分类结果。





（2）Progressive UPsampling-PUP





但是通过一步上采样回去也许会引入噪声，因此本文还设计了一种渐进式的上采样策略替换常规的卷积-上采样操作。为了减少引入的噪声，将上采样倍数限制为2倍，因此需要 执行4次上采样才能恢复到原始分辨率。这一过程可以参考Fig 1(b)






（3）Multi-Level feature Aggregation-MLA






多尺度特征聚合类似特征金字塔模型，但是区别在于本文所使用的的所有feature map分辨率都是一样的，没有经过下采样操作。







从M层中选择feature map,有几层就有几流；对于每一stream都会将特征维度从2D变为3D，搭建一个3层的网络，分别使用1x1,3x3,3x3的卷积核，第一层和第三层分别将通道数减半，第三层还会经过双线性插值将空间分辨率提升4X。为了增强stream之间的信息交互，本文还在第一层之后引入了自上而下的多尺度特征聚合操作，会将特征进行element-wise addition；相加之后的特征还会经过一个3x3的卷积；第三层输出后会级联所有通道的输出然后恢复到原始图像分辨率。

Section IV Experiments

Part 1实验设定

Cityscapes


包含19类目标检测对象，training:testing:validation = 2975:500:1525 图像分辨率 2048x 1525



ADE20k




包含150类细粒度分割图像，trainging:validation:testing = 20210:2000:3352




PASCAL context




本文选取最常见的59类+背景共60类 training:validation = 4998:5105









Baselines




对比的基准模型是dialted FCN和Semantic FPN以及目前的一些SOTA模型
SETR 及其变体
本文设置了3种不同的decoder结构，encoder也设置了2种，分别是T-Base和T-Large，区别在于Transformer层数分别是12层和24层，详情参见Table 1.





此外还设置了一种混合模型，即FCN+Transformer，将基于ResNet-50的FCN提取的结果送入SETR中。
评价指标
评价指标为mIoU，ADE20K细粒度分割还额外比较了像素级别的accuracy.












消融实验






Table 2和Table 3分别展示了不同SETR变体和SETR模型与FCN模型预训练之后的性能比较。






从Table 2可以看到：PUP+SETR+Deit获得了最佳精度。






本文认为MLA性能没有达到最佳的原因可能是不同Transformer层输出的特征不像空间金字塔那样具有不同尺度的信息，详情可以参考Fig 5；但是MLA已经比原始SETR性能提升了很多。
而SETR Base性能比Hybrid-Base好也证明了可以将FCN编码部分在语义分割中进行替换。
第三点就是预训练之后确实会带来精度的进一步提升。






为了进一步探究预训练的作用，本文进一步进行了消融实验，详情参见Table3.






可以看到在更大规模数据集上预训练后FCN本身性能已经有明显提升；而本文的性能比FCN更加，也证实了本文结构的有效性来自于提出的seq2seq的建模策略，而不是预训练数据集的规模。















Comparison to SOTA








Table 4，5展示的是目前与SOTA网络的对比结果。Table 4是ADE20K上的结果，可以看到仅通过单尺度的推理就获得了48.64%的精度，采用多尺度特征后精度提升到50.28%.Fig 2展示的是ADE20K上的一个分割对比结果，右侧是SETR左侧是FCN，可以看到在一些细粒度特征的地方FCN会有明显的误分割现象，而SETR就没有。
Table 5是Pascal数据集的对比结果，可以看到SETR明显优于APCnet，Fig 3则可视化了SETR与Dialted FCN的对比结果，可以看到SETR可以更好的关注到具有语义的前景区域，说明可以学到更丰富的特征表述用于分割。








Table 6和Table 7则是Pascal数据集的结果，可以看到SETR+PUP的效果优于FCN及FCN+Attention类的网络。尽管Axial Deeplab训练更多epoch输入分辨率也更大但性能还是略逊于SETR。










Section V Conclusion

本文将语义分割任务转换为序列预测任务，不同于当前主流的扩大FCN有效感受野的改进策略，本文则是完全改变了FCN的基础架构，来解决感受野受限的问题。
通过Transformer在特征学习的每一层都可以学习全局上下文关联，结合不同复杂程度的decoder的设计，可以进行语义分割。在ADE20K，Cityscapes,Pascal数据集上均取得了优异的性能，并且目前在ADE20K比赛中排名第一。

Appendix
-Visualizations

Fig 7展示了position embedding的可视化结果，展示的是当前位置与其他所有patch之间的余弦相似度。



Fig 9展示的则是不同Transformer layer输出的结果。


Fig 10则可视化了SETR-PUP每一transformer层的注意力图谱；Fig 8则是选取了不同位置（红点）对应的attention map。