轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer

点击下方卡片,关注“自动驾驶之心”公众号

ADAS巨卷干货,即可获取

点击进入→自动驾驶之心【模型部署】技术交流群

后台回复【模型部署工程】获取基于TensorRT的分类、检测任务的部署源码!

摘要

随着视觉Transformers(ViTs)在计算机视觉任务中的成功,最近的技术试图优化ViT的性能和复杂性,以实现在移动设备上的高效部署。研究人员提出了多种方法来加速注意力机制,改进低效设计,或结合mobile-friendly的轻量级卷积来形成混合架构。然而,ViT及其变体仍然比轻量级的CNNs具有更高的延迟或更多的参数,即使对于多年前的MobileNet也是如此。实际上,延迟和大小对于资源受限硬件上的高效部署都至关重要。在这项工作中,论文研究了一个中心问题,ViT模型是否可以像MobileNet一样快速运行并保持类似的大小?论文重新审视了ViT的设计选择,并提出了一种具有低延迟和高参数效率的改进型超网络。论文进一步引入了一种细粒度联合搜索策略,该策略可以通过同时优化延迟和参数量来找到有效的架构。所提出的模型EfficientFormerV2在ImageNet-1K上实现了比MobileNetV2和MobileNetV1高约4%的top-1精度,具有相似的延迟和参数。论文证明,适当设计和优化的ViT可以以MobileNet级别的大小和速度实现高性能。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第1张图片

总结来说,本文的贡献如下:

  • 论文提供了一项全面的研究,以验证和改进移动友好的设计选择,这是获得超高效视觉主干的实用指南;

  • 论文提出了一种细粒度联合搜索算法,该算法同时优化了模型大小和速度,实现了卓越的Pareto最优;

  • EfficientFormerV2系列模型实现了超快的推理和极小的模型大小,大大超过了以前的技术,并且可以在各种下游任务中充当强大的骨干。

相关工作

Vaswani等人[72]提出了对NLP任务中的序列建模的注意力机制,这形成了transformer架构。Dosovitskiy等人[18]和Carion等人[6]后来将transformer用于视觉任务。DeiT[68]通过蒸馏训练提高了ViT,并实现了与传统CNN比肩的性能。随后的研究通过结合分层设计[70,73]、借助卷积[15,22,23,25,64]注入局部性,或探索不同类型的token混合机制,如局部注意力[17,48]、空间MLP混合器[66,67]和非参数化池化混合器[79],进一步改进了ViT。通过适当的改进,视觉transformer在下游视觉任务中也表现出最先进的性能[19,38,39,77,80,82,83]。为了从有利的性能中获益,ViT的有效部署已成为研究热点,尤其是对于移动设备[12,43,53,54,56]。为了降低视觉transformer的计算复杂性,许多工作提出了新的模块并改进了架构设计[10,20,26,37,41,61],而其他工作消除了注意力机制中的冗余[9,14,29,44,60,71,74]。类似于对CNN的优化,也对ViT进行了架构搜索[8,11,21,46,85]、修剪[81]和量化[50]。

论文总结了有效ViT研究的两个主要缺点。首先,许多优化不适合移动部署。例如,通过正则化注意力机制的跨度或模式,注意力机制的二次计算复杂性可以降低为线性[10,17,48]。然而,在资源受限的设备上甚至不支持复杂的reshaping和索引操作[43]。重新思考移动友好型设计至关重要。第二,尽管最近的混合设计和网络搜索方法展示了具有强大性能的高效ViT[43,46,54],但它们主要优化了一个度量的Pareto curve,而在其他度量中竞争力较弱。例如,MobileViT[54]是参数高效的,同时比轻量级的CNN[62,65]慢几倍。EfficientFormer[43]在移动设备上使用超快的速度,但模型尺寸巨大。LeViT[22]和MobileFormer[12]以冗余参数为代价实现了有利的FLOPs。

重新思考混合Transformer网络

