EfficientViT | 边缘设备上实时语义分割

EfficientViT: Lightweight Multi-Scale Attention for On-Device Semantic Segmentation

• paper https://arxiv.org/pdf/2205.14756.pdf
• code https://github.com/mit-han-lab/efficientvit

demo

Abstract

EfficientViT is a new family of vision models for efficient high-resolution dense prediction. The core building block of EfficientViT is a new lightweight multi-scale linear attention module that achieves global receptive field and multi-scale learning with only hardware-efficient operations.

• 在针对高分辨率视觉应用时，ViT 不如卷积神经网络 (CNN)。ViT 的关键计算瓶颈是 softmax 注意力模块，其计算复杂度与输入分辨率成二次方。降低 ViT 的成本以将其部署在边缘设备上至关重要。现有方法（例如 Swin、PVT）将 softmax 注意力限制在局部窗口内或降低键/值张量的分辨率以降低成本，从而牺牲了 ViT 在全局特征提取方面的核心优势。
• EfficientViT是一种用于高分辨率低计算视觉识别的高效 ViT 架构。我们建议用线性注意力代替softmax注意力，而不是限制softmax注意力，同时通过深度卷积增强其局部特征提取能力。EfficientViT 保持全局和局部特征提取能力，同时享受线性计算复杂度。
• EfficientViT是一个全新的轻量级多尺度注意力的新语义分割模型系列。与之前的语义分割模型不同，这些模型依赖于沉重的自我关注、硬件效率低下的大核卷积或复杂拓扑结构来获得良好性能，而我们的轻量级多尺度注意力只需轻量级和高效的硬件操作，就能实现全局感受野和多尺度学习（语义分割模型的两个关键特征）。
• 在基准数据集上，EfficientViT 比以往最先进的语义分割模型有显著的性能提升，并在移动平台上显著提速。在 Cityscapes 上，我们的 EfficientViT 在不损失性能的情况下，在移动平台上的性能分别提高了 15 倍和 9.3 倍。和 SegNeXt 相比，移动延迟分别减少了 15 倍和 9.3 倍。和 SegNeXt 相比，移动延迟分别减少了 15 倍和 9.3 倍。在保持相同移动延迟，EfficientViT 在 ADE20K 上比 SegNeXt 增加了 7.4 mIoU

Contributions

• 将注意力机制softmax平方复杂度降低到线性复杂度，同时引入深度卷积增强局部特征提取能力

image-20230919102753047

Related Work

• ViT-based model
• 设备端高分辨率语义分割

Methodology

Overview

Lightweight Multi-Scale Attention

image-20230919102314552

• 虽然线性注意力在计算复杂度和硬件延迟方面优于softmax注意力，但线性注意力有局限性。以前的研究 表明，在 NLP 中，线性注意力和 softmax 注意力之间通常存在显著的性能差距。对于视觉任务，之前的工作也表明线性注意力不如 softmax 注意力。
• 线性注意力的劣势主要是由于局部特征提取能力的损失。如果没有 softmax 注意力中使用的非线性分数归一化，线性注意力很难像 softmax 注意力那样集中其注意力分布。给定相同的原始注意力分数，使用 softmax 比不使用 softmax 更能集中注意力。因此，线性注意力不能有效地集中在局部模式产生的高注意力分数，削弱了其局部特征提取能力
• 我们的想法是通过卷积来增强线性注意力，这在局部特征提取中非常有效。这样，我们就不需要依赖线性注意力来捕获局部特征，而可以专注于全局特征提取。具体来说，为了保持线性注意力的效率和简单性，我们建议在每个 FFN 层中插入一个深度卷积（DWConvGconv），这会产生很少的计算开销，同时大大提高线性注意力的局部特征提取能力。
• 展示了增强型线性注意力的详细架构，它由一个线性注意力层和一个 FFN 层组成。将深度卷积插入 FFN 的中间。与之前的方法 不同，我们在 EfficientViT 中不使用相对位置偏差。虽然相对位置偏差可以提高性能，但它使模型对分辨率变化很脆弱。多分辨率训练或新分辨率下的测试在检测和分割中很常见。去除相对位置偏差使 EfficientViT 对输入分辨率更加灵活。与以前的低计算 CNN中的设计不同，我们为下采样块添加了额外的下采样shortcuts。每个下采样shortcut由一个平均池化和一个 1x1 卷积组成。在我们的实验中，这些额外的下采样shortcuts可以稳定 EfficientViT 的训练并提高性能。

EfficientViT Architecture

image-20230919102354273

• 它由输入stem和 4 个阶段组成。最近的研究 表明，在早期阶段使用卷积对 ViT 效果更好。我们遵循这个设计，并在第 3 阶段开始使用增强的线性注意力。为了突出高效的骨干本身，我们使用相同的扩展比 e 为 MBConv 和 FFN (e = 4) ，以保持超参数简单，所有深度卷积的核大小 k 相同（k = 5，除了输入stem），以及所有层的相同激活函数（hard swish）。P2、P3 和 P4 表示阶段 2、3 和 4 的输出，形成特征图金字塔。我们按照惯例将 P2、P3 和 P4 馈送到检测头。我们使用 YoloX 进行检测。对于分割，我们融合了 P2 和 P4。在 Fast-SCNN之后，融合的特征被馈送到包含多个卷积层的轻量级头部。对于分类，我们将 P4 馈送到轻量级头部，与 MobileNetV3相同。

Experiments

image-20230919104023799

Conclusions

• 主要是为语义分割任务设计的模型，但在图像分类任务ImageNet上不仅精度高，而且速度快 （不是SOTA，只是比较高而已，Top1 acc排名100左右了）
• 针对高分辨率场景，降低了attention的复杂度
• 针对边缘设备端分割模型的推理性能优化，达到实时推理的目的
• 金标准（Conv+Pseudo-MHSA 卷积加多头自注意力机制）的问题，至少可以说明两点，一就是卷积的局部性+等变性的归纳偏置加MHSA的长距离表示能力；二是硬件和编译软件和运行软件对CNN支持的较好，对transformer还需要更多的支持

References