边缘计算(一)：EdgeNeXt，面向移动开发的CNN-Transformer混合架构

单位：MBZUAI（位于阿布扎比的默罕默德人工智能大学），IIAI（起源研究院，邵岭团队）
ArXiv: https://arxiv.org/abs/2206.10589
Github: https://github.com/mmaaz60/EdgeNeXt

导读：CNN和Transformer单独玩出的花样层出不穷。你中有我我中有你的思想，交替出现，例如Large Kernel CNN试图去模仿Transformer的远程建模能力，Local Transformers试图去模仿CNN分层建模的能力。而这篇文章，设计了一个混合框架，将CNN和Transformer结合起来，既快又好。其中CNN模块来自Res2Net，Transformer模块参考协方差注意力机制，组合创新，但取得了SOTA的性能，绝了！

---

---

摘要

为了追求不断增长的精度，通常会设计大而复杂的模型。这样的模型需要大量计算资源，因此不能在边缘设备上部署。由于在不少应用领域上具有实用性，构建资源高效型的通用模型受到欢迎。这篇文章试图有效地结合CNN和Transformer的各自优势，提出了一个混合架构，即EdgeNeXt。具体的，作者引入了可拆分的深度转置注意力编码器（split depth-wise transpose attention,SDTA），它将输入的张量拆分为多个通道组，并利用深度卷积网络和跨通道的自注意来隐式地增加感受野，和编码多尺度特征。在多个任务如识别、检测与分割上进行了广泛的实验，揭示该方法的先进性，它取得超过了SOTA的精度并具有更低的运算量。例如，在ImageNet-1K上取得了71.2%的top1精度，超过MobileViT2.2%但减少了28%的FLOPs。

---

一、介绍

计算效率对于资源有限的设备如移动平台等至关重要。很多实际应用如机器人、自动驾驶汽车等，识别处理过程既要注重识别精度，也要兼顾低延迟。当前很多处理方法，如设计高效的CNN变体去实现速度和精度的平衡；也有一些方法采用硬件感知的神经架构搜索（neural architecture search，NAS）去构建低延迟的高精度模型。尽管这些CNN模型容易训练，在编码局部图像细节方面很高效，但不能明确地建模像素之间的全局关系。ViT中的自注意力能够显式地建模这种全局关系，但因其运算量大，常以缓慢的推理为代价。针对移动视觉应用，设计轻量级的ViT很有挑战。

基于CNN的高效推理模型成为主流，但有两个固有的限制：1、它具有局部视觉感受野，全局建模上下文有限；2、学习的权重是固定的，使得CNN无法灵活适应输入内容。尽管这两个问题能够被Transformer解决，但他们是计算密集型的。有一些工作开始调研设计轻量级的架构，通过借助CNN和Transformer的优势来构建针对移动平台的任务。但这些工作关注参数优化而引入更多的乘加操作。作者认为设计有效地将CNN和ViT优势结合起来的统一架构，能够针对资源有效的设备时减少模型尺寸、参数量和乘加操作量。

二、EdgeNeXt

整体架构

整体架构主要由两部分组成：1、自适应的$N\times N$的Conv Encoder；2、可分离的深度转置注意力编码器即SDTA。EdgeNeXt遵从ConvNeXt的设计原则，在四个阶段中提出四个不同尺度的层次化特征。输入图像$H\times W\times 3$，首先会经过一个分块流层，它是通过一个无重叠的$4\times$卷积+Layer Norm即层归一化，输出$\frac{H}{4}\times \frac{W}{4}\times C_1$特征图。然后它会经过$3\times 3$卷积层去提取局部信息。在第二阶段，首先通过步长为2的卷积层去下采样，将空间尺度减半并增加通道数，输出$\frac{H}{8}\times \frac{W}{8}\times C_2$特征图，然后跟着叠加两个连续的$5\times 5$卷积层。此时，位置编码也在SDTA模块前添加进来，仅在阶段2添加。虽然位置编码对于密集的预测任务如分割和检测等非常敏感，在所有阶段都添加位置编码会影响推理速度。然后，在第三第四阶段分别产生$\frac{H}{16}\times \frac{W}{16}\times C_3$和$\frac{H}{32}\times \frac{W}{32}\times C_4$的输出特征图。

卷积编码器 Convolution Encoder

这个模块采用深度可分离的带有自适应核大小的卷积操作。我们可以将之定义为两部分：1、带有自适应核$N\times N$的DW-CNN，我们采用$K=3,5,7,9$对应不同的四个阶段。2、Point-wise CNN，为了丰富局部表征，后面通过会跟随layer Norm和GELU即Gaussian Error Linear Unit激活函数，为了非线性的特征映射。最后，一个Skip连接被添加到网络中让信息演着层次化的网络流动。这个模块与ConvNeXt模块相似，但采用核大小是动态的。作者发现动态大小的的卷积编码器比固定核尺度的表现要好。卷积编码器能够归纳为：
$$x_{i+1}=x_i+Linea{r}_G(Linear(LN(D_w(x_i))))$$
其中，$x_i$表示$H\times W\times 3$的输入特征图，$Linea{r}_G$表示point-wise CNN层与GELU一起作用的函数表示。

