作者丨小馬
編輯丨極市平臺(tái)

【寫在前面】

注意力機(jī)制是一個(gè)被廣泛應(yīng)用在各種CV任務(wù)中的方法。注意力機(jī)制根據(jù)施加的維度大致可以分為兩類：通道注意力和空間注意力。對(duì)于通道注意力機(jī)制，代表性的工作有SENet[2]、ECANet[3]；對(duì)于空間注意力機(jī)制，代表性的工作有Self-Attention[4]。隨著空間和通道注意力機(jī)制的提出，很自然的，結(jié)合空間和通道兩個(gè)維度的雙重注意力機(jī)制也被提出，代表工作有CBAM[1]，DANet[5]。

基于雙重注意力機(jī)制，本文針對(duì)Pixel-wise regression的任務(wù)，提出了一種更加精細(xì)的雙重注意力機(jī)制——極化自注意力（Polarized Self-Attention）。作為一個(gè)即插即用的模塊，在人體姿態(tài)估計(jì)和語(yǔ)義分割任務(wù)上，作者將它用在了以前的SOTA模型上，并達(dá)到了新的SOTA性能，霸榜COCO人體姿態(tài)估計(jì)和Cityscapes語(yǔ)義分割。

1. 論文和代碼地址

Polarized Self-Attention: Towards High-quality Pixel-wise Regression

官網(wǎng)代碼：https://github.com/DeLightCMU/PSA （暫未開源）

核心代碼：https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/attention/PolarizedSelfAttention.py

2. Motivation

細(xì)粒度的像素級(jí)任務(wù)（比如語(yǔ)義分割）一直都是計(jì)算機(jī)視覺中非常重要的任務(wù)。不同于分類或者檢測(cè)，細(xì)粒度的像素級(jí)任務(wù)要求模型在低計(jì)算開銷下，能夠建模高分辨率輸入/輸出特征的長(zhǎng)距離依賴關(guān)系，進(jìn)而來(lái)估計(jì)高度非線性的像素語(yǔ)義。CNN中的注意力機(jī)制能夠捕獲長(zhǎng)距離的依賴關(guān)系，但是這種方式比較復(fù)雜并且是對(duì)噪聲比較敏感的。

對(duì)于這類任務(wù)，通常采用的是encoder-decoder的結(jié)構(gòu)，encoder（e.g.,ResNet）用來(lái)降低空間維度、提高通道維度；decoder通常是轉(zhuǎn)置卷積或者上采樣，用來(lái)提高空間的維度、降低通道的維度。因此連接encoder和decoder的tensor通常在空間維度上比較小，雖然這對(duì)于計(jì)算和顯存的使用比較友好，但是對(duì)于像實(shí)例分割這樣的細(xì)粒度像素級(jí)任務(wù)，這種結(jié)構(gòu)顯然會(huì)造成性能上的損失。

基于此，作者提出了一個(gè)即插即用的模塊——極化自注意力機(jī)制（ Polarized Self-Attention(PSA)），用于解決像素級(jí)的回歸任務(wù)，相比于其他注意力機(jī)制，極化自注意力機(jī)制主要有兩個(gè)設(shè)計(jì)上的亮點(diǎn)：

1）極化濾波（ Polarized filtering）：在通道和空間維度保持比較高的resolution（在通道上保持C/2的維度，在空間上保持[H,W]的維度 ），這一步能夠減少降維度造成的信息損失；

2）增強(qiáng)（Enhancement）：采用細(xì)粒度回歸輸出分布的非線性函數(shù)。

3. 方法

3.1. Overview

本文提出極化自注意力結(jié)構(gòu)如上圖所示，分為兩個(gè)分支，一個(gè)分支做通道維度的自注意力機(jī)制，另一個(gè)分支做空間維度的自注意力機(jī)制，最后將這兩個(gè)的分支的結(jié)果進(jìn)行融合，得到極化自注意力結(jié)構(gòu)的輸出。

如上表所示，相比于其他注意力的方法，本文在空間維度和通道維度都沒有進(jìn)行很大程度的壓縮（在空間維度上保持[H,W]的大小，在通道維度上使用了C/2的大小），并將復(fù)雜度也保持在了一個(gè)比較小的水平（CWH）。

