RepViT（CVPR 2024）论文解读

paper：RepViT: Revisiting Mobile CNN From ViT Perspective

official implementation：https://github.com/THU-MIG/RepViT

third-party implementation：https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/repvit.py

背景

近年来，轻量级视觉Transformer（ViTs）在资源受限的移动设备上表现出优越的性能和更低的延迟，优于轻量级卷积神经网络（CNNs）。尽管研究人员发现了轻量级ViTs和轻量级CNNs之间的许多结构上的联系，但在block结构、宏观和微观设计方面的显著架构差异尚未得到充分研究。

出发点

该研究旨在从ViT的角度重新审视轻量级CNN的高效设计，并强调其在移动设备上的前景。具体来说，研究团队通过集成轻量级ViTs的高效架构设计，逐步增强标准轻量级CNN（即MobileNetV3）的移动友好性，最终开发出一个新的轻量级纯CNN架构，即RepViT。

创新点

1. RepViT架构：通过将轻量级ViT的高效架构设计整合到标准轻量级CNN（如MobileNetV3）中，逐步增强了移动友好性，最终形成了新的轻量级纯CNN系列模型RepViT。
2. 宏观和微观设计优化：对网络的宏观架构（如stem、下采样层、分类器和整体阶段比例）和微观架构（如内核大小选择和SE层放置）进行了优化，以适应移动设备的需求。
3. 结构重参数化（SR）：在RepViT block中采用了结构重参数化技术，以提高模型在训练期间的学习效率，并在推理期间减少计算和内存成本。
4. 跨块SE层放置：采用跨块方式放置挤压-激励（SE）层，以最大化准确性收益，同时最小化延迟增加。
5. 性能与延迟的权衡：RepViT在各种视觉任务中展现出比现有的轻量级ViT更高的性能，并且在延迟方面也具有优势。特别是在ImageNet数据集上，RepViT在iPhone 12上的延迟仅为1.0毫秒，准确率达到80%以上，这在轻量级模型中尚属首次。

方法介绍

这篇文章的结构和ConvNeXt比较类似（具体介绍可参考ConvNeXt（CVPR 2022）论文解读-CSDN博客），都是以一个模型作为baseline，然后从另一个框架或结构中借鉴好的结构设计（ConvNeXt中是Transformer架构，本文是轻量级Transformer架构）逐步添加到baseline中，包括宏观设计和微观设计。

本文的目标是构建一个新的轻量级网络，因此首先需要确定评价指标。和之前的许多框架采用参数量或FLOPs不同，本文直接采用在移动设备上的延迟作为基准度量。实践中baseline模型选择MobileNetV3-L（可参考MobileNet v3-CSDN博客），使用iPhone 12作为测试设备，使用Core ML Tools作为编译器，为了避免Core ML Tools中不支持的函数，在MobileNetV3-L中使用GeLU激活函数。

  我们测量了MobileNetV3-L的延迟为1.01ms。

训练策略和最近的轻量ViTs一样都采用DeiT的训练策略，具体包括用AdamW optimizer和cosine learning rate schedule训练300个epochs。用RegNetY-16GF作为教师模型进行蒸馏。数据增强采用Mixup、auto-augmentation和random erasing。采用Label Smoothing作为正则方法。为了公平比较，作者将MobileNetV3-L和其它轻量ViTs的训练策略保持一致，除了暂时不用蒸馏。最终用这种训练策略MobileNetV3-L得到了71.5%的top-1精度。

  现在默认使用这个训练策略。

整个改进过程如图2所示，接下来具体介绍细节。

Block design

Separate token mixer and channel mixer. MetaFormer（具体介绍见MetaFormer（CVPR 2022，Sea）-CSDN博客）这篇文章指出ViTs的有效性主要源于通用的token mixer和channel mixer结构即MetaFormer结构，而不是token mixer具体的结构。因此作者希望分离MobileNetV3-L中的token mixer和channel mixer来模拟现有轻量ViTs的设计。

如图3(a)所示，MobileNetV3-L中的两个1x1卷积可以起到channel mixer的作用，而3x3深度卷积起到token mixer的作用，这种设计将两者耦合到一起了。为了解耦它们，作者将3x3 DW移到最上面，此外可选的SE模块也随着一起移到上面，这样就将token mixer和channel mixer分离开了。作者还使用了结构重参数化技术（具体介绍见结构重参数化之二：RepVGG-CSDN博客）来增强训练过程中模型的学习能力，修改后的结果如图3(b)所示，其中右侧是推理时的结构可以通过结构重参数化技术将3x3 DW和1x1 DW融合到一起，加快推理速度。作者将这个新的结构称为RepViT block，这一修改将延迟降低到了0.81ms，同时精度也降低到了68.3%。

Reducing expansion ratio and increasing width. 在原始的ViT中，channel mixer中的expansion ratio通常设置为4，这使得FFN消耗了大量的计算资源，同时也占据了相当一部分的推理时间。为了缓解这一瓶颈，最近的一些工作采用了更窄的FFN，例如LV-ViT的expansion ratio为3，LeViT的expansion ratio为2。此外Nvit这篇文章指出FFN中存在大量的通道冗余，因此使用较小的expansion ratio是合理的。

在MobileNetV3-L中，expansion ratio从2.3到6，最后两个stage的expansion ratio为6导致通道数量过多。对于RepViT block，作者设置所有stage中channel mixer的expansion ratio都为2，这将延迟降低到了0.65ms。在较小的expansion ratio下，我们可以增加网络的宽度来弥补参数量的下降，作者将每个stage的通道数翻倍，得到48、96、192、384，这一修改将精度提升到了73.5%，延迟为0.89ms。

