作者丨Ziyang Li

編輯丨極市平臺

寫在前面的話

本文從減輕計算資源消耗以部署於邊緣裝置入手，提出EdgeNeXt混合架構，該架構引入分割深度轉置注意力（SDTA）編碼器，透過通道組分割與合理的注意力機制來提高資源利用率。該架構憑藉1。3M的引數量在ImageNet-1K上表現超過MobileViT，擴容後憑藉5。6M的引數量在ImageNet-1K上達到79。4%的最高精度。

論文地址：https：//arxiv。org/abs/2206。10589

程式碼連結：https：//github。com/mmaaz60/EdgeNeXt

一、為什麼會想到分割深度轉置注意力？

EdgeNeXt是一種集CNN與Transformer於一體的混合架構，鑑於現有大多流行方法，在追求更高精度的同時並不能很好的保持更快推理速度，以至於在資源受限裝置上的部署表現不足。

因此EdgeNeXt希望可以在保持較小模型容量的同時提升模型效能。而CNN作為捕捉區域性特徵的經典代表，自然被作為架構設計的核心。然而隨著現有研究的進行，一些經典的

CNN有如下兩個主要限制

：

首先，CNN是使用區域性感受野來進行特徵捕獲，因此無法直接對全域性環境進行建模；

其次，CNN學習到的權重在推理時是靜止的，這使得CNN不能靈活地適應不同輸入內容。

現有方法試圖將Vision Transformers（ViTs）的優良特性與CNN結合，以解決上述限制。然而大部分結合方式都不可避免的帶來了一個新問題：

由於自注意力計算方式的特點，架構推理速度被嚴重限制。

因此，EdgeNeXt試圖在將CNN與ViTs結合的同時，兼顧模型效能與推理速度。作者等人透過引入一個分割深度轉置注意力（SDTA）編碼器，在做到高效結合的同時，不增加額外的引數量和乘法-加法運算量（MAdds）。一份簡單的架構比較圖如下所示：

圖1：EdgeNeXt模型與SOTA ViTs和混合架構設計的比較，資料集為ImageNet-1K。

圖中1。3M引數量的EdgeNeXt在ImageNet-1K上的效能超過同樣輕量級的混合架構MobileViT，而5。6M引數量的版本憑藉適中的MAdds成為對比架構中的SOTA模型。

EdgeNeXt是如何結合CNN與ViTs的？又是如何做到在引入SDTA模組後不增加額外引數量的？

下面筆者帶領大家一探究竟~

二、EdgeNeXt剖析

圖2：EdgeNeXt架構整體示意圖。

如圖，EdgeNeXt同樣採用標準的“四階段”金字塔式設計規範，其主要包括卷積編碼器和SDTA編碼器兩個核心模組。

現有研究工作表明，CNN淺層的更多特徵體現為紋理、區域性邊緣等，深層則體現為抽象語義，更偏向全域性資訊。因此在淺層應用較小核，深層應用較大核是符合CNN結構特性的。

卷積編碼器和SDTA編碼器作為架構核心模組，將分成兩個單獨的小節進行講解~

卷積編碼器（Convolution Encoder）

圖3：卷積編碼器示意圖。

如圖，編碼器起始部分的N×N DWConv代表具有自適應核大小的Depth-wise卷積，如上所述，N在EdgeNeXt四個階段中的數值分別為{3，5，7，9}。經過層歸一化後使用逐點卷積（即圖中的Linear）來擴充套件通道維度至原來的4倍，並輔以GELU來增加非線性，最後透過逐點卷積恢復通道維度後，新增殘差連線來保障網路的正常學習。一份簡單的pytorch程式碼示例如下所示：