Token Mixers vs. Feed Forward Network

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第2张图片

结合局部信息可以提高性能,并使ViT在缺少显式位置嵌入的情况下更加鲁棒[5]。PoolFormer[79]和EfficientFormer[43]使用3×3平均池化层(图2(a))作为local token混合器。用相同内核大小的depth-wise卷积(DWCONV)替换这些层不会引入耗时开销,而使用可忽略的额外参数(0.02M),性能提高了0.6%。此外,最近的工作[5,21]表明,在ViT中的前馈网络(FFN)中注入局部信息建模层也有利于以较小的开销提高性能。值得注意的是,通过在FFN中放置额外的depth-wise 3×3卷积来捕获局部信息,复制了原始局部混合器(池或卷积)的功能。基于这些观察,论文移除了显式残差连接的local token混合器,并将 depth-wise 3×3 CONV移动到FFN中,以获得locality enabled的统一FFN(图2(b))。论文将统一的FFN应用于网络的所有阶段,如图2(a,b)所示。这种设计修改将网络架构简化为仅两种类型的block(local FFN和global attention),并在相同的耗时(见表1)下将精度提高到80.3%,参数开销较小(0.1M)。更重要的是,该修改允许我们直接使用模块的确切数量搜索网络深度,以提取局部和全局信息,尤其是在网络的后期阶段。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第3张图片

搜索空间优化

通过统一的FFN和删除残差连接的token mixer,论文检查来自EfficientFormer的搜索空间是否仍然足够,特别是在深度方面。论文改变了网络深度(每个阶段中的block数)和宽度(通道数),并发现更深和更窄的网络会带来更好的精度(0.2%的改进)、更少的参数(0.3M的减少)和更少的耗时(0.1ms的加速),如表1所示。因此,论文将此网络设置为新的基线(精度80.5%),以验证后续的设计修改,并为架构搜索提供更深入的超网络。

此外,具有进一步缩小的空间分辨率(1/64)的5阶段模型已广泛用于有效的ViT工作[12,22,46]。为了证明是否应该从一个5阶段超网络中搜索,论文在当前的基线网络中添加了一个额外的阶段,并验证了性能增益和开销。值得注意的是,尽管考虑到小的特征分辨率,计算开销不是一个问题,但附加阶段是参数密集型的。因此需要缩小网络维度(深度或宽度),以将参数和延迟与基线模型对齐,以便进行公平比较。如表1所示,尽管节省了MACs(0.12G),但5阶段模型的最佳性能在更多参数(0.39M)和延迟开销(0.2ms)的情况下意外降至80.31%。这符合我们的直觉,即五阶段计算效率高,但参数密集。鉴于5阶段网络无法在现有的规模和速度范围内引入更多潜力,论文坚持4阶段设计。这一分析也解释了为什么某些ViT在MACs精度方面提供了出色的Pareto curve,但在大小上往往非常冗余[12,22]。作为最重要的一点,优化单一度量很容易陷入困境。

多头注意力改进

然后,论文研究了在不增加模型大小和耗时的额外开销的情况下提高注意力模块性能的技术。如图2(c)所示,论文研究了MHSA的两种方法。首先通过添加depth-wise 3×3 CONV将局部信息注入到Value矩阵(V)中,[21,64]也采用了这种方法。其次通过在head维度[63]上添加全连接层来实现注意力头之间的通信,如图2(c)所示。通过这些修改,进一步将性能提高到80.8%,与基线模型相比,具有相似的参数和延迟。

在更高分辨率上的注意力

注意机制有利于性能。然而,将其应用于高分辨率特征会损害部署效率,因为它具有与空间分辨率相对应的二次时间复杂度。论文研究了将MHSA有效应用于更高分辨率(早期阶段)的策略。回想一下,在当前基线网络中,MHSA仅在输入图像空间分辨率为1/32的最后阶段使用。论文将额外的MHSA应用于具有1/16特征大小的倒数第二个阶段,并观察到准确度提高了0.9%。另一方面,推理速度减慢了几乎2.7倍。因此,有必要适当降低注意力模块的复杂性。

