深度学习论文: 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}，N表示Scale的数量。
- 然后，将Token平均池化到目标大小，例如R^{\frac{H}{64} * \frac{W}{64}}。
- 最后，将来自不同尺度的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：
