[Transformer] TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation

TopFormer：打造Arm端实时分割与检测模型，完美超越MobileNet!

CVPR2022

TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation

Paper: http://arxiv.org/pdf/2204.05525

Code: https://github.com/hustvl/TopFormer

1 Introduction

为了使ViT适应各种密集的预测任务，最近的ViTs，如PVT、CvT、LeViT以及MobileViT都采用了分层结构。这些ViTs将全局自注意力及其变体应用到高分辨率Token上，由于Token数量的二次复杂度，这带来了巨大的计算成本。

为了提高效率，Swin-Transformer、Shuffle-Transformer、Twins和HR-Former，都在计算在局部/窗口区域内的自注意力。然而，窗口分区在移动设备上是非常耗时的。此外，Mobile-Former通过减少Token的数量而降低了计算能力，但也牺牲了它们的精度。

在这些ViTs中，MobileViT和Mobile-Former是专门为移动设备设计的，但在移动设备上的实际延迟方面并没有显示出优势。

是否有可能设计出移动友好型网络，在移动语义分割任务上获得更好的性能？

在本文中，作者提出了一种移动端友好的架构，名为Token Pyramid Vision Transformer(TopFormer)。

构建了一个基于CNN的模块，称为Token Pyramid Module，用于处理高分辨率图像，以快速生成局部特征金字塔。

为了获得丰富的语义和较大的感受野，构建了一个基于ViT的模块Semantics Extractor，以Token作为输入。

将来自不同尺度（阶段）的Token池化到非常小的数字（分辨率），并沿着通道维度进行拼接。

然后，新的Token被输入到Transformer Block中以产生全局语义。

由于Transformer Block中的残差连接学习到的语义与Token的尺度有关，因此该模块被表示为Scale-aware Global Semantics。

所提出的Semantics Injection Module可以将Scale-aware Global Semantics注入到相应的Token中，构建强大的层次特征。

为了证明方法的有效性，在具有挑战性的分割数据集上进行了实验：ADE20K，Pascal和COCOStuff。

实验结果表明，TopFormer在多个语义分割数据集上显著优于基于CNN和ViT的网络，并在准确性和实时性之间取得了良好的权衡。

TopFormer的小版本在基于ARM的移动设备上实现实时推理，具有竞争性的结果。

2 Architecture

TopFormer的整体网络架构如图2所示。

网络由几个部分组成：

• Token Pyramid Module
• Semantics Extractor
• Semantics Injection Module
• Segmentation Head

Token Pyramid Module

由堆叠的MobileNet blocks组成，将一个图像作为输入，并生成Token Pyramid。

首先，通过一些MobileNetV2 Block产生一系列Token。

然后，将Token平均池化到目标大小。

最后，将来自不同尺度的Token沿着通道维度连接起来，产生新的Token。

新的Token将被输入Vision Transformer，以产生具有尺度感知的语义特征。

由于新的Token的数量较小，因此即使新的Token具有较大的通道，Vision Transformer也可以以非常低的计算成本运行。

Scale-aware Semantics Extractor

由几个堆叠的Transformer Block组成。Transformer Block数为L。

Transformer Block由Multi-head Attention module、FFN和残差连接组成。

将线性层替换为1×1的卷积层。

所有的非线性激活都是ReLU6，而不是GELU。

对于Multi-head Attention module，遵循LeViT的配置，将key K和query Q的Head尺寸设置为D=16，value V的head 设置为2D=32通道。

在计算Attention Map和输出时，减少K和Q的通道将降低计算成本。

同时，还去掉了Layer Normalization Layer，并向每个卷积添加了Batch Normalization。

在推理过程中，Batch Normalization可以与前面的卷积融合。

对于FFN，通过在2个1×1卷积层之间插入一个Depth-wise卷积，来增强Vision Transformer的局部连接。

将FFN的扩展系数设为2，以降低计算成本。

Vision Transformer将来自不同尺度的Token作为输入。

