全新ViT Backbone | PLG-ViT 同时具有并行局部和全局自注意力的轻量化视觉Transformer-阿里云开发者社区

近年来，Transformer架构在计算机视觉领域占据主导地位。然而，这些模型不容易部署在具有有限资源的硬件上，特别是针对需要实时性能的自动驾驶任务。它们的计算复杂度和内存需求限制了它们的使用，尤其是在高分辨率输入的应用中。

在本文的研究中，本文重新设计了功能强大的最先进的Vision Transformer PLG-ViT，将其转化为更紧凑高效的架构，适用于这类任务。

本文识别出原始PLG-ViT架构中计算代价高昂的模块，并提出了多种重新设计方案，旨在减少参数量和浮点运算。通过本文的重新设计，本文能够将PLG-ViT的大小缩小5倍，性能略有下降。

本文提出了2个变体，针对参数数量与运行时间以及参数数量与准确性之间的最佳平衡进行优化。仅使用500万参数，在ImageNet-1K分类基准测试中，本文实现了79.5%的top-1准确率。本文的网络在像COCO实例分割这样的通用视觉基准测试中表现出色。

此外，本文进行了一系列实验，展示了本文的方法在解决专门针对自动驾驶和交通挑战的各种任务方面的潜力。

1、简介

视觉Transformer（ViT）作为卷积神经网络（CNNs）的一种可行替代方案的出现，源于多头自注意力机制的成功应用。与标准卷积相比，多头自注意力机制提供了全局感受野。

然而，虽然ViT在各种视觉识别任务中展示了潜力，但自注意力的计算复杂性在将这些方法应用到资源受限的应用中带来了挑战。对于具有有限资源的硬件上的图像处理，特别是在自动驾驶领域，具有大量参数和高计算复杂度的Transformer模型不适用于需要实时或接近实时性能的任务。通常，在资源受限的平台上使用轻量级CNN进行视觉识别任务。

最近，已经出现了一些轻量级Transformer设计的初步尝试，例如「MobileViT」和「EdgeViT」。这些网络采用混合架构，将CNN与Transformer结合起来，以减少复杂性和内存，并展现了在移动应用中使用Transformer的巨大潜力。此外，对MobileViT的持续改进显示出这种潜力尚未得到充分利用。

在本文中，本文提出了一种轻量级的PLG-ViT（Light-Weight PLG-ViT），它是一种计算效率高但功能强大的特征提取器，适用于资源受限的硬件。

本文基于通用的PLG-ViT构建，因为其并行的局部-全局注意力使其成为适合硬件有限平台的理想选择。

具体而言，利用分辨率自适应窗口进行全局注意力的独特概念在低计算复杂性和高准确性之间保持了良好的平衡。尽管这种特定的设计选择是有益的，但PLG-ViT的其他部分仍然需要大量资源。

在一项综合研究中，本文通过创新的方法识别和替代计算效率低下的部分，以提高网络的计算效率，同时减少参数数量和内存需求。

具体而言，这些创新包括在Transformer阶段之间进行高效降采样操作以及Transformer层内部的新型、更强大的前向网络。最后，本文通过对通用视觉基准以及智能交通领域的4个基准进行全面评估，展示了「LW PLGViT」的卓越性能（见图1）。

2、本文方法

2.1 原始PLG-ViT

原始PLG-ViT架构中的分层编码器结构由5个阶段组成，其中第一阶段包括一个卷积层，作为重叠的 Patch 嵌入。这个分词层将图像分辨率降低了4倍，并将输入投影到一个C维嵌入空间中。

接下来的四个Transformer阶段包括PLG自注意力力（PLG-SA）和一个卷积前馈网络，每个阶段之前都有一个层归一化处理。每个Transformer阶段之间通过卷积下采样分隔，它将特征的分辨率减半，同时增加深度。

PLG-SA通过沿特征深度分割输入，生成局部和全局特征。通过减少特征图，可以节省计算复杂性和参数，从而减少计算量。局部特征经过窗口自注意力力处理，而全局特征则通过自适应 Patch 采样操作进行处理，该操作结合了最大池化和平均池化。

这个操作创建一个具有固定 Token-size 的单个全局窗口，用于计算自注意力。自适应 Patch 采样的属性使得PLG-ViT在增加输入分辨率时具有计算效率。在对全局特征应用自注意力之后，使用双线性插值恢复原始分辨率。最后，将局部和全局特征连接起来，并在自注意力层后面添加一个前馈神经网络(FFN)。

2.2 模型设计

