SA-Net解读
簡介
這篇文章是南京大學Yu-Bin Yang等人于2021年初開放的一篇文章,已經被收錄于ICASSP2021,文章提出了一種新的視覺注意力機制,稱為Shuffle Attention(置換注意力),它通過置換單元組合空間注意力和通道注意力,相比此前的混合注意力更加高效,是一種非常輕量的注意力結構。實驗表明,在ImageNet、COCO等benchmark上,超越了當前的SOTA注意力模型如SE、SGE等,且擁有更低的計算復雜度和參數量。
-
論文標題
SA-Net: Shuffle Attention for Deep Convolutional Neural Networks
-
論文地址
http://arxiv.org/abs/2102.00240
-
論文源碼
https://github.com/wofmanaf/SA-Net
介紹
注意力機制如今已經被廣泛用于卷積神經網絡中,大大提升了很多任務上的性能表現。目前視覺中的注意力機制主要有兩種,如下圖所示,分別是通道注意力和空間注意力(我在之前的文章介紹了視覺中一些比較有名的注意力方法,可以訪問查看)。
通道注意力著重于捕獲通道間的依賴而空間注意力則關于像素間的關系捕獲,不過它們都是通過不同的聚合策略、轉換方法和強化函數來從所有位置聚合相同的特征來強化原始特征。CBAM和GCNet同時處理空間和通道信息,獲得了較好的精度提升,然而它們通常存在收斂困難、計算負擔大等問題。也有一些工作關于簡化注意力結構,如ECA-Net將SE模塊中通道權重的計算改用了1D卷積來簡化。SGE沿著通道維度對輸入進行分組形成表示不同語義的子特征,在每個子特征上進行空間注意力。遺憾的是,這些方法都沒用很好地利用空間注意力和通道注意力之間的相關性,所以效率偏低,所以自然而然引出這篇文章的出發點:能否以一個高效輕量的方式融合不同的注意力模塊?
不妨先回顧一下輕量級網絡的代表之一的ShuffleNetv2,它構建了一個可并行的多分支結構,如下圖所示,在每個單元的入口,含有ccc個通道的輸入被分割為c?c′c-c'c?c′和c′c'c′兩個分支,接著幾個卷積層用來提取關于輸入更高級的信息,然后結果concat到一起保證和輸入通道數相同,最后channel shuffle操作用來進行兩個分支之間的信息通訊。類似的,SGE將輸入沿著通道分組,所有子特征并行增強。
基于這些前人的工作,論文作者提出了一種更加輕量但是更加高效的Shuffle Attention(SA)模塊,它也是將輸入按照通道進行分組,對每組子特征,使用Shuffle Unit來同時構建通道注意力和空間注意力。對每個注意力模塊,論文還設計了針對每個位置的注意力掩碼來抑制噪聲加強有效的語義信息。論文主要的貢獻為設計了一個輕量但是有效的注意力模塊,SA,它將輸入特征圖按照通道分組并對每個分組用Shuffle Unit實現空間和通道注意力.
Shuffle Attention
所謂一圖勝千言,這篇文章其實忽略掉一些細節上圖就已經展示了Shuffle Attention模塊(SA模塊)的設計了,不過我這里還是按照論文的思路來逐步理解作者的設計,整個SA模塊其實分為三個步驟,分別為特征分組、混合注意力、特征聚合。
Feature Grouping
首先來看特征分組這個操作,它將輸入特征圖分為多組,每組為一個子特征(sub-feature)。具體來看,輸入特征圖X∈RC×H×WX \in \mathbb{R}^{C \times H \times W}X∈RC×H×W被沿著通道維度分為GGG組,表示為X=[X1,?,XG],Xk∈RC/G×H×WX=\left[X_{1}, \cdots, X_{G}\right], X_{k} \in \mathbb{R}^{C / G \times H \times W}X=[X1?,?,XG?],Xk?∈RC/G×H×W,其中每個子特征XkX_kXk?隨著訓練會逐漸捕獲一種特定的語義信息。這部分對應上圖最左邊的Group標注的部分。
Channel Attention and Spatial Attention
接著,XkX_kXk?會被分為兩個分支,依然是沿著通道維度劃分,兩個子特征表示為Xk1,Xk2∈RC/2G×H×WX_{k 1}, X_{k 2} \in \mathbb{R}^{C / 2 G \times H \times W}Xk1?,Xk2?∈RC/2G×H×W,如上圖中間Split標注后的部分,兩個分支分別是綠色和藍色表示,上面的綠色分支實現通道注意力開采通道間的依賴,下面的藍色分支則捕獲特征之間的空間依賴生成空間注意力圖,這樣,模型同時完成了語義和位置信息的注意。
具體來看通道注意力這個分支,這里可以作者沒用采用SE模塊的設計,主要是考慮到輕量化設計的需求(SE的參數還是比較多的),也沒有使用ECA-Net的設計采用一維卷積(ECA-Net要想精度高,對卷積核尺寸要求比較大),而是采用最簡單的GAP+Scale+Sigmoid的單層變換,公式如下。
s=Fgp(Xk1)=1H×W∑i=1H∑j=1WXk1(i,j)s=\mathcal{F}_{g p}\left(X_{k 1}\right)=\frac{1}{H \times W} \sum_{i=1}^{H} \sum_{j=1}^{W} X_{k 1}(i, j) s=Fgp?(Xk1?)=H×W1?i=1∑H?j=1∑W?Xk1?(i,j)
Xk1′=σ(Fc(s))?Xk1=σ(W1s+b1)?Xk1X_{k 1}^{\prime}=\sigma\left(\mathcal{F}_{c}(s)\right) \cdot X_{k 1}=\sigma\left(W_{1} s+b_{1}\right) \cdot X_{k 1} Xk1′?=σ(Fc?(s))?Xk1?=σ(W1?s+b1?)?Xk1?
上面的式子只有兩個變換參數,即W1∈RC/2G×1×1W_{1} \in \mathbb{R}^{C / 2 G \times 1 \times 1}W1?∈RC/2G×1×1和b1∈RC/2G×1×1b_{1} \in \mathbb{R}^{C / 2 G \times 1 \times 1}b1?∈RC/2G×1×1。
下面來看空注意力這個分支,思路也很簡單,先是對輸入特征圖進行Group Norm,然后也是通過一個變換Fc(?)\mathcal{F}_{c}(\cdot)Fc?(?)來增強輸入的表示,具體公式如下。這里的參數也只有W2W_2W2?和b2b_2b2?,它們的尺寸都是RC/2G×1×1\mathbb{R}^{C / 2 G \times 1 \times 1}RC/2G×1×1。
Xk2′=σ(W2?GN(Xk2)+b2)?Xk2X_{k 2}^{\prime}=\sigma\left(W_{2} \cdot G N\left(X_{k 2}\right)+b_{2}\right) \cdot X_{k 2} Xk2′?=σ(W2??GN(Xk2?)+b2?)?Xk2?
最后,兩種注意力的結果被concat到一起Xk′=[Xk1′,Xk2′]∈RC/G×H×WX_{k}^{\prime}=\left[X_{k 1}^{\prime}, X_{k 2}^{\prime}\right] \in \mathbb{R}^{C / G \times H \times W}Xk′?=[Xk1′?,Xk2′?]∈RC/G×H×W,此時它已經和該組的輸入尺寸一致了。
Aggregation
最后一步的聚合也很簡單,通過ShuffleNetv2采用的channel shuffle操作來保證各組子特征之間的交互,最后得到和輸入XXX同維的注意力圖。這是上圖中最右邊的部分。
至此,SA模塊的構建已經完成,其實W1W_1W1?、b1b_1b1?、W2W_2W2?和b2b_2b2?就是整個SA模塊所有的參數,它可以通過PyTorch輕易的實現。SA模塊可以取代SE模塊,因此替換SENet的網絡稱為SA-Net。
為了驗證SA模塊對語義信息的提取能力,作者在ImageNet上訓練了沒用channel-shuffle的SA-Net50B和有channel-shuffle的SA-Net50,給出精度對比如下圖。一方面在SA采用之后,top1精度表示出了提升,這也就意味著特征分組可以顯著提升特征的語義表達能力;另一方面,不同類別的分布在前面的層中非常相似,這也就意味著在前期特征分組的重要性被不同類別共享,而隨著深度加深,不同的特征激活表現出了類別相關性。
為了驗證SA有效性,基于GradCAM進行可視化,得到下圖的結果,可以發現,SA模塊使得分類模型關注于目標信息更相關的區域,進而有效的提升分類精度。
實驗
為了驗證性能,作者也在ImageNet、COCO上進行了分類、檢測和分割任務的實驗,分別如下。
實驗表明,SA模塊相比SE等注意力方法在參數量更少的情況下,達到了更高的精度,是非常高效的輕量級注意力機制。
總結
從ShuffleNet得到啟發,作者設計了輕量級網絡上適用的高效注意力機制Shuffle Attention,推動了注意力在輕量級網絡上的發展,在多個任務上達到SOTA表現。
總結
- 上一篇: MMDetection-运行时
- 下一篇: TrackFormer解读