CCNet: Criss-Cross Attention for Semantic Segmentation
Abstract
上下文資訊對于語義分割和目标檢測任務都很重要，這裡提出CCNet。對于每個像素，criss-cross attention子產品能獲得其交叉路徑上所有像素的上下文資訊，通過進一步的遞歸操作，每個像素最終可以捕獲全圖像的依賴關系。此外，提出類别一緻損失使得criss-cross attention子產品生成更具判别性的特征。CCNet有以下優點：(1)GPU顯存友好，比non-local block少11倍顯存消耗 (2)高計算效率，比non-local block少85% (3)最先進性能，Cityscapes可達81.9%
開源位址：https://github.com/speedinghzl/CCNet
這裡推薦一篇很有意思的工作，通過注意力機制處理預測head關聯多目标跟蹤領域的檢測和ReID特征：Rethinking the competition between detection and ReID in Multi-Object Tracking，該論文不但構思了一種簡單有效的注意力機制，還巧妙地利用注意力機制來交叉關聯兩個任務，避免了檢測和ReID的競争，并聯合提升了彼此分支的性能。
Introduction
FCN是固定的幾何結構(卷積的網格效應)，局部感受野隻能提供短距離上下文資訊。為了彌補FCN的缺點，deeplab系列提出多尺度空洞卷積結構的ASPP子產品聚合上下文資訊，PSPNet引入金字塔池化子產品捕獲上下文資訊。然而，基于空洞卷積的方法從一些周圍像素收集資訊，不能準确産生稠密的上下文資訊；基于池化的方法以非自适應的相同上下文提取政策處理所有像素，不能滿足不同像素需要不同上下文依賴的要求。Non-local的空間和時間複雜度高，需要改進。
Related work
UNet，Deeplabv3+，MSCI，SPGNet，RefineNet，DFN采取encoder-decoder結構，融合低層次和高層次資訊做出稠密預測。Scale-adaptive Convolutions(SAC)和Deformable Convolutional Network(DCN)改善标準卷積處理目标形變和各種尺寸目标…
Approach
1.Network Architecture
上圖是Non-local block和Criss-Cross Attention block結構的簡化示意圖，Non-local block通過計算任意兩個位置之間的互動直接捕獲遠距離依賴，而不用局限于相鄰點，但是由于每個位置對應的向量(共 H × W H\times W H×W個)都要和 H × W H\times W H×W個向量相乘，帶來的計算量偏大；而Criss-Cross Attention block一次隻用考慮"十字交叉"的同行同列的向量，即每個位置對應的向量( H × W H\times W H×W個)都要和 ( H + W − 1 ) (H+W-1) (H+W−1)個向量相乘，這樣捕獲的是一個位置和”十字交叉“路徑上其他位置的依賴，但是将輸出結構果再次送入Criss-Cross Attention block即可獲得一個位置與全局位置的依賴，進一步的了解可以看下文的Criss-cross Attention子產品細節。
Backbone是全卷積網絡，移去最後兩個下采樣操作，并且随後的卷積層都采取空洞卷積(帶空洞卷積的FCN)，輸出特征圖 X X X為輸入圖像的1/8。 X X X通過卷積層降低通道次元輸出 H H H， H H H送入criss-cross attention子產品聚合每個像素交叉路徑上的上下文資訊得到 H ′ H' H′， H ′ H' H′再次送入criss-cross attention子產品輸出 H ′ ′ H'' H′′，則 H ′ ′ H'' H′′的每個像素聚合了所有像素的資訊。兩個criss-cross attention子產品共享參數，取名為recurrent Criss-Cross Attention(RCCA)子產品。然後， H ′ ′ H'' H′′和特征 X X X進行concat，接着是一個或幾個帶BN的卷積層和激活層用于特征融合，最後融合的特征送入分割層預測最終的分割結果。
2.Criss-cross Attention
考慮一個局部特征圖 H ∈ R C × W × H H\in \mathbb{R}^{C\times W\times H} H∈RC×W×H，首先通過兩個 1 × 1 1\times1 1×1卷積生成兩個特征圖 Q Q Q和 K K K， { Q , K } ∈ R C ′ × W × H \{Q,K\} \in{\mathbb{R}^{C'\times W\times H}} {Q,K}∈RC′×W×H， C ′ C' C′是比 C C C小的通道數，形狀為 C ′ × H × W C'\times H\times W C′×H×W的三維特征圖可以很容易reshape成二維的 C ′ × ( H × W ) C'\times (H \times W) C′×(H×W)的矩陣。通過Affinity操作生成注意力圖 A A A，對于特征圖 Q Q Q的每一個位置 u u u，拉出一條次元為 C ′ C' C′的向量 Q u ∈ R C ′ Q_u\in \mathbb{R}^{C'} Qu∈RC′。同時，從特征圖 K K K拉出同屬于 u u u位置的同行或同列的 ( H + W − 1 ) (H+W-1) (H+W−1)條( H + W H+W H+W會包括 u u u位置兩次)次元均為 C ’ C’ C’的向量集 Ω u ∈ R ( H + W − 1 ) × C ′ \Omega_{u}\in\mathbb{R}^{(H+W-1)\times C'} Ωu∈R(H+W−1)×C′， Ω i , u ∈ R C ′ \Omega_{i,u}\in\mathbb{R}^{C'} Ωi,u∈RC′是 Ω u \Omega_{u} Ωu的第 i i i個元素(向量)，則Affinity操作可用公式表達如下： d i , u = Q u Ω i , u T d_{i,u}=Q_u\Omega_{i,u}^T di,u=QuΩi,uT d i , u ∈ D d_{i,u}\in D di,u∈D是特征 Q u Q_u Qu和 Ω i , u \Omega_{i,u} Ωi,u之間的關聯度， i = [ 1 , . . . , H + W − 1 ] i=[1,...,H+W-1] i=[1,...,H+W−1]，且 D ∈ R ( H + W − 1 ) × ( W × H ) D\in \mathbb{R}^{(H+W-1)\times (W\times H)} D∈R(H+W−1)×(W×H)，然後對 D D D在通道次元上添加softmax層，輸出注意力圖A
輸入特征圖H通過另外一個 1 × 1 1\times1 1×1卷積生成特征圖 V ∈ R C × W × H V \in{\mathbb{R}^{C\times W\times H}} V∈RC×W×H，對于特征圖 V V V的每一個位置 u u u，同理拉出次元為 C C C的向量 V u ∈ R C V_u\in \mathbb{R}^{C} Vu∈RC和向量集 Φ u ∈ R ( H + W − 1 ) × C \Phi_u \in \mathbb{R}^{(H+W-1)\times C} Φu∈R(H+W−1)×C，然後給出聚合(Aggregation)操作的公式如下： H u ′ = ∑ i = 1 H + W − 1 A i , u Φ i , u + H u H_u'=\sum_{i=1}^{H+W-1}A_{i,u}\Phi_{i,u}+H_u Hu′=i=1∑H+W−1Ai,uΦi,u+Hu
其中， H u ′ H_u' Hu′是 H ′ ∈ R C × W × H H' \in \mathbb{R}^{C\times W \times H} H′∈RC×W×H中位置 u u u的特征向量， A i , u A_{i,u} Ai,u是注意力圖A中位置 u u u對應的第 i i i個數值。最後是以殘差的形式輸出 H ′ H' H′，增強了像素級的表達能力，并聚合了全局上下文資訊，提升了語義分割的性能。
Recurrent Criss-Cross Attention(RCCA)子產品包含兩個Criss-Cross Attention子產品，且是共享參數的，RCCA可獲得一個位置與全局位置的依賴，能夠獲得稠密豐富的下文資訊。
3.Learning Category Consistent Features
對于語義分割任務，同一類像素應該有相似的特征，不同類像素應該有差别大的特征，這被稱作類别一緻性。論文認為，RCCA子產品聚合的特征可能會存在過度平滑的問題，這是圖神經網絡的常見問題，是以除了使用交叉熵損失 l s e g l_{seg} lseg監督外，還提出了類别一緻損失。RCCA子產品輸出後接 1 × 1 1\times1 1×1卷積降低特征圖通道數(這裡設定16)，在這個低通道數特征圖 M M M上添加類别一緻性損失。假定 C C C是mini-batch images裡存在的類别數， N c N_c Nc是屬于類别 c ∈ C c\in C c∈C的有效元素數目， h i ∈ H h_i \in H hi∈H是特征圖M空間位置 i i i對應的特征向量( i i i是屬于類别 c c c的，是 N c N_c Nc中的一個元素)， u c u_c uc是類别 c c c的平均特征向量(聚類中心)， φ v a r ( h i , u c ) \varphi_{var}(h_i,u_c) φvar(hi,uc)計算兩者之間的距離進行懲罰，希望同類别像素對應的特征向量具有相似性，靠近該類聚類中心最好； u c a u_{c_a} uca和 u c b u_{c_b} ucb是兩個不同類别的聚類中心（特征向量）， φ d i s ( u c a , u c b ) \varphi_{dis}(u_{c_a},u_{c_b}) φdis(uca,ucb)計算兩個類别中心之間的距離進行懲罰，兩兩類别計算，希望不同類别像素的聚類中心越遠越好。 l r e g l_{reg} lreg是聚類中心向量的正則項損失，最終損失是所有損失的權重和： l = l s e g + α l v a r + β l d i s + γ l r e g l=l_{seg}+\alpha l_{var}+\beta l_{dis}+\gamma l_{reg} l=lseg+αlvar+βldis+γlreg其中，設定 δ v = 0.5 , δ d = 1.5 , α = β = 1. γ = 0.001 \delta_v=0.5,\delta_d=1.5,\alpha=\beta=1.\gamma=0.001 δv=0.5,δd=1.5,α=β=1.γ=0.001。總之，類别一緻損失是從特征上，希望同類别像素特征具有相似性，不同類相似特征具有差異性。損失具體公式如下：