【文献阅读】Asymmetric Non-local Neural Networks for Semantic Segmentation

原文链接：https://arxiv.org/abs/1908.07678
代码：https://github.com/MendelXu/ANN.git

---

Non-local 是一种特别有用的语义分割技术，但也因其难以进行计算和占用GPU内存而受到批评。本文提出了Asymmetric Non-local Neural Network，其中有两个突出的组成部分：Asymmetric Pyramid Non-local Block（APNB） 和 Asymmetric Fusion Non-local Block（AFNB）。APNB利用金字塔采样模块，在不牺牲性能的前提下，极大地减少了计算和内存消耗；AFNB是由APNB演化而来的，在充分考虑了长期相关性的前提下，融合了不同层次的特征，从而大大提高了性能。

核心思路是：
只要query分支和key分支的输出保持相同的大小，Non-local的输出大小就保持不变。考虑到这一点，如果只从key分支和value分支中选取几个有代表性的点作为样本，就有可能在不牺牲性能的情况下显著降低时间复杂度。如下图所示，其中，$S \ll N$S≪N。

---

Asymmetric Non-local Neural Network

提出了APNB和AFNB，其中APNB旨在减少Non-local的计算开销，AFNB提高了Non-local的学习能力，从而提高了分割性能。

重识Non-local

