深度学习论文: TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation及其PyTorch实现

深度学习论文: TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation及其PyTorch实现
TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation
PDF: https://arxiv.org/pdf/2204.05525.pdf
PyTorch代码: https://github.com/shanglianlm0525/CvPytorch
PyTorch代码: https://github.com/shanglianlm0525/PyTorch-Networks

1 概述

提出的Token Pyramid Vision Transformer(TopFormer)最优算法以不同尺度的Token作为输入，产生尺度感知的语义特征，然后将其注入到相应的Token中，以增强表征。

2 TopFormer

TopFormer的整体网络架构如图

2-1 Token Pyramid Module

Token Pyramid Module将一个图像作为输入，并生成Token Pyramid。Token Pyramid Module的目标并不是获得丰富的语义和较大的感受野，而是使用更少的块来构建Token Pyramid。

• 首先，通过一些MobileNetV2 Block产生一系列Token, T 1 , T 2 , . . . , T N T^{1},T^{2},...,T^{N} T1,T2,...,TN，N表示Scale的数量。
• 然后，将Token平均池化到目标大小，例如 R H 64 ∗ W 64 R^{\frac{H}{64} * \frac{W}{64}} R64H​∗64W​。
• 最后，将来自不同尺度的Token沿着通道维度连接起来，产生新的Token。新的Token将被输入Vision Transformer，以产生具有尺度感知的语义特征。

2-2 Scale-aware Semantics Extractor

Scale-aware Semantics Extractor由几个堆叠的Transformer Block组成。Transformer Block数为L。

• Transformer Block由Multi-head Attention module、FFN和残差连接组成。
• 为了保持Token的空间形状和减少重塑的数量，这里将线性层替换为1×1的卷积层。
• 此外，在ViT中，所有的非线性激活都是ReLU6，而不是GELU。

Multi-head Attention module：

• 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：

• 对于FFN，通过在2个1×1卷积层之间插入一个Depth-wise卷积，来增强Vision Transformer的局部连接。将FFN的扩展系数设为2，以降低计算成本。

2-3 Semantics Injection Module

Semantics Injection Module 用来缓解在融合不同尺度Token 之间的语义差距。

• 局部Token通过1×1卷积层，然后进行批归一化，生成要注入的特征。
• 全局语义输入1×1卷积层 + 批归一化层 + sigmoid层产生语义权重，同时全局语义也通过1×1卷积层 + 批归一化。
  

2-4 Segmentation Head

Segmentation Head首先将低分辨率Token上采样到与高分辨率Token相同的大小，并按元素方式对所有尺度的Token进行sum up。最后，将该特征通过2个卷积层，生成最终的分割图。

3 Experiments

ADE20K:

Cityscapes：