EdgeFormer: 向視覺 Transformer 學習，構建一個比 MobileViT 更好更快的卷積網路

阿新 • • 發佈：2022-05-07

前言本文主要探究了輕量模型的設計。通過使用 Vision Transformer 的優勢來改進卷積網路，從而獲得更好的效能。

歡迎關注公眾號CV技術指南，專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀、CV招聘資訊。

論文：https://arxiv.org/abs/2203.03952

程式碼：https://github.com/hkzhang91/EdgeFormer

核心內容

本文主要探究了輕量模型的設計。通過使用 Vision Transformer 的優勢來改進卷積網路，從而獲得更好的效能。

本文提出的核心運算元，即 global circular convolution (GCC)，是一個卷積操作，但是會引入位置嵌入，同時還具有全域性的感受野。另一個改進點是使用提出的 GCC 和 SE 操作構建了類似於 Vision Transformer 的基礎操作單元。藉助於 SE 引入了樣本相關的注意力機制。

提出的結構組合起來可以作為一個基礎的單元用於替換卷積網路或者 Transformer 架構中的相關模組。

基於這些提出的結構構建的輕量卷積結構在更少的引數量和推理速度的同時，獲得了比輕量卷積網路和 Vision Transformer 更好的效能。

對於 ImageNet-1K 的分類，相較於基線模型 MobileViT，EdgeFormer 以約 5M 引數實現 78.6％的 TOP-1 準確性，節省 11％的引數和 13％的計算成本，但準確性提高了 0.2％，並且推理速度更快（在基於 ARM 的 Rockchip RK3288 上）。僅使用了 DeIT 0.5 倍引數量，但與 DeiT 相比，準確率提升了 2.7％。在 MS-CCOCO 目標檢測和 Pascal VOC 分割任務上，EdgeFormer 還顯示出更好的效能。

ViT 和 ConvNet

二者都是必不可少的。

從實用角度而言，二者各有優缺點：
- ViT 模型逐漸體現出更好的效能，但是通常面領著高昂的計算成本，並且難以訓練。
- 相較於 ViT 類結構，卷積網路顯示出了稍差一點的效能，但是仍然有著獨特的優勢。例如卷積網路有著更好的硬體支援，容易訓練。另外，卷積網路在針對移動或者小模型上仍然大有可為。
從資訊處理的視角而言，二者有著獨特的特性：
- ViT 架構擅長於提取全域性資訊，並且在資料驅動下，使用注意力機制從不同位置上提取資訊。
- 卷積網路專注於建模區域性關係，並且有著更強的歸納偏置帶來的先驗資訊。

經過這樣的對比，作者們思考，是否可以從 ViT 的設計上學習一些東西來提升針對移動或者邊緣裝置的卷積網路的效能。這也正是本文的目的所在。

文中，作者們想要設計一個新的輕量的卷積網路，來進一步增強其在移動裝置和邊緣計算裝置友好模型領域中的力量。純卷積是更加移動裝置友好的，因為卷積操作被現有的工具鏈已經高度優化，並且被廣泛使用於部署模型到資源受限裝置上。

甚至，由於這些年來卷積網路的流行，一些現有的神經網路加速器圍繞著卷積風格的操作被設計出來，而並沒有有效支援例如 softmax 這樣複雜的非線性操作以及大矩陣乘法所需要的資料匯流排頻寬資料。這些硬體和軟體上的限制使得純卷積的輕量模型相對於那些即使在其他方面具有同樣競爭力的 ViT 類的模型仍然是更好的選擇。

為了設計這樣一個卷積網路，作者們比較並總結了 ViT 相較於卷積網路的三個主要不同：

ViT 更擅長於提取全域性特徵
ViT 應用了 meta-former block
在 ViT 中，資訊整合是資料驅動的（但是這種資料驅動的形式過於耗費記憶體，於是大量基於先驗知識對此整合區域進行約束，壓縮空間，從而提升 Attention 操作的效率）