考虑输入特征图$X \in \mathbb{R}^{C \times H \times W}$X∈RC×H×W，三个$1 \times 1$1×1卷积$W_{\phi}$Wϕ​、$W_{\theta}$Wθ​和$W_{\gamma}$Wγ​用来将$X$X变换到不同的嵌入式空间$\phi \in \mathbb{R}^{C' \times H \times W}$ϕ∈RC′×H×W、$\theta \in \mathbb{R}^{C' \times H \times W}$θ∈RC′×H×W、$\gamma \in \mathbb{R}^{C' \times H \times W}$γ∈RC′×H×W
$\phi = W_{\phi}(X), \space \theta = W_{\theta}(X), \space \gamma = W_{\gamma}(X)$ϕ=Wϕ​(X), θ=Wθ​(X), γ=Wγ​(X)
接着将他们三个$flatten$flatten为$C' \times N$C′×N，$N$N为空间中总像素数，$N = H \times W$N=H×W。相似度矩阵$V \in \mathbb{R}^{N \times N}$V∈RN×N可由矩阵乘法计算
$V = \phi^T \times \theta$V=ϕT×θ
接着，将$V$V进行归一化
$\vec{V}=f(V)$V=f(V)
$f$f有三种形式，分别为$softmax$softmax、$rescaling$rescaling和$none$none，这里选择用$softmax$softmax。对于$\gamma$γ中的每个位置，都可以计算输出为
$O=\vec{V} \times \gamma^T$O=V×γT
其中，$O \in \mathbb{R}^{N \times C'}$O∈RN×C′。通过参考Non-local的设计，最终的输出为
$Y = W_o(O^T) + X \space or \space Y = cat(W_o(O^T),X)$Y=Wo​(OT)+X or Y=cat(Wo​(OT),X)
其中，$W_o$Wo​使用$1 \times 1$1×1卷积，作为一个加权参数来调整Non-local的重要性，并且将通道维度从$C'$C′还原为$C$C。

Asymmetric Pyramid Nonlocal Block（APNB）

动机和分析

Non-local在两个矩阵相乘的时间复杂度为$O(C'N^2)=O(C'H^2W^2)$O(C′N2)=O(C′H2W2)。在语义分割任务中，网络的输出通常有较大的分辨率来保持细节的语义信息，$N$N通常很大。
一个更直接的$pipeline$pipeline如下：
$\mathbb{R}^{N \times C'} \times \mathbb{R}^{C' \times N}\rightarrow \mathbb{R}^{N \times N} \times \mathbb{R}^{N \times C'} \rightarrow \mathbb{R}^{N \times C'}$RN×C′×RC′×N→RN×N×RN×C′→RN×C′
通过将$N$N变为$S$S（$S \ll N$S≪N），输出尺寸保持不变
$\mathbb{R}^{N \times C'} \times \mathbb{R}^{C' \times S}\rightarrow \mathbb{R}^{N \times S} \times \mathbb{R}^{S \times C'} \rightarrow \mathbb{R}^{N \times C'}$RN×C′×RC′×S→RN×S×RS×C′→RN×C′
将$N$N变为更小的$S$S相当于从图中采样几个有代表性的点，而不是考虑空间中所有的点，从而将计算量大大减少。

方法

在$\theta$θ和$\gamma$γ后，通过添加采样模块$P_\theta$Pθ​和$P_\gamma$Pγ​，来采样一些稀疏的锚点，得到$\theta_P \in \mathbb{R}^{C' \times S}$θP​∈RC′×S和$\gamma_P \in \mathbb{R}^{C' \times S}$γP​∈RC′×S，其中$S$S为采样锚点的个数。数学公式为
$\theta_P = P_{\theta}(\theta),\space \gamma_P = P_{\gamma}(\gamma)$θP​=Pθ​(θ), γP​=Pγ​(γ)
$\phi$ϕ与锚点$\theta_P$θP​间的相似度矩阵$V_P$VP​如下
$V_P=\phi^T \times \theta_P$VP​=ϕT×θP​
其中，$V_P$VP​是一个不对称的矩阵，大小为$N \times S$N×S。接下来与Non-local一样进行归一化，得到$\vec{V_P}$VP​​。$attention$attention输出为
$O_P=\vec{V_P} \times \gamma_P^T$OP​=VP​​×γPT​
最终的输出$Y_P$YP​为
$Y_P = cat(W_o(O_P^T),X)$YP​=cat(Wo​(OPT​),X)
时间复杂度仅为$O(C'NS)$O(C′NS)，比标准Non-local中的$O(C'N^2)$O(C′N2)小很多。

前面的工作表明，全局和多尺度对语义分割很有帮助，因此，在Non-local块中嵌入金字塔池化来增强全局表示。

这四个池化结果被平铺并连接起来，作为被采样的点。设置$n \subseteq \{1,3,6,8\}$n⊆{1,3,6,8}，则锚点的总量为
$S = 110 = \sum_{n \subseteq \{1,3,6,8\}}n^2$S=110=n⊆{1,3,6,8}∑​n2

Asymmetric Fusion Nonlocal Block（AFNB）

不同层级的特征融合对语义分割有帮助，将特征融合加入到Non-local中，得到Fusion Non-local Block。
标准的Non-local只有一个输入源，而FNB有两个输入源：一个高级特征图$X_h\in \mathbb{R}^{C_h \times N_h}$Xh​∈RCh​×Nh​和一个低级特征图$X_l \in \mathbb{R}^{C_l \times N_l}$Xl​∈RCl​×Nl​。同样的，$1 \times 1$1×1卷积$W_{\phi}^h$Wϕh​、$W_{\theta}^l$Wθl​和$W_{\gamma}^l$Wγl​被用来将$X_h$Xh​和$X_l$Xl​变换到嵌入式空间$\phi_h \in \mathbb{R}^{C' \times N_h}$ϕh​∈RC′×Nh​、$\theta_l \in \mathbb{R}^{C' \times N_l}$θl​∈RC′×Nl​和$\gamma_l \in \mathbb{R}^{C' \times N_l}$γl​∈RC′×Nl​
$\phi_h=W_{\phi}^h(X_h),\space \theta_l=W_{\theta}^l(X_l),\space \gamma_l=W_{\gamma}^l(X_l)$ϕh​=Wϕh​(Xh​), θl​=Wθl​(Xl​), γl​=Wγl​(Xl​)
$\phi_h$ϕh​和$\theta_l$θl​间的相似度矩阵$V_F \in \mathbb{R}^{N_h \times N_l}$VF​∈RNh​×Nl​为
$V_F=\phi_h^T \times \theta_l$VF​=ϕhT​×θl​
同样的，将$V_F$VF​进行归一化得到$\vec{V_F} \in \mathbb{R}^{N_h \times N_l}$VF​​∈RNh​×Nl​。同理输出为
$O_F=\vec{V_F} \times \gamma_l^T$OF​=VF​​×γlT​

$Y_P = cat(W_o(O_P^T),X)$YP​=cat(Wo​(OPT​),X)

网络结构

---

实验结果