【2024 CVPR-Backbone】RepViT: Revisiting Mobile CNN From ViT Perspective
摘要
近期,轻量级视觉Transformer(ViT)在资源受限的移动设备上表现出比轻量级卷积神经网络(CNN)更优异的性能和更低的延迟。研究人员已发现轻量级ViT与轻量级CNN之间存在许多结构关联,但二者在模块结构、宏观和微观设计上的显著架构差异尚未得到充分研究。本研究从ViT视角重新审视轻量级CNN的高效设计,并强调其在移动设备上的应用前景。具体而言,我们通过整合轻量级ViT的高效架构设计,逐步提升标准轻量级CNN(即MobileNetV3)的移动友好性,最终形成一系列纯轻量级CNN新架构,命名为RepViT。大量实验表明,RepViT在各项视觉任务中超越了现有轻量级ViT,并展现出良好的延迟表现。值得注意的是,在ImageNet数据集上,RepViT在iPhone 12上实现了超过80%的top-1准确率且延迟仅1.0毫秒,据我们所知这是轻量级模型首次达到该性能。此外,当RepViT与SAM结合时,我们的RepViT-SAM推理速度比先进的MobileSAM快近10倍。代码与模型已开源:https://github.com/THU-MIG/RepViT。
1. 引言
在计算机视觉领域,设计轻量级模型一直是实现高性能与低计算成本平衡的核心研究方向,这对资源受限的移动设备部署视觉模型至关重要。过去十年间,研究者主要聚焦轻量级卷积神经网络(CNN),并取得显著进展:提出了可分离卷积[27]、倒残差瓶颈[53]、通道混洗[44,75]和结构重参数化[13,14]等高效设计原则,催生了MobileNets[26,27,53]、ShuffleNets[44,75]和RepVGG[14]等代表性模型。
近年来,视觉Transformer(ViT)[18]作为CNN的替代方案崭露头角,在图像分类[39,63]、语义分割[6,66]和目标检测[4,34]等任务中展现出超越CNN的性能。然而,ViT通过增加参数量提升性能的趋势导致模型体积庞大、延迟高昂[11,40],难以适配移动设备[36,46]。直接缩小ViT模型尺寸虽能满足资源限制,但性能往往不及轻量级CNN[5]。因此,研究者开始探索ViT的轻量化设计,目标是在移动设备上实现超越轻量级CNN的性能。
针对ViT的轻量化设计,研究者提出了多项高效原则[5,35,46,49]:包括结合卷积与ViT的混合架构[5,46]、线性复杂度的自注意力操作[47],以及维度一致的设计原则[35,36]等。这些研究表明,轻量级ViT[35,47,49]在移动设备上可实现更低延迟并超越轻量级CNN[26,53,60](如图1所示)。
尽管轻量级ViT取得进展,仍面临硬件与计算库支持不足[60]、高分辨率输入导致延迟剧增[3]等实际挑战。相比之下,CNN凭借高度优化的卷积操作(输入线性复杂度)更适合边缘部署[54,73]。这促使我们重新审视轻量级CNN的设计,通过系统对比轻量级ViT与CNN的差异来寻求突破。
轻量级ViT与CNN存在结构相似性:两者均使用卷积模块学习局部表征[46,47,49,61],而CNN通过扩大卷积核感受野[73]、ViT通过多头自注意力[46,47]学习全局表征。但二者在模块结构、宏观/微观设计上的显著差异尚未被充分研究:例如ViT偏好MetaFormer模块[69],而CNN倾向倒残差瓶颈[53]。这引发关键问题:ViT的架构设计能否提升CNN性能?为此,本研究从ViT视角重构轻量级CNN,旨在弥合二者差距并证明CNN在移动部署中的潜力。
基于[41]的方法,我们以MobileNetV3-L[26]为起点,逐步融入轻量级ViT的高效设计[35,36,38,46],最终得到纯轻量级CNN新家族——RepViT。该模型完全由重参数化卷积构成,采用类似ViT的MetaFormer结构[36,69,70]。实验表明,RepViT在ImageNet分类、COCO目标检测/实例分割、ADE20k语义分割等任务中,性能与效率均超越现有轻量级ViT[35,49]。值得注意的是,RepViT在iPhone 12上实现80%以上top-1准确率且延迟仅1.0ms(据我们所知为轻量级模型首次达成);其最大模型RepViT-M2.3达83.7%准确率,延迟仅2.3ms。与SAM[33]结合后,RepViT-SAM推理速度比MobileSAM[71]快近10倍,且零样本迁移能力显著提升。我们希望RepViT能成为轻量级模型研究的强基准,并推动边缘部署的进一步探索。
2. 相关工作
过去十年间,卷积神经网络(CNN)凭借其固有的局部归纳偏置和平移等效性,成为计算机视觉任务的主流方法[16,23,24,43,62,67]。然而标准CNN的高计算量使其难以部署在资源受限的移动设备上。为克服这一挑战,研究者提出了多种轻量化技术,包括可分离卷积[27]、倒残差瓶颈[53]、通道混洗[44,75]和结构重参数化[14]等,催生了MobileNets[26,27,53]、ShuffleNets[44,75]和RepVGG[14]等广泛使用的轻量级CNN。
随后出现的视觉Transformer(ViT)[18]将Transformer架构适配至大规模图像识别任务,其性能超越CNN[18,58]。后续研究通过引入空间归纳偏置[10,22]、设计高效自注意力机制[17,79]及扩展ViT至多种视觉任务[19,74]进一步提升性能。尽管ViT在多项任务中表现优异,但其庞大的参数量和计算/内存需求[39,58]使其难以适配资源有限的移动设备[46,49]。
为此,研究者致力于探索ViT轻量化技术,例如MobileViT[46]通过结合MobileNet模块与多头自注意力(MHSA)模块的混合架构,以及EfficientFormer[36]提出的维度一致性设计范式来优化延迟-性能权衡。这些轻量级ViT在移动设备上实现了超越传统轻量级CNN[53,60]的最新性能-延迟平衡。
轻量级ViT的成功通常归因于MHSA模块学习全局表征的能力,但轻量级CNN与ViT在模块结构、宏观/微观设计上的显著差异往往被忽视。与现有工作不同,我们的核心目标是通过整合轻量级ViT的架构设计来重新审视轻量级CNN,旨在弥合二者差距并突出CNN在移动设备部署中的优势。
3. 方法
图2展示了从MobileNetV3-L到RepViT的逐步现代化过程。我们主要关注移动设备延迟和ImageNet-1K top-1准确率两个指标。最终得到纯轻量化CNN新家族RepViT,实现了更低的延迟和更高的性能。这些结果都是在不使用蒸馏技术的情况下获得的。
在这一节中,我们首先使用一个标准的轻量级CNN,即MobileNetV3-L,然后通过融入轻量化ViT的架构设计,逐步从多个粒度对其进行现代化改造。我们首先介绍一种在移动设备上测量延迟的指标,然后在3.1节中与现有轻量化ViT的训练方案保持一致。基于一致的训练设置,我们在3.2节中探索最优的块设计。我们进一步在3.3节中从宏观架构元素(即stem、下采样层、分类器和整体阶段比例)优化MobileNetV3-L在移动设备上的性能。然后在3.4节中通过逐层微观设计对轻量级CNN进行微调。图2展示了我们在每一步中实现的整个过程和结果。最后,我们在3.5节中获得了一个专为移动设备设计的纯轻量级CNN新家族,名为RepViT。所有模型都在ImageNet-1K上进行训练和评估。
3.1 初步准备
- 延迟指标
传统方法[5,57]使用FLOPs或模型大小优化推理速度,但这些指标与真实移动设备延迟相关性差[36]。我们采用实际设备延迟作为基准(iPhone 12 + CoreMLTools[1],参考[35,36,46,60]),测得MobileNetV3-L延迟1.01ms。为避免CoreMLTools不支持的函数,使用GeLU激活函数(参考[36,60])。 - 训练方案对齐
轻量化ViT[35,36,46,49]普遍采用DeiT[58]的训练方案:AdamW优化器[42] + 余弦学习率调度,300轮训练(RegNetY-16GF[51]教师模型蒸馏),Mixup[72]、自动增强[8]、随机擦除[77]数据增强,Label Smoothing[56]正则化。我们将MobileNetV3-L训练方案与之对齐(暂不使用知识蒸馏),获得71.5% top-1准确率。此方案将作为后续默认设置。
3.2 模块设计
图3展示了模块设计。图(a)是一个带有可选挤压-激励(SE)层的MobileNetV3块。图(b)是设计的RepViT块,它通过结构重参数化技术分离了token mixer和channel mixer。SE层在RepViT块中也是可选的。为了简洁起见,省略了归一化层和非线性激活函数。
- 分离token mixer和channel mixer
轻量级ViT的重要特征是分离token mixer和channel mixer[70]。根据MetaFormer架构理论[69],我们拆分MobileNetV3-L的token mixer(3×3深度卷积)与channel mixer(1×1卷积),并将SE层移至深度卷积后。采用结构重参数化技术[7,14]消除推理时跳跃连接的计算/内存成本,形成RepViT模块(延迟降至0.81ms,准确率68.3%)。 - 降低扩展比并增加宽度
ViT的FFN模块扩展比通常为4,导致高计算开销[76]。参考LV-ViT[30]和LeViT[21],我们将RepViT模块的扩展比统一设为2,并将各阶段通道数翻倍至48/96/192/384,使准确率提升至73.5%(延迟0.89ms)。
3.3 宏观设计
图4. 宏观设计。(a)与(b)、(c)与(d)、(e)与(f)分别表示 stem、下采样层和分类器的设计。
RepViT 具有四个阶段,分辨率分别为 H/4×W/4、H/8×W/8、H/16×W/16 和 H/32×32,其中 H 和 W 分别表示输入图像的宽度和高度。C 代表通道维度,B 表示批次大小。图中省略了归一化层和非线性激活函数。
- 早期卷积作为stem
ViT通常使用16×16非重叠卷积作为stem,而MobileNetV3-L采用复杂stem(含3×3卷积、深度可分离卷积和倒残差)。我们改用两个3×3卷积(通道数24/48,步长2)作为stem,使延迟降至0.86ms,准确率提升至73.9%。 - 更深的下采样
ViT通过独立patch合并层实现下采样。我们采用深度可分离卷积+1×1卷积+RepViT块组成更深的下采样层(含FFN模块),使准确率提升至75.4%(延迟0.96ms)。 - 简单分类器
ViT通常使用全局平均池化+线性层的简单分类器。我们移除MobileNetV3-L分类器中的额外1×1卷积,采用相同简单设计,使延迟降至0.77ms(准确率76.3%)。 - 整体阶段比例
参考ViT的1:1:7:1阶段比例,我们将网络深度调整为2:2:14:2,使准确率提升至76.9%(延迟0.91ms)。
3.4 微观设计
-
卷积核尺寸选择:全部采用3×3卷积(避免大核卷积的计算复杂度),维持76.9%准确率(延迟0.89ms)。
-
SE层放置策略:采用跨模块放置策略(每阶段第1、3、5…个模块使用SE层),使准确率提升至77.4%(延迟0.87ms)。
3.5 网络架构
开发RepViT-M0.9/M1.0/M1.1/M1.5/M2.3系列模型(后缀"-MX"表示iPhone12上的延迟为X毫秒),通过通道数和模块数量区分不同变体。详细信息见补充材料。