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

picture.image

虽然ViT在计算机视觉方面取得了巨大的成功，但巨大的计算成本阻碍了它们在密集的预测任务上的应用，如在移动设备上的语义分割。

在本文中，作者提出了一种移动端友好的架构，名为Token Pyramid Vision Transformer ( TopFormer)。所提出的最优算法以不同尺度的 Token作为输入，产生尺度感知的语义特征，然后将其注入到相应的 Token中，以增强表征。

实验结果表明， TopFormer在多个语义分割数据集上显著优于基于 CNN和 ViT的网络，并在准确性和实时性之间取得了良好的权衡。在ADE20K数据集上， TopFormer的mIoU比 MobileNetV3的延迟更高5%。此外， TopFormer的小版本在基于ARM的移动设备上实现实时推理，具有竞争性的结果。

开源地址：https://github.com/hustvl/TopFormer

背景

为了使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的计算核心。

picture.image 图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的移动设备上进行实时分割。

本文方法

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

Token Pyramid Module
Semantics Extractor
Semantics Injection Module
Segmentation Head

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

picture.image 图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。

picture.image 图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之前的语义差距。

picture.image 图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为目标通道数。各个版本的模型配置如下：

picture.image 模型配置表

实验

4.1 消融实验

1、Token Pyramid的影响

picture.image

如表所示，将来自不同尺度的堆叠 Token作为语义提取器的输入，并将最后一个 Token分别作为语义提取器的输入。为了公平的比较，附加了一个1×1卷积层来扩展与堆叠的 Token一样的通道。实验结果证明了使用Token Pyramid作为输入的有效性。

picture.image

如表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的影响

picture.image

结果如表所示。在这里使用不带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的影响

picture.image

如表所示，将局部 Token与 Sigmoid层之后的语义相乘，表示为“SigmoidAttn”。将语义提取器中的语义添加到相应的局部 Token中，称为“SemInfo”。与“SigmoidAttn”和“SemInfo”相比，同时添加“SigmoidAttn”和“SemInfo”通过一点额外的计算可以带来很大的改进。

在这里还讨论了 Segmentation Head的设计。将特征传递到 Semantic Injection Module后，输出的层次特征具有较强的语义和丰富的空间细节。提出的 Segmentation Head简单地将它们相加，然后使用2个 1×1卷积层来预测分割图。

picture.image