一文詳解self-attention機制在語義分割中的應用(含論文解析)

Table of Contents

背景

Self-Attention Mechanism

論文解析

DANet

CCNet

ISSA


關於1×1卷積的作用:1×1 卷積

背景

語義分割經歷多年的發展,提出了FCN、U-Net、SegNet、DeepLab等一大批優秀的語義分割網絡。但是FCN等結構限制了局部感受野的範圍和短距離上下文信息,傳統的深度卷積神經網絡主要通過疊加多個卷積來模擬長距離依賴關係。爲了捕獲長距離的依賴關係,Chen等[1]人提出帶有多尺度空洞卷積的ASPP模塊集成上下文信息;Zhao等[2]人進一步提出帶有金字塔池化模塊的PSPNet捕獲上下文信息。但是基於空洞卷積的方法仍然是從少數的周圍點中獲取信息而不能形成密集的上下文信息。同時,基於池化的方法以非適應的方式獲得上下文信息並且對圖像所有像素獲得同質的上下文信息,這不能滿足不同像素需要不同上下文依賴的需求。

爲了獲得密集的像素級的上下文信息,PSANet[3]通過預測注意力圖中學習彙總每個位置的上下文信息,Non-local[4]網絡利用自注意力機制使任何位置的單一特徵能夠感知所有其他位置的特徵,能夠產生更強大的像素級的表徵能力。

Self-Attention機制能夠捕獲特徵圖中任意兩個位置的空間依賴關係,獲得長距離上下文依賴信息。Ulku等[5]人認爲全局上下文信息決定了最終的性能。2020年來自EPFL的ICLR 2020的一篇論文闡述到自注意力機制可以表達任何卷積濾波層[6]。那麼什麼是self-attention機制呢?

Self-Attention Mechanism

首先介紹Self-Attention機制。Self-Attention是從NLP中借鑑過來的思想,因此仍然保留了Query, Key和Value等名稱。下圖是self-attention的基本結構,feature maps是由基本的深度卷積網絡得到的特徵圖,如ResNet、Xception等,這些基本的深度卷積網絡被稱爲backbone,通常將最後ResNet的兩個下采樣層去除使獲得的特徵圖是原輸入圖像的1/8大小。

Self-attention結構自上而下分爲三個分支,分別是query、key和value。計算時通常分爲三步:

  • 第一步是將query和每個key進行相似度計算得到權重,常用的相似度函數有點積,拼接,感知機等;
  • 第二步一般是使用一個softmax函數對這些權重進行歸一化;
  • 第三步將權重和相應的鍵值value進行加權求和得到最後的attention。

 

下面我們通過實例代碼講述self-attention的原理。

class Self_Attn(nn.Module):
    """ Self attention Layer"""
    def __init__(self,in_dim,activation):
        super(Self_Attn,self).__init__()
        self.chanel_in = in_dim
        self.activation = activation

        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
        self.gamma = nn.Parameter(torch.zeros(1))

        self.softmax  = nn.Softmax(dim=-1) 
    def forward(self,x):
        """
            inputs :
                x : input feature maps( B X C X W X H)
            returns :
                out : self attention value + input feature 
                attention: B X N X N (N is Width*Height)
        """
        m_batchsize,C,width ,height = x.size()
        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
        energy =  torch.bmm(proj_query,proj_key) # transpose check
        attention = self.softmax(energy) # BX (N) X (N) 
        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N

        out = torch.bmm(proj_value,attention.permute(0,2,1) )
        out = out.view(m_batchsize,C,width,height)

        out = self.gamma*out + x
        return out,attention

假設feature maps的大小是Batch_size×Channels×Width×Height

 

在初始化函數中,定義了三個1×1卷積,分別是query_conv , key_conv 和 value_conv。

  • 在query_conv卷積中,輸入爲B×C×W×H,輸出爲B×C/8×W×H;
  • 在key_conv卷積中,輸入爲B×C×W×H,輸出爲B×C/8×W×H;
  • 在value_conv卷積中,輸入爲B×C×W×H,輸出爲B×C×W×H。

 

在forward函數中,定義了self-attention的具體步驟。

proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1)