Macro design

这一部分作者关注于优化网络的宏观结构设计。

Early convolutions for stem. 原始的ViT采用patchify操作作为stem，即将输入图片切分成不重叠的patch。具体通过一个stride=16的16x16大核卷积实现。Early Convolutions Help Transformers See Better（NeurIPS 2021, Meta）-CSDN博客这篇文章指出，这种操作会导致模型的可优化性变差并且对训练策略的设置比较敏感，为了缓解这个问题，它们提出使用堆叠多个步长为2的3x3卷积作为替代，即early convolutions。这提高了优化的稳定性和模型性能。

MobileNetV3-L原始的stem如图4(a)所示，结构复杂可能会成为移动设备上的速度瓶颈，为了平衡延迟将初始卷积数量减少到了16，这反过来又限制了stem的表示能力。本文提出采用early convolution的方法，即两个步长为2的3x3卷积，如图4(b)所示，其中第一个卷积数量为24第二个为48。这一修改将延迟降低到了0.86ms，同时精度提升到了73.9%。

  我们现在使用early convolutions作为stem。

Deeper downsampling layers. MobileNetV3-L中的下采样层采用了一个inverted bottleneck block其中包含一个步长为2的深度卷积，如图4(c)所示。这种设计可能缺乏足够的网络深度，导致信息丢失并对模型性能产生负面影响。所以为了获得一个分离且更深的下采样层，作者使用了一个步长为2的深度卷积和一个1x1卷积如图4(d)所示。并在前面加上一个RepViT block来增加下采样层的深度，在后面放置一个FFN来保存更多的信息。这一修改后，精度提升到了75.4%，延迟为0.96ms。

  我们现在将使用更深的下采样层。

Simple classifier. 在轻量ViTs中，分类器由一个全局平均池化层和一个线性层构成，这样一个简单的分类器对延迟是友好的，尤其是对于移动设备。相反MobileNetV3-L采用了一个复杂的分类器，包括一个额外的1x1卷积和一个额外的线性层将特征映射到更高的维度，如图4(e)所示。这种设计对于MobileNetV3-L生成丰富的预测特征至关重要，特别是考虑到最后阶段的输出通道数较小。但反过来又对延迟造成了影响。而在本文对block的设计修改之后，最后一个阶段现在有了更多的通道，因此就可以用一个简单的分类器来替代了，即一个全局平均池化层和一个线性层，如图4(f)所示。这一修改后导致精度下降了0.6%，但延迟也降低到了0.77ms。

  我们现在使用简单的分类器。

Overall stage ratio. stage ratio表示不同阶段block数量的比值，之前的工作指出在第三个stage放置更多的block可以得到一个更好的精度-速度的平衡。本文采用1:1:7:1的stage ratio，然后将网络深度增加到2:2:14:2，得到一个更深的网络。这一步修改将精度提升到了76.9%，延迟为0.91ms。

  我们现在使用这个stage ratio。

Micro design

这一部分作者主要关注轻量CNN的微观架构设计，包括卷积kernel的大小选择和SE层的放置。

Kernel size selection. CNN的性能和延迟往往受到kernel大小的影响，为了捕捉到像MHSA的长期依赖关系，ConNeXt采用大核卷积，并获得了性能的提升。但是由于计算复杂和内存访问成本，大核卷积对移动设备并不友好。此外与3x3卷积相比，大核卷积没有被编译器或计算库进行专门优化。MobileNetV3-L主要采用3x3卷积，并在某些block中使用5x5卷积。为了保证在移动设备上的推理效率，本文在所有block中统一使用3x3卷积，这一修改后，精度保持不变还是76.9%，但延迟降低到了0.89ms。

  我们现在使用3x3卷积。

Squeeze-and-excitation layer placement. 与卷积相比，自注意力的一个优势是可以根据输入调整权重，称之为数据驱动属性。作为一个通道注意力模块，SE（具体介绍见SENet: Squeeze-and-Excitation Networks-CSDN博客）层可以弥补缺乏数据驱动属性的卷积的局限性，带来更好的性能。MobileNetV3-L在某些层使用了SE，主要是最后两个阶段。但是Tresnet这篇文章指出，与高分辨率的stage相比，低分辨率stage使用SE获得的精度收益较小。同时随着精度提升，SE也带来了不可忽略的计算成本。因此本文使用了一种cross-block使用SE的策略，即在每个stage的第1、3、5...个block中使用SE，以最小的延迟增加来最大化提升精度。这一步将精度提升到了77.4%，延迟为0.87ms。

  我们现在使用cross-block SE。这一步后就得到了我们的最终模型，即RepViT。

实验结果

作者设计了五种不同大小的RepViT，即RepViT-M0.9/M1.0/M1.1/M1.5/M2.3，其中后缀-MX表示在移动设备如iPhone12上的延迟，不同variants的区别在于每个stage中的通道数量和block数量。

在ImageNet上的结果如表1所示，可以看到，RepViT在不同的模型尺寸下都取得了SOTA的结果。

在下游任务上的表现如表5所示，其中包括在COCO数据集上的目标检测和实例分割任务，在ADE20K数据集上的语义分割任务。其中目标检测和实例分割采用Mask R-CNN模型，语义分割采用Semantic FPN。可以看到RepViT在各个任务上的效果都超越了同等范围内的其它轻量模型。