尽管一些工作提出了基于窗口的注意力[17,48],或下采样Key和Value[40]来缓解这个问题,但论文发现它们不是最适合移动部署的选项。由于复杂的窗口划分和重新排序,基于窗口的注意力很难在移动设备上加速。对于[40]中的下采样Key(K)和Value(V),需要全分辨率查询(Q)来保持注意力矩阵乘法后的输出分辨率(Out):

db75a289b18792f8b18967e7cfb824e4.png

根据测试该模型的耗时仅下降到2.8ms,仍然比基线模型慢2倍。因此,为了在网络的早期阶段执行MHSA,论文将所有Query、Key和Value降采样到固定的空间分辨率(1/32),并将注意力的输出插值回原始分辨率,以馈送到下一层,如图2所示((d)和(e))。我们称这种方法为“Stride Attention”。如表1所示,这种简单的近似值将延迟从3.5ms显著降低到1.5ms,并保持了具有竞争力的准确性(81.5%对81.7%)。

注意力降采样

大多数视觉主干利用跨步卷积或池化层来执行静态和局部下采样,并形成分层结构。最近的一些研究开始探索注意力下采样。例如,LeViT[22]和UniNet[46]建议通过注意力机制将特征分辨率减半,以实现全局感受野的上下文感知下采样。具体而言,Query中的token数量减少一半,以便对注意力模块的输出进行下采样:

2810e2151c7ac71acaaa74e702b94436.png

然而,决定如何减少Query中的token数量是非常重要的。Graham等人根据经验使用池化对Query进行下采样[22],而Liu等人建议搜索局部或全局方法[46]。为了在移动设备上实现可接受的推理速度,将注意力下采样应用于具有高分辨率的早期阶段是不利的,这限制了以更高分辨率搜索不同下采样方法的现有工作的价值。

相反,论文提出了一种同时使用局部性和全局依赖性的组合策略,如图2(f)所示。为了获得下采样的Query,论文使用池化层作为静态局部下采样,使用3×3 DWCONV作为可学习的局部下采样。此外,注意力下采样模块残差连接到regular strided CONV,以形成local-global方式,类似于下采样bottlenecks[28]或inverted bottlenecks[62]。如表1所示,通过略微增加参数和耗时开销,论文进一步将注意力下采样的准确率提高到81.8%。

EfficientFormerV2

EfficientFormerV2的设计

如前文所述,论文采用了四阶段分层设计,其获得的特征尺寸为输入分辨率的{1/4,1/8,1/16,1/32}。与其前身[43]类似,EfficientFormerV2从一个小的内核卷积stem开始嵌入输入图像,而不是使用非重叠patch的低效嵌入,

7dc59559111b2874f3a013c4fb149cf6.png

前两个阶段捕获高分辨率的局部信息;因此论文只使用统一的FFN(FFN,图2(b)):

b7f33bd84b04b215445103d88e5a664a.png

在最后两个阶段,使用局部FFN和全局MHSA块。因此式4,global blocks定义为:

befb0fe09d75916941faf047e6d9b638.png 9388eb720e3a24483cdf1dfb865897d1.png

联合优化模型大小和速度

尽管基线网络EfficientFormer[43]是通过耗时驱动搜索发现的,并且在移动设备上具有快速的推理速度,但搜索算法有两个主要缺点。首先,搜索过程仅受速度限制,导致最终模型是参数冗余的,如图1所示。其次,它仅搜索深度(每个阶段的blocks数)和阶段宽度,这是一种粗粒度的方式。事实上,网络的大多数计算和参数都在FFN中,并且参数和计算复杂度与其扩展比线性相关。可以针对每个FFN独立地指定,而不必相同。因此,搜索实现了更细粒度的搜索空间,其中计算和参数可以在每个阶段内灵活且非均匀地分布。这是最近ViT NAS艺术[21,43,46]中缺失的属性,其中在每个阶段保持相同。论文提出了一种搜索算法,该算法实现了灵活的per-block配置,并对大小和速度进行了联合约束,并找到了最适合移动设备的视觉主干。

