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

虽然ViT在计算机视觉方面取得了巨大的成功，但巨大的计算成本阻碍了它们在密集的预测任务上的应用，如在移动设备上的语义分割。
在本文中，作者提出了一种移动端友好的架构，名为Token Pyramid Vision Transformer(TopFormer)。所提出的最优算法以不同尺度的Token作为输入，产生尺度感知的语义特征，然后将其注入到相应的Token中，以增强表征。
实验结果表明，TopFormer在多个语义分割数据集上显著优于基于CNN和ViT的网络，并在准确性和实时性之间取得了良好的权衡。在ADE20K数据集上，TopFormer的mIoU比MobileNetV3的延迟更高5%。此外，TopFormer的小版本在基于ARM的移动设备上实现实时推理，具有竞争性的结果。
开源地址：https://github.com/hustvl/Top...

1 背景

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

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

在这些ViTs中，MobileViT和Mobile-Former是专门为移动设备设计的。它们都结合了CNN和ViT的优势。在图像分类方面，MobileViT比与参数数量相似的MobileNets具有更好的性能。Mobile-Former在比MobileNets更少的FLOPs的情况下取得了更好的性能。然而，与MobileNets相比，它们在移动设备上的实际延迟方面并没有显示出优势。这提出了一个问题：是否有可能设计出移动友好型网络，在移动语义分割任务上获得更好的性能？

MobileViT和Mobile-Former的启发，作者也利用了CNN和ViT的优势。构建了一个基于CNN的模块，称为Token Pyramid Module，用于处理高分辨率图像，以快速生成局部特征金字塔。考虑到在移动设备上非常有限的计算能力，在这里使用一些堆叠的轻量级MobileNetV2 Block和Fast Down-Sampling策略来构建一个Token Pyramid。

为了获得丰富的语义和较大的感受野，作者还构建了一个基于ViT的模块，即Semantics Extractor，并将Token作为输入。为了进一步降低计算成本，使用Average Pooling Operator将Token减少到一个非常小的数字，例如，输入大小的1/(64×64)。

与ViT不同，T2T-ViT和LeViT使用嵌入层的最后一个输出作为输入Token，而TopFormer将来自不同尺度（阶段）的Token池化到非常小的数字（分辨率），并沿着通道维度进行拼接。然后，新的Token被输入到Transformer Block中以产生全局语义。由于Transformer Block中的残差连接学习到的语义与Token的尺度有关，因此该模块被表示为Scale-aware Global Semantics。

为了获得密集预测任务的强大层次特征，将尺度感知的全局语义通过不同尺度的Token通道进行分割，然后将标度感知的全局语义与相应的Token融合，以增强表示。增强的Token被用作分割Head的输入。

为了证明方法的有效性，在具有挑战性的分割数据集上进行了实验：ADE20K，Pascal上下文和COCOStuff。并测试了硬件上的延迟，即一个现成的基于Arm的计算核心。

图1

如图1所示，TopFormer比延迟较低的MobileNets获得了更好的结果。为了证明方法的泛化性,还在COCO数据集上进行了目标检测实验。

综上所述，本文的贡献如下：

所提出的最优预测算法以不同尺度的Token作为输入，并将Token池化到非常小的尺寸，以获得计算代价非常轻的Scale-aware Global Semantics;
所提出的Semantics Injection Module可以将Scale-aware Global Semantics注入到相应的Token中，构建强大的层次特征;
与MobileNetV3相比，TopFormer可以实现5%的mIoU提升，在ADE20K数据集上基于Arm的移动设备上的延迟更低。TopFormer-Tiny可以在基于Arm的移动设备上进行实时分割。

2 相关工作

2.1 轻量化Vision Transformers

其实以及有很多工作对Vision Transformers结构在图像识别中的应用进行了探索。ViT是第一个将纯Transformer应用于图像分类的工作，实现了最先进的性能。随后，DeiT引入了基于Token的蒸馏，以减少训练Transformer所需的数据量。T2T-ViT通过递归地将相邻Token聚合为一个Token来减少Token长度。Swin-Transformer在每个局部窗口内计算自注意力，带来了输入Token数量的线性计算复杂度。然而，这些Vision Transformers和后续工作往往是大量的参数和沉重的计算复杂度。

为了构建一个轻量级的Vision Transformer，LeViT设计了一个混合架构，使用堆叠的标准卷积层和stride=2来减少Token的数量，然后附加一个改进的Vision Transformer来提取语义。在分类任务中，LeViT在CPU上的性能明显优于EfficientNet。

MobileViT也采用了相同的策略，并使用MobilenetV2 Block而不是标准的卷积层来对特征图进行降采样。Mobile-Former采用并行结构与双向桥，利用了MobileNet和Transformer的优势。然而，有研究表明MobileViT和其他ViT-based的网络在移动设备上明显慢于MobileNets。

对于分割任务，输入的图像总是高分辨率的。因此，ViT-based网络比MobileNets的执行速度更具挑战性。本文中的目标是设计一个轻量级的视Vision Transformer，同时改模型可以超过MobileNets，以实现实时的分割任务。

2.2 高效的卷积神经网络