為此，作者們針對這些點設計了 EdgeFormer 模組，來改進卷積網路：

提出一個全域性迴圈卷積（Global Circular Convolution，GCC）操作來提取全域性特徵從而改善傳統卷積操作的有限的感受野。其中，基於例項的卷積核（這裡想說的應該是卷積核的引數與輸入相關）和位置嵌入策略分別被用來應對輸入尺寸的變化，並且向輸出特徵圖中注入位置資訊。
基於提出的 GCC，構建了一個純卷積風格的 meta-former 模組作為基礎結構。
通過對 FFN 新增通道注意力（SE Block）來引入輸入相關的核權重。

對於最終的模型 EdgeFormer，則是借鑑 CoatNet 和 MobileViT 的結構風格，採用了雙叉形式的結構（bifurcate structure）。根據這些設計構建的最終模型在三個主流的視覺任務上的輕量模型的對比中獲得了不錯的效能表現。

從相關工作看論文

NLP 中的 Transformer 被 ViT 引入視覺任務後，出現了大量工作對其結構進行優化，以提升其計算效率和訓練效率。

從提升訓練策略的角度來看，DeiT（Training data-efficient image transformers & distillation through attention）使用知識蒸餾策略來訓練 ViT，從而使用更少的預訓練資料實現了極具競爭力的效能。
從提升模型架構的角度來看：
- LeViT（Levit: a vision transformer in convnet’s clothing for faster inference）混合了卷積網路和 Transformer，實現了很好的速度和準確率的權衡。
- BoTNet（Bottleneck transformers for visual recognition）則替換 ResNet 深層的卷積層為多頭自注意力結構。
- ViT-C（Early convolutions help transformers see better）則將原始的 ViT 的 patchify stem 結構（核大小等於跨步的卷積）替換為 convolutional stem 結構從而有效提升了 ViT 優化的穩定性和峰值效能。
- ConViT（Convit: Improving vision transformers with soft convolutional inductive biases）則將通過門控位置自注意力合併了 soft convolutional inductive biases。
- CMT（Cmt: Convolutional neural networks meet vision transformers）則應用了基於深度卷積的區域性感知單元和一個輕量的 Transformer 模組（使用跨步等於卷積核的深度卷積縮減 K 和 V）。
- CoatNet（Coatnet: Marrying convolution and attention for all data sizes）合併了卷積和自注意力來設計新的 transformer 模組（淺層使用卷積，深層使用額外引入全域性靜態權重的 Attention），這可以同時關注區域性和全域性資訊。
- PVT（Pyramid vision transformer: A versatile backbone for dense prediction without convolutions）和 CVT（Cvt: Introducing convolutions to vision transformers）在不同階段中插入卷積操作來減少 token 的數量，構建了分層的多階段結構。
- Swin（Swin transformer: Hierarchical vision transformer using shifted windows）則是將 Attention 的計算約束在了固定尺寸的區域性視窗之中。
- PiT（Rethinking spatial dimensions of vision transformers）則是通過聯合使用池化層和深度卷積層來實現特徵圖通道的擴張和空間的縮減。
- 一些研究者試圖通過借鑑卷積網路的思想來優化 ViT，他們將卷積網路的一些操作整合到 ViT 中。
- 而另一些工作則是通過直接組合兩種風格來設計新的混合架構。他們都相似地在模型起始階段使用卷積 stem 來提取區域性特徵，而在後續的結構中使用 transformer 風格的結構提取“全域性”或者“區域性-全域性”特徵。本文提出的純卷積架構的工作同樣使用類似的思想。

針對本文著重探究的輕量模型的設計的領域，作者們也詳細的回顧了卷積神經網路和視覺 Transformer 的一些工作。

