[论文解读] EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers

Junting Pan, Adrian Bulat+6 more|arXiv (Cornell University)|May 6, 2022

Advanced Neural Network Applications被引用 14

一句话总结

EdgeViTs 提出了一类新型轻量级视觉变换器，通过整合一种结合自注意力与深度可分离卷积的局部-全局-局部（LGL）注意力机制，在移动设备上超越了当前最先进的高效CNN模型。该模型在真实硬件上的图像分类、目标检测和语义分割任务中，实现了帕累托最优的准确率-延迟与准确率-能耗权衡，甚至超越了此前唯一专为移动设备设计的ViT架构——MobileViT。

ABSTRACT

Self-attention based models such as vision transformers (ViTs) have emerged as a very competitive architecture alternative to convolutional neural networks (CNNs) in computer vision. Despite increasingly stronger variants with ever-higher recognition accuracies, due to the quadratic complexity of self-attention, existing ViTs are typically demanding in computation and model size. Although several successful design choices (e.g., the convolutions and hierarchical multi-stage structure) of prior CNNs have been reintroduced into recent ViTs, they are still not sufficient to meet the limited resource requirements of mobile devices. This motivates a very recent attempt to develop light ViTs based on the state-of-the-art MobileNet-v2, but still leaves a performance gap behind. In this work, pushing further along this under-studied direction we introduce EdgeViTs, a new family of light-weight ViTs that, for the first time, enable attention-based vision models to compete with the best light-weight CNNs in the tradeoff between accuracy and on-device efficiency. This is realized by introducing a highly cost-effective local-global-local (LGL) information exchange bottleneck based on optimal integration of self-attention and convolutions. For device-dedicated evaluation, rather than relying on inaccurate proxies like the number of FLOPs or parameters, we adopt a practical approach of focusing directly on on-device latency and, for the first time, energy efficiency. Specifically, we show that our models are Pareto-optimal when both accuracy-latency and accuracy-energy trade-offs are considered, achieving strict dominance over other ViTs in almost all cases and competing with the most efficient CNNs. Code is available at https://github.com/saic-fi/edgevit.

研究动机与目标

• 弥合轻量级视觉变换器（ViTs）与当前最先进的高效CNN在移动和边缘设备上的性能差距。
• 通过关注实际效率指标（如设备端延迟和能耗）而非代理指标（如FLOPs），解决现有ViTs在移动部署中的局限性。
• 设计一种计算高效且可使用标准深度学习框架（无需自定义算子）直接部署的ViT架构。
• 在保持移动硬件低推理成本的同时，实现在密集预测任务（目标检测、语义分割）上的竞争力表现。

提出的方法

• 提出一种结构化且成本高效的局部-全局-局部（LGL）信息交换瓶颈，将深度可分离卷积与自注意力机制有机结合。
• 将自注意力分解为局部、全局和重新组合三个阶段，以降低计算复杂度，同时保留长距离建模能力。
• 采用受CNN启发的分层多阶段架构，通过空间分辨率下采样控制序列长度，减少二次方复杂度。
• 引入深度可分离卷积以增强局部特征建模能力，并提升参数效率。
• 设计模型时确保可直接使用标准深度学习框架（如ONNX、TensorRT）部署，以保障实现友好性与高设备端推理速度。
• 在ImageNet-1K上训练模型，并在COCO和ADE20K上使用标准协议进行微调，以评估其在多样化视觉任务中的性能。

实验结果

研究问题

• RQ1 视觉变换器能否被优化到足够高效，使其在准确率和设备端推理效率方面超越或匹配最佳轻量级CNN？
• RQ2 将局部卷积整合到自注意力机制中，是否能显著提升效率而不牺牲移动硬件上的准确率？
• RQ3 在移动设备部署中，设备端延迟和能耗等实际指标，与FLOPs相比，作为模型效率预测指标的优劣如何？
• RQ4 基于ViT的架构能否在无需框架特定优化的情况下，既保持高准确率又可使用标准深度学习框架直接部署？
• RQ5 与CNN和通用ViT相比，视觉变换器在准确率-延迟与准确率-能耗权衡上，能达到何种程度的帕累托最优？

主要发现

• EdgeViT-XXS在ADE20K语义分割任务上达到39.7%的mIoU，优于PVTv2-B0 2.5个百分点，且在相似GFLOPs和参数量下表现更优。
• EdgeViT-S在ADE20K上达到45.9%的mIoU，较PVTv2-B1高出3.4个百分点，展现出在密集预测任务中的强大性能。
• 在COCO目标检测任务中，使用RetinaNet时，EdgeViT-XXS较PVTv2-B0高出1.5 AP；使用Mask R-CNN时，其掩码分割性能超出PVTv2-B0 1.7 AP。
• EdgeViT-XXS在三星Galaxy S21上推理耗时仅33.3ms，快于EfficientNet-B0的52.1ms，尽管其GFLOPs更高，表明其设备端效率更优。
• EdgeViTs在准确率-延迟与准确率-能耗权衡上全面超越其他ViT模型，能耗仅略高于EfficientNet-B0。
• 所有评估指标下，模型均达到帕累托最优，全面超越此前唯一专为移动设备设计的ViT家族——MobileViTs，在所有基准测试和效率指标上表现更优。