摘要
https://arxiv.org/pdf/2206.10589
为了追求更高的准确性,通常会开发大型且复杂的神经网络。这些模型需要高计算资源,因此无法部署在边缘设备上。构建资源高效通用网络在多个应用领域都非常有用,因此备受关注。在本研究中,我们努力有效地结合卷积神经网络(CNN)和Transformer模型的优势,并提出了一种新的高效混合架构EdgeNeXt。在EdgeNeXt中,我们特别引入了拆分深度转置注意力(STDA)编码器,该编码器将输入张量拆分为多个通道组,并利用深度卷积以及跨通道维度的自注意力来隐式增加感受野并编码多尺度特征。我们在分类、检测和分割任务上进行了大量实验,揭示了所提方法的优点,其在相对较低的计算需求下优于最先进的方法。我们的EdgeNeXt模型具有130万个参数,在ImageNet-1K上实现了71.2%的Top-1准确率,优于MobileViT,绝对增益为2.2%,同时浮点运算次数(FLOPs)减少了28%。此外,我们的EdgeNeXt模型具有560万个参数,在ImageNet-1K上实现了79.4%的Top-1准确率。代码和模型可在https://t.ly/_Vu9上获取。
关键词:边缘设备、混合模型、卷积神经网络、自注意力、Transformer、图像分类、目标检测和分割
1 引言
卷积神经网络(CNN)和最近引入的视觉Transformer(ViT)在包括目标识别、检测和分割在内的几个主流计算机视觉任务中显著推进了最先进的技术[37,20]。普遍趋势是为了追求不断提高的准确性,使网络架构更深、更复杂。然而,在追求更高准确性的同时,大多数现有的基于CNN和ViT的架构忽略了计算效率(即模型大小和速度)方面,这对于在资源受限的设备(如移动平台)上运行至关重要。在许多实际应用中,例如机器人和自动驾驶汽车,期望识别过程在资源受限的移动平台上既准确又具有低延迟。
大多数现有方法通常利用精心设计的卷积高效变体,以在资源受限的移动平台上实现速度和准确性之间的权衡[19,28,36]。除了这些方法之外,少数现有工作[16,39]采用硬件感知的神经架构搜索(NAS)来为移动设备构建低延迟准确模型。上述轻量级CNN易于训练,在编码局部图像细节方面高效,但它们没有显式建模像素之间的全局交互。
视觉Transformer(ViT)[8]中自注意力的引入使得显式建模这种全局交互成为可能,然而,这通常以自注意力计算的慢推理为代价[24]。这为设计用于移动视觉应用的轻量级ViT变体带来了重要挑战。
大多数现有工作在开发高效模型时采用基于CNN的设计。然而,CNN中的卷积运算存在两个主要限制:首先,它具有局部感受野,因此无法建模全局上下文;其次,学习到的权重在推理时是固定的,使得CNN无法灵活适应输入内容。虽然这两个问题都可以通过Transformer来缓解,但它们通常计算密集。最近的一些工作[46,29]研究了通过结合CNN和ViT的优势来为移动视觉任务设计轻量级架构。然而,这些方法主要关注优化参数,并产生了更高的乘加(MAdds)运算,这限制了移动设备上的高速推理。由于注意力块的复杂度与输入大小成二次方关系[29],因此MAdds更高。在网络架构中有多个注意力块时,这一问题进一步加剧。在这里,我们认为在设计一个有效结合CNN和ViT互补优势的统一移动架构时,模型大小、参数和MAdds对于资源受限设备而言都应尽可能小(见图1)。
贡献。我们提出了一种新的轻量级架构,名为EdgeNeXt,它在模型大小、参数和MAdds方面都很高效,同时在移动视觉任务上的准确性也很优越。具体来说,我们引入了拆分深度转置注意力(SDTA)编码器,该编码器有效地学习了局部和全局表示,解决了CNN中感受野有限的问题,而没有增加参数数量和MAdd运算。与最先进的移动网络相比,我们的提出架构在各种任务(包括图像分类、目标检测和语义分割)的准确性和延迟方面都表现出良好的性能。我们的EdgeNeXt主干网络具有560万个参数和13亿次MAdds,实现了79.4%的ImageNet-1K分类Top-1准确率,优于最近引入的MobileViT对应网络[29],同时MAdds减少了35%。对于目标检测和语义分割任务,与文献中发表的所有轻量级模型相比,提出的EdgeNeXt实现了更高的平均精度(mAP)和平均交并比(mIoU),同时MAdds更少,参数数量相当。
2 相关工作
近年来,为移动视觉任务设计轻量级硬件高效卷积神经网络已在文献中得到了充分研究。当前的方法专注于为低功耗边缘设备设计高效的卷积版本[19,17]。在这些方法中,MobileNet[17]是最广泛使用的架构,它采用了深度可分离卷积[5]。另一方面,ShuffleNet[47]使用通道混洗和低成本分组卷积。MobileNetV2[36]引入了具有线性瓶颈的倒残差块,在各种视觉任务上取得了有前景的性能。ESPNetv2[31]利用深度膨胀卷积来增加网络的感受野,而不增加网络复杂性。还探索了硬件感知的神经架构搜索(NAS),以在移动设备上找到速度和准确性之间更好的权衡[16,39]。尽管这些CNN在移动设备上训练和推理更快,但它们缺乏像素之间的全局交互,从而限制了其准确性。
最近,Dosovitskiy等人[8]基于自注意力机制[41]引入了一种用于视觉任务的视觉Transformer架构。他们提出的架构利用大规模预训练数据(如JFT-300M)、广泛的数据增强和更长的训练计划来获得具有竞争力的性能。随后,DeiT[40]提出在该架构中集成蒸馏标记,并且仅在ImageNet-1K[35]数据集上进行训练。自那时起,文献中提出了几种视觉Transformer(ViT)的变体和混合架构,为ViT增加了图像特定的归纳偏置,以在不同视觉任务上获得更好的性能[38,9,42,48,11]。
ViT模型在多个视觉识别任务中取得了具有竞争力的结果[8,24]。然而,由于多头自注意力(MHA)的高计算成本,这些模型很难部署在资源受限的边缘设备上。最近有研究致力于设计用于移动视觉任务的轻量级混合网络,这些网络结合了卷积神经网络(CNN)和Transformer的优点。MobileFormer[4]采用了MobileNetV2[36]和ViT[8]的并行分支,并通过一个桥接结构连接两个分支以实现局部-全局交互。Mehta等人[29]将Transformer视为卷积,并提出了一种MobileViT块用于局部-全局图像上下文融合。他们的方法在图像分类上取得了优越的性能,超过了以前使用相似参数量的轻量级CNN和ViT。
尽管MobileViT[29]主要关注优化参数和延迟,但MHA仍然是该模型中的主要效率瓶颈,特别是对于乘加运算(MAdds)的数量和边缘设备上的推理时间。在MobileViT-S模型中,由于其现有的九个注意力块,MHA的复杂度与输入大小呈二次方关系,这是主要的效率瓶颈。在本工作中,我们努力为移动设备设计一种新的轻量级架构,该架构在参数和MAdds方面都是高效的,同时在移动视觉任务上具有更高的准确性。我们提出的架构EdgeNeXt是基于最近引入的CNN方法ConvNeXt[25]构建的,该方法遵循ViT的设计选择对ResNet[14]架构进行了现代化改造。在我们的EdgeNeXt中,我们引入了一个SDTA块,该块以高效的方式结合了具有自适应核大小的深度卷积和转置注意力,从而获得了最佳的准确性-速度权衡。
3 EdgeNeXt
本工作的主要目标是开发一种轻量级混合设计,有效地融合ViT和CNN的优点,用于低功耗边缘设备。ViT(如MobileViT[29])中的计算开销主要来自于自注意力操作。与MobileViT相比,我们模型中的注意力块对于输入空间维度的复杂度是线性的,为 O ( N d 2 ) \mathcal{O}\left(Nd^{2}\right) O(Nd2),其中 N N N是块的数量, d d d是特征/通道维度。我们模型中的自注意力操作是跨通道维度应用的,而不是空间维度。此外,我们证明,即使使用数量少得多的注意力块(与MobileViT中的9个相比,我们使用了3个),我们也可以超过其性能水平。这样,所提出的框架可以用有限数量的MAdds来建模全局表示,这是确保边缘设备上低延迟推理的基本标准。为了论证我们提出的架构,我们展示了两个理想的特性。
a) 高效地编码全局信息。自注意力学习全局表征的内在特性对于视觉任务至关重要。为了有效地继承这一优势,我们使用交叉协方差注意力,在相对较少的网络块内,将注意力操作纳入特征通道维度,而非空间维度。这将原始自注意力操作的复杂度从二次方降低到线性(就标记数量而言),并有效地隐式编码全局信息。
b) 自适应核大小。众所周知,大核卷积的计算成本很高,因为随着核大小的增加,参数数量和浮点运算次数(FLOPs)会二次方增加。虽然较大的核大小有助于增加感受野,但在整个网络层次结构中使用如此大的核既昂贵又非最优。我们提出了一种自适应核大小机制来降低这种复杂度,并在网络中捕获不同级别的特征。受卷积神经网络(CNN)层次结构的启发,我们在卷积编码器块的早期阶段使用较小的核,而在后期阶段使用较大的核。这种设计选择是最佳的,因为CNN的早期阶段通常捕获低级特征,而较小的核适合这一目的。然而,在网络的后期阶段,需要较大的卷积核来捕获高级特征[45]。接下来,我们将解释我们的架构细节。
整体架构。图2展示了所提出的EdgeNeXt架构的概览。主要组成部分有两个:(1)自适应 N × N N \times N N×N卷积编码器,和(2)分割深度转置注意力(SDTA)编码器。我们的EdgeNeXt架构基于ConvNeXt[25]的设计原则,并在四个阶段中提取四个不同尺度的层次特征。大小为 H × W × 3 H \times W \times 3 H×W×3的输入图像在网络开始时通过一个补丁化主干层,该层使用 4 × 4 4 \times 4 4×4非重叠卷积,后跟一层归一化,从而得到 H 4 × W 4 × C 1 \frac{H}{4} \times \frac{W}{4} \times C1 4H×4W×C1的特征图。然后,将输出传递给 3 × 3 3 \times 3 3×3卷积编码器以提取局部特征。第二阶段以一个下采样层开始,该层使用 2 × 2 2 \times 2 2×2步长卷积实现,将空间尺寸减半并增加通道数,从而得到 H 8 × W 8 × C 2 \frac{H}{8} \times \frac{W}{8} \times C2 8H×8W×C2的特征图,然后是两个连续的 5 × 5 5 \times 5 5×5卷积编码器。仅在第二阶段的SDTA块之前添加位置编码(PE)。我们观察到,PE对于密集预测任务(如目标检测和分割)很敏感,并且在所有阶段中添加它会增加网络的延迟。因此,我们只在网络中添加一次PE来编码空间位置信息。输出特征图进一步传递到第三和第四阶段,分别生成 H 16 × W 16 × C 3 \frac{H}{16} \times \frac{W}{16} \times C3 16H×16W×C3和 H 32 × W 32 × C 4 \frac{H}{32} \times \frac{W}{32} \times C4 32H×32W×C4维特征。
卷积编码器。该块由具有自适应核大小的深度可分离卷积组成。我们可以将其定义为两个单独的层:(1)具有自适应
N
×
N
N \times N
N×N核的深度卷积。对于阶段1、2、3和4,我们分别使用
k
=
3
,
5
,
7
k=3,5,7
k=3,5,7和9。然后,(2)使用两个逐点卷积层来丰富局部表征,同时采用标准的层归一化[2](LN)和高斯误差线性单元[15](GELU)激活函数进行非线性特征映射。最后,添加一个跳跃连接以使信息在网络层次结构中流动。该块与ConvNeXt块类似,但核大小是动态的,并且根据阶段而变化。我们观察到,与静态核大小相比,卷积编码器中的自适应核大小表现更好(表8)。卷积编码器可以表示如下:
x
i
+
1
=
x
i
+
Linear
G
(
Linear
(
L
N
(
D
w
(
x
i
)
)
)
)
\boldsymbol{x}_{i+1}=\boldsymbol{x}_{i}+\text{Linear}_{G}(\text{Linear}(LN(Dw(\boldsymbol{x}_{i}))))
xi+1=xi+LinearG(Linear(LN(Dw(xi)))),
其中, x i \boldsymbol{x}_{i} xi表示形状为 H × W × C H \times W \times C H×W×C的输入特征图, Linear G \text{Linear}_{G} LinearG是由GELU激活函数跟随的逐点卷积层, D w Dw Dw是 k × k k \times k k×k深度卷积, L N LN LN是归一化层,而 x i + 1 \boldsymbol{x}_{i+1} xi+1表示卷积编码器的输出特征图。
SDTA编码器。所提出的分裂深度转置注意力(SDTA)编码器有两个主要组件。第一个组件致力于通过学习输入图像内的各种空间层次来获取自适应多尺度特征表示,而第二个组件则隐式地编码全局图像表示。我们的编码器的第一部分受到Res2Net[12]的启发,我们通过将层次表示开发为单个块来采用多尺度处理方法。这使得输出特征表示的空间感受野更加灵活和自适应。与Res2Net不同的是,我们的SDTA编码器的第一个块不使用 1 × 1 1 \times 1 1×1逐点卷积层,以确保在参数和MAdds数量受限的情况下网络保持轻量。此外,我们在每个阶段使用自适应数量的子集,以实现有效且灵活的特征编码。在我们的STDA编码器中,我们将输入张量 H × W × C H \times W \times C H×W×C拆分为 s s s个子集,每个子集用 x i \boldsymbol{x}_{i} xi表示,并具有与 C / s C / s C/s通道相同的空间大小,其中 i ∈ { 1 , 2 , … , s } i \in\{1,2, \ldots, s\} i∈{1,2,…,s}, C C C是通道数。每个特征图子集(第一个子集除外)都传递给 3 × 3 3 \times 3 3×3深度卷积,用 d i d_{i} di表示,其输出用 y i \boldsymbol{y}_{i} yi表示。此外, d i − 1 d_{i-1} di−1的输出,用 y i − 1 \boldsymbol{y}_{i-1} yi−1表示,被加到特征子集 x i \boldsymbol{x}_{i} xi上,然后输入到 d i d_{i} di中。子集的数量 s s s根据阶段数 t t t(其中 t ∈ { 2 , 3 , 4 } t \in\{2,3,4\} t∈{2,3,4})是自适应的。我们可以将 y i \boldsymbol{y}_{i} yi表示如下:
y i = { x i i = 1 d i ( x i ) i = 2 , t = 2 d i ( x i + y i − 1 ) 2 < i ≤ s , t \boldsymbol{y}_{i}=\left\{\begin{array}{ll} \boldsymbol{x}_{i} & i=1 \\ d_{i}\left(\boldsymbol{x}_{i}\right) & i=2, t=2 \\ d_{i}\left(\boldsymbol{x}_{i}+\boldsymbol{y}_{i-1}\right) & 2<i \leq s, t \end{array}\right. yi=⎩ ⎨ ⎧xidi(xi)di(xi+yi−1)i=1i=2,t=22<i≤s,t
在图2所示的SDTA编码器中,每个深度操作 d i d_{i} di都会接收来自所有先前拆分 { x j , j ≤ i } \left\{\boldsymbol{x}_{j}, j \leq i\right\} {xj,j≤i}的特征图输出。
如前所述,对于边缘设备上的视觉任务而言,变压器自注意力层的开销是不可行的,因为它会以更高的MAdds和延迟为代价。为了缓解这个问题并有效地编码全局上下文,我们在SDTA编码器[1]中使用了转置的查询和键注意力特征图。此操作通过在通道维度而不是空间维度上应用MSA的点积操作,具有线性复杂度,这使我们能够计算跨通道的互协方差,从而生成具有关于全局表示的隐式知识的注意力特征图。给定形状为 H × W × C H \times W \times C H×W×C的归一化张量 Y \boldsymbol{Y} Y,我们使用三个线性层计算查询( Q \boldsymbol{Q} Q)、键( K \boldsymbol{K} K)和值( V \boldsymbol{V} V)投影,得到 Q = W Q Y \boldsymbol{Q}=\boldsymbol{W}^{Q} \boldsymbol{Y} Q=WQY, K = W K Y \boldsymbol{K}=\boldsymbol{W}^{K} \boldsymbol{Y} K=WKY和 V = W V Y \boldsymbol{V}=\boldsymbol{W}^{V} \boldsymbol{Y} V=WVY,维度为 H W × C H W \times C HW×C,其中 W Q \boldsymbol{W}^{Q} WQ、 W K \boldsymbol{W}^{K} WK和 W V \boldsymbol{W}^{V} WV分别是 Q \boldsymbol{Q} Q、 K \boldsymbol{K} K和 V \boldsymbol{V} V的投影权重。然后,在计算互协方差注意力之前,对 Q \boldsymbol{Q} Q和 K \boldsymbol{K} K应用L2范数,以稳定训练。我们不是沿着空间维度(即 ( H W × C ) ⋅ ( C × H W ) (H W \times C) \cdot(C \times H W) (HW×C)⋅(C×HW))在 Q \boldsymbol{Q} Q和 K T \boldsymbol{K}^{T} KT之间应用点积,而是在 Q T \boldsymbol{Q}^{T} QT和 K \boldsymbol{K} K之间沿着通道维度(即 ( C × H W ) ⋅ ( H W × C ) (C \times H W) \cdot(H W \times C) (C×HW)⋅(HW×C))应用点积,生成 C × C C \times C C×C的softmax缩放注意力得分矩阵。为了获得最终的注意力图,我们将得分乘以 V \boldsymbol{V} V并将它们相加。转置注意力操作可以表示如下:
X ^ = Attention ( Q , K , V ) + X s.t., Attention ( Q , K , V ) = V ⋅ softmax ( Q T ⋅ K ) \begin{array}{c} \hat{\boldsymbol{X}}=\operatorname{Attention}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})+\boldsymbol{X} \\ \text{s.t., } \operatorname{Attention}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\boldsymbol{V} \cdot \operatorname{softmax}\left(\boldsymbol{Q}^{T} \cdot \boldsymbol{K}\right) \end{array} X^=Attention(Q,K,V)+Xs.t., Attention(Q,K,V)=V⋅softmax(QT⋅K)
其中,
X
\boldsymbol{X}
X是输入,
X
^
\hat{\boldsymbol{X}}
X^是输出特征张量。之后,使用两个
1
×
1
1 \times 1
1×1逐点卷积层、LN和GELU激活函数来生成非线性特征。表1显示了Conv.和STDA编码器的序列,以及每层对应的输入大小和关于极小、小和小型模型的更多设计细节。
4 实验
在本节中,我们在ImageNet-1K分类、COCO目标检测和Pascal VOC分割基准上评估了我们的EdgeNeXt模型。
4.1 数据集
在所有分类实验中,我们使用ImageNet-1K [35]数据集。该数据集为1000个类别提供了约128万张训练图像和5万张验证图像。根据文献[17,29],我们在所有实验中报告验证集上的Top-1准确率。对于目标检测,我们使用COCO [22]数据集,该数据集分别提供了约11.8万张训练图像和5千张验证图像。对于分割任务,我们使用Pascal VOC 2012数据集[10],该数据集提供了近1万张带有语义分割掩码的图像。按照[29]中的标准做法,我们还使用了来自[22]和[13]的额外数据和标注。
4.2 实现细节
我们以 256 × 256 256 \times 256 256×256的输入分辨率和4096的有效批处理大小来训练我们的EdgeNeXt模型。所有实验均使用AdamW [27]优化器运行300个周期,学习率和权重衰减分别设置为 6 e − 3 6 \mathrm{e}-3 6e−3和0.05。我们使用余弦学习率调度[26],并在前20个周期进行线性预热。训练期间使用的数据增强方法包括随机调整大小的裁剪(RRC)、水平翻转和RandAugment [6],其中RandAugment仅用于EdgeNeXt-S模型。我们还在训练期间使用了多尺度采样器[29]。此外,仅对EdgeNeXt-S模型使用速率为0.1的随机深度[18]。在训练期间,我们使用动量为0.9995的指数移动平均(EMA)[32]。对于推理,图像被调整大小为 292 × 292 292 \times 292 292×292,然后以 256 × 256 256 \times 256 256×256的分辨率进行中心裁剪。为了与以前的方法进行公平比较,我们还训练并报告了EdgeNeXt-S模型在 224 × 224 224 \times 224 224×224分辨率下的准确率。分类实验在八张A100 GPU上进行,EdgeNeXt-S模型的平均训练时间约为30小时。
对于检测和分割任务,我们按照与[29]类似的设置对EdgeNeXt进行微调,并分别报告在IOU为0.50-0.95时的平均精度(mAP)和平均交并比(mIOU)。实验在四张A100 GPU上进行,检测和分割的平均训练时间分别约为36小时和7小时。
我们还报告了我们的模型在NVIDIA Jetson Nano 1 { }^{1} 1和NVIDIA A100 40GB GPU上的延迟。对于Jetson Nano,我们将所有模型转换为TensorRT 2 { }^{2} 2引擎,并在FP16模式下使用批处理大小为1进行推理。对于A100,与[25]类似,我们使用PyTorch v1.8.1和批处理大小为256来测量延迟。
4.3 图像分类
表2将我们提出的EdgeNeXt模型与以前的最先进全卷积(ConvNets)、基于变换器(ViTs)和混合模型进行了比较。总体而言,与所有这三类方法相比,我们的模型在计算(参数和MAdds)与准确率之间的权衡方面表现更好(见图1)。
与ConvNets的比较。在参数相似的情况下,EdgeNeXt在Top-1准确率方面以相当大的优势超越了轻量级ConvNets(表2)。通常,由于没有注意力计算,ConvNets与变换器和混合模型相比具有更少的MAdds,但它们缺乏全局感受野。例如,与MobileNetV2 [36]相比,EdgeNeXt-S具有更高的MAdds,但它在Top-1准确率上获得了4.1%的提升,且参数数量更少。此外,我们的EdgeNeXt-S在参数数量相当的情况下,分别比ShuffleNetV2 [28]和MobileNetV3 [16]高出了4.3%和3.6%。
与ViTs的比较。我们的EdgeNeXt在ImageNet1K数据集上超越了最近的ViT变体,且参数和MAdds更少。例如,EdgeNeXt-S获得了78.8%的Top-1准确率,分别比T2T-ViT [44]和DeiT-T [40]高出了2.3%和6.6%的绝对优势。
与混合模型的比较。我们提出的EdgeNeXt在参数数量和MAdds更少的情况下,超越了MobileFormer [4]、ViT-C [42]、CoaT-Lite-T [7](表2)。为了与MobileViT [29]进行公平比较,我们以
256
×
256
256 \times 256
256×256的输入分辨率训练我们的模型,并显示不同模型大小(即S、XS和XXS)在边缘设备上具有更少的MAdds和更快的推理速度的一致增益(表3)。例如,我们的EdgeNeXt-XXS模型仅以1.3M参数就达到了71.2%的Top-1准确率,比相应的MobileViT版本高出了2.2%。最后,我们的EdgeNeXt-S模型在ImageNet上仅以5.6M参数就达到了79.4%的准确率,比相应的MobileViT-S模型高出了1.0%。这证明了我们的设计的有效性和泛化性。
此外,我们还按照[34]的方法使用知识蒸馏训练了我们的EdgeNeXt-S模型,并实现了81.1%的Top-1 ImageNet准确率。
4.4 ImageNet-21K 预训练
为了进一步探索EdgeNeXt的能力,我们设计了拥有1850万参数和380万次加法的EdgeNeXt-B模型,并首先在ImageNet-21K [35]数据集的一个子集上进行预训练,然后在标准的ImageNet-1K数据集上进行微调。ImageNet-21K('21年冬季版)包含约1300万张图像和1.9万个类别。我们遵循[33]的方法对预训练数据进行预处理,移除样本较少的类别,并将其分为训练集和验证集,分别包含约1100万张和52.2万张图像,涵盖10450个类别。我们将此数据集称为ImageNet-21K-P。我们严格遵循[25]中的训练方案对ImageNet-21K-P进行预训练。此外,我们使用ImageNet-1K预训练模型来初始化ImageNet 21K-P的训练,以加快收敛速度。最后,我们在ImageNet-1K上对ImageNet-21K模型进行30个周期的微调,学习率设为
7.5
×
1
0
−
5
7.5 \times 10^{-5}
7.5×10−5,有效批量大小设为512。结果总结在表4中。
4.5 边缘设备上的推理
我们在NVIDIA Jetson Nano边缘设备上计算了EdgeNeXt模型的推理时间,并将其与最先进的MobileViT [29]模型进行了比较(表3)。所有模型都转换为TensorRT引擎,并以FP16模式进行推理。我们的模型在边缘设备上实现了低延迟,具有相似的参数数量、更少的MAdds和更高的Top-1准确率。表3还列出了MobileViT和EdgeNeXt模型在A100 GPU上的推理时间。可以观察到,与在Jetson Nano上仅快约8%相比,我们的EdgeNeXtXXS模型在A100上比MobileViT-XSS模型快约34%,这表明与MobileViT相比,EdgeNeXt能更好地利用先进硬件。
4.6 目标检测
我们将EdgeNeXt作为SSDLite的主干网络,并在输入分辨率为
320
×
320
320 \times 320
320×320的COCO 2017数据集[22]上对模型进行微调。SSD [23]和SSDLite的区别在于,SSD头部中的标准卷积被替换为可分离卷积。结果报告在表5中。EdgeNeXt始终优于MobileNet主干网络,并且与MobileVit主干网络相比具有竞争性的性能。在MAdds数量更少且参数数量相当的情况下,EdgeNeXt实现了最高的27.9框AP,比MobileViT少约38%的MAdds。
4.7 语义分割
我们将EdgeNeXt作为DeepLabv3 [3]的主干网络,并在输入分辨率为
512
×
512
512 \times 512
512×512的Pascal VOC [10]数据集上对模型进行微调。DeepLabv3使用级联设计中的膨胀卷积以及空间金字塔池化来编码多尺度特征,这对于编码多尺度对象非常有用。我们的模型在验证数据集上获得了80.2的mIOU,比MobileViT提高了1.1个百分点,同时MAdds减少了约36%。
5 消融实验
在本节中,我们对提出的EdgeNeXt模型中的不同设计选择进行了消融实验。SDTA编码器和自适应核大小。表7说明了SDTA编码器和自适应核大小在我们提出的架构中的重要性。将SDTA编码器替换为卷积编码器会使准确率降低1.1%,这表明SDTA编码器在我们的设计中的有效性。当我们将网络中所有四个阶段的核大小固定为7时,准确率会进一步降低0.4%。总体而言,我们提出的设计提供了最佳的速度-准确率权衡。
我们还在表7中消融了SDTA组件(如自适应分支和位置编码)的贡献。移除自适应分支和位置编码会略微降低准确率。
混合设计。表8对我们EdgeNeXt模型的不同混合设计选择进行了消融实验。受MetaFormer [43]的启发,我们将最后两个阶段中的所有卷积模块替换为SDTA编码器。结果表明,当最后两个阶段的所有块都是SDTA块时,性能更优,但会增加延迟(第2行与第3行对比)。我们的混合设计中,建议在最后三个阶段的最后一个块中使用SDTA模块,这提供了最佳的速度-准确率权衡。
表9展示了在网络不同阶段使用SDTA编码器的重要性消融实验。值得注意的是,逐步将SDTA编码器作为最后三个阶段中的最后一个模块加入,可以在略微牺牲推理延迟的情况下提高准确性。然而,在第4行中,当SDTA编码器作为网络最后三个阶段中的最后一个模块加入时,我们获得了准确性和速度之间的最佳权衡。此外,我们注意到,在网络的第一阶段添加全局SDTA编码器并没有帮助,因为此时的特征尚不成熟。
我们还对在每个阶段开始与结束时使用SDTA模块进行了消融实验。表10显示,在每个阶段结束时使用全局SDTA编码器更有益。这一观察结果与最近的工作[21]一致。
激活与归一化。EdgeNeXt在整个网络中使用了GELU激活函数和层归一化。我们发现,当前PyTorch中GELU和层归一化的实现对于高速推理来说并不理想。为此,我们将GELU替换为Hard-Swish,将层归一化替换为批归一化,并重新训练了我们的模型。图3表明,这略微降低了准确性,但大幅减少了延迟。
6 定性结果
图4和图5分别展示了EdgeNeXt检测和分割模型的定性结果。我们的模型可以检测并分割不同视图中的物体。
7 结论
与卷积神经网络(CNN)相比,Transformer模型的成功带来了更高的计算开销。自注意力操作是这一开销的主要来源,这使得视觉Transformer在与基于CNN的移动架构相比时,在边缘设备上运行较慢。在本文中,我们提出了一种混合设计,该设计结合了卷积和基于高效自注意力的编码器,以有效地联合建模局部和全局信息,同时在参数和乘加运算(MAdds)方面表现高效,并且在视觉任务上的性能优于最先进的方法。我们的实验结果表明,EdgeNeXt的不同变体展现出了令人鼓舞的性能,这证明了所提出模型的有效性和泛化能力。