霸榜COCO和Cityscapes!南理工CMU提出极化自注意力,更精细的双重注意力建模结构
作者丨小馬
編輯丨極市平臺
【寫在前面】
注意力機制是一個被廣泛應用在各種CV任務中的方法。注意力機制根據施加的維度大致可以分為兩類:通道注意力和空間注意力。對于通道注意力機制,代表性的工作有SENet[2]、ECANet[3];對于空間注意力機制,代表性的工作有Self-Attention[4]。隨著空間和通道注意力機制的提出,很自然的,結合空間和通道兩個維度的雙重注意力機制也被提出,代表工作有CBAM[1],DANet[5]。
基于雙重注意力機制,本文針對Pixel-wise regression的任務,提出了一種更加精細的雙重注意力機制——極化自注意力(Polarized Self-Attention)。作為一個即插即用的模塊,在人體姿態估計和語義分割任務上,作者將它用在了以前的SOTA模型上,并達到了新的SOTA性能,霸榜COCO人體姿態估計和Cityscapes語義分割。
1. 論文和代碼地址
Polarized Self-Attention: Towards High-quality Pixel-wise Regression
官網代碼:https://github.com/DeLightCMU/PSA (暫未開源)
核心代碼:https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/attention/PolarizedSelfAttention.py
2. Motivation
細粒度的像素級任務(比如語義分割)一直都是計算機視覺中非常重要的任務。不同于分類或者檢測,細粒度的像素級任務要求模型在低計算開銷下,能夠建模高分辨率輸入/輸出特征的長距離依賴關系,進而來估計高度非線性的像素語義。CNN中的注意力機制能夠捕獲長距離的依賴關系,但是這種方式比較復雜并且是對噪聲比較敏感的。
對于這類任務,通常采用的是encoder-decoder的結構,encoder(e.g.,ResNet)用來降低空間維度、提高通道維度;decoder通常是轉置卷積或者上采樣,用來提高空間的維度、降低通道的維度。因此連接encoder和decoder的tensor通常在空間維度上比較小,雖然這對于計算和顯存的使用比較友好,但是對于像實例分割這樣的細粒度像素級任務,這種結構顯然會造成性能上的損失。
基于此,作者提出了一個即插即用的模塊——極化自注意力機制( Polarized Self-Attention(PSA)),用于解決像素級的回歸任務,相比于其他注意力機制,極化自注意力機制主要有兩個設計上的亮點:
1)極化濾波( Polarized filtering):在通道和空間維度保持比較高的resolution(在通道上保持C/2的維度,在空間上保持[H,W]的維度 ),這一步能夠減少降維度造成的信息損失;
2)增強(Enhancement):采用細粒度回歸輸出分布的非線性函數。
3. 方法
3.1. Overview
本文提出極化自注意力結構如上圖所示,分為兩個分支,一個分支做通道維度的自注意力機制,另一個分支做空間維度的自注意力機制,最后將這兩個的分支的結果進行融合,得到極化自注意力結構的輸出。
如上表所示,相比于其他注意力的方法,本文在空間維度和通道維度都沒有進行很大程度的壓縮(在空間維度上保持[H,W]的大小,在通道維度上使用了C/2的大小),并將復雜度也保持在了一個比較小的水平(CWH)。
個人覺得,極化自注意力機制能夠在細粒度的像素級任務上保持比較好的性能,很大程度上依賴于它在空間和通道維度上都沒有進行很大程度的壓縮,使得信息損失比較小。
除此之外,以前的attention方法采用非線性函數進行概率估計時,通常只采用Softmax或者Sigmoid。為了能夠擬合出細粒度回歸結果的輸出分布,極化自注意力機制在通道和空間分支都在采用了Softmax和Sigmoid相結合的函數。
3.2. Polarized Self-Attention (PSA) Block
為了解決同時對空間和通道建模時,如果不進行維度縮減,就會導致計算量、顯存爆炸的問題。作者在PSA中采用了一種極化濾波(polarized filtering)的機制。類似于光學透鏡過濾光一樣,每個SA的作用都是用于增強或者削弱特征。(在攝影時,所有橫向的光都會進行反射和折射。極化濾波的作用就是只允許正交于橫向方向的光通過,以此來提高照片的對比度。 由于在濾波過程中,總強度會損失,所以濾波后的光通常動態范圍較小,因此需要額外的提升,用來以恢復原始場景的詳細信息。 )
基于上面的思想,作者提出了Polarized Self-Attention (PSA)機制,同上面的思想一樣,作者也是現在一個方向上對特征進行壓縮,然后對損失的強度范圍進行提升,具體可分為兩個結構:
1)濾波(Filtering):使得一個維度的特征(比如通道維度)完全坍塌,同時讓正交方向的維度(比如空間維度)保持高分辨率。
2)High Dynamic Range(HDR):首先在attention模塊中最小的tensor上用Softmax函數來增加注意力的范圍,然后再用Sigmoid函數進行動態的映射。
3.3. Channel-only branch
如上圖所示,PSA分為兩個分支,通道分支和空間分支,通道分支的權重計算公式如下:
可以看出,作者先用了1x1的卷積將輸入的特征X轉換成了Q和V,其中Q的通道被完全壓縮,而V的通道維度依舊保持在一個比較高的水平(也就是C/2)。因為Q的通道維度被壓縮,如上面所說的那樣,就需要通過HDR進行信息的增強,因此作者用Softmax對Q的信息進行了增強。然后將Q和K進行矩陣乘法,并在后面接上1x1卷積、LN將通道上C/2的維度升為C。最后用Sigmoid函數使得所有的參數都保持在0-1之間。代碼如下:
class Channel_only_branch(nn.Module):def __init__(self, channel=512):super().__init__()self.ch_wv=nn.Conv2d(channel,channel//2,kernel_size=(1,1))self.ch_wq=nn.Conv2d(channel,1,kernel_size=(1,1))self.softmax=nn.Softmax(1)self.ch_wz=nn.Conv2d(channel//2,channel,kernel_size=(1,1))self.ln=nn.LayerNorm(channel)self.sigmoid=nn.Sigmoid()def forward(self, x):b, c, h, w = x.size()#Channel-only Self-Attentionchannel_wv=self.ch_wv(x) #bs,c//2,h,wchannel_wq=self.ch_wq(x) #bs,1,h,wchannel_wv=channel_wv.reshape(b,c//2,-1) #bs,c//2,h*wchannel_wq=channel_wq.reshape(b,-1,1) #bs,h*w,1channel_wq=self.softmax(channel_wq)channel_wz=torch.matmul(channel_wv,channel_wq).unsqueeze(-1) #bs,c//2,1,1channel_weight=self.sigmoid(self.ln(self.ch_wz(channel_wz).reshape(b,c,1).permute(0,2,1))).permute(0,2,1).reshape(b,c,1,1) #bs,c,1,1channel_out=channel_weight*xreturn channel_out3.4. Spatial-only branch
空間分支計算權重的公式如下:
可以看出,與Channel-only branch相似,作者先用了1x1的卷積將輸入的特征轉換為了Q和V,其中,對于Q特征,作者還用了GlobalPooling對空間維度壓縮,轉換成了1x1的大小;而V特征的空間維度則保持在一個比較大的水平(HxW)。由于Q的空間維度被壓縮了,所以作者就用了Softmax對Q的信息進行增強。然后將Q和K進行矩陣乘法,然后接上reshape和Sigmoid使得所有的參數都保持在0-1之間。代碼如下:
class Spatial_only_branch(nn.Module):def __init__(self, channel=512):super().__init__()self.sigmoid=nn.Sigmoid()self.sp_wv=nn.Conv2d(channel,channel//2,kernel_size=(1,1))self.sp_wq=nn.Conv2d(channel,channel//2,kernel_size=(1,1))self.agp=nn.AdaptiveAvgPool2d((1,1))def forward(self, x):b, c, h, w = x.size()#Spatial-only Self-Attentionspatial_wv=self.sp_wv(x) #bs,c//2,h,wspatial_wq=self.sp_wq(x) #bs,c//2,h,wspatial_wq=self.agp(spatial_wq) #bs,c//2,1,1spatial_wv=spatial_wv.reshape(b,c//2,-1) #bs,c//2,h*wspatial_wq=spatial_wq.permute(0,2,3,1).reshape(b,1,c//2) #bs,1,c//2spatial_wz=torch.matmul(spatial_wq,spatial_wv) #bs,1,h*wspatial_weight=self.sigmoid(spatial_wz.reshape(b,1,h,w)) #bs,1,h,wspatial_out=spatial_weight*xreturn out(完整代碼見https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/attention/PolarizedSelfAttention.py,如有任何代碼歡迎在issue中指出。)
3.5. Composition
對于兩個分支的結果,作者提出了兩種融合的方式:并聯和串聯(先進行通道上的注意力,再進行空間上的注意力):
從組合方式上看,這篇文章的處理比較簡單,只是對通道分支和空間分支的結果進行了串聯或者并聯,這一點和CBAM[1]一樣。
不同的是CBAM等雙重注意力的結構,大多都是用全連接層、卷積層來獲取注意力權重,所以在挖掘信息的時候可能不是非常有效。而本文采用的是Self-Attention來獲取注意力權重,充分了利用了Self-Attention結構的建模能力,并且作者對Q也進行了特征降維,所以在保證計算量的情況下,實現了一種非常有效的長距離建模。
4.實驗
對于baseline網絡,作者在3x3的卷積后面接上了PSA模塊,來檢驗PSA帶來的性能提升。
4.1. PSA vs Baselines
4.1.1. Top-Down 2D Human Pose Estimation
在人體姿態估計任務上,通過少量計算開銷和參數數量,PSA在所有baseline網絡上能夠增加2.6到4.3 AP。
4.1.2. Semantic Segmentation
在語義分割任務上,通過少量計算開銷和參數數量,PSA在所有baseline網絡上能夠增加1.8到2.6 mIoU。
4.2. Comparing with State-of-the Arts
4.2.1. Top-Down 2D Human Pose Estimation
在MS-COCO keypoint數據集上,PSA基于UDP-Pose結構,提升了1.7個點的性能,達到了79.5 AP的新SOTA性能。
4.2.2. Semantic Segmentation
在Cityscapes語義分割數據集上,PSA基于HRNet-OCR結構,達到了86.95 mIoU的新SOTA性能。
4.3. Ablation Study
從消融實驗中,可以看出PSA相比于其他attention模塊,具有非常明顯的優越性。
4.4. 可視化
從可視化結果中可以看出,在關鍵點檢測和分割任務上,PSA的結果比其他baseline結果更加精準。
5. 總結
這篇文章提出了一種極化自注意力機制,通過在空間和通道維度內部保持比較高的resolution,并采用了Softmax-Sigmoid聯合的非線性函數,使得帶有極化自注意力機制的模型,能夠在pixel-wise的任務上獲得更好的性能。
個人覺得,本質上這篇文章的結構還是和Self-Attention比較相似,但是如果直接用Self-Attention進行通道和空間維度上的建模,就會導致計算量非常大。所以作者提出了Filtering和HDR兩種設計方式,使得這個結構在比較合理的計算復雜度下,也能達到比較高的準確率。
另外[5]這篇文章也是借鑒了Self-Attention的結構(網絡結構如下圖所示),也是在空間和通道兩個維度上建模,也是用在分割任務上,有興趣的同學也可以比較一下兩篇文章的異同點。
參考文獻
[1]. Woo, Sanghyun, et al. “Cbam: Convolutional block attention module.” Proceedings of the European conference on computer vision (ECCV) . 2018.
[2]. Hu, Jie, Li Shen, and Gang Sun. “Squeeze-and-excitation networks.” Proceedings of the IEEE conference on computer vision and pattern recognition . 2018.
[3]. Wang, Qilong, et al. “ECA-Net: efficient channel attention for deep convolutional neural networks, 2020 IEEE.” CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE . 2020.
[4]. Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems . 2017.
[5]. Fu, Jun, et al. “Dual attention network for scene segmentation.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2019.
與50位技術專家面對面20年技術見證,附贈技術全景圖總結
以上是生活随笔為你收集整理的霸榜COCO和Cityscapes!南理工CMU提出极化自注意力,更精细的双重注意力建模结构的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: CVPR 2021 Short-vide
- 下一篇: 轻量高效!清华智能计算实验室开源基于Py