为了进一步减少计算量，使用平均池化操作将不同尺度的Token数量减少到输入大小的1/64。

concat之后作为Vision Transformer的输入。

Vision Transformer可以获得全图像的感受野和丰富的语义。

在通过几个Transformer Block后，获得了尺度感知语义。

Semantics Injection Module & Segmentation Head

Semantics Injection Module 来缓解在融合这些Token之前的语义差距。

以Token Pyramid Module的局部Token和Vision Transformer的全局语义作为输入。

局部Token通过1×1卷积层，然后进行批归一化，生成要注入的特征。

全局语义输入1×1卷积层 + 批归一化层 + sigmoid层产生语义权重，同时全局语义也通过1×1卷积层 + 批归一化。

这3个输出的大小相同。

通过哈达玛积将全局语义注入到局部标记中，并在注入后添加全局语义和特征。

经过Semantics Injection Module后，来自不同尺度的增强Token同时捕获了丰富的空间信息和语义信息，这对语义分割至关重要。

Segmentation Head首先将低分辨率Token上采样到与高分辨率Token相同的大小，

并按元素方式对所有尺度的Token进行sum up。最后，将该特征通过2个卷积层，生成最终的分割图。

消融实验

1、Token Pyramid的影响

如表所示，将来自不同尺度的堆叠Token作为语义提取器的输入，并将最后一个Token分别作为语义提取器的输入。

为了公平的比较，附加了一个1×1卷积层来扩展与堆叠的Token一样的通道。实验结果证明了使用Token Pyramid作为输入的有效性。

如表3所示，使用{1/4,1/8,1/16,1/32}的Token可以在最重的计算下获得最佳性能。

使用{1/16,1/32}的Token在最轻的计算下获得较差的性能。

为了在精度和计算成本之间实现良好的权衡，作者选择在所有其他实验中使用{1/8,1/16,1/32}的Token。

2、Scale-aware Semantics Extractor的影响

在这里使用不带SASE的Topformer作为基线。

加入SASE将带来约10%的mIoU收益，这是一个显著的改善。

为了验证Transformer Block中的多头自注意力模块(MHSA)，删除了所有的MHSA模块，并添加了更多的ffn，以进行公平的比较。

结果表明，在精心的架构设计下是一个高效有效的模块中MHSA可以获得约2.4%的mIoU收益。

同时，将SASE与流行的上下文模型进行了比较，如ASPP和PPM。

如表4所示，“+SASE”比“+PSP”和“+ASPP”可以以更低的计算成本获得更好的性能。

实验结果表明，SASE更适合用于移动设备。

PSP

ASPP

3、Semantic Injection Module和Segmentation Head的影响

将局部Token与Sigmoid层之后的语义相乘，表示为“SigmoidAttn”。

将语义提取器中的语义添加到相应的局部Token中，称为“SemInfo”。

与“SigmoidAttn”和“SemInfo”相比，同时添加“SigmoidAttn”和“SemInfo”通过一点额外的计算可以带来很大的改进。

在这里还讨论了Segmentation Head的设计。

将特征传递到Semantic Injection Module后，输出的层次特征具有较强的语义和丰富的空间细节。

提出的Segmentation Head简单地将它们相加，然后使用2个1×1卷积层来预测分割图。

作者还设计了另外2个分割头，如图所示。

Sum Head等同于只在SIM中添加SemInfo。

Concat Head使用1×1卷积层来减少SIM输出的通道，然后将特征拼接在一起。

4、SIM宽度的影响

如表所示，M=256,192,128通过非常接近的计算实现了类似的性能。

因此，在tiny, small和base模型中分别设置M=128,192,256。

5、output stride的影响

不同分辨率的结果如表所示。s32、s64、s128表示集合分辨率为输入大小的1/x。

考虑到计算量和精度的权衡性, 选择s64作为语义提取器的输入 Token 的输出 stride。

6、参数量与实时性 figure3

如图所示，虽然语义提取器具有大部分参数（74%），但语义提取器的FLOPs和实际延迟相对较低（约10%）。