搜索目标

首先,论文介绍了指导联合搜索算法的度量。考虑到在评估移动友好模型时网络的大小和延迟都很重要,我们考虑了一个通用的、公平的度量标准,它可以更好地理解移动设备上网络的性能。在不失一般性的情况下,我们定义了Mobile Efficiency Score(MES):

65c06c760059503118e7a842d42516c2.png

分数是为简单起见设置为100的预定义基础分数。模型大小是通过参数的数量来计算的,延迟是在设备上部署模型时的运行时间。由于本文专注于移动部署,因此以MobileNetV2的大小和速度为单位。具体来说,将和定义为使用CoreMLTools部署的iPhone 12(iOS 16)上的1ms延迟[1]。为了强调速度,,。减小大小和耗时可以获取更高的MES,论文在MES精度上搜索Pareto optimality。MES的形式是通用的,可以扩展到其他感兴趣的度量,例如推理时间内存占用和能耗。此外,通过适当定义,可以轻松调整每个度量的重要性。

搜索空间和SuperNet

搜索空间组成:

  • 网络的深度,由每个阶段的blocks数测量;

  • 网络的宽度,即每阶段的通道维度;

  • 每个FFN的膨胀率。

MHSA的数量可以在深度搜索期间无缝地确定,这控制了超网络中block的保存或删除。因此,论文在超网络的最后两个阶段将每个block设置为MHSA,然后是FFN,并通过深度搜索获得具有所需数量的全局MHSA的子网络。

SuperNet是通过使用在弹性深度和宽度上执行的可精简网络[78]来构建的,以实现纯基于评估的搜索算法。弹性深度可以通过stochastic drop path augmentation实现[32]。关于宽度和扩展比,论文遵循Yu等人[78]构建具有共享权重但独立规正则化层的可切换层,使得相应层可以在预定义集合的不同通道数量上(即16或32的倍数)执行。具体而言,膨胀比由每个FFN中的depth-wise 3×3 Conv的通道确定,并且stage宽度通过对齐FFN和MHSA block的最后投影(1×1 Conv)的输出通道确定。可切换执行可以表示为:

bc8e1bd76480723309ee2fba8d7695d3.png

通过在每次迭代中训练最大、最小和随机采样的两个子网(论文在算法1中将这些子网表示为max、min、rand-1和rand-2),使用Sandwich Rule[78]对超网络进行预训练。

搜索算法

现在,搜索目标、搜索空间和超网络集合都已公式化,论文提出了搜索算法。由于超网络在弹性深度和可切换宽度下可执行,论文可以通过分析每个瘦身动作的效率增益和精度下降来搜索具有最佳Pareto曲线的子网络。论文将行动池定义为:

ad3b346f45e814d8fa6e1f32472bfdab.png

使用全深度和全宽度(最大子网)初始化状态,论文评估ImageNet-1K验证集上每个前沿操作的精度结果(),这只需要大约4 GPU分钟。同时,参数缩减()可以直接从层属性(即kernel大小、输入通道和输出通道)计算。论文通过使用CoreMLTools在iPhone 12上测量的预先构建的耗时查找表来获得耗时减少()。有了这些度量就可以通过和计算,并选择每个MES精度下降最小的动作。值得注意的是,尽管动作组合非常庞大,但只需要在每个step评估前,这在复杂性上是线性的。详细信息可在Alg 1中找到。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第4张图片

实验

ImageNet-1K分类

