1.3ms耗时！清华最新开源移动端神经网络架构 RepViT

论文地址：https://arxiv.org/abs/2307.09283

代码地址：https://github.com/THU-MIG/RepViT

RepViT 在移动端 ViT 架构中表现出色，展现出显著的优势。接下来，我们将探讨本研究的贡献所在。

1. 文中提到，轻量级 ViTs 通常比轻量级 CNNs 在视觉任务上表现得更好，这主要归功于它们的多头自注意力模块(MSHA)可以让模型学习全局表示。然而，轻量级 ViTs 和轻量级 CNNs 之间的架构差异尚未得到充分研究。
2. 在这项研究中，作者们通过整合轻量级 ViTs 的有效架构选择，逐步提升了标准轻量级 CNN（特别是 MobileNetV3 的移动友好性。这便衍生出一个新的纯轻量级 CNN 家族的诞生，即RepViT。值得注意的是，尽管 RepViT 具有 MetaFormer 结构，但它完全由卷积组成。
3. 实验结果表明，RepViT 超越了现有的最先进的轻量级 ViTs，并在各种视觉任务上显示出优于现有最先进轻量级ViTs的性能和效率，包括 ImageNet 分类、COCO-2017 上的目标检测和实例分割，以及 ADE20k 上的语义分割。特别地，在ImageNet上，RepViT 在 iPhone 12 上达到了近乎 1ms 的延迟和超过 80% 的Top-1 准确率，这是轻量级模型的首次突破。

好了，接下来大家应该关心的应该时“如何设计到如此低延迟但精度还很6的模型”出来呢？

方法

再 ConvNeXt 中，作者们是基于 ResNet50 架构的基础上通过严谨的理论和实验分析，最终设计出一个非常优异的足以媲美 Swin-Transformer 的纯卷积神经网络架构。同样地，RepViT也是主要通过将轻量级 ViTs 的架构设计逐步整合到标准轻量级 CNN，即MobileNetV3-L，来对其进行针对性地改造（魔改）。在这个过程中，作者们考虑了不同粒度级别的设计元素，并通过一系列步骤达到优化的目标。

训练配方的对齐

在论文中，新引入了一个用于衡量移动设备上延迟的指标，并确保训练策略与目前流行的轻量级 ViTs 保持一致。这一举措的目的是为了保证模型训练的一贯性，其中涉及到延迟度量和训练策略的调整两个关键概念。

延迟度量指标

为了更准确地衡量模型在真实移动设备上的性能，作者选择了直接测量模型在设备上的实际延迟，以此作为基准度量。这个度量方法不同于之前的研究，它们主要通过FLOPs或模型大小等指标优化模型的推理速度，这些指标并不总能很好地反映在移动应用中的实际延迟。

训练策略的对齐

这里，将 MobileNetV3-L 的训练策略调整以与其他轻量级 ViTs 模型对齐。这包括使用 AdamW 优化器【ViTs 模型必备的优化器】，进行 5 个 epoch 的预热训练，以及使用余弦退火学习率调度进行 300 个 epoch 的训练。尽管这种调整导致了模型准确率的略微下降，但可以保证公平性。

块设计的优化

接下来，基于一致的训练设置，作者们探索了最优的块设计。块设计是 CNN 架构中的一个重要组成部分，优化块设计有助于提高网络的性能。

分离 Token 混合器和通道混合器

这块主要是对 MobileNetV3-L 的块结构进行了改进，分离了令牌混合器和通道混合器。原来的 MobileNetV3 块结构包含一个 1x1 扩张卷积，然后是一个深度卷积和一个 1x1 的投影层，然后通过残差连接连接输入和输出。在此基础上，RepViT 将深度卷积提前，使得通道混合器和令牌混合器能够被分开。为了提高性能，还引入了结构重参数化来在训练时为深度滤波器引入多分支拓扑。最终，作者们成功地在 MobileNetV3 块中分离了令牌混合器和通道混合器，并将这种块命名为 RepViT 块。

降低扩张比例并增加宽度

在通道混合器中，原本的扩张比例是 4，这意味着 MLP 块的隐藏维度是输入维度的四倍，消耗了大量的计算资源，对推理时间有很大的影响。为了缓解这个问题，我们可以将扩张比例降低到 2，从而减少了参数冗余和延迟，使得 MobileNetV3-L 的延迟降低到 0.65ms。随后，通过增加网络的宽度，即增加各阶段的通道数量，Top-1 准确率提高到 73.5%，而延迟只增加到 0.89ms！