卷積神經網路：ShuffleNet 系列，MobileNet 系列，MicroNet，GhostNet，EfficientNet，TinyNet 和 MnasNet。相較於標準卷積網路，他們具有更少的引數，更低的計算成本，更快的推理速度。另外，輕量的卷積神經網路可以被應用到各種裝置上。但是儘管有這樣的優勢，但是這些輕量模型效能稍差於更重的模型。
卷積神經網路與視覺 Transformer 的組合：
- Mobile-Former（Mobile-former: Bridging mobilenet and transformer）並行組合了 MobileNet 和 Transformer，利用前者提取區域性特徵，利用後者捕獲全域性資訊。
- MobileViT（Mobilevit: light-weight, general-purpose, and mobilefriendly vision transformer）中使用提出的 MobileViT 模組來替換 MobileNetV2 的深層模組。提出的模組中，通過卷積提取的區域性表徵，與全域性表徵被拼接來獲得“區域性-全域性”表徵。

就本文而言，提出的結構與 Mobile-Former 和 MobileViT 比較相關。不同點在於這兩個方法仍然去保留 Transformer 塊，而本文結構卻是一個純的卷積網路，這使得提出的結構對移動端更加友好。論文針對低算力平臺的的實驗也證實了這一點。在通過向視覺 Transformer 學習來設計一個純卷積網路這一點上，本文和 ConvNeXt 有相近之處。主要有兩點不同：

想法和架構並不相同。ConvNeXt 通過將一系列增量但是有效的設計引入，從而將一個標準的 ResNet 朝著視覺 Transformer 的方向來改動。而本文則直接從三個卷積網路和視覺 Transformer 的差異的角度入手來設計網路，這從一個巨集觀的層面上填補了缺口。由於想法不同，對應的結構也就有所差異。
提出的目的不同。本文主要針對移動裝置。相較於 ConvNeXt，在輕量模型的設定下，本文的方法更有優勢。

提出的方法

EdgeFormer Block

圖中顯示了卷積網路和 ViT 類結構主要差異。而圖中下半部分則是提出的模組的結構。

使用全域性迴圈卷積提取全域性特徵

卷積操作中，每個位置上的輸出特徵是通過對以其為中心的區域性鄰域上的輸入特徵進行加權聚合得到的。而 ViT 類結構中，自注意力模組提取特徵則是通過對全域性空間範圍上的特徵進行聚合後得到的，而且聚合權重是基於特徵相似性生成的動態引數。相較於卷積的區域性感受野，後者可以從整個空間維度上學習全域性特徵。

針對這點不足，本文提出了全域性迴圈卷積 GCC。提出的 GCC 通過使用水平和垂直方向進行分解從而得到了一個單一來看足夠輕量的操作。二者各自可以同時覆蓋到向同行和相同列的所有輸入位置。

具體而言，以垂直方向的 GCC 為例，即 GCC-V，其步驟如下：

輸入特徵為 CxHxW
引入一個基礎位置嵌入資訊為 CxBx1，通過擴充套件它，使其完全等於輸入特徵的形狀，從而可以靈活地適應不同尺寸的輸入特徵。
1. 通過雙線性插值，變為 CxHx1
2. 軸向複製擴充套件，變為 CxHxW
3. 加到輸入上
將得到的特徵言 H 方向堆疊，得到 Cx(2H-1)xW，這步操作非常重要，由此便可以實現具有全域性感受野且仍然引數共享的標準卷積操作了。
接下來開始構造卷積核引數，其同樣由一個大小為 CxBx1 的基礎量插值到 CxHx1 而得到。此時卷積核的空間感受野為 Hx1。搭配著 GCC-H 中的 Cx1xW，便近似得到了全域性的覆蓋。當然，這裡同樣可以引入偏置，具體的實現也同樣是基於一個基礎量擴充套件而成。
使用插值後的卷積核和嵌入位置資訊的輸入張量計算標準卷積操作，得到大小為 CxHxW 的輸出。

