Asymmetric Non-local Neural Networks

论文：Asymmetric Non-local Neural Networks for Semantic Segmentation,ICCV,2019.

链接：paper

代码：github

本文是在Non-local Neural Networks（cvpr,2018) 这篇论文的基础上改进得来的，non-local block Non-local模块作为语义分割任务中很有用的技术，但因为计算量较大，需要较大的显存开销，这阻碍了non-local network在实际应用中的使用。本文提出了APNB来减少non-local block的计算量和显存开销,AFNB通过提升分割性能增强non-local block的学习能力。我们着重分析这篇论文的改进点。

Revisiting Non-local Block

输入 $X \in \mathbb{R}^{C*H*W}$

1.经过3个1x1的卷积Query，Key和Value变换之后，分别得到 : $\varnothing \in \mathbb{R}^{\hat{C}HW}==\mathbb{R}^{\hat{C}N}$ ， $\theta \in \mathbb{R}^{\hat{C}*H*W}==\mathbb{R}^{\hat{C}*N}$ ， $\gamma \in \mathbb{R}^{\hat{C}*H*W}==\mathbb{R}^{\hat{C}*N}$

2.计算所有空间位置特征向量的相似度： $V=\varnothing ^{T}*\theta$

3.归一化： $\overrightarrow{V}=softmax(V)$

4.根据所有空间位置特征向量的相似度加权求和： $O=\overrightarrow{V}*\gamma^ {T}$

5.新特征经过1*1卷积变换后与原输入特征拼接： $Y = cat\left(W_{o}\left(O^{T} \right ),X\right)$

计算复杂度： $O\left ( \hat{C}N^{2} \right )=O\left ( \hat{C}H^{2}W^{2} \right )$

Asymmetric Pyramid Non-local Block(APNB)

Non-local 的计算复杂度主要来自第2步和第4步：

可以看出经N（红色）改变为S，不会改变输出的size。

将N 减小到S ,相当于从θ 和 γ 采样一些特征点，代替使用HxW空间范围的点，如上图1(b)所示，增加两个sample模块Pθ和PγP_θ 和 P_γPθ和Pγ，采样一些稀疏的锚点 $\theta_{P}\in\mathbb{R}^{\hat{C}*S}$ 和 $\gamma _{P}\in\mathbb{R}^{\hat{C}*S}$ ，S就是锚点的数量。

1.经过3个1x1的卷积Query，Key和Value变换之后，分别得到 $\varnothing \in \mathbb{R}^{\hat{C}*H*W}==\mathbb{R}^{\hat{C}*N}$ ， $\theta \in \mathbb{R}^{\hat{C}*H*W}==\mathbb{R}^{\hat{C}*N}$ ， $\gamma \in \mathbb{R}^{\hat{C}*H*W}==\mathbb{R}^{\hat{C}*N}$

2.采样： $\theta_{P} = P_{\theta}\left ( \theta\right )$ , $\gamma _{P} = P_{\gamma }\left ( \gamma \right )$

3.计算所有空间位置与S个采样点的特征相似度： $V_{P}=\varnothing ^{T}*\theta_{P}$

4.归一化： $\overrightarrow{V_{P}}=softmax(V_{P})$

5.根据相似度加权求和： $O_{P}=\overrightarrow{V_{P}}*\gamma_{P}^ {T}$

5.新特征经过1X1卷积变换后与原输入特征拼接： $Y_{P} = cat\left(W_{o}\left(O_{P}^{T} \right ),X\right)$

计算复杂度： $O\left ( \hat{C}NS\right )< < O\left ( \hat{C}H^{2}W^{2} \right )$

可以看出sample之后，S(110)<<N，则该模型就大大减少了计算复杂度和内存开销。

分割网络框架

由于作者将传统的non-local模型中的N经过Sample操作之后变成了较小的S，这样就不能捕捉到每个像素点之间的相关性信息，可能会造成效果不是很好，为了使模型性能进一步的提升，作者又在考虑长依赖性的情况下融合不同级别下的特征。

作者提出的非对称non-local网络的框架如下图所示：

首先作者采用ResNet - 101 网络作为模型的backbone,并且去掉了原模型中最后两个下采样层，使用空洞卷积来保持最后两个阶段得到的特征分辨率是原图像的1/8,作者将stage4和stage5的输出用AFNB模块进行融合，作者将Stage5得到的特征称为high-level 特征，将stage4得到的特征称为low-level特征，将这两个特征输入AFNB模块进行融合，其具体的过程如下图：

该模块将stage5产生的高级特征作为Query,将stage4得到的低级特征作为Key 和 Value,首先对Query,Key和Value 分别使用1x1 的卷积核进行降维，根据Asymmetric non-local 网络的特征需要进行Pyramid Pooling 操作，来对Key 和 Value 进行进一步的特征抽取，从而得到较小的特征表示。

APNB锚点采样策略

通过金字塔池化（何凯明在fast-rcnn 中提出的）对Key路径和Value路径得到的特征图 $\theta \in \mathbb{R}^{\hat{C}*N}$ 和 $\gamma \in \mathbb{R}^{\hat{C}*N}$ 进行下采样，得到 $\theta _{P}\in\mathbb{R}^{\hat{C}*S}$ 和 $\gamma _{P}\in\mathbb{R}^{\hat{C}*S}$

使用4中不同的池化操作：{1,3,6,8},然后将池化后的特征图展开成一维向量，拼接起来，一共得到 110 个锚点。经过Pyramid Pooling操作后原本CN大小的特征变成了大小为CS的特征矩阵，然后计算Query和Key的相关性矩阵计算，就是将Query转置后与Key相乘最终得到一个N*S的相关性矩阵，最后，将这个相关性矩阵与Value 特征进行转置相乘最终得到了不同leve下的特征融合表示OFO_FOF,最后将这个表示与Stage5输出的特征拼接后得到原始图像最新的特征表示YFY_FYF.

然后对得到的最新的特征表示进行非对称金字塔non-local 模块，其目的与self-attention的目的是相似的，但是其过程是不相同的。其具体的流程如下图所示：

该过程和Non-local block 过程相似，只不过是key 和 Value 进行了统一的下采样操作。

该过程对YFY_FYF使用两次1x1 的卷积操作，得到non-local 模块中的query,key 和value.为了减少计算量使key=value，所以针对key 和 value 只进行一次Pyramid Pooling 操作，得到了新的特征表示，其他的过程与AFNB的过程相同，都是根据Query和Key计算出相关性矩阵，然后将相关性矩阵与Value的特征转置相乘得到最终的特征，再将这个特征与YFY_FYF拼接，得到作者较为满意的图像特征表示，最后在采用双线性插值的方法将特征分辨率还原，从而得到了分割结果。

实验结果

可以看出效果还是不错的。