PLG-ViT的架构已经具有分辨率无关的高效性，这也是本文重新设计以适应较弱硬件的决定性论据。

然而，PLG-ViT的一些结构并没有考虑到硬件受限应用。需要重新设计的第一个网络部分是卷积 Patch 嵌入，在每个Transformer阶段之前应用。原始的 Patch 嵌入是Fused-MBConv层，可以描述为：

其中，LN代表Layer Normalization，GeLU代表Gaussian Error Linear Units，SE代表Squeeze and Excitation。表示带有 Stride 2的3×3卷积层。

在本文重新设计的轻量级 Patch 嵌入中，保留了该层的第一部分，只是进行了些许修改。仅将GeLU操作替换为性能更好的SiLU操作，并且将SE层的通道减少从1/4增加到1/8。

在公式1的第二部分中，本文将 Stride 卷积替换为一个1×1的层，以调整空间（ Stride 为2）和通道维度，并使用一个深度可分离的3×3层生成高级特征。所有这些改变可以节省近80%的原始块参数。本文的新 Patch 嵌入的描述如下：

本文的第二个重大改变发生在卷积前馈网络（FFN）中，该网络紧跟在PLG-SA层之后，如图2所示。

在图3（a）中，展示了原始PLG-ViT的CCF-FFN。在原始FFN的开头，特征按照默认值α进行了深度增加，其中α默认为4。然后是一个深度可分离的3×3层和一个全连接层（FC），恢复了原始的深度。

在本文的重新设计过程中，本文为CCF-FFN添加了一些新的特性，生成了本文的新CCF-FFN+（参见图3（b）），一方面略微减少了参数数量，另一方面也提高了性能。

本文的CCF-FFN+保留了原始FFN的一般结构，并将来自初始逐点卷积的特征馈送到后续的深度可分离卷积中。然而，这些特征还与从深度可分离的3×3卷积中得到的特征进行了拼接。这种融合可以在减小α值（从4减小到3）的同时实现更丰富的特征，从而减少参数数量。

此外，本文在深度可分离卷积之后利用ConvNeXt-v2的全局响应归一化（GRN）来提高模型的特征表示能力，减少训练的计算成本。

本文的最后一次重新设计针对的是PLG-SA层全局分支中的 Patch 采样。原始的采样通过添加自适应最大池化和平均池化来获得一个全局窗口，为后续的全局自注意力提供了固定数量的 Token。

对于较大的网络变体，这种池化特征的添加与准确性的提高相关。然而，本文观察到，对于本文优化后的变体，这种情况并不成立，因此本文发现省略「平均池化有利于减少内存消耗和浮点运算次数」（FLOPs）。

2.3 架构变体

原始的PLG-ViT提出了3个变体，它们在参数数量和浮点运算次数（FLOPs）上有所不同。由于本文的目标是开发一个更小、更轻量的变体，本文将最小的模型PLG-ViT Tiny作为baseline。

首先，参数数量将会减少。这主要通过减小每个阶段的特征图深度来实现。在第一个Transformer阶段，PLG-ViT Tiny的特征深度为，并且在接下来的每个阶段中翻倍。由于特征不断翻倍，参数数量迅速增加。

为了解决这个问题，本文的轻量级变体中的特征不会翻倍，而是可以单独设置，类似于MobileNet。在最后一个阶段，本文添加一个3×3的层来进行通道扩展，以获得更丰富的特征。

在原始的PLG-ViT中，局部自注意力力和全局自注意力头的数量总是相同的。然而，本文观察到由于固定的全局窗口，全局自注意力的性能更好。

因此，本文决定解耦局部自注意力和全局自注意力，并允许不同数量的并行 Head 。这样，本文可以通过使用更多的全局 Head 来优先执行全局自注意力。这也可以在图2中看到。LW PLG-SA的输入形状为，其中H和W表示特征的空间维度，C表示特征深度。与原始实现不同，本文沿着特征深度将输入z分割，并生成局部特征和全局特征。这里的描述了局部窗口和全局窗口之间的比例，可以单独设置。在原始的PLG-SA中，始终设置为0.5。

在本论文中，本文提出了2种网络配置：精确度优化的轻量级PLG-ViT-A和运行时优化的轻量级PLG-ViT-R。这两个网络总共有500万个参数，但在FLOPs的数量上有所不同。每个模型的配置可以在表格I中找到。

在表格中，本文列出了 Stride s，通道深度C以及每个阶段Transformer设置的更多信息。这里和表示全局自注意力和局部自注意力。大括号中的第一个数字表示窗口大小，第二个数字表示Head的数量。方括号后面的数字表示块的重复次数。