EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications
Abstract
-
为了追求不断提高的精度,通常会开发大型复杂的神经网络。这种模型需要很高的计算资源,因此不能部署在边缘设备上。由于资源节约型通用网络在多个应用领域的有用性,因此构建资源节约型通用网络具有极大的兴趣。
-
在这项工作中,努力有效地结合CNN和Transformer模型的优点,并提出了一种新的高效混合架构EdgeNeXt。具体来说,在EdgeNeXt中,引入了分裂深度转置注意(STDA)编码器,该编码器将输入张量分割为多个通道组,并利用深度卷积和跨通道维度的自注意来隐式增加接受域并编码多尺度特征。我们在分类、检测和分割任务上的大量实验,揭示了所提出方法的优点,优于最先进的方法,计算需求相对较低。
-
本文的EdgeNeXt模型具有1.3M参数,在ImageNet-1K上达到71.2%的top-1精度,优于MobileViT,绝对增益2.2%,FLOPs减少28%。此外,本文的EdgeNeXt模型具有5.6M参数,在ImageNet-1K上达到了79.4%的top-1精度。
-
EdgeNeXt: 使用 Channel Attention 来提升 CNN 模型,将 Attention 机制引入 CNN 模型,不过本文使用的并非传统的空间注意力机制,而是使用的通道助力机制来减少 Attention计 算量。从 CNN+Attention 的角度探索了轻量级模型设计的新思路,不过总体上讲,相比已有的 CNN 轻量级模型,在速度上仍没有特别明显的优势。由MBZUAI(位于阿布扎比的默罕默德人工智能大学),IIAI(起源研究院,邵岭团队)发表在ECCV2022
-
为了达到轻量化的ViT,首先需要对self-attention操作进行优化,使得复杂度从关于图像二次方下降为线性,其次需要减少transformer模块的堆叠次数,而采用轻量化卷积如DW卷积,来减少参数量,同时不减少图像表征能力。
Introduction
-
卷积神经网络(cnn)和最近引入的视觉变压器(ViTs)在包括物体识别、检测和分割[Deep learning in neural networks: An overview,Transformers in vision: A survey]在内的几个主流计算机视觉任务中显著提高了最先进水平。总的趋势是使网络架构更加深入和复杂,以追求不断提高的准确性。在追求更高精度的同时,大多数现有的基于CNN和vit的架构忽略了计算效率(即模型大小和速度)方面,这对于在资源受限的设备(如移动平台)上运行至关重要。在许多现实应用中,比如机器人和自动驾驶汽车,在资源受限的移动平台上,希望识别过程既准确又具有低延迟。
-
大多数现有方法通常利用精心设计的高效卷积变体,在资源受限的移动平台上实现速度和准确性之间的权衡。除了这些方法之外,很少有现有的工作使用硬件感知神经结构搜索(NAS)来为移动设备构建低延迟的精确模型。虽然这些轻量级cnn在编码局部图像细节方面易于训练和有效,但它们并没有显式地模拟像素之间的全局相互作用。
-
在(ViTs)中引入自我注意使显式建模这种全局相互作用成为可能,然而,这通常是以自我注意计算为代价的缓慢推理。这就成为了为移动视觉应用设计轻量级ViT变体的一个重要挑战。
-
现有的大部分工作采用基于cnn的设计来开发高效的模型。然而,cnn中的卷积运算继承了两个主要的局限性:第一,它具有局部的接受域,因此不能对全局上下文建模;其次,学习到的权重在推理时间是固定的,使得cnn不灵活地适应输入内容。虽然这两个问题都可以通过transformer得到缓解,但它们通常是计算密集型的。最近很少有研究[Edgeformer,Mobilevit]通过结合cnn和vit的优势来设计移动视觉任务的轻量级架构。
-
然而,这些方法主要集中在优化参数和引起较高的乘加运算,这限制了移动设备上的高速推理。MAdds更高,因为注意力块的复杂度是关于输入大小的二次函数。由于网络架构中的多个注意力块,这变得更加成问题。在这里,我们认为,在设计一个统一的移动架构时,模型大小、参数和MAdds都希望相对于资源受限的设备较小,该架构有效地结合了cnn和vit的互补优势(见下图)。
-
-
我们提出的EdgeNeXt模型与SOTA vit和混合架构设计的比较。x轴显示乘法和加法(MAdd)操作,y轴显示前1 ImageNet-1K分类精度。图中每个对应的点都提到了参数的数量。与最近的方法相比,我们的EdgeNeXt显示了更好的计算(参数和MAdds)与精度权衡。
-
Contributions.
-
本文提出了一种新的轻量级架构,名为EdgeNeXt,它在模型尺寸、参数和MAdds方面都很有效,同时在移动视觉任务上的精度更高。具体来说,引入了分裂深度转置注意(SDTA)编码器,该编码器有效地学习了局部和全局表示,以解决cnn中接受阈有限的问题,而不增加参数和MAdd操作的数量。
-
与最先进的移动网络相比,本文出的架构在包括图像分类、对象检测和语义分割在内的各种任务上都表现出了良好的准确性和延迟性能。本文的EdgeNeXt骨干具有5.6M参数和1.3G MAdds,实现了79.4%的前1 ImageNet-1K分类精度,优于最近推出的MobileViT对应工作,同时所需的MAdds减少了35%。对于目标检测和语义分割任务,与文献中所有发表的轻量级模型相比,所提出的EdgeNeXt以更少的MAdds和相当数量的参数实现了更高的mAP和mIOU。
-
Related Work
-
近年来,针对移动视觉任务设计轻量级、硬件高效的卷积神经网络已被广泛研究。目前的方法主要集中在为低功耗边缘设备设计高效的卷积版本。在这些方法中,MobileNet是使用最广泛的架构,它采用了深度可分离卷积。另一方面,ShuffleNet使用信道变换和低成本的组卷积。MobileNetV2引入了带有线性瓶颈的反向残差块,在各种视觉任务上取得了良好的性能。ESPNetv2利用深度扩张卷积来增加网络的接受场,而不增加网络的复杂性。硬件感知神经架构搜索(NAS)也被用于在移动设备上的速度和准确性之间找到更好的平衡点[。尽管这些cnn在移动设备上的训练和推断速度更快,但它们缺乏像素之间的全局交互,这限制了它们的准确性。
-
最近,Desovitskiy等人[An image is worth 16x16 words: Transformers for image recognition at scale]介绍了一种基于自注意机制的视觉转换器架构,用于视觉任务。他们提出的体系结构利用了大规模的预训练数据(例如JFT-300M)、广泛的数据增强和更长的训练计划来实现有竞争力的性能。随后,DeiT提出在该架构中集成蒸馏令牌,只在ImageNet-1K数据集上进行训练。此后,文献中提出了几种vit和混合架构的变体,为vit添加特定于图像的感应偏差,以获得在不同视觉任务上的改进性能。
-
ViT模型在一些视觉识别任务中取得了具有竞争力的结果[vits,Swin transformer]。然而,由于多头自注意(MHA)的计算成本较高,这些模型很难部署在资源受限的边缘设备上。最近的工作是为移动视觉任务设计轻量级混合网络,它结合了cnn和变压器的优势。MobileFormer采用了MobileNetV2和ViTs的并行分支,并通过桥接连接两个分支以实现本地-全局交互。Mobilevit将变压器视为卷积,并提出了一种用于局部全局图像上下文融合的MobileViT块。他们的方法在图像分类上取得了优于以前的轻量级cnn和vit的性能,使用类似的参数预算。
-
虽然MobileViT主要关注优化参数和延迟,但MHA仍然是该模型的主要效率瓶颈,特别是在MHA数量和边缘设备上的推理时间方面。MobileViT中MHA的复杂度是输入大小的二次函数,考虑到MobileViT- s模型中存在9个注意块,这是主要的效率瓶颈。在这项工作中,努力为移动设备设计一种新的轻量级架构,该架构在参数和MAdds方面都是有效的,同时在移动视觉任务中具有更高的准确性。本文提出的架构EdgeNeXt建立在最近引入的CNN方法ConvNeXt之上,该方法根据ViT设计选择对ResNet架构进行了现代化改造。在我们的EdgeNeXt中,我们引入了一个SDTA块,以有效的方式将深度卷积与自适应内核大小以及转置注意结合起来,获得最佳的精度-速度权衡。
EdgeNeXt
-
这项工作的主要目标是开发一种轻量级的混合设计,有效地融合了低功耗边缘设备的ViTs和cnn的优点。ViTs(如MobileViT)的计算开销主要是由于自注意操作。与MobileViT相比,本文模型中的注意力块相对于输入空间维度O(N d2)具有线性复杂度,其中N是补丁的数量,d是特征/通道维度。模型中的自注意操作应用于跨通道维度,而不是空间维度。此外,我们证明,在注意力块数量更少的情况下(在本文中是3个,而在MobileViT中是9个),我们可以超过他们的性能标志。通过这种方式,所提出的框架可以用有限数量的MAdds建模全局表示,这是确保边缘设备上低延迟推理的基本标准。为了激发我们提出的体系结构,我们提出了两个理想的特性。
-
a)有效编码全局信息。
- 自我注意学习全局表征的内在特征对视觉任务至关重要。为了有效地继承这一优势,我们使用交叉协方差注意在相对较少的网络块内跨特征通道维度而不是空间维度合并注意操作。该方法将原有的自注意操作在令牌数量上的复杂度从二次型降低为线性型,并对全局信息进行了有效的隐式编码。
-
b)自适应内核大小。
- 众所周知,大核卷积的计算成本很高,因为参数和flop的数量会随着核大小的增长而成倍增加。尽管较大的内核大小有助于增加接受域,但在整个网络层次结构中使用如此大的内核代价高昂,而且不是最优的。我们提出了一种自适应内核大小机制来降低这种复杂性,并捕获网络中不同级别的特征。受cnn的层次结构的启发,在卷积编码器块的早期阶段使用较小的内核,而在后期阶段使用较大的内核。这种设计选择是最优的,因为CNN的早期阶段通常捕获低级别的特征,较小的内核适合这个目的。然而,在网络的后期阶段,需要大的卷积内核来捕获高级特征。接下来我们将解释架构细节。
-
总体的结构
-
下图说明了所提议的EdgeNeXt架构的概述。主要成分有两方面:
-
(1)自适应N×N convc .编码器,
-
(2)分割深度转置注意(SDTA)编码器。
-
-
本文的EdgeNeXt架构建立在ConvNeXt的设计原则之上,并在四个阶段中以四个不同的规模提取分层特征。大小为H×W ×3的输入图像在网络的开始通过一个patchify stem层,使用一个4×4非重叠卷积和一个层范数来实现,结果是h/4 × w/4 ×C1特征映射。然后,输出被传递给3×3 convc . encoder来提取局部特征。第二阶段从使用2×2跨步卷积实现的下采样层开始,该层将空间大小减少了一半并增加了通道,从而得到H /8 × W /8 ×C2特征映射,然后是两个连续的5×5 convs编码器。位置编码(PE)也只在第二阶段的SDTA块之前添加。观察到PE对于密集的预测任务(例如,对象检测和分割)是敏感的,并且在所有阶段添加它会增加网络的延迟。因此,只在网络中添加一次,以编码空间位置信息。将输出的特征图进一步传递到第三和第四阶段,分别生成H /16× W /16 ×C3和H /32× W /32 ×C4维度特征。
-
-
第一行:我们框架的整体架构是分阶段设计的。
-
在这里,第一阶段使用4 × 4跨步卷积将输入图像降采样到1/4分辨率,然后使用三个3×3卷积(Conv.)编码器。
-
在阶段2-4中,开始时使用2×2跨步卷积进行下采样,随后使用N×N卷积和分割深度转置注意(SDTA)编码器。
-
下面一行:我们展示了convs编码器(左)和SDTA编码器(右)的设计。convl .编码器使用N×N深度卷积进行空间混合,然后使用两个点卷积进行信道混合。SDTA编码器将输入张量分成B通道组,并应用3 × 3深度卷积进行多尺度空间混合。分支之间的跳跃连接增加了网络的整体接受域。分支B3和B4在第3和第4阶段逐渐被激活,增加了网络更深层次的整体接受域。在提出的SDTA中,我们利用转置注意力,然后使用轻量级MLP,将注意力应用于特征通道,并与输入图像具有线性复杂性。
-
-
模型构造了如上图的架构,主要包括 Conv Encoder 和 SDTA Encoder,其中 Conv Encoder 是 ConvNeXt 中的 block 机制,如下图所示,DWConv 代表 Depth-wise Conv。
-
SDTA Encoder 的基本结构,其中左侧是一个级联的卷积结构,将特征按通道分为 s 组,按图中的方式依次进行级联, 3x3 代表的是 depth-wise 卷积,中间为通道自注意力机制,最右侧为一个 FFN 结构。本模型相当于在 ConvNeXt 的基础上在每个阶段的最后,添加了一个轻量级的注意力机制。
-
Convolution Encoder .
-
该块由具有自适应核大小的深度可分离卷积组成。可以通过两个单独的层来定义它:
-
(1)使用自适应N×N内核的深度卷积。分别用k = 3、5、7和9来表示阶段1、2、3和4。
-
(2)使用两个点卷积层与标准层归一化(LN)和高斯误差线性单元(GELU)激活来丰富局部表示,用于非线性特征映射。
-
最后,添加一个跳过连接,使信息在网络层次结构中流动。这个块类似于ConvNeXt块,但是内核大小是动态的,并且根据阶段而变化。我们观察到,与静态内核大小相比,convl .编码器中的自适应内核大小表现得更好。convl .编码器可以表示如下:
-
x i + 1 = x i + L i n e a r G ( L i n e a r ( L N ( D w ( x i ) ) ) ) , ( 1 ) x_{i+1}=x_i+Linear_G(Linear(LN(D_w(x_i)))),(1) xi+1=xi+LinearG(Linear(LN(Dw(xi)))),(1)
-
-
其中xi表示形状为H×W ×C的输入特征图,LinearG是一个点级卷积层,GELU是一个点级卷积层,Dw是k×k深度级卷积层,LN是一个归一化层,xi+1表示convl .编码器的输出特征图。
-
-
SDTA Encoder .
-
在所提出的分割深度转置注意编码器中有两个主要组件。第一个组件通过编码输入图像中的各种空间级别来努力学习自适应多尺度特征表示,第二部分隐式编码全局图像表示。编码器的第一部分受到Res2Net的启发,其中我们采用了多尺度处理方法,将层次表示开发到单个块中。这使得输出特征表示的空间感受场更加灵活和自适应。与Res2Net不同,我们的SDTA编码器中的第一个块不使用1×1点卷积层,以确保具有有限数量的参数和MAdds的轻量级网络。此外,使用自适应数量的子集每个阶段允许有效和灵活的特征编码。在我们的STDA编码器中,我们将输入张量H×W ×C分成s个子集,每个子集用xi表示,并且具有与C/s通道相同的空间大小,其中i∈{1,2,…, s}, C为通道数。每个特征映射子集(除了第一个子集)被传递给3×3深度卷积,用di表示,输出用yi表示。同样,di−1的输出,用yi−1表示,被添加到特征子集xi,然后馈送给di。子集的数量s是基于阶段数t自适应的,其中t∈{2,3,4}。我们可以这样写yi:
-
f ( x ) = { x i if i=1 d i ( x i ) if i=2,t=2 d i ( x i + y i − 1 ) if 2<i<=s,t ( 2 ) f(x) = \begin{cases} x_i &\text{if i=1}\\ d_i(x_i) &\text{if i=2,t=2}\\ d_i(x_i+y_{i-1}) &\text{if 2<i<=s,t}\\ \end{cases}(2) f(x)=⎩ ⎨ ⎧xidi(xi)di(xi+yi−1)if i=1if i=2,t=2if 2<i<=s,t(2)
-
每个深度操作di,如图2中的SDTA编码器所示,接收之前所有分割{xj, j≤i}的特征映射输出。
-
-
如前所述,变压器自注意层的开销对于边缘设备上的视觉任务是不可用的,因为它以更高的MAdds和延迟为代价。为了缓解这个问题并有效地编码全局上下文,我们在SDTA编码器中使用了转置查询和关键注意特征映射。通过跨通道维度而不是空间维度应用MSA的点积操作,该操作具有线性复杂性,这允许计算跨通道的交叉协方差,以生成具有关于全局表示的隐式知识的注意力特征图。给定一个形状为H×W ×C的归一化张量Y,我们使用以下方法计算查询(Q)、键(K)和值(V)投影三个线性层,得到Q=W QY, K=W KY, V =W V Y,维度HW ×C,其中W Q,W K和W V分别是Q, K和V的投影权值。然后,对Q和K采用L2范数,计算训练稳定时的交叉方差注意。
-
我们不再沿着空间维度应用Q和KT之间的点积,即(HW ×C)·(C×HW),而是在QT和K之间的通道维度上应用点积,即(C×HW)·(HW ×C),生成C×C softmax缩放注意评分矩阵。为了得到最终的注意力地图,我们将分数乘以V并将它们相加。注意转置运算可表示为:
-
X ^ = A t t e n t i o n ( Q , K , V ) + X , ( 3 ) s . t . , A t t e n t i o n ( Q , K , V ) = V ⋅ s o f t m a x ( Q T ⋅ K ) , ( 4 ) \hat{X}=Attention(Q,K,V)+X,(3)\\ s.t.,Attention(Q,K,V)=V·softmax(Q^T·K),(4) X^=Attention(Q,K,V)+X,(3)s.t.,Attention(Q,K,V)=V⋅softmax(QT⋅K),(4)
-
其中X是输入,而 X ^ \hat{X} X^是输出特征张量。之后,使用两个1×1点卷积层,LN和GELU激活来生成非线性特征。下表给出了convc .和STDA编码器在每一层对应输入尺寸的顺序,并给出了关于超小、超小和小型号的更多设计细节。
-
-
EdgeNeXt架构。关于输出大小、内核大小和输出通道的模型层的描述,重复n次,以及模型的MAdds和参数。选择小模型、超小模型和超超小模型的输出通道数量,使参数数量与对应的MobileViT模型相匹配。在Conv. Encoder中使用自适应内核大小来降低模型的复杂性并捕获不同级别的特征。此外,我们还填充了最后一个阶段的输出大小,以便能够应用9×9过滤器。
-
对于一张输入图像,分为四个stage,通过不同的stage后,特征层的分辨率分布变成原来的四分之一、八分之一、十六分之一和三十二分之一,这种分层结构可以很方便地嵌入到图像分割和目标检测模型上面。
-
-
self-attention的计算复杂度较高,有多种减少计算量的方法。一:使用小窗口分割,仅在小窗口内计算self-attention,通过信息交互获得其他窗口的信息,信息交互的方式可以通过移窗策略(swin-transformer)或者如LightViT添加一个全局信息表征进行信息交互;二:同样使用小窗口分割,但是self-attention是在不同窗口的对应位置进行计算,然后通过融合卷积获得全局感受野,如mobilevit。
-
本文使用了一种新的信息交互方式且不使用小窗口分割方法。一般在ViT中,是在空间维度上做self-attention。如,对于一个HWC的特征图,要通过K、Q然后计算相关系数,首先通过K、Q后的特征图都是H*W*C,然后计算相关矩阵 Q H W ∗ C ⋅ K C ∗ H W Q^{HW*C}·K^{C*HW} QHW∗C⋅KC∗HW ,得到HW*HW维度的矩阵,这一步计算量是 ( H W ) 2 ⋅ C (HW)^2·C (HW)2⋅C ,是关于HW成二次方复杂度。这个矩阵再做scale,与V矩阵做矩阵乘法。
-
为了将复杂度将为关于图像分辨率成线性的。因此,本文在通道维度上做self-attention,即 Q C ∗ H W ⋅ K H W ∗ C Q^{C*HW}·K^{HW*C} QC∗HW⋅KHW∗C ,得到C*C的矩阵,这一步骤的计算量是 C 2 ( H W ) C^2(HW) C2(HW) ,是关于通道数成二次方的,因此将复杂度降下来了。然后再与V做矩阵乘法。经过self-attention后,通过一个残差连接,然后连续两次线性层获得输出特征,再与STDA的输入做跳跃连接获得多层信息表征。
Experiments
-
在本节中,我们将在ImageNet-1K分类、COCO对象检测和Pascal VOC分割基准上评估我们的EdgeNeXt模型
-
Dataset
- 所有分类实验均使用ImageNet-1K数据集。该数据集为1000个类别提供了大约1.28M的训练图像和50K的验证图像。本文报告了所有实验验证集的top-1精度。对于目标检测,我们使用COCO数据集,该数据集分别提供大约118k的训练图像和5k的验证图像。对于分割,使用Pascal VOC 2012数据集,该数据集提供了近10k张带有语义分割掩码的图像。遵循mobilevit中的标准实践,我们也使用额外数据和注释。
实现细节
-
以256×256的输入分辨率训练EdgeNeXt模型,有效批大小为4096。所有实验均使用AdamW优化器运行300个epoch,学习率为6e-3,权值衰减为0.05。使用余弦学习率计划和20个周期的线性预热。训练期间使用的数据增强是随机调整大小裁剪(RRC)、水平翻转和RandAugment,其中RandAugment仅用于EdgeNeXt-S模型。在训练过程中,还使用了多尺度采样器。进一步的随机深度与0.1的速率仅用于EdgeNeXt-S模型。在训练中使用EMA,动量为0.9995。为了进行推断,将图像大小调整为292×292,然后以256×256分辨率进行中心裁剪。还训练并报告了我们的EdgeNeXt-S模型在224×224分辨率下的准确性,以便与以前的方法进行公平的比较。分类实验在8个A100 gpu上运行,EdgeNeXt-S模型的平均训练时间接近30小时。
-
对于检测和分割任务,我们按照与[mobilevit]相似的设置微调EdgeNeXt,并分别报告IOU为0.50-0.95的平均平均精度(mAP)和平均交叉除以联合(mIOU)。实验在4个A100图形处理器上运行,检测和分割的平均训练时间分别为~ 36和~ 7小时。
-
我们还报告了我们的模型在NVIDIA Jetson Nano和NVIDIA A100 40GB GPU上的延迟。对于Jetson Nano,我们将所有型号转换为TensorRT开发板并使用批量大小为1的FP16模式执行推理。对于A100,类似于[A convnet for the 2020s.],我们使用批处理大小为256的PyTorch v1.8.1来测量延迟。
图像分类
-
下表将我们提出的EdgeNeXt模型与之前最先进的全卷积(ConvNets)、基于变压器(ViTs)和混合模型进行了比较。总体而言,与所有三类方法相比,我们的模型显示出更好的精度(参数和MAdds)权衡。
-
-
与ConvNets的比较。在参数相似的情况下,EdgeNeXt在top-1精度方面超过了轻量级ConvNets(下表)。通常情况下,由于没有注意力计算,ConvNets与变压器和混合模型相比具有更少的MAdds,但是,它们缺乏全局接受场。例如,与MobileNetV2相比,EdgeNeXtS具有更高的MAdds,但它以较少的参数数量获得4.1%的top-1精度增益。此外,本文的EdgeNeXt-S在参数数量相当的情况下,性能分别优于ShuffleNetV2和MobileNetV3 4.3%和3.6%。
-
与ViTs的比较。我们的EdgeNeXt在ImageNet1K数据集上优于最近的ViT变体,具有更少的参数和MAdds。例如,EdgeNeXt-S的top-1准确率为78.8%,绝对裕度分别超过T2T-ViT和DeiT-T 2.3%和6.6%。
-
与混合模型的比较。提出的EdgeNeXt在参数数量和MAdds数量更少的情况下优于MobileFormer , vitc,CoaT- lite - t(上表)。为了与MobileViT进行公平的比较,在256×256的输入分辨率下训练我们的模型,并在不同模型尺寸(即S, XS和XXS)下显示一致的增益,MAdds更少,边缘设备推理更快(下表)。例如,我们的EdgeNeXt-XXS模型仅1.3M参数就实现了71.2%的top-1精度。超过相应的MobileViT版本2.2%。最后,我们的EdgeNeXt-S模型在ImageNet上仅使用5.6M参数就达到了79.4%的精度,与相应的MobileViT-S模型相比,误差为1.0%。这说明了我们设计的有效性和泛化性。
-
-
EdgeNeXt不同变体与MobileViT对应模型的比较。最后两列分别列出了NVIDIA Jetson Nano和A100设备上以毫秒和µs为单位的延迟。我们的EdgeNext模型为每个模型尺寸提供更高的精度和更低的延迟,这表明我们设计的灵活性,可缩小至1.3M参数。
-
此外,我们还使用[for training any backbone to top results]之后的知识蒸馏来训练我们的EdgeNeXt-S模型,并达到了81.1%的top-1 ImageNet精度。
ImageNet-21K Pretraining
-
为了进一步探索EdgeNeXt的容量,本文设计了具有18.5M参数和3.8MAdds的EdgeNeXt- b模型,并在ImageNet-21K数据集的子集上对其进行预训练,然后在标准ImageNet-1K数据集上进行微调。ImageNet-21K(21年冬季发布)包含大约13M张图像和19K个类。我们遵循[Imagenet-21k pretraining for the masses]对预训练数据进行预处理,去除样本较少的类,并将其分为训练集和验证集,分别包含约11M张和522K张图像,超过10450个类。我们将此数据集称为ImageNet-21K-P。我们严格按照[A convnet for the 2020s]的训练配方进行ImageNet-21K-P预处理。此外,我们使用ImageNet-1K预训练模型初始化ImageNet21K-P训练,以实现更快的收敛。最后,我们在ImageNet-1K上微调ImageNet-21K模型30个epoch,学习率为7.5e−5,有效批处理大小为512。结果汇总在下表中。
-
-
EdgeNeXt-B模型的大规模ImageNet-21K-P预训练。与SOTA ConvNeXt和MobileViT-V2相比,本文的模型获得了更好的精度和计算权衡。
Inference on Edge Devices
- 本文计算EdgeNeXt模型在NVIDIA Jetson Nano边缘设备上的推断时间,并将其与最先进的MobileViT模型进行比较。所有模型都转换为TensorRT引擎,并在FP16模式下执行推理。我们的模型在具有相似参数的边缘设备上实现了低延迟,更少的MAdds和更高的top-1精度。还列出了MobileViT和EdgeNeXt模型在A100 GPU上的推断时间。可以观察到,我们的EdgeNeXtXXS模型比A100上的MobileViT- xss模型快约34%,而在Jetson Nano上仅快约8%,这表明EdgeNeXt比MobileViT更好地利用了先进的硬件。
对象检测
-
我们使用EdgeNeXt作为SSDLite的骨干,并在COCO 2017数据集上以320×320的输入分辨率微调模型。SSD和SSDLite之间的区别在于标准卷积被SSD磁头中的可分离卷积所取代。结果如下表所示。与MobileVit主干相比,EdgeNeXt始终优于MobileNet主干,并提供具有竞争力的性能。由于MAdds数量较少,但参数数量相当,EdgeNeXt获得了最高的27.9 box AP,比MobileViT少38%。
-
-
与SOTA在COCO对象检测上的比较。EdgeNeXt比以前的方法有所改进。
语义分割
-
我们在DeepLabv3中使用EdgeNeXt作为骨干,并在Pascal VOC数据集上以512×512的输入分辨率微调模型。DeepLabv3在级联设计中使用扩展卷积以及空间金字塔池来编码多尺度特征,这在编码多尺度对象时非常有用。我们的模型在验证数据集上获得了80.2 mIOU,比MobileViT提供了1.1点的增益,MAdds减少了36%。
-
-
VOC语义分割与SOTA的比较。我们的模型提供了合理的收益。
Ablations
-
在本节中,我们将在EdgeNeXt模型中删除不同的设计选择。
-
SDTA编码器和自适应内核大小。下表说明了SDTA编码器和自适应内核大小在我们提议的体系结构中的重要性。用卷积编码器代替SDTA编码器,精度降低了1.1%,这表明SDTA编码器在我们的设计中是有用的。当我们在网络的所有四个阶段都将内核大小固定为7时,它进一步降低了0.4%的准确性。总的来说,本文提出的设计提供了一个最佳的速度-精度权衡。我们还在下表中删除了SDTA组件的贡献(例如,自适应分支和位置编码)。去除自适应分支和位置编码会略微降低准确性。
-
-
对EdgeNeXt和SDTA编码器的不同组件进行烧蚀。结果表明,SDTA编码器和自适应内核在我们的设计中的优势。此外,SDTA模块还需要自适应分支和位置编码(PE)。
Hybrid design.
-
下表列出了EdgeNeXt模型的不同混合设计选择。从MetaFormer的激励下,我们用SDTA编码器替换了最后两个阶段的所有卷积模块。当后两个阶段的所有块都是SDTA块时,结果显示出更好的性能,但它增加了延迟(第2行vs第3行)。在我们的混合设计中,我们建议使用SDTA模块作为最后三个阶段的最后一个模块,这提供了一个最佳的速度-精度权衡。
-
-
混合结构上的消融。使用一个SDTA编码器作为最后三个阶段的最后一个块提供了最佳的精度-延迟权衡。
-
下表提供了在网络的不同阶段使用SDTA编码器的重要性。值得注意的是,逐步添加SDTA编码器作为最后三个阶段的最后一个块提高了精度,但减少了推理延迟。然而,在第4行中,我们获得了精度和速度之间的最佳权衡,其中SDTA编码器被添加为网络最后三个阶段的最后一个块。此外,我们注意到在网络的第一阶段添加一个全局SDTA编码器在功能还不太成熟的情况下是没有帮助的。
-
-
利用SDTA编码器在网络的不同阶段进行消融。在后三个阶段中包含SDTA编码器可以提高性能,而在第一阶段中使用全局SDTA编码器没有帮助。
-
我们还提供了在每个阶段开始时使用SDTA模块与在结束时使用SDTA模块的区别。下表显示,在每个阶段的末尾使用全局SDTA编码器更有益。这一观察结果与最近的工作[Improved multiscale vision transformers for classification and detection]一致。
-
-
在EdgeNeXt的每个阶段的开始和结束使用SDTA进行消融。结果表明,在每个阶段的末尾使用SDTA通常是有益的。
-
激活和归一化。EdgeNeXt在整个网络中使用GELU激活和层规范化。我们发现,目前PyTorch对GELU和层规范化的实现对于高速推理并不是最优的。为此,我们用Hard-Swish代替GELU,用batch-norm代替layer-norm,并重新训练我们的模型。从下图可以看出,它稍微降低了精度,但是却大大降低了延迟。
-
-
消融使用不同的激活函数和归一化层对我们的网络变体的准确性和延迟的影响。使用Hard Swish激活和批处理归一化而不是GELU和层归一化显著地提高了延迟,但代价是准确性的损失。
-
定性结果
-
下图分别显示了EdgeNeXt检测模型和分割模型的定性结果。我们的模型可以在不同的视图中检测和分割对象。
-
本文的EdgeNeXt检测模型在COCO验证数据集中的定性结果。该模型在具有80个检测类的COCO数据集上进行训练。该模型能够有效地对不同场景下的物体进行定位和分类。
-
-
我们的EdgeNeXt分割模型在未见COCO验证数据集上的定性结果。该模型在Pascal VOC数据集上训练,该数据集包含20个分割类。(a)为预测的语义分割掩码,其中“黑色”表示背景像素。(b)在原始图像上显示预测的掩码。©表示所显示分割掩码的所有Pascal VOC类的颜色编码。我们的模型在未见过的COCO图像上提供了高质量的分割掩码。
-
-
Conclusion
- 与cnn相比,transformer模型的成功带来了更高的计算开销。自注意操作是造成这种开销的主要原因,与基于cnn的移动架构相比,这使得边缘设备上的视觉transformer速度较慢。在本文中,我们介绍了一种由卷积和高效的基于自我注意力的编码器组成的混合设计,以有效地联合建模局部和全局信息,同时在视觉任务的参数和MAdds方面具有比最先进的方法更好的性能。实验结果表明,不同变体的EdgeNeXt具有良好的性能,证明了所提出模型的有效性和泛化能力。