個(gè)人覺得，極化自注意力機(jī)制能夠在細(xì)粒度的像素級(jí)任務(wù)上保持比較好的性能，很大程度上依賴于它在空間和通道維度上都沒有進(jìn)行很大程度的壓縮，使得信息損失比較小。

除此之外，以前的attention方法采用非線性函數(shù)進(jìn)行概率估計(jì)時(shí)，通常只采用Softmax或者Sigmoid。為了能夠擬合出細(xì)粒度回歸結(jié)果的輸出分布，極化自注意力機(jī)制在通道和空間分支都在采用了Softmax和Sigmoid相結(jié)合的函數(shù)。

3.2. Polarized Self-Attention (PSA) Block

為了解決同時(shí)對(duì)空間和通道建模時(shí)，如果不進(jìn)行維度縮減，就會(huì)導(dǎo)致計(jì)算量、顯存爆炸的問題。作者在PSA中采用了一種極化濾波（polarized filtering）的機(jī)制。類似于光學(xué)透鏡過濾光一樣，每個(gè)SA的作用都是用于增強(qiáng)或者削弱特征。（在攝影時(shí)，所有橫向的光都會(huì)進(jìn)行反射和折射。極化濾波的作用就是只允許正交于橫向方向的光通過，以此來(lái)提高照片的對(duì)比度。 由于在濾波過程中，總強(qiáng)度會(huì)損失，所以濾波后的光通常動(dòng)態(tài)范圍較小，因此需要額外的提升，用來(lái)以恢復(fù)原始場(chǎng)景的詳細(xì)信息。 ）

基于上面的思想，作者提出了Polarized Self-Attention (PSA)機(jī)制，同上面的思想一樣，作者也是現(xiàn)在一個(gè)方向上對(duì)特征進(jìn)行壓縮，然后對(duì)損失的強(qiáng)度范圍進(jìn)行提升，具體可分為兩個(gè)結(jié)構(gòu)：

1）濾波（Filtering）：使得一個(gè)維度的特征（比如通道維度）完全坍塌，同時(shí)讓正交方向的維度（比如空間維度）保持高分辨率。

2）High Dynamic Range（HDR）：首先在attention模塊中最小的tensor上用Softmax函數(shù)來(lái)增加注意力的范圍，然后再用Sigmoid函數(shù)進(jìn)行動(dòng)態(tài)的映射。

3.3. Channel-only branch

如上圖所示，PSA分為兩個(gè)分支，通道分支和空間分支，通道分支的權(quán)重計(jì)算公式如下：

可以看出，作者先用了1x1的卷積將輸入的特征X轉(zhuǎn)換成了Q和V，其中Q的通道被完全壓縮，而V的通道維度依舊保持在一個(gè)比較高的水平（也就是C/2）。因?yàn)镼的通道維度被壓縮，如上面所說的那樣，就需要通過HDR進(jìn)行信息的增強(qiáng)，因此作者用Softmax對(duì)Q的信息進(jìn)行了增強(qiáng)。然后將Q和K進(jìn)行矩陣乘法，并在后面接上1x1卷積、LN將通道上C/2的維度升為C。最后用Sigmoid函數(shù)使得所有的參數(shù)都保持在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_out

3.4. Spatial-only branch

空間分支計(jì)算權(quán)重的公式如下：

可以看出，與Channel-only branch相似，作者先用了1x1的卷積將輸入的特征轉(zhuǎn)換為了Q和V，其中，對(duì)于Q特征，作者還用了GlobalPooling對(duì)空間維度壓縮，轉(zhuǎn)換成了1x1的大小；而V特征的空間維度則保持在一個(gè)比較大的水平（HxW）。由于Q的空間維度被壓縮了，所以作者就用了Softmax對(duì)Q的信息進(jìn)行增強(qiáng)。然后將Q和K進(jìn)行矩陣乘法，然后接上reshape和Sigmoid使得所有的參數(shù)都保持在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

