《TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation》

期刊：CVPR

年份：2022

代码：https://github.com/hustvl/TopFormer

摘要

尽管视觉Transformer(ViTs)在计算机视觉领域取得了巨大的成功，但沉重的计算成本阻碍了它们在密集预测任务中的应用，如移动设备上的语义分割。在本文中，我们提出了一种名为 TokenPyramid Vision Transformer (TopFormer) 的移动友好架构。提出的TopFormer以不同尺度的令牌作为输入，生成尺度感知的语义特征，然后将其注入到相应的令牌中以增强表示。实验结果表明，我们的方法在多个语义分割数据集上显着优于基于 CNN 和 ViT 的网络，并在准确性和延迟之间取得了良好的平衡。在ADE20K数据集上，TopFormer在基于ARM的移动设备上的延迟较低的情况下，mIoU的准确率比MobileNetV3高5%。此外，TopFormer 的微小版本在基于 ARM 的移动设备上实现了实时推理，并取得了有竞争力的结果。

Introduce

主要贡献：

所提出的 TopFormer 将不同尺度的标记作为输入，并将标记汇集到非常小的数字中，以获得具有非常轻计算成本的尺度感知语义。
所提出的语义注入模块可以将尺度感知语义注入到相应的标记中，以构建强大的层次特征，这对密集预测任务至关重要。
所提出的基础模型可以比 MobileNetV3 获得更好的 5% mIoU，在 ADE20K 数据集上基于 ARM 的移动设备上的延迟较低。微小的版本可以在基于arm的移动设备上执行实时分割，并具有具有竞争力的结果。

Related Work

1.轻量级Vision Transformer：

讨论了将Transformer结构应用于图像识别的多种探索，包括ViT、DeiT、T2T-ViT、Swin Transformer和LeViT等。
指出这些Transformer通常参数量大，计算复杂度高，不适合移动设备。

2.高效的卷积神经网络（CNN）：

强调了在移动和嵌入式设备上部署视觉模型的需求，促进了对高效CNN设计的探索。
提到了MobileNet、IGCNet、ShuffleNet、GhostNet、AdderNet、MobileNeXt和EfficientNet等网络，它们通过不同的结构优化来提高效率。

3.移动语义分割：

讨论了在移动设备上进行语义分割的挑战，以及如何通过不同的方法来加速分割过程并降低计算成本。
提到了ICNet、DFANet、SwiftNet、BiSeNet、AlignSeg、SFNet、ESPNets、AutoML技术和NRD等方法。

4.MobileViT和Mobile-Former：

特别提到了MobileViT和Mobile-Former，这两种架构是为移动设备特别设计的，结合了CNN和ViT的优势。
指出尽管MobileViT在图像分类任务上表现优于MobileNets，但在移动设备上的实际延迟并没有显示出优势。

5.TopFormer的设计灵感：

作者们受到MobileViT和Mobile-Former的启发，利用CNN和ViT的优势，提出了Token Pyramid Module和Semantics Extractor。
描述了如何使用轻量级MobileNetV2块和快速下采样策略构建Token Pyramid，以及如何使用Semantics Extractor获取丰富的语义和大的感受野。

Method

网络由几个部分组成：令牌金字塔模块（Token Pyramid Module）、语义提取器（Semantics Extractor）、语义注入模块（Semantics Injection）和分割头（Segmentation Head）。

Token Pyramid 模块将图像作为输入并生成令牌金字塔。
Vision Transformer 被用作语义提取器，它将令牌金字塔作为输入并产生尺度感知语义。
语义被注入到相应尺度的标记中，以增强语义注入模块的表示。
Segmentation Head 使用增强的令牌金字塔来执行分割任务。

3.1 Token Pyramid Module

Token Pyramid Module (TPM) 的主要作用是将输入图像转换成一系列不同尺度的Tokens，形成一个Token金字塔结构，为后续的语义提取和注入提供基础。以下是TPM的详细介绍：

1. 结构组成：
- TPM 由多个堆叠的轻量级 MobileNetV2 块组成，这些块被设计为具有不同的输出尺寸，以便快速生成不同分辨率的特征表示。

2. 处理流程：
- 输入图像首先通过 MobileNetV2 块进行处理，每个块的输出将作为下一个块的输入，这样逐层深入，逐步降低图像的空间分辨率，同时增加特征通道数。
- 随着处理的深入，生成的Tokens尺寸逐渐减小，但同时能够捕捉到更抽象的特征表示。

3. 尺度生成：
- 通过不同配置的MobileNetV2块，TPM能够产生多个尺度的特征图（Tokens），例如，可能包括 1/4、1/8、1/16 和 1/32 原始图像尺寸的特征表示。