import torchimport torch。nn as nnclass ConvEncoder（nn。Module）： def __init__（self， n， channels）： super（ConvEncoder， self）。__init__（） self。identity = nn。Identity（） self。stem_conv = nn。Conv2d（channels， channels， kernel_size=（n， n）， padding=（n // 2， n // 2）， groups=channels） self。norm = nn。LayerNorm（channels） # 這裡可以直接呼叫，在forward中進行維度調整，也可以另寫一個層歸一化處理的函式再呼叫 self。linear1 = nn。Conv2d（channels， channels * 4， kernel_size=1） self。gelu = nn。GELU（） self。linear2 = nn。Conv2d（channels * 4， channels， kernel_size=1） def forward（self， x）： identity = self。identity（x） x = self。stem_conv（x） x = self。norm（x。permute（0， 2， 3， 1））。permute（0， 3， 1， 2） x = self。linear2（self。gelu（self。linear1（x））） out = x + identity return out

分割深度轉置注意編碼器（SDTA Encoder）

圖4：SDTA編碼器示意圖。

如圖，SDTA編碼器包含兩個主要元件：特徵編碼模組和自注意計算模組。

在第一個元件中，作者等人參考Res2Net的架構設計，希望能獲得具有更靈活和自適應空間感受野的輸出特徵。首先輸入特徵

透過直接的通道切分被劃分為s個子集（圖中s=4），每個子集的尺寸均為

，接著每個子集的計算方式都是將上一個子集的輸出特徵融合後再經過3×3的depth-wise卷積。以圖4中的4個子集為例，子集B1由於處在起始位置，故不做任何處理；子集B2將B1的輸出進行特徵融合，再經過3×3的depth-wise卷積進行特徵編碼；子集B3將B2的輸出進行特徵融合，再經過3×3的depth-wise卷積進行特徵編碼；子集B4將B3的輸出進行特徵融合，再經過3×3的depth-wise卷積進行特徵編碼。最終4個子集的輸出特徵被拼接後得到具有多尺度感受野的輸出特徵。

需要注意的是，特徵通道的切分也可以使用逐點卷積來近似的完成，但作者等人希望進一步減少引數量，故採用直接的通道切分。

第二個元件是本文的核心，如圖，輸入特徵

透過reshape被修改為

，輔以三個線性層得到Q、K、V，在計算交叉協方差注意力之前，透過對Q、K應用 L2 範數來穩定訓練。之後，以往的計算在Q、K之間沿空間維度進行，即

（HW × C）和

（C × HW）運算後得到（HW × HW）。然而本文方法希望在通道維度上計算注意力，即僅對Q進行轉置，因此（C×HW） ·（HW ×C）可得到（C × C），計算結果輔以softmax得到注意力得分矩陣，應用於V得到最終的注意力圖。

為什麼要在通道維度進行點積運算？作者等人的解釋是：傳統自注意力層的開銷對於邊緣裝置上的視覺任務是不可行的，因為它是以更高的 MAdd 和延遲為代價的。為了緩解這個問題並有效地編碼全域性上下文，提出在通道維度上應用點積。

模型整體的結構與引數量示意如下所示：

圖5：EdgeNeXt 架構詳細示意圖。

三、效能對比

在實驗部分，作者等人對影象分類（ImageNet-1K）、目標檢測（COCO）、影象分割（Pascal VOC 2012）等下游任務均進行詳細的效能對比。其對比結果如下所示：

影象分類：

圖6：影象分類對比結果。

目標檢測：

圖7：目標檢測對比結果。

影象分割：

圖8：影象分割對比結果。

四、思考與總結

本文從輕量化角度出發，設計了卷積與transformer的混合架構EdgeNeXt，兼顧了模型效能與模型大小/推理速度。

整體架構採取標準的“四階段”金字塔正規化設計，其中包含卷積編碼器與SDTA編碼器兩個重要的模組。在卷積編碼器中，自適應核大小的設計被應用，這與SDTA中的多尺度感受野的思想相呼應。而在SDTA編碼器中，特徵編碼部分使用固定的3×3卷積，但透過層次級聯實現多尺度感受野的融合，而此處若使用不同尺寸的卷積核是否會帶來更好的效果有待考證。在自注意計算部分，透過將點積運算應用於通道維度，得到了兼顧計算複雜度與全域性注意力的輸出，是支撐本文的一個核心點。

從分類效能來看，效果確實很好，但結合檢測、分割任務來看，供對比模型略少，僅提供了部分輕量級網路對比，希望以後能看到更多優秀工作的產出~