非局部注意力：NL论文笔记——Non-local Neural Networks

综述

论文题目：《Non-local Neural Networks》

会议时间：IEEE Conference on Computer Vision and Pattern Recognition 2018 (CVPR, 2018)

论文地址：https://openaccess.thecvf.com/content_cvpr_2018/papers/Wang_Non-Local_Neural_Networks_CVPR_2018_paper.pdf

简介

  传统的卷积运算和循环运算都是处理局部领域的运算，以卷积运算为例，经过$3\times 3$卷积运算得到的特征图中，每个特征数据只与运算前的9个数据有关，即感受野为3，如下图所示，9个深蓝数据经过某种运算得到数据A，A可以同时反应9个相邻的特征数据：

卷积运算的一大优点就是局部感知，只对局部区域做特征提取，降低了参数量，但这种特征提取方式不利于构建远程依赖关系（long-range dependencies），比如在上图中，单次卷积只能将相邻的9个特征数据加以关联，但是如果想让不相邻的数据做关联，如上图中的B和C，则需要经过多次卷积运算，即堆积多个卷积层，逐步扩大感受野，当感受野增大到一定程度时，B和C就可以关联起来，此时网络可以挖掘B和C之间共有的特性，利用这种共性可以辅助网络做判断。但堆叠多个卷积层会带来参数过多，计算量过大等问题，不利于网络的优化。

  直观地来讲，单看下面第一张图，足球的位置不仅仅由足球本身所反映，人的朝向，眼睛所注视的方向同样也可以间接反映足球的位置，因此这些特征之间具有某种联系（对应空间联系），单纯使用卷积运算很难让网络挖掘这种联系。扩展到视频中，不同帧之间同样也有一定的联系（对应时间联系），帧之间的图像特点也可以相互反映：

想要建立这种联系，核心就是利用某种运算将每个特征与所有特征数据都关联起来，对此，作者参考非局部均值操作（non-local mean operation），定义了一种通用的非局部操作，利用不同位置之间的联系（relationship）来计算响应值，可以用如下公式表示：
$$y_i=\frac{1}{\mathcal{C}(x)}\sum _{\forall j}f(x_i,x_j)g(x_j)$$
其中，$i$表示当前像素位置，$j$表示其他所有的像素位置，$x$是输入的特征数据，$f(\cdot )$表示计算位置$i$和$j$之间关系的函数，$g(\cdot )$表示计算$j$处输入信号的表示（representation），$\mathcal{C}(x)$用于归一化，形式因$f(\cdot )$而异，$y_i$可以表示其他所有位置上的特征数据对当前位置$i$上特征数据的响应，也就是$y_i$可以表示其他位置与当前位置的关系。

$f(\cdot )$的选择：

• Gaussian：$f(x_i,x_j)=e^{x_i^T{x}_j}$，$\mathcal{C}(x)={\sum }_{\forall j}f(x_i,x_j)$，$x_i^T{x}_j$表示点积运算
• Embedded Gaussian：$f(x_i,x_j)=e^{\theta (x_i{)}^T\varphi (x_j)}$，$\mathcal{C}(x)={\sum }_{\forall j}f(x_i,x_j)$

其中$\theta (x_i)=W_{\theta }{x}_i$，$\varphi (x_j)=W_{\varphi }{x}_j$，可以利用$1\times 1$的卷积运算实现，前两类加一个$e$的指数，相当于经过了一次softmax函数，即Embedded Gaussian可以写为$y=softmax(x^T{W}_{\theta }^T{W}_{\varphi }x)$

• Dot product：$f(x_i,x_j)=\theta (x_i{)}^T\varphi (x_j)$，$\mathcal{C}(x)=N$
• Concatenation：$f(x_i,x_j)=ReLU(w_f^T[\theta (x_i),\varphi (x_j)])$，$\mathcal{C}(x)=N$

注意：$f(\cdot )$的形式不重要，重要的是非局部关系的计算，即$y_i$的计算形式；

  为了将非局部操作嵌入到卷积神经网络中，作者进一步定义了非局部块（non-local block），可以表示为：
$$z_i=W_z{y}_i+x_i$$
由于残差连接$+x_i$的存在，使得非局部块可以嵌入到任何预训练网络中，不需要改变初始参数（此时可以将$W_z$初始化为0），模块结构如下图所示：

其中，$W_g、W_{\theta }、W_{\varphi }$将每个位置上的特征数目压缩一半，也就是将整张特征图的通道数压缩一半，可以利用$1\times 1$的卷积运算实现，同理，$W_z$将$y_i$的通道数再扩充回来，这种瓶颈式的操作可以将计算量减小一半。

  同时，还可以利用下采样操作进一步降低计算量，$y_i$的计算方式可以改为：
