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

Wenqiang Zhang, Zilong Huang+6 more|arXiv (Cornell University)|Apr 12, 2022

Advanced Neural Network Applications被引用 16

一句话总结

TopFormer 提出了一种面向移动设备的视觉变换器架构，用于语义分割。该模型采用令牌金字塔模块提取多尺度特征，并通过语义注入机制将具备尺度感知能力的全局表征融合到局部令牌中。在移动设备上，其 mIoU 比 MobileNetV3 高出 5%，同时延迟更低；小型版本可在基于 ARM 的硬件上实现实时推理。

ABSTRACT

Although vision transformers (ViTs) have achieved great success in computer vision, the heavy computational cost hampers their applications to dense prediction tasks such as semantic segmentation on mobile devices. In this paper, we present a mobile-friendly architecture named extbf{To}ken extbf{P}yramid Vision Trans extbf{former} ( extbf{TopFormer}). The proposed extbf{TopFormer} takes Tokens from various scales as input to produce scale-aware semantic features, which are then injected into the corresponding tokens to augment the representation. Experimental results demonstrate that our method significantly outperforms CNN- and ViT-based networks across several semantic segmentation datasets and achieves a good trade-off between accuracy and latency. On the ADE20K dataset, TopFormer achieves 5\% higher accuracy in mIoU than MobileNetV3 with lower latency on an ARM-based mobile device. Furthermore, the tiny version of TopFormer achieves real-time inference on an ARM-based mobile device with competitive results. The code and models are available at: https://github.com/hustvl/TopFormer

研究动机与目标

• 设计一种轻量化、高效的视觉变换器架构，专为移动语义分割而优化。
• 解决标准变换器在资源受限设备上执行密集预测任务时计算成本过高的问题。
• 在基于 ARM 的移动 CPU 上，实现超越 MobileNetV3 的精度，同时降低延迟。
• 在不使用模型量化或 GPU 加速的情况下，实现在移动设备上的实时推理。
• 将该架构推广至其他密集预测任务，如目标检测。

提出的方法

• 令牌金字塔模块通过堆叠轻量级 MobileNetV2 块并快速下采样，处理高分辨率图像，生成多尺度特征令牌。
• 语义提取器通过平均池化将不同阶段的令牌聚合到极低分辨率（例如输入尺寸的 1/64），然后将其输入变换器块，以学习具备尺度感知能力的全局语义。
• 语义注入模块按通道分割具备尺度感知能力的全局语义，并将其重新融合到对应尺度的令牌中，以增强表征能力。
• 注入后的令牌被用作分割头的输入，实现具有最小计算开销的层次化特征学习。
• 变换器块中使用残差连接，以在全局语义中保留尺度感知能力。
• 模型采用交叉熵损失和 Dice 损失进行端到端训练，用于语义分割任务，并在 ADE20K、Pascal Context 和 COCO-Stuff 数据集上进行评估。

实验结果

研究问题

• RQ1 能否设计一种视觉变换器，在移动 CPU 上实现语义分割的高精度与低延迟？
• RQ2 如何高效地提取并注入具备尺度感知能力的全局语义到多尺度特征中，以提升表征能力？
• RQ3 混合 CNN-Transformer 架构能否在移动设备上同时超越 MobileNetV3 的精度与推理速度？
• RQ4 所提方法是否可泛化至其他密集预测任务，如目标检测？
• RQ5 面向移动部署，模型大小、FLOPs 与推理延迟之间的最优权衡是什么？

主要发现

• 在 ADE20K 验证集上，TopFormer-B 达到 45.28% 的 mIoU，仅需 1.25 GFLOPs，相比 MobileNetV3 提高 5% mIoU，且延迟更低。
• TopFormer-T 的小型版本在基于 ARM 的 Snapdragon 865 CPU 上实现实时推理，延迟为 110ms，计算量为 0.44 GFLOPs。
• 在 Pascal Context 数据集上，TopFormer-S 在 60 个类别上达到 43.68% mIoU，主干网络仅需 0.80 GFLOPs，头部网络仅需 0.18 GFLOPs，优于所有先前的 CNN 和 ViT 基础方法。
• 在 COCO-Stuff 数据集上，TopFormer-B 达到 33.43% mIoU，仅需 1.38 GFLOPs，比 MobileNetV3 高出 8%，且计算量相近。
• 在 COCO 数据集的目标检测任务中，使用 TopFormer-S 作为主干的 RetinaNet 达到 30.4% mAP，仅需 3.7 GFLOPs，优于 MobileNetV3 和 ShuffleNet，且 FLOPs 更低。
• 尽管语义提取器占总参数量的 74%，但其实际延迟仅占 10%，表明其计算效率极高。