SETR 论文笔记

Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers

SETR

论文链接: https://arxiv.org/abs/2012.15840

一、 Problem Statement

传统的CNN网络在语义分割方面通常是encoder-decoder结构，通过encoder逐渐降低分辨率的大小，增大感知域来获取特征信息。而decoder用于对encoder生成的特征表示进行像素级分类。 这样的设计存在两个好处就是translation equivariance 和 locality。前者能够加强模型的泛化能力，而后者控制了模型的复杂度。但是存在一个局限性就是: 学习对无约束场景图像中的语义分割至关重要的long-range dependency信息，由于仍然有限的感受野而变得具有挑战性。

二、 Direction

提出SEgmentation TRansformer(SETR)，即使用纯粹的transformer来替由代堆叠卷积层组成的encoder。 也就是encoder-decoder变成了sequence-to-sequence。

三、 Method

先来看一下整体的网络结构:

可以看出来，网络结构整体分为三个部分: Image to sequence，encoder(transformer backbone) 和 decoder。Decoder作者提供了三种选择，分别是Naive upsampling， Progressive UPsampling和Multi-Level feature Aggregation(MLA)。

提供了两种模型，区别在于transformer layers的层数，hidden layer size 和 attention head数量。

1. Image to sequence

既然需要输入到transformer，需要把输入图像$x\in {\mathbb{R}}^{H\times W\times 3}$ 转换到 feature embeddings $Z\in {\mathbb{R}}^{L\times C}$。其中$L$是序列的长度，$C$是 hidden channel size。 通常，CNN会降采样输入图像16倍，作者遵循这个原则，设计了序列长度为$L=\frac{H}{16}\times \frac{W}{16}=\frac{HW}{256}$。所以输入图像，会被平均分割成$\frac{H}{16}\times \frac{W}{16}$个patches。然后使用linear projection將维度变为$C$。同时，添加一个position embeddings，最后形成了以下输入:

E = { e 1 + p 1 , e 2 + p 2 , . . . , e L + p L } E = \{e_1 + p_1, e_2+p_2,...,e_L+p_L\} E={e1​+p1​,e2​+p2​,...,eL​+pL​}

2. Transformer Backbone

Transformer encoder由$L_e$个transformer层组成，每个都是由multi-head self-attention(MSA)和Multilayer Perceptron(MLP) blocks组成。对于第$l$层的输入的(query, key, value)，都是来自于$Z^{l-1}\in {\mathbb{R}}^{L\times C}$层。
$$\text{query}=Z^{l-1}{W}_Q,\text{key}=Z^{l-1}{W}_K,\text{value}=Z^{l-1}{W}_V$$

其中$W_Q,W_K,W_V\in {\mathbb{R}}^{C\times d}$都是需要学习的参数。$d$是的(query, key, value)的维度。

self-attention计算为:
$$\text{SA}(Z^{l-1})=Z^{l-1}+\text{softmax}(\frac{Z^{l-1}{W}_Q(Z{W}_K{)}^T}{\sqrt{d}})(Z^{l-1}{W}_V)$$

Multi-head self attention为:
$$\text{MSA}(Z^{l-1})=[S{A}_1(Z^{l-1});S{A}_2(Z^{l-1});...;S{A}_m(Z^{l-1})]W_O$$
其中$W_O\in {\mathbb{R}}^{md\times C}$。$d$通常设置为$C/m$。

MSA的输出会传入到MLP中，

$$Z^l=\text{MSA}(Z^{l-1})+\text{MLP}(\text{MSA}(Z^{l-1}))\in {\mathbb{R}}^{L\times C}$$

最后，每一个transformer layers的输出特征表示如下:
{ Z 1 , Z 2 , . . . , Z L e } \{Z^1, Z^2, ... , Z^{L_e}\} {Z1,Z2,...,ZLe​}

3. Decoder

对于transformer的输出特征$Z$，作者使用了三种不同的decoder进行pixel-level的分割。对于一张2D图像的输入$(H\times W)$，我们需要reshape 特征$Z\in {\mathbb{R}}^{\frac{HW}{256}\times C}$到$\frac{H}{16}\times \frac{W}{16}\times C$。

1. Naive upsampling

最简单的decoder就是使用一个简单的网络:1x1 conv + sync batch norm(w/ ReLU) + 1x1 conv，把transformer feature $Z^{L_e}$ 映射到为目标类别数量的维度，比如说Cityscapes的是19。 之后，简单的做bilinearly upsample到full image resolution。然后使用pixel-wise cross-entropy loss进行分类。

2. Progressive UPsampling

第一种会引入noisy prediction。那第二种decoder就是逐步的进行上采样。这里面首先会把transformer feature $Z^{L_e}$进行reshape，然后通过四个卷积层，最后得到维度为$H\times W\times 19$的输出。

3. Multi-Level feature Aggregation

第三种类似于FPN的结构。由于所有的feature representations$Z^l$输出的维度都一样，无法像金字塔似的形状。因此，作者取相隔$\frac{L_e}{M}$层的输入进行融合，如上图所示。
{ Z m } ( m ∈ { L e M , 2 L e M , . . . , M L e M } ) \{Z^m\}(m \in \{\frac{L_e}{M}, 2\frac{L_e}{M},...,M\frac{L_e}{M}\}) {Zm}(m∈{MLe​​,2MLe​​,...,MMLe​​})
提出之后，把维度为$\frac{HW}{256}\times C$ reshape成维度为$\frac{H}{16}\times \frac{W}{16}\times C$。然后分别通过1x1, 3x3, 3x3的卷积层后，进行4X 的bilinear 操作得到feature maps。为了增强信息的流通，引入了top-down aggregation。每一个element-wise addition之后，都有一个3x3的卷积。最后每一层再进行一次4X的插值，得到full image resolution后进行concatenation。

四、 Conclusion

把encoder变成了transformer backbone， 表现SOTA。但是参数量增加挺多的，表现和DeepLab-v3(ResNet-101)差不多。

Reference