论文阅读笔记：EfficientViT: Lightweight Multi-Scale Attention for On-Device Semantic Segmentation

HollowKnightZ

已于 2024-05-30 11:21:39 修改

阅读量1.5k

点赞数 35

分类专栏：语义分割深度学习论文阅读文章标签：论文阅读深度学习计算机视觉

于 2024-05-30 09:45:37 首次发布

本文链接：https://blog.csdn.net/Z960515/article/details/139311699

版权

深度学习论文阅读同时被 2 个专栏收录

32 篇文章 3 订阅

订阅专栏

语义分割

13 篇文章 0 订阅

订阅专栏

论文阅读笔记：EfficientViT: Lightweight Multi-Scale Attention for On-Device Semantic Segmentation

1 背景
2 创新点
3 方法
4 模块
- 4.1 基于ReLU的全局注意力
- 4.2 多尺度Tokens

代码：https://github.com/mit-han-lab/efficientvit

论文：https://arxiv.org/pdf/2205.14756

1 背景

SOTA语义分割模型所需要的计算成本和资源受限的边缘设备之间存在巨大差距。这使得在边缘设备上部署这些模型变得不切实际。同时，语义分割是一个密集的鱼粉任务，需要高分辨率的图像和强大的上下文信息提取能力才能很好的工作，因此直接从图像分类中移植搞笑的模型架构不适用于语义分割。

本文介绍了一种新的设备语义分割模型系列EfficientVIT。EfficientViT的核心是一个新颖的轻量级多尺度注意力模块，能够实现全局感受野核多尺度学习，并具有高效的硬件效率。本文提出的模块是收到先前的SOTA语义分割模型的启发，即多尺度学习和全局感受野在提高语义分割性能方面起关键作用，然而，先前的方法没有考虑硬件效率。

因此，模块的设计原则是启用这两个关键功能，同时避免硬件低效问题。具体来说，为了拥有一个全局的感受野，作者提出轻量级的基于ReLU的全局注意力来代替低效的自注意力。通过利用矩阵乘法的结合率，基于ReLU的全局注意力可以保持功能性的同时，将计算复杂度从二次降低到线性。此外，它避免了softmax等硬件低效的操作，使其更适用于设备上的语义分割。

此外，作者基于ReLU-based全局注意力提出了一种新的轻量级多尺度注意力模块，将附近的tokens与小核卷积聚合生成多尺度tokens，并对多尺度token执行基于ReLU的全局注意力，将全局感受野核多尺度学习相结合。在表1中总结了本文工作和之之前SOTA语义分割模型之间的比较。可以看到本文提出的模型比之前的模型更适合于设备上的语义分割。
在这里插入图片描述
在Cityscapes和ADE20K上对本文方法进行了评估。与之前的SOTA方法相比，EfficientViT提供了显著的性能提升。同时在图像分类任务上也表现了强大的性能。在ImageNet上没有精度损失的情况下，比EfficientNet实现了7.9倍的延迟降低。如图1。
在这里插入图片描述

2 创新点

提出了一种新颖的轻量级多尺度注意力机制用于设备上的语义分割。它实现了全局感受野和多尺度学习并在边缘设备上保持了良好的效率。
基于提出的轻量级多尺度注意力模块设计了一个新的模型系列EfficientViT。
在流行的语义分割基准数据集和ImageNet上，EfficientViT在移动端上比之前的SOTA语义分割模型显示出显著的加速比。

3 方法

本文提出的EfficientViT的模块由一个轻量级的多尺度注意力模块和一个MBConv组成。轻量级的MSA模块负责捕获上下文信息，而MBConv负责捕获局部信息。在多尺度注意力模块中先通过线性投影层得到Q/K/V的tokens后，再使用轻量级小核心卷积聚合附近的tokens来生成多尺度tokens。将基于ReLU的全局注意力应用于多尺度tokens，并将输出串联后送入最终的线性投影层进行特征融合。如图2。
在这里插入图片描述

MBConv是MobileNetV2中提出的块，是具有深度可分离卷积的导致线性瓶颈层。如下图。

4 模块

4.1 基于ReLU的全局注意力

给定输入 $x∈R^{N×f}$ ，自注意力的一般形式可以写为：
在这里插入图片描述

其中， $Q=xW_Q，K=xW_K，V=xW_v$ ， $W_Q/W_K/W_V∈R^{f×d}$ 是可学习的线性映射矩阵。 $O_i$ 表示矩阵 $O$ 的第 $i$ 行， $S im (\cdot,\cdot)$ 是相似函数。当相似函数 $Sim(Q,K)=exp(\frac{QK^T}{\sqrt{d}})$ 时，式（1）就变成了原始的自注意力。

除了 $exp(\frac{QK^T}{\sqrt{d}})$ ，我们可以使用其他的相似度函数。本文使用基于ReLU的全局注意力来实现全局感受野的线性计算复杂度。在基于ReLU的全局注意力中，相似度函数定义为：
在这里插入图片描述

当相似度函数定义为式（2）时，式（1）可被重写为：
在这里插入图片描述

然后，我们可以利用矩阵乘法结合律，在不改变功能的情况下，将计算复杂度和内存占用从平方减少到线性：
在这里插入图片描述

在式（3）中，只需要计算一次 $(\sum_{j=1}^NReLU(K_j)^TV_j)∈R^{d×d}$ 和 $(\sum_{j=1}^NReLU(K_j)^T)∈R^{d×1}$ ，就可以为每个查询重用他们，从而只需要 $O (N)$ 的计算成本和 $O (N)$ 的内存成本。