proj_query本質上就是卷積,只不過加入了reshape的操作。首先是對輸入的feature map進行query_conv卷積,輸出爲B×C/8×W×H;view函數是改變了輸出的維度,就單張feature map而言,就是將W×H大小拉直,變爲1×(W×H)大小;就batchsize大小而言,輸出就是B×C/8×(W×H);permute函數則對第二維和第三維進行倒置,輸出爲B×(W×H)×C/8。proj_query中的第i行表示第i個像素位置上所有通道的值。

 

proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height)

proj_key與proj_query相似,只是沒有最後一步倒置,輸出爲B×C/8×(W×H)。proj_key中的第j行表示第j個像素位置上所有通道的值。

 

energy =  torch.bmm(proj_query,proj_key)

這一步是將batch_size中的每一對proj_query和proj_key分別進行矩陣相乘,輸出爲B×(W×H)×(W×H)。Energy中的第(i,j)是將proj_query中的第i行與proj_key中的第j行點乘得到。這個步驟的意義是energy中第(i,j)位置的元素是指輸入特徵圖第j個元素對第i個元素的影響,從而實現全局上下文任意兩個元素的依賴關係。

 

attention = self.softmax(energy)

這一步是將energe進行softmax歸一化,是對行的歸一化。歸一化後每行的之和爲1,對於(i,j)位置即可理解爲第j位置對i位置的權重,所有的j對i位置的權重之和爲1,此時得到attention_map。

 

proj_value = self.value_conv(x).view(m_batchsize,-1,width*height)

proj_value和proj_query與proj_key一樣,只是輸入爲B×C×W×H,輸出爲B×C×(W×H)。從self-attention結構圖中可以知道proj_value是與attention_map進行矩陣相乘,即下面兩行代碼。

 

out = torch.bmm(proj_value,attention.permute(0,2,1) )
out = out.view(m_batchsize,C,width,height)

在對proj_value與attention_map點乘之前,先對attention進行轉置。這是由於attention中每一行的權重之和爲1,是原特徵圖第j個位置對第i個位置的權重,將其轉置之後,每一列之和爲1;proj_value的每一行與attention中的每一列點乘,將權重施加於proj_value上,輸出爲B×C×(W×H)。

out = self.gamma*out + x

這一步是對attention之後的out進行加權,x是原始的特徵圖,將其疊加在原始特徵圖上。Gamma是經過學習得到的,初始gamma爲0,輸出即原始特徵圖,隨着學習的深入,在原始特徵圖上增加了加權的attention,得到特徵圖中任意兩個位置的全局依賴關係。

 

以上是self-attention的原理,下面針對一些論文的具體網絡進行分析。

 

論文解析

DANet

DANet[7]是一種經典的應用self-Attention的網絡,它引入了一種自注意力機制來分別捕獲空間維度和通道維度中的特徵依賴關係。

 

從其結構圖中可以看到,它由兩個並列的attention module組成,第一個就是前文所述的原理,得到的是特徵圖中任意兩個位置的依賴關係,稱爲Position Attention Module(PAM);第二個是任意兩個通道間的依賴關係,稱爲Channel Attention Module(CAM)。

從其具體的模塊中來看,PAM中的attention_map的大小爲B×(W×H)×(W×H),而CAM中的attention_map大小爲B×C×C,這就是PAM與CAM的區別,他們所代表的一個是任意兩個位置之間的依賴關係,一個代表的是任意兩個通道之間的依賴關係。

其具體代碼地址:https://github.com/junfu1115/DANet/blob/master/encoding/nn/attention.py

 

CCNet

CCNet[8]是針對self-attention佔用GPU內存大和計算量大提出來的,它能減少11倍的內存佔用和85% FLOPs的同時獲得長距離依賴關係。從下面示意圖中可以看出來,若要計算輸入feature map中藍色點與其他像素點之間的依賴關係,(a)是一般的Non-local模塊,其得到的權重是W×H個;(b)是論文提出中Criss-Cross Attention模塊,它只關注每個像素點所在行所在列的權重,得到的權重是H+W-1個,循環兩次(爲什麼要循環兩次後文解釋)後便能達到與self-attention一樣的效果,還能節省大量的計算與內存。

 