宏观架构元素的优化

在这一步，本文进一步优化了MobileNetV3-L在移动设备上的性能，主要是从宏观架构元素出发，包括 stem，降采样层，分类器以及整体阶段比例。通过优化这些宏观架构元素，模型的性能可以得到显著提高。

浅层网络使用卷积提取器

图片

ViTs 通常使用一个将输入图像分割成非重叠补丁的 "patchify" 操作作为 stem。然而，这种方法在训练优化性和对训练配方的敏感性上存在问题。因此，作者们采用了早期卷积来代替，这种方法已经被许多轻量级 ViTs 所采纳。对比之下，MobileNetV3-L 使用了一个更复杂的 stem 进行 4x 下采样。这样一来，虽然滤波器的初始数量增加到24，但总的延迟降低到0.86ms，同时 top-1 准确率提高到 73.9%。

更深的下采样层

在 ViTs 中，空间下采样通常通过一个单独的补丁合并层来实现。因此这里我们可以采用一个单独和更深的下采样层，以增加网络深度并减少由于分辨率降低带来的信息损失。具体地，作者们首先使用一个 1x1 卷积来调整通道维度，然后将两个 1x1 卷积的输入和输出通过残差连接，形成一个前馈网络。此外，他们还在前面增加了一个 RepViT 块以进一步加深下采样层，这一步提高了 top-1 准确率到 75.4%，同时延迟为 0.96ms。

更简单的分类器

在轻量级 ViTs 中，分类器通常由一个全局平均池化层后跟一个线性层组成。相比之下，MobileNetV3-L 使用了一个更复杂的分类器。因为现在最后的阶段有更多的通道，所以作者们将它替换为一个简单的分类器，即一个全局平均池化层和一个线性层，这一步将延迟降低到 0.77ms，同时 top-1 准确率为 74.8%。

整体阶段比例

阶段比例代表了不同阶段中块数量的比例，从而表示了计算在各阶段中的分布。论文选择了一个更优的阶段比例 1:1:7:1，然后增加网络深度到 2:2:14:2，从而实现了一个更深的布局。这一步将 top-1 准确率提高到 76.9%，同时延迟为 1.02 ms。

微观设计的调整

接下来，RepViT 通过逐层微观设计来调整轻量级 CNN，这包括选择合适的卷积核大小和优化挤压-激励（Squeeze-and-excitation，简称SE）层的位置。这两种方法都能显著改善模型性能。

卷积核大小的选择

众所周知，CNNs 的性能和延迟通常受到卷积核大小的影响。例如，为了建模像 MHSA 这样的远距离上下文依赖，ConvNeXt 使用了大卷积核，从而实现了显著的性能提升。然而，大卷积核对于移动设备并不友好，因为它的计算复杂性和内存访问成本。MobileNetV3-L 主要使用 3x3 的卷积，有一部分块中使用 5x5 的卷积。作者们将它们替换为3x3的卷积，这导致延迟降低到 1.00ms，同时保持了76.9%的top-1准确率。

SE 层的位置

自注意力模块相对于卷积的一个优点是根据输入调整权重的能力，这被称为数据驱动属性。作为一个通道注意力模块，SE层可以弥补卷积在缺乏数据驱动属性上的限制，从而带来更好的性能。MobileNetV3-L 在某些块中加入了SE层，主要集中在后两个阶段。然而，与分辨率较高的阶段相比，分辨率较低的阶段从SE提供的全局平均池化操作中获得的准确率提升较小。作者们设计了一种策略，在所有阶段以交叉块的方式使用SE层，从而在最小的延迟增量下最大化准确率的提升，这一步将top-1准确率提升到77.4%，同时延迟降低到0.87ms。【这一点其实百度在很早前就已经做过实验比对得到过这个结论了，SE 层放置在靠近深层的地方效果好】

网络架构

最终，通过整合上述改进策略，我们便得到了模型RepViT的整体架构，该模型有多个变种，例如RepViT-M1/M2/M3。同样地，不同的变种主要通过每个阶段的通道数和块数来区分。

实验

图像分类

检测与分割

总结

本文通过引入轻量级 ViT 的架构选择，重新审视了轻量级 CNNs 的高效设计。这导致了 RepViT 的出现，这是一种新的轻量级 CNNs 家族，专为资源受限的移动设备设计。在各种视觉任务上，RepViT 超越了现有的最先进的轻量级 ViTs 和 CNNs，显示出优越的性能和延迟。这突显了纯粹的轻量级 CNNs 对移动设备的潜力。