4. Token 池化：
- 一旦生成了不同尺度的Tokens，TPM 将这些Tokens通过平均池化操作进一步减少到目标尺寸，例如，将所有Tokens池化到 1/64 的输入图像尺寸。

5. 通道拼接：
- 池化后的Tokens沿通道维度进行拼接，形成一个统一的特征张量，这个张量将作为Semantics Extractor的输入。

6. 设计优势：
- 由于TPM使用轻量级的MobileNetV2块，它在计算上非常高效，适合于移动设备上的应用。
- 通过构建Token金字塔，TPM能够为后续的Transformer模块提供多尺度的特征信息，有助于捕捉到不同层次的语义内容。

Token Pyramid Module 是 TopFormer 能够实现高效语义分割的关键，它通过多尺度特征提取和轻量化设计，为移动设备上的深度学习应用提供了一个实用的解决方案。

3.2 Semantics Extractor

主要任务是提取和生成具有尺度意识的语义特征。

1.组成结构：

SE由多个Transformer块堆叠而成，每个Transformer块包含多头自注意力（Multi-head Self-Attention, MHSA）模块、前馈网络（Feed-Forward Network, FFN）和残差连接。

MHSA：对于多头注意力模块，我们遵循 LeViT的设置，并将键 K 和查询 Q 的头部维度设置为 D = 16，值 V 的头部有 2D = 32 个通道。在计算注意力图和输出时，减小 K 和 Q 的通道将降低计算成本。同时，我们还删除了层归一化层，并在每个卷积上附加批量归一化。在推理过程中，批归一化可以与前面的卷积合并，比层归一化运行得更快。

FFN：通过在两个1×1卷积层之间插入深度卷积层来增强Vision Transformer的局部连接。FFN 的扩展因子设置为 2 以降低计算成本。

Vision Transformer 可以获得全图像感受野和丰富的语义。更具体地说，全局自注意力在空间维度的令牌之间交换信息。1×1卷积层将在不同尺度的令牌之间交换信息。在每个 Transformer 块中，残差映射是在从所有尺度交换标记信息后学习的，然后将残差映射添加到标记中以增强表示和语义。最后，通过几个Transformer 块后获得尺度感知语义。

3.3 Semantics Injection Module

作用是将 Semantics Extractor (SE) 提取的尺度意识语义信息有效地注入到 Token Pyramid Module (TPM) 生成的局部 Tokens 中，以此来增强这些 Tokens 的语义表示，为最终的语义分割任务提供丰富的特征支持。以下是 SIM 的详细介绍：

主要目标：SIM 的目标是缩小局部 Tokens 和全局语义之间的语义差距，并通过融合这些信息来增强 Tokens 的表示能力。

处理流程：
- 局部 Tokens 通过 1x1 卷积层和批量归一化层（BatchNorm）处理，以生成待注入的特征。
- 全局语义特征同样通过 1x1 卷积层和 BatchNorm 层，然后通过 Sigmoid 函数生成语义权重。

语义注入：
- SIM 通过 Hadamard 积（逐元素乘积）将语义权重与局部 Tokens 结合，实现语义注入。
- 注入后的语义信息不仅与局部 Tokens 融合，还通过加法操作进一步与原始的全局语义特征相结合。

通过将全局语义信息有效地注入到局部 Tokens 中，SIM 增强了模型对不同尺度特征的表示能力，从而在保持计算效率的同时，提高了语义分割的性能。

3.4 Segmentation Head

在语义注入之后，不同尺度的增强标记同时捕获丰富的空间和语义信息，这对于语义分割至关重要。此外，语义注入减轻了令牌之间的语义差距。所提出的分割头首先将低分辨率标记上采样到与高分辨率标记相同的大小，并对来自所有尺度的标记进行元素求和。最后，特征通过两个卷积层来生成最终的分割图。

3.5 Architecture and Variants

为了定制各种复杂性的网络，分别介绍了TopFormer-Tiny(TopFormer-T)、TopFormer-Small(TopFormer-S)和TopFormer-Base(TopFormer-B)。

base模型、small模型和tiny模型的模型尺寸和FLOPs如表所示。base模型、small模型和tiny模型在每个多头自关注模块中分别有8头、6头和4头，目标通道数分别为M = 256、M = 192和M = 128。

Conclusion

本文提出了一种新的移动视觉任务架构。结合 CNN 和 ViT 的优势，所提出的 TopFormer 在准确性和计算成本之间取得了很好的平衡。TopFormer 的微小版本可以在具有竞争性结果的基于 ARM 的移动设备上产生实时推理。实验结果表明了所提出方法的有效性。TopFormer的主要限制是目标检测的微小改进。我们将继续提高目标检测的性能。此外，我们将在未来的工作中探索 TopFormer 在密集预测中的应用。