TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation论文解读

前言

  本文试图在移动平台上实现低成本的transfomer计算来进行语义分割，并且实现多尺度融合。对于类似语义分割这种密集预测任务来说多尺度交互是很重要的，其他文章的多尺度融合一般是将编码器中的transfomer实现类似CNN的分层结构，然后将不同scale的特征图拿来进行融合。而这篇文章的创新点是通过金字塔结构将token多尺度化，然后将不同尺度的token进行融合。这样做既融合了多尺度的图像又因为H,W减小减小了计算量。再通过原始token与经过transfomer处理后的特征图进行融合，在分割的任务表现出了不错的效果。
论文地址：论文PDF地址
代码地址：github代码地址

0.摘要

  本文试图通过topfomer来打造一个移动友好型的基于transfomer的语义分割框架。通过token金字塔来显著减小输入的数据量，通过多尺度融合来保持语义分割的效果，从下图来看，蓝点是topfomer,红点是其他模型，圆的大小代表模型的大小。总的来说，topfomer在延时率和准确性之间取得了较好的平衡。（测试平台是移动平台上的高通骁龙865，没有其他类如GPU之类的加速器）

1.introduction

  视觉transfomer在一些视觉任务（类似image classification, object detection, semantic segmentation）上取得了不错的效果，但是attention机制带来的计算量实在是太大了，这大大超出了移动平台或者是嵌入式设备的计算能力。swin在窗口内计算注意力，但窗口分区在移动设备上也够耗时的。像是MobileViT 和 Mobile-Former是专门为移动平台打造的，但是他们的延时效率没有什么优势。本文的topfomer希望可以达到更好的延时率与准确度之间的平衡。
  topfomer吸收了CNN与transfomer结构上的优势，通过CNN常用的金字塔结构来把高分辨率图像快速生成特征，金字塔分为4层，使用下采样来构建金字塔，代码里面是用1x1Conv,步长取2来逐步实现H,W的成倍减小。再使用减小后融合的token作为trnasfomer模块的输入，这样计算量就大大减小了。transfomer模块负责提取出丰富的语义特征。

  本文认为从金字塔里面提取出来的特征都是局部特征，从transfomer里面出来的特征是全局语义。把两者进行融合是全局语义和局部语义的融合。
  最后看看本文的contribution：
1.将不同尺度的token作为输入，并将它们融合到非常小的scale,以极低的成本来获得语义。
2.Semantics Injection Module可以把语义注入token,实现局部特征和全局特征的融合
3.topformer的 base model 在 ARM移动设备上测试ADE20K 数据集，比MobileNetV3高 5% mIoU, tiny模型可以实现实时分割，分割效果也是比较有竞争力的。

2.related work

。。。

3.Architecture

3.1 Token Pyramid Module

  金字塔部分用的是MobileNets块的堆叠，特点就是一步stride=1,一步stride=2这样叠了4块相当于H,W下降16倍。因为之前conv_stem将H,W降为1/2，维度从3升到16，所以在融合前H,W降为1/32，通过平均池化融合后再降1/2，所以输入transfomer的时候H,W是1/64.因为H,W实在太小了，所以通道数即使拼接起来计算成本也很小。
  我们可以看到在输入transfomer的时候，H,W已经减小到原来的1/64，C根据config文件给出的参数来看，tiny模型4个阶段C分别为16，32，64，96.在送入transfomer前拼接起来通道数C为208。原始图像的通道数是3，那么数据量是从1 * 1 * 3=3到1/64 * 1/64 * 208=0.05078125，相当于数据量减小到1.69%，这对平方时间复杂度的attention可以说计算效率大大提升。

3.2 Vision Transformer as Scale-aware Semantics Extractor

  本部分是transfomer编码器部分，本文没有对attention结构进行修改，attention计算的时间复杂度也是平方。所谓语义提取器就是将数个transfomer块堆叠在一块，图上数量为L。q,k的shapeNxC维度C取16,v取32。attention计算中shape是这样变化的：N * C @ C * N -> N * N,N * N @ N * 2C->N * 2C（文中指定C=16)
  本文transfomer模块中另外几个改变是1x1卷积代替linear，ReLU6代替GELU，每个卷积后跟 batch normalization而不是Layer normalization 。

3.3 Semantics Injection Module and Segmentation Head

  从transfomer块里面出来的全局语义要与金字塔中的局部语义融合，但是它们存在着明显的语义差距，本文使用Semantics Injection Module来缓解它们之间的语义差距。
语义注入模块 (SIM) 将token金字塔模块的局部特征token和transfomer的全局特征图作为输入。 局部token我们把它暂时叫T，transfomer的全局特征图我们暂时叫F。
  1.T通过 1×1 卷积层，BN以生成特征。
  2.F通过一边 1×1 卷积层，BN， sigmoid 层，
  3.F通过 1×1 卷积层，BN。
  总共三个输出，具有相同的大小。
  这三个输出怎么融合呢，先是1和2做Hadamard 积，文中把他叫做把全局语义注入局部特征。。。然后是在注入之后将积与3相加。 文中也说语义注入减轻了token之间的语义差距。

  上图的黄×就是Semantics Injection Module，下方是它的实现细节：

  经过处理后局部语义和全局语义相乘再加上全局语义。这就是局部语义和全局语义融合的方法。这样我们获得了4张特征图。
  将多尺度特征图转化为分割结果也很简单。将4张特征图上采样到统一规格后相加，在经过1x1conv通道取物体class种类的数量即可。实际上就是将不同分辨率的token通过bilinear插值到与高分辨率token相同的大小，然后对所有特征图求和。

4.experiment

  本文使用mmsegmentation框架和pytorch，先是在imgenet上做了预训练，然后又在ade20k上做了160K次迭代训练，batch size设置为16，初始学习率设置为 0.00012，权重衰减为 0.01。 使用因子为 1.0 的“poly”方法。本文按照head[8,6,4]和金字塔最底层token channel[256,192,128]不同大小设计了3个模型，TopFormer-Tiny (TopFormer-T)、TopFormer-Small (TopFormer-S) 和 TopFormer-Base (TopFormer-B)。 可以看出topformer的准确率与延时是具有竞争力的。