$$y_i=\frac{1}{\mathcal{C}(\hat{x})}\sum _{\forall j}f(x_i,{\hat{x}}_j)g({\hat{x}}_j)$$
其中$\hat{x}$表示$x$的下采样数据，利用采样数据去计算特征之间的联系可以降低$1/4$的计算量，可以在$\varphi$和$g$之后添加一个最大池化层来实现。

  作者在论文中以ResNet网络为基础网络做实验，实验结果显示非局部模块加在res2、res3、res4阶段较好，每个阶段加在最后一个残差块前面，并且之前提到的$f(\cdot )$计算方法效果均类似（除了Gaussian提升不大），论文后半部分的实验均用Embedded Gaussian来计算特征联系。

补充：

  在神经网络中，全连接层也可以将所有特征数据加以关联（加权求和），下一组特征中每个特征数据均由上一组所有的特征数据构成，但是全连接层中的加权求和操作忽略了特征位置之间的联系，这也是NL这篇文章一直在强调的内容，NL更注重像素点之间的联系，这里的联系以一种函数关系呈现。举例来说，在第一张图中，由于非局部模块的引入，可以让网络知道B点的存在对C点的影响，在第二张图中，也可以让网络发现身体朝向与足球位置之间的联系，利用身体朝向辅助预测足球的位置。全连接层只是用于像素之间简单的求和相加，不能让网络挖掘这种隐含联系。（并且全连接层会打混特征的空间分布，因此不能加在卷积运算中间做辅助运算，全连接层常适用于加在最后做分类。）

代码实现

参考代码：

• https://github.com/AlexHex7/Non-local_pytorch

以Embedded Gaussian为例：

class _NonLocalBlockND(nn.Module):
    def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):
        """
        :param in_channels:
        :param inter_channels:
        :param dimension:
        :param sub_sample:
        :param bn_layer:
        """

        super(_NonLocalBlockND, self).__init__()

        assert dimension in [1, 2, 3]

        self.dimension = dimension
        self.sub_sample = sub_sample

        self.in_channels = in_channels
        self.inter_channels = inter_channels

        # 默认压缩一次通道，将原始通道数压缩一半
        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1

        # 判断用哪类卷积运算，对应不同的计算机视觉任务
        if dimension == 3:
            conv_nd = nn.Conv3d
            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
            bn = nn.BatchNorm3d
        elif dimension == 2:
            conv_nd = nn.Conv2d
            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
            bn = nn.BatchNorm2d
        else:
            conv_nd = nn.Conv1d
            max_pool_layer = nn.MaxPool1d(kernel_size=(2))
            bn = nn.BatchNorm1d

        # 定义g运算，对应论文公式1中的g(xj)
        self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0)

        # 定义最后的卷积层，用于还原通道数量
        if bn_layer:
            self.W = nn.Sequential(
                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                        kernel_size=1, stride=1, padding=0),
                bn(self.in_channels)
            )
            nn.init.constant_(self.W[1].weight, 0)
            nn.init.constant_(self.W[1].bias, 0)
        else:
            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                             kernel_size=1, stride=1, padding=0)
            nn.init.constant_(self.W.weight, 0)
            nn.init.constant_(self.W.bias, 0)

        # 定义θ运算
        self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                             kernel_size=1, stride=1, padding=0)
        # 定义φ运算
        self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                           kernel_size=1, stride=1, padding=0)
        # 是否下采样，如果执行下采样的话，参与g运算和φ运算的特征图之后会经过一次下采样
        if sub_sample:
            self.g = nn.Sequential(self.g, max_pool_layer)
            self.phi = nn.Sequential(self.phi, max_pool_layer)

    def forward(self, x):
        """
        :param x: (b, c, t, h, w)
        :param return_nl_map: if True return z, nl_map, else only return z.
        :return:
        """

        batch_size = x.size(0)
        # 特征图经过g(·)运算，并且做一次转置，将特征图尺寸转为(b, t×h×w, c//2)，便于后续矩阵相乘，得到g(xj)
        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)
        # 特征图经过θ运算，并且再转置，尺寸同样为(b, t×h×w, c//2)，得到θ(xi)
        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)
        # 特征图经过φ运算，得到φ(xj)，尺寸为(b, c//2, t×h×w)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        # θ(xi)与φ(xj)做矩阵相乘，相当于元素相乘并且按j求和，得到特征图尺寸为(b, t×h×w, t×h×w)
        f = torch.matmul(theta_x, phi_x)
        # 传入softmax运算(沿最后一个维度)
        f_div_C = F.softmax(f, dim=-1)
        
        # 将得到的结果再和g(xj)做矩阵相乘，得到特征图尺寸为(b, t×h×w, c//2)
        y = torch.matmul(f_div_C, g_x)
        # 之后将结果经过转置、调整维度，调整为(b, c//2, t, h, w)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        # 传入最后一层卷积，通道数还原为c
        W_y = self.W(y)
        # 与输入x相加，得到最后的结果
        z = W_y + x

        return z

注：以上仅是笔者个人见解，若有错误，欢迎指正。