对在移动和嵌入式设备上部署视觉模型的需求不断增加，鼓励了对高效卷积神经网络设计的研究。MobileNet提出了一种inverted bottleneck的结构，该结构主要是叠加了Depth-wise和Point-wise卷积。IGCNet和ShuffleNet使用通道Shuffle/Permutation操作，为多个组卷积层进行跨Group信息流。GhostNet使用更简单的操作符，即Depth-wise卷积，来生成更多的特性。AdderNet利用add来替换大量的乘法。MobileNeXt翻转了反向残差块的结构，并呈现了一个连接高维表示的构建块。EfficientNet和TinyNet研究了深度、宽度和分辨率的复合尺度。

2.3 移动端语义分割

最精确的分割网络通常需要数十亿个FLOPs的计算，这可能会超过移动设备和嵌入式设备的计算能力。为了加快分割速度和降低计算成本，ICNet使用多尺度图像作为输入，并使用级联网络来提高计算效率。DFANet利用一个轻量级的Backbone来加速其网络，并提出了一种跨级特征聚合来提高精度。SwiftNet使用横向连接作为经济有效的解决方案，在保持速度的同时恢复预测分辨率。BiSeNet引入了Spatial path和Semantic path来减少计算。AlignSeg和SFNet对齐了来自相邻level的特征映射，并使用特征金字塔框架进一步增强了特征映射。ESPNets通过将标准卷积分解为Point-wise convolution和Spatial pyramid of dilated convolution来节省计算。

3 本文方法

TopFormer的整体网络架构如图2所示。网络由几个部分组成：

Token Pyramid Module
Semantics Extractor
Semantics Injection Module
Segmentation Head

Token Pyramid Module将一个图像作为输入，并生成Token Pyramid。

图2.TopFormer的整体架构

3.1 Token Pyramid Module

受MobileNets的启发，所提出的Token Pyramid Module由堆叠的MobileNet blocks组成。与MobileNets不同，Token Pyramid Module的目标并不是获得丰富的语义和较大的感受野，而是使用更少的块来构建Token Pyramid。

如图2所示，把一个图像作为输入，其中3，H，W分别表示RGB通道，高度，宽度；

Token Pyramid Module:

首先，通过一些MobileNetV2 Block产生一系列Token，N表示Scale的数量。
然后，将Token平均池化到目标大小，例如，。
最后，将来自不同尺度的Token沿着通道维度连接起来，产生新的Token。新的Token将被输入Vision Transformer，以产生具有尺度感知的语义特征。

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

3.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，遵循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，以降低计算成本。Transformer Block的数量是L。

图3.Scale-aware Semantics Extractor

如图3所示，Vision Transformer将来自不同尺度的Token作为输入。为了进一步减少计算量，使用平均池化操作将不同尺度的Token数量减少到输入大小的。来自不同尺度的集合Token具有相同的分辨率，它们被连接在一起作为Vision Transformer的输入。Vision Transformer可以获得全图像的感受野和丰富的语义。

更具体地说，全局自注意力在空间维度的Token之间交换信息。1×1卷积层将在来自不同尺度的Token之间交换信息。在每个Transformer Block中，在交换来自所有尺度的Token信息后学习残差映射，然后将残差映射添加到Token中，以增强表示和语义。最后，在通过几个Transformer Block后，获得了尺度感知语义。

3.3 Semantics Injection Module and Segmentation Head

在获得尺度感知语义，直接将它们与其他Token 相加。然而，在Token和尺度感知语义之间存在着显著的语义差距。为此，引入了Semantics Injection Module来缓解在融合这些Token之前的语义差距。

图4.Semantics Injection Module和Segmentation Head

如图4所示，Semantics Injection Module(SIM)以Token Pyramid Module的局部Token和Vision Transformer的全局语义作为输入。
局部Token通过1×1卷积层，然后进行批归一化，生成要注入的特征。
全局语义输入1×1卷积层 + 批归一化层 + sigmoid层产生语义权重，同时全局语义也通过1×1卷积层 + 批归一化。

这3个输出的大小相同。然后，通过阿达玛生产将全局语义注入到局部标记中，并在注入后添加全局语义和特征。几个sim的输出共享相同的通道数，记为M。

经过Semantics Injection Module后，来自不同尺度的增强Token同时捕获了丰富的空间信息和语义信息，这对语义分割至关重要。此外，Semantics Injection Module还缓解了Token之间的语义差距。所提出的Segmentation Head首先将低分辨率Token上采样到与高分辨率Token相同的大小，并按元素方式对所有尺度的Token进行sum up。最后，将该特征通过2个卷积层，生成最终的分割图。

3.4 架构及其变体

为了定制各种复杂的网络，作者设计了TopFormer-Tiny (TopFormer-T)和TopFormer-Small(TopFormer-S)和TopFormer-Base(TopFormer-Base)。

下表给出了Base、Small以及Tiny模型的尺寸和FLOPs。Base、Small以及Tiny模型在每个Multi-Head self-attention module中分别有8、6和4个Head，以M=256、M=192和M=128为目标通道数。各个版本的模型配置如下：

模型配置表

4 实验

4.1 消融实验

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更适合用于移动设备。

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卷积层来预测分割图。