实施细节:论文通过PyTorch 1.12[58]和Timm库[75]实现了该模型,并使用16个NVIDIA A100 GPU来训练的模型。使用AdamW[51]优化器在ImageNet-1K[16]上从头开始训练模型300和450个epoch。使用三个测试台来测试耗时:

  • iPhone 12-NPU:通过在神经引擎(NPU)上运行模型来获得iPhone 12(iOS 16)上的耗时。使用CoreML[1]编译模型(batch大小为1);

  • Pixel 6 - CPU:在Pixel 6(Android)CPU上测试模型耗时。为了获得大多数正在比较的作品的耗时,论文将所有模型的激活替换为ReLU,以获得公平的比较。使用XNNPACK[4]编译模型(batch大小为1);

  • Nvidia GPU:论文还提供了高端GPU Nvidia A100的延迟。这些模型(batch大小为64)部署在ONNX[2]中,由TensorRT[3]执行。

单一指标评估

论文在表2中显示了比较结果。EfficientFormerV2系列在单个指标(即参数数量或耗时)上实现了最先进的结果。在模型大小方面,EfficientFormerV2-S0比EdgeViT XXS[56]的精度高出1.3%,参数甚至少了0.6M,比MobileNetV2×1.0[62]的精度高3.9%,参数数量相似。对于大模型,EfficientFormerV2-L模型实现了与最近的EfficientFormer-L7[43]相同的精度,同时小3.1倍。在速度方面,EfficientFormerV2-S2具有可比或更低的延迟,分别以0.8%、0.6%和2.4%的top-1精度优于UniNet-B1[46]、EdgeViT-S[56]和EfficientFormer-L1[43]。论文希望这些结果能够为未来的架构设计提供实用见解:现代深度神经网络对架构排列具有鲁棒性,通过联合约束(如耗时和模型大小)优化架构不会损害单个度量。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第5张图片
联合优化大小和速度

此外,当考虑模型大小和速度时,论文证明了EfficientFormerV2的优越性能。在这里使用MES作为评估移动效率的更实用的指标,而不是仅使用大小或耗时。EfficientFormerV2-S1分别以4.2%、4.6%和1.5%的优势超过MobileViT-XS[54]、EdgeViT-XXS[56]和EdgeViT-XS[56]。凭借1.8倍更高的MES,EfficientFormerV2-L比MobileFormer-508M[12]高出4.0%的top-1精度。MES与准确度的可视化如图3所示。评估结果回答了开头提出的中心问题:使用拟议的移动效率基准,我们可以避免陷入一个陷阱,在一个指标上实现看似良好的性能,同时为其他指标牺牲太多。相反,我们可以获得既轻便又快速的高效移动ViT主干。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第6张图片

下游任务

目标检测和实例分割:选择Mask-RCNN [27]作为自己先,实验结果如表3所示。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第7张图片
搜索算法的消融实验

论文将所提出的搜索算法与EfficientFormer[43]的普通算法进行了比较。如表4所示,论文的搜索算法获得的模型参数和耗时与EfficientFormer[43]相似,但精度更高,证明了细粒度搜索和耗时和大小联合优化的有效性。

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第8张图片

讨论和结论

在这项工作中,论文全面研究了混合视觉主干,并验证了移动友好的设计选择。进一步提出了一种基于大小和速度的细粒度联合搜索,并获得了EfficientFormerV2模型家族,该模型族在推理速度上既轻量级又超快。由于论文专注于大小和速度以实现简单,未来的一个方向是将联合优化方法应用于后续研究,探索其他关键指标,如内存占用和二氧化碳排放。

往期回顾

手把手教学!TensorRT部署实战:YOLOv5的ONNX模型部署

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第9张图片

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第10张图片

自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区,聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多传感器融合、SLAM、光流估计、深度估计、轨迹预测、高精地图、NeRF、规划控制、模型部署落地、自动驾驶仿真测试、硬件配置、AI求职交流等方向;

轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer_第11张图片

添加汽车人助理微信邀请入群

备注:学校/公司+方向+昵称

你可能感兴趣的:(轻量ViT王者 | EfficientFormerV2:全新轻量级视觉Transformer)