這個流程中，GCC-V 中大小為 CxHx1 的卷積核的輸出特徵座標 i，實際上對應著輸入特徵上的座標範圍 (i+t)%h，這裡的 t 是一個 0 到 H-1 之間的值，不同的 t 對應著卷積核覆蓋的區域性鄰域中的相對座標。

也就是說，通過這樣一個取模的計算，作者們構建出了一個“迴圈”的概念。可以說是非常會變通了。但是作者們的實現，卻是更考驗著思維上的轉化了。因為針對這樣卷積操作計算方式的改變，一般有兩種實現策略，一種是從卷積運算的本身著手，另一種則是從被卷積的特徵的角度去變換。這一點在我自己的一份關於動態濾波器的工作中也有所體現。

一般而言，前者雖然直接但是實現卻非常複雜，可能需要重新寫運算元。而後者則更考驗思維變通的能力。本文的實現思路就是後者。所以更具參考和啟發意義。

這部分結構主要實現可見：

https://github.com/hkzhang91/EdgeFormer/blob/df5ab09c7f60187d022a891ad083f4596838ff54/cvnets/modules/edgeformer_block.py#L534-L677

具體而言，作者們直接將兩個特徵沿著指定軸向維度進行拼接（注意這裡後面的特徵只取了 H-1 行的資料，不然輸出的特徵形狀就比輸入大 1 了）：torch.cat((x_1_r, x_1_r[:, :, :-1, :]), dim=2)

相較於傳統卷積操作，提出的 GCC 引入了三處修改：

感受野擴大到了全域性空間。要注意和直接擴大傳統區域性卷積核到全域性範圍的設定區分開。因為傳統卷積需要引入 padding 操作來維持特徵尺寸，如果直接增大卷積核，絕大多數位置上的處理過程中，仍然只有部分畫素被覆蓋到。尤其是在邊緣區域。大部分割槽域都是 0 值。
位置嵌入被用來保證輸出特徵對於空間位置的敏感性。GCC 可以提取全域性特徵，但是也打亂了原始輸入的空間結構（因為迴圈的存在，所以也就無法嚴格按照區域性鄰域的本身的相對順序來處理了）。對於分類任務而言，空間結構的保持可能不是一個大的問題，但是文中的實驗中也展現了分割檢測這樣的密集預測任務對於保持空間結構的需求還是有的。也就是說位置資訊的嵌入對於這樣的任務而言有一定的益處。
卷積核和位置嵌入是根據輸入尺寸動態插值生成的。GCC 中，卷積核的尺寸和位置嵌入和輸入是一致的，這增強了對於不同尺寸輸入的適應能力。

使用 GCC 構建 EdgeFormer 塊

從卷積網路到 ViT 架構，一個相當重要的修改就是 meta-former block 替換了殘差塊。meta-former block 一般包含著兩個核心元件：token mixer 和 channel mixer。前者用於在不同空間位置上的 token 之間交換資訊，後者則是重點關注通道的互動。而且二者都會應用殘差結構。

受此啟發，GCC 被插入 meta-former 風格的結構中，從而構建了 EdgeFormer block。其中替換自注意力模組為 GCC，從而作為一個新的空間處理單元。這主要有兩點原因：

GCC 可以提取全域性特徵，並在不同畫素之間進行資訊互動，這滿足 token mixer 的要求。
自注意力模組的計算複雜度是平方級別的，通過改用 GCC，可以顯著降低計算成本，有助於實現設計輕量的卷積網路。

向 Channel Mixer 中新增通道注意力

ViT 架構中，自注意力模組可以根據輸入調整權重，這造就了 ViT 這來資料驅動的模型。通過引用注意力機制，資料驅動模型可以專注於重要的特徵，並且抑制不必要的資訊，這帶來了更好的效能。曾經的一些工作（例如 SENet 等)也強調了保持資料驅動的重要性。

通過替換自注意力操作為 GCC 結構，我們得到了一個可以提取全域性特徵的純卷積網路。但是替換後的模型不再是資料驅動的結構了，為此，文中向 Channel Mixer 中引入了通道注意力。這裡參考了 SENet 的結構。