全局建模

自注意力学习全局表征对视觉任务至关重要。为了有效继承这一优势，作者设计了相对较少的网络模块（3层叠加，MobileViT采用9层），采用交叉协方差注意力将注意力操作整合到特征通道维度上而不是在空间维度上。这将原始的自注意力操作从相对token的数量的平方复杂度，降到线性复杂度$O(N{d}^2)$，并有效地隐式编码全局信息。$N$是图像块的数量，$d$是特征通道维度。

自适应的核大小

众所周知，大核卷积的运算量很大，因为随着核大小增加，参数量与运算量会平方增加。尽管大核尺度有利于增加视觉感受野，但在整个层次的网络架构中使用如此大的卷积核是高代价的和次优的。作者提出了一个自适应的核大小机制去减少复杂度，和捕捉不同级别的特征。作者提出在第一阶段使用小卷积核，而后续阶段使用大卷积核。这个设计是优化的，因为早起阶段提出的是低级别的特征，使用小卷积核很合适。

SDTA Encoder，可分离深度转置注意力

SDTA主要有两个组成：第一部分尝试去学习自适应的多尺度特征表征，通过在输入图像上编码不同的空间级别；第二部分直观地编码全局图像表征。第一部分启发于Res2Net网络模块，作者设计了一个多尺度的处理方法通过在单个模块内学习分层的特征表征，这让空间感受野更加灵活与自适应。但与Res2Net不同，作者没有采用$1\times 1$point-wise CNN 确保该网络具有有限参数和乘加运算。并且，作者采用自适应数量的子集允许特征编码的灵活性。作者将输入张量$H\times W\times C$分成$s$个分组，每一组用$x_i$表示和有相同通道数$C/s$的空间尺度，其中$i\in 1,2,\cdot \cdot \cdot ,s$，$C$表示通道数。每一个特征图子集（第一个除外）将会通过一个$3\times 3$的DW-CNN层，用$d_i$表示，其输出表示为$y_i$。并且，$d_{i-1}$的输出即$y_{i-1}$，也被加入到特征子集$x_i$中，然后输入到$d_i$中。子集的数量$s$基于不同阶段t，其中$t\in 2,3,4$。那么，$y_i$可以表述为：

即DW-CNN操作$d_i$接受前一分割 { x j , j ≤ i } \{x_j,j \le i\} {xj​,j≤i}的特征图输出。
所以，第一部分参考Res2Net的思想，有很小的改动就是分割不再固定

第二部分，为了消除自注意力机制的高复杂度运算量问题，作者借鉴了转置QK注意力特征图。它通过在通道维度上而不是空间维度上使用MSA的点乘操作来达到线性复杂度，这允许作者在通道上计算互协方差去产生注意力的特征图。它具有全局表征的隐式知识。给出一个归一化的张量$Y\in H\times W\times C$，我们使用三个线性层去计算$Q,K,V\in HW\times C$投影，产生：
$$Q=W^{Q}Y,K=W^{K}Y,V=W^{V}Y$$
然后，在计算互协方差注意力前，使用L2归一化Q和K，因为它能稳定训练。不同于标准的自注意力机制在空间维度上计算$Q$和$K^T$之间的点乘，即
$$(HW\times C)\cdot (C\times HW),$$
作者在通道维度上计算$Q^T$和$K$的点乘，即$$(C\times HW)\cdot (HW\times C)$$产生$C\times C$softmax后尺度化的注意力分数矩阵。为了得到最终的注意力图，将V乘以这个分数矩阵并累加，转置注意力操作可以表示为：
$$\hat{X}=Attention(Q,K,V)+X,$$
$$s_{.}{t}_{.},Attention(Q,K,V)=V\cdot softmax(Q^T\cdot K)$$其中$X$输入张量，$\hat{X}$是输出特征张量。然后两个$1\times 1$的point-wise操作，LN和GELU激活层被用于产生非线性的特征。

第二部分参考协方差注意力的机制

实验

Table 2. Classification performance comparison of our proposed EdgeNeXt model with state-of-the-art lightweight fully convolutional, transformer-based, and hybrid models on the ImageNet-1K validation set. Our model outperforms the SOTA models and achieves a better tradeoff between accuracy and compute (i.e., parameters and multiplication-addition (MAdds) operations).
Table展示了当前轻量级的CNN，Transformer和混合架构在ImageNet-1K上的性能。EdgeNeXt超过了SOTA的性能，并在精度和运算量上达到更大的平衡。

---