對(duì)于兩個(gè)分支的結(jié)果，作者提出了兩種融合的方式：并聯(lián)和串聯(lián)（先進(jìn)行通道上的注意力，再進(jìn)行空間上的注意力）：

從組合方式上看，這篇文章的處理比較簡(jiǎn)單，只是對(duì)通道分支和空間分支的結(jié)果進(jìn)行了串聯(lián)或者并聯(lián)，這一點(diǎn)和CBAM[1]一樣。

不同的是CBAM等雙重注意力的結(jié)構(gòu)，大多都是用全連接層、卷積層來(lái)獲取注意力權(quán)重，所以在挖掘信息的時(shí)候可能不是非常有效。而本文采用的是Self-Attention來(lái)獲取注意力權(quán)重，充分了利用了Self-Attention結(jié)構(gòu)的建模能力，并且作者對(duì)Q也進(jìn)行了特征降維，所以在保證計(jì)算量的情況下，實(shí)現(xiàn)了一種非常有效的長(zhǎng)距離建模。

4.實(shí)驗(yàn)

對(duì)于baseline網(wǎng)絡(luò)，作者在3x3的卷積后面接上了PSA模塊，來(lái)檢驗(yàn)PSA帶來(lái)的性能提升。

4.1. PSA vs Baselines

4.1.1. Top-Down 2D Human Pose Estimation

在人體姿態(tài)估計(jì)任務(wù)上，通過少量計(jì)算開銷和參數(shù)數(shù)量，PSA在所有baseline網(wǎng)絡(luò)上能夠增加2.6到4.3 AP。

4.1.2. Semantic Segmentation

在語(yǔ)義分割任務(wù)上，通過少量計(jì)算開銷和參數(shù)數(shù)量，PSA在所有baseline網(wǎng)絡(luò)上能夠增加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數(shù)據(jù)集上，PSA基于UDP-Pose結(jié)構(gòu)，提升了1.7個(gè)點(diǎn)的性能，達(dá)到了79.5 AP的新SOTA性能。

4.2.2. Semantic Segmentation

在Cityscapes語(yǔ)義分割數(shù)據(jù)集上，PSA基于HRNet-OCR結(jié)構(gòu)，達(dá)到了86.95 mIoU的新SOTA性能。

4.3. Ablation Study

從消融實(shí)驗(yàn)中，可以看出PSA相比于其他attention模塊，具有非常明顯的優(yōu)越性。

4.4. 可視化

從可視化結(jié)果中可以看出，在關(guān)鍵點(diǎn)檢測(cè)和分割任務(wù)上，PSA的結(jié)果比其他baseline結(jié)果更加精準(zhǔn)。

5. 總結(jié)

這篇文章提出了一種極化自注意力機(jī)制，通過在空間和通道維度內(nèi)部保持比較高的resolution，并采用了Softmax-Sigmoid聯(lián)合的非線性函數(shù)，使得帶有極化自注意力機(jī)制的模型，能夠在pixel-wise的任務(wù)上獲得更好的性能。

個(gè)人覺得，本質(zhì)上這篇文章的結(jié)構(gòu)還是和Self-Attention比較相似，但是如果直接用Self-Attention進(jìn)行通道和空間維度上的建模，就會(huì)導(dǎo)致計(jì)算量非常大。所以作者提出了Filtering和HDR兩種設(shè)計(jì)方式，使得這個(gè)結(jié)構(gòu)在比較合理的計(jì)算復(fù)雜度下，也能達(dá)到比較高的準(zhǔn)確率。

另外[5]這篇文章也是借鑒了Self-Attention的結(jié)構(gòu)（網(wǎng)絡(luò)結(jié)構(gòu)如下圖所示），也是在空間和通道兩個(gè)維度上建模，也是用在分割任務(wù)上，有興趣的同學(xué)也可以比較一下兩篇文章的異同點(diǎn)。

參考文獻(xiàn)

[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位技術(shù)專家面對(duì)面20年技術(shù)見證，附贈(zèng)技術(shù)全景圖

總結(jié)

以上是生活随笔為你收集整理的霸榜COCO和Cityscapes！南理工CMU提出极化自注意力，更精细的双重注意力建模结构的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。