EdgeFormer

通過為提出的 EdgeFormer block 選擇一個外框架（outer frame），構建了本文的模型 EdgeFormer。

現存的混合架構基本可以分為三類，包括序列結構、並行結構和雙叉式（bifurcate）結構。其中雙叉式結構目前獲得了最好的效果，典型如 CoatNet，而對於輕量結構，MobileViT 也採用了這一形式。因此本文也延續了這一設計，基於 MobileViT 構建了最終的外框架。具體而言，針對 baseline MobileViT，本文做出瞭如下改進：

MobileViT 包含兩個主要型別的模組，ViT 和 MobileNetV2。本文保持了所有 MobileNetV2 模組，替換了 ViT 結構。這一替換使得最終模型從混合架構變成了純卷積架構。
適當地增加了 EdgeFormer 塊的寬度。即便如此，替換的模型仍然具有較少的引數和較少的計算成本。
現有的雙叉結構包含著一些互動模組，這負責本地和全域性特徵模組之間的資訊互動。原始 MobileViT 中，ViT 塊是最重的結構，經過替換後，這些互動結構反而成了大頭。因此，本文將分組卷積和點卷積引入這些模組來減少了引數數量還不會傷害效能。

實驗結果

效能對比

作者們在分類檢測分割上都驗證了提出模型和元件的效能。

在輕量模型的對比中，可以看到混合架構超越了以往的純卷積架構和 ViT 類架構，而本文的升級版的純卷積架構獲得了整體更好的效果。

消融實驗

這裡的消融實驗在三個任務上驗證了提出元件的效果，這裡的 baseline 為 MobileViT。主要做了以下幾個方面的對比：

用於替換自注意力的（對比消融 GCC）卷積核：普通卷積、兩種尺度的大核卷積、GCC
- 這裡一開始就引入了通道注意力，似乎不太滿足控制變數的要求
- 2、3、6 的對比可見 GCC 的優勢。
- 這裡在大核卷積的實驗中（2、3），隨著卷積核的增大效能區域飽和（倒是再增大是否會下降就很難說，只能說可能有這個趨勢。這一點現象據說在 ConvNeXt 中也有反映）。而直接覆蓋了全域性的 GCC 卻可以獲得更好的表現。
meta-former 架構的應用
- 這裡為了替換 meta-former 結構，作者們將 GCC 整合到 ResNeXt block 中作為替代品。
- 4、7 的對比可以看到使用 meta-former 構建純卷積網路的潛力。
通道注意力的應用
- 5、7 的對比可以看到通道注意力的效能增益。
位置嵌入的引入
- 位置嵌入在檢測和分割這樣的任務中有著積極的意義。這可能也是因為二者相較於分類任務，對於空間位置更加敏感。

低功率平臺的推理速度

作者們在 Rockchip RK3288 上對 EdgeFormer-S 和基線模型 MobileViT-S 模型進行了對比。通過使用 ONNX 和 MNN 將這些模型遷移到 RK3288 上，對他們迭代一百次從而測試平均推理速度。

對比中可以看到，EdgeFormer 快了 23%。除了更少的 Madd 操作，作者們認為速度提升可能還因為：

卷積已經被現有的廣泛用於部署模型到這些資源有限的裝置上的工具鏈高度優化。
相較於卷積，Transformer 需要更多的資料頻寬來計算 attention map，這涉及到兩個大的矩陣 K 和 Q，然而在卷積中，核相對於輸入特徵而言，是一個很小的矩陣。如果頻寬需求超過了晶片設計，那麼 CPU 將閒置來等待資料，從而導致 CPU 利用率較低和總體推理速度較慢。

歡迎關注公眾號CV技術指南，專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀、CV招聘資訊。

CV技術指南建立了一個免費的知識星球。關注公眾號新增編輯的微訊號可邀請加入。