TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation 论文笔记

TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation

论文链接： https://arxiv.org/abs/2204.05525

一、 Problem Statement

是否能够设计出一个mobile-friendly Vision-Transformer的网络，能够比mobileNets有更低的延迟性，但能有更好的性能。

二、 Direction

结合MobileNet和vision transformer。

三、 Method

先来看一下网络结构:

主要包括以下几个部分:

• Token Pyramid Module (backbone)
• Scale-aware Semantics Extractor (Vision Transformer)
• Semantics Injection Module
• Segmentation Head

1. Token Pyramid Module(backbone)

先来看一下主要的backbone。这里主要是借鉴于mobileNetv2中的blocks。可以看到，输入一张图片$I\in {\mathbb{R}}^{3\times H\times W}$，会生成一系列的特征图$\{T^1,...,T^N\}$，这里的$N$代表尺度的数量。之后，这些不同尺度的特征图会通过average pooling，变成目标的大小${\mathbb{R}}^{\frac{H}{64}\times \frac{W}{64}}$。最后，这些不同层级的特征图，会拼接在一起，生成一个新的特征图。
这里作者做了几个比较实验:

• Token pyramid 和Token of the last scale
  

  可以看出多尺度的特征会带来性能的提升。

• 下采样的大小
  

  为了trade-off情况下，作者选择了stride=64。

2. SASE (Vision Transformer)

这里的SASE就是transformer blocks的堆叠了，数量为L。Transformer包括multi-head attention, Feed-Forward Network和residual connections。为了保持特征图的空间大小，同时减少reshape操作的数量，作者将linear layers用1x1卷积进行替换。同时用ReLU6替换GELU。

对于Multi-head attention模块，将keys K 和 queries Q设置为16通道，values V 维度为32通道。减少K和Q的通道数，可以减少计算attention maps时候的消耗。同时，使用batch normalization替换layer normalization。因为batch normalization可以在推理的时候和卷积一起合并，加速推理。

对于feed-forward network，在两个1x1卷积中，插入一个depth-wise convolution layer，提升vision transformer局部的关联。expansion factor设置为2，来减少计算消耗。

作者认为Vision Transformer可以获得full-image的感知域和丰富的语义信息，而不是在backbone中获得。

1. the Token Pyramid Module does not aim to obtain rich semantics and large receptive field, but uses fewer blocks to build a token pyramid.

2. The Vision Transformer can obtain full-image receptive field and rich semantics.

理由如下:

• global self-attention 在特征图中间交换了信息。
• 1x1卷积会交换不同尺度的特征图信息。
• 每一个transformer block，在来自所有尺度特征图交换信息后，residual mapping会被学习到，然后又相加到特征图，来加强了特征表示和语义信息。
• Scale-aware 语义可以通过几个transformer blocks后获得。

3. Semantics Injection Module

获得了scale-aware 语义之后，我们直接与局部的特征图相加。但是，这里会出现一个称之为 “semantic gap” 的现象。因此作者引入了Semantics Injection Module，来消除这个现象。

可以看到，每个尺度特征图会经过一个1x1卷积和batch normlization之后，与通过1x1 conv, BN和sigmoid激活函数后的global semantics相乘 (Hadamard production)。相乘后又相加。每个尺度特征图经过这个模块之后，都又同样的通道数。

4. Segmentation Head

Segmentation head首先会把低分辨率的特征图上采样到同样大小的特征图，然后进行element-wise相加。最后通过两个1x1卷积生成最后的segmentation map。

四、 Performance

五、 Conclusion

结合了MobileNet和tranformer blocks。效果能在ARM-based的设备上跑到实时。

六、 Reference