基于ReLU的全局注意力的另一个关键优点是它不涉及想softmax这样的硬件不友好操作。图3展示了softmax注意力和基于ReLU的线性注意力之间的延迟对比。在移动端上，相似的计算量下，基于ReLU的线性注意力明显快于softmax注意力。
在这里插入图片描述

4.2 多尺度Tokens

只使用基于ReLU的注意力，模型的容量是有限的。为了增强基于ReLU的全局注意力机制的多尺度学习能力，作者提出将邻近的Q/K/V tokens信息进行聚合，得到多尺度tokens。聚合过程如图5所示。这个信息聚合对于每个头中的Q，K，V都是独立的，只使用小尺寸卷积核进行信息聚合，以避免损害硬件效率。
在这里插入图片描述

在实际实现中，独立执行这些聚合操作在GPU上是低效的，因此作者使用组卷积来减少总的运算次数。将所有的DWConv融合成单个DWConv，将所有1×1卷积合并成为单个1×1组卷积，如下图，其中组数为 3 × #heads，每组通道数为d。

这里3指的是Q，K，V这3个输入，#heads指的是Q，K，V的head数。Q，K，V共3×Dims的通道，先用DWConv进行group=3×Dims的通道可分离卷积，然后进行group= 3 × #heads的组卷积，相当于把Q，K，V在通道上分成了#heads组，每组 Dims/#heads 的维度。所以这里的多尺度是是指如下图用不同大小的卷积核进行不同尺寸的局部信息交互。

这里的具体实现可见代码：https://github.com/mit-han-lab/efficientvit/blob/master/efficientvit/models/nn/ops.py#L371
self.aggreg = nn.ModuleList(
	[
		nn.Sequential(
			nn.Conv2d(
				3 * total_dim,
				3 * total_dim,
				scale,
				padding=get_same_padding(scale),
				groups=3 * total_dim,
				bias=use_bias[0],
			),
			nn.Conv2d(3 * total_dim, 3 * total_dim, 1, groups=3 * heads, bias=use_bias[0]),
		)
		for scale in scales
	]
)


![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/9fd41259ba334496a4743e5db1257190.png#pic_center)

得到多尺度特征后，对其进行全局关注以提取多尺度全局特征。最后，将不同尺度的特征沿着头部维度进行拼接，并送到最终的线性投影层进行融合。

## 4.3 EfficientViT结构

基于提出的轻量级MSA模块，作者构建了一个模型系列。其中核心块EfficientViT Module如下图。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/c4daabbbde7d482db37dc7c922515763.png#pic_center)

具体来说，EfficientViT模块包括一个轻量级的MSA模块和一个MBConv 。轻量级的MSA模块用于上下文信息的提取，而MBConv用于局部信息的提取。

EfficientViT的宏观架构如图4所示，采用标准的backbone-head/encoder-decoder设计。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/fd8e540e832744af8113fa0660d75f1c.png#pic_center)

**Backbone**：由输入层和四个阶段组成，特征图大小逐渐减小，通道数量逐渐增加。在阶段3和4中插入EfficientViT模块。对于下采样，使用步幅为2的MBConv。

**Head**：P2、P3和P4表示阶段2、3和4的输出，形成特征图的金字塔。为了简单和高效，使用1x 1卷积和标准上采样操作（例如，双线性/双三次上采样）以匹配它们的空间和信道大小并经由加法来融合它们。简单的头部设计，其包括若干MBConv块和输出层（即预测和上采样）。

除了语义分割外，模型还可以通过将主干和任务特定的头相结合，应用于其他视觉任务，如图像分类。遵循相同的宏观架构，作者设计了一系列不同尺寸的模型来满足效率约束，如表2。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/6b173ae347d24ea98b8bf7d49dbd515b.png#pic_center)
# 5 效果

## 5.1 消融实验

作者在Cityscapes上进行了消融研究实验，以研究轻量级MSA模块的两个关键设计组件的有效性，即多尺度学习和全局注意力。为了消除预训练的影响，实验从随机初始化中训练所有模型。结果汇总于表3。可以看到，去除全局注意力或多尺度学习都会显著地损害性能。这表明，所有这些对于在性能和效率之间取得更好的折衷是至关重要的。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/39855135f2db42578f5a9db001d13ee6.png#pic_center)

## 5.2 和SOTA方法对比

为了理解EfficientViT的骨干网络在图像分类中的有效性，在ImageNet上训练对比模型效果并将模型与SOTA图像分类模型进行了比较，如表4。耗时为CPU运行耗时。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/53f42b3b80444fc5b970ae458e802349.png#pic_center)
图6为各SOTA方法之间指标和乘加累积操作数Macs的对比。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/866994946f874118b3a58bde26ca0a03.png#pic_center)

表5为CityScapes上的比较结果，1024\*2048的输入尺寸下，EfficientVIT-B0比SeformerB0快了整整14倍，并且指标稍高。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/f7200e3b10dc44d38b3ef405ebd321eb.png#pic_center)


表6为在ADE20K上的比较结果，720\*960的输入尺寸下，EfficientViT-B1比Segformer B1要快6.5倍，并且指标稍高。
![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/72bdee0ef369471d8b209e642448cb9b.png#pic_center)
# 6 结论

EfficientViT可用作基础网络和其他方法相结合。