在其具體的criss-cross attention module中,與self-attention不同的是通過Query和Key獲得的attention_map中權重的個數是H+W-1,只獲得了水平和垂直方向的上下文信息,而這些信息不足以進行語義分割。

 

 

那麼如何獲得全局上下文信息呢?論文中認爲criss-cross attention module循環兩次即可獲得任意兩個像素點的依賴關係,即全局特徵。

比如我們要獲得下圖右上角藍色點與左下角綠色兩個像素點的依賴關係,在Loop1中,右上角(ox,oy)藍色點將信息傳遞到(ux,oy)和(ox,uy),還不能傳播到左下角點(ux,uy),在Loop2中,左下角點(ux,uy)能夠從左上角點(ux,oy)和右下角點(ox,uy)中得到信息,這時已經包含了藍色點的信息,所以右上角點信息傳播到左下角點。同理,任何不能一次遍歷的位於十字位置的點只需兩次就能完全遍歷。這就是爲什麼論文中循環使用兩次criss-cross attention module的原因,此結構稱之爲Recurrent Criss-Cross Attention module(RCCA)。

ISSA

ISSA[9]同樣是針對self-attention需要大量計算與內存而提出來的一種模型,是將interlace機制與self-attention機制結合從而獲得任意兩個位置的依賴關係的,其主要思想是將密集相似矩陣分解爲兩個稀疏相似矩陣的乘積。

 

我們將所有的輸入位置分爲Q個大小相等的子集(上圖中Q=4),每個子集中包含P個位置(N=P*Q,上圖中P=4)。對於長距離注意模塊,我們從每一個子集中採樣一個位置構建一個含有Q個位置的新子集(因爲原先劃分爲Q個子集,每個子集拿出一個位置,就有了Q個位置),根據這樣採樣策略能夠獲得P個這樣的子集。這樣每個構造子集中的位置是長空間間隔距離的位置。在每個子集上採用自注意力機制計算稀疏相似性矩陣AL。對於短距離注意模塊,直接在原始的Q個子集上(相當於在long-range模塊permute之後再permute回來)利用self-attention計算稀疏相似性矩陣AS。融合這兩個機制,便可以將信息從每個輸入位置傳播到所有輸出位置。

 

上圖是ISSA的僞代碼,原理相對簡單,是對self-attention的一種改進,節省了大量的計算與內存佔用情況。

 

以上就是對self-attention機制的一些解釋,當然還有一些其他優秀論文,語義分割也有其他方向的發展,等閱讀到了一些優秀論文再更新。

 

參考文獻:

[1] Chen L C, Papandreou G, Kokkinos I, et al. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs[J]. IEEE transactions on pattern analysis and machine intelligence, 2017, 40(4): 834-848.

[2] Zhao H, Shi J, Qi X, et al. Pyramid scene parsing network[C].Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 2881-2890.

[3] Zhao H, Zhang Y, Liu S, et al. Psanet: Point-wise spatial attention network for scene parsing[C].Proceedings of the European Conference on Computer Vision (ECCV). 2018: 267-283.

[4] Wang X, Girshick R, Gupta A, et al. Non-local neural networks[C].Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 7794-7803.

[5] Ulku I, Akagunduz E. A Survey on Deep Learning-based Architectures for Semantic Segmentation on 2D images[J]. arXiv preprint arXiv:1912.10230, 2019.

[6] Cordonnier J B, Loukas A, Jaggi M. On the Relationship between Self-Attention and Convolutional Layers[J]. arXiv preprint arXiv:1911.03584, 2019.

[7] Fu J, Liu J, Tian H, et al. Dual attention network for scene segmentation[C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019: 3146-3154.

[8] Huang Z, Wang X, Huang L, et al. Ccnet: Criss-cross attention for semantic segmentation[C]//Proceedings of the IEEE International Conference on Computer Vision. 2019: 603-612.

[9] Huang L, Yuan Y, Guo J, et al. Interlaced sparse self-attention for semantic segmentation[J]. arXiv preprint arXiv:1907.12273, 2019.

 

 

參考博客:

Self-Attention GAN 中的 self-attention 機制

關於attention機制在nlp中的應用總結

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章