Non-local Neural Networks

Non-local Neural Networks 网络详解

论文：Non - local Neural Networks，cvpr,2017.

链接：paper

代码：github

摘要

卷积神经网络和循环神经网络都是处理局部区域的操作，是典型的local operations。受经典的非局部均值（non-local means ）的启发，文中提出了一种non-local operations 用于捕获长距离依赖（long-range dependencies),(即如何建立图像上两个有一定距离的像素之间的联系，如何建立视频里两帧的联系，如何建立一段话中不同词的联系等) non-local operations 将某个位置的响应计算为所有位置特征的加权和。所有位置可以是空间的，时间的，时空的。这个结构可以被插入到很多计算机视觉结构中，在视频分类的任务上，non-local模型在Kinetics和Charades上都达到了最好的结果。在图像识别的任务上，non-local模型提高了COCO上物体检测／物体分割／姿态估计等任务的结果。

其实就是一种自注意力机制，和DANet 的自注意力机制有点相似，输入和输出大小相同，可以可容易的加入到网络中

引言

使用非局部操作的几个有点：（1）非局部操作通过计算任意两个位置之间的交互，直接捕获长期依赖关系，而不考虑它们的位置距离。（2）非局部操作是有效的，即使有几个层，也可以获得很好的效果。（3）文中的非局部操作保持了输入输出的大小，可以很容易的与其它操作结合。

Non-local Neural Networks

1.定义

按照非局部均值的定义，定义深度神经网络中的non-local 如下：

其中$x$ 表示输入信号，$y$ 表示输出信号，输入和输出的大小相同。$f(x_i,x_j)$用来计算$i$ 和所有可能关联的位置$j$ 之间的关系（文中其实计算的就是位置$i$ 和位置 $j$ 之间的相似度），这个关系可以是比如$i$ 和$j$ 的位置距离越远，f 值越小，表示j位置对i影响越小。$g(x_j)$ 用于计算输入信号在j位置的特征值。$C(x)$ 是归一化参数。

从上式中可以看出任意位置上的j，可以是时间的、空间的、时空的任意位置，都可以影响到i的输出值。作为对比，conv的操作是对一个局部邻域内的特征值进行加权求和，比如kernel size=3时，$i-1\le j\le i+1$。 recurrent的操作则是i时刻的值仅基于当前时刻或前一时刻（j=i or i-1）。

non-local的操作也和fc层不同。公式(1)计算的输出值受到输入值之间的关系的影响（因为要计算pairwise function），而fc则使用学习到的权重计算输入到输出的映射，在fc中$x_j$和$x_i$的关系是不会影响到输出的，这一定程度上损失了位置的相关性。另外，non-local能够接受任意size的输入，并且得到的输出保持和输入size一致。而fc层则只能有固定大小的输入输出。

2.表示形式

讨论几个$f$ 和 $g$ 的表示的几种形式，但实验显示不同的表现形式其实对non-local 的结果并没有太大影响，这表明 non-local 这个行为才是主要的提升因素。

为了简化，只考虑$g$ 是线性的情况，即 $g(x_i)=W_g{x}_j$,其中$W_g$ 是一个可学习的权重矩阵，实际上通过空间域的 1x1 的conv 或 时空域 的 1x1x1 的conv 实现的。

$f$ 的不同形式：

• **Gaussian:**常见的$f$是高斯函数。

其中$x_i^T{x}_j$是点乘相似度（dot-product similarity）。也可以用欧式距离，但是点乘在深度学习平台上更好实现。此时归一化参数$C(x)=\sum \forall jf(xi,xj)$

• **Embedded Gaussian:**高斯函数的一个简单的变种，在embedding space 中计算相似度。

其中 $\theta (xi)=W_{\theta }xi$ 和$\varphi (xj)=W_{\varphi }xj$是两个embedding。归一化参数和之前一致，为$C(x)=\sum \forall jf(xi,xj)$。

• Dot product: f也可以定义成点乘相似度.

这里使用embedded version。在这里，归一化参数设为C(x)=N (N=H*W)，其中N是x的位置的数目，而不是f的和，这样可以简化梯度的计算。

• Concatenation:

这里[.,.]表示的是concat，$w_f$是能够将concat的向量转换成一个标量的权重向量。这里设置C(x)=N。

文中发现self-attention模块其实就是non-local的embedded Gaussian版本的一种特殊情况。

3.Non-local Block

将(1)式中的non-local操作变形成一个non-local block，以便其可以被插入到已有的结构中。我们定义一个non-local block为:

中$y_i$已经在(1)式中给出了，$+x_i$则表示的是一个residual connection。residual connection 的结构使得我们可以在任意的pretrain模型中插入一个新的non-local block而不需要改变其原有的结构（如$W_z=0$ 作为初始化则完全和原始结构一致）。

non-local Block 模块图：

non-local block的pairwise的计算可以是非常light-weight的。

Non-local Blocks的高效策略：我们设置$W_g,W_{\theta },W_{\varphi }$的channel的数目为x的channel数目的一半，这样就形成了一个bottleneck，能够减少一半的计算量。$W_z$再重新放大到x的channel数目，保证输入输出维度一致。
还有一个subsampling的trick可以进一步使用，就是将(1)式变为：
其中x^是x下采样得到的（比如通过pooling），我们将这个方式在空间域上使用，可以减小1/4的pairwise function的计算量。这个trick并不会改变non-local的行为，而是使计算更加稀疏了。这个可以通过在图2中的ϕ和g后面增加一个max pooling层实现。本文中的所有non-local模块中都使用了上述的高效策略。

细节图

实验

• a 使用non-local对baseline结果是有提升的，但是不同相似度计算方法之间差距并不大
• b non-local加入网络的不同stage下性能都有提升,但是对较小的feature map提升不大
• c 添加越多的non-local 模块，效果提升越明显，但是会增大计算量
• d 同时在时域和空域上加入non-local 操作效果会最好
• 

pytorch 代码

import torch
import torch.nn as nn

class NonLocalBlock(nn.Module):
    def __init__(self,channel):
        super(NonLocalBlock, self).__init__()
        self.inter_channel = channel // 2
        self.conv_phi = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1,padding=0, bias=False)
        self.conv_theta = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1, padding=0, bias=False)
        self.conv_g = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1, padding=0, bias=False)
        self.softmax = nn.Softmax(dim=1)
        self.conv_mask = nn.Conv2d(in_channels=self.inter_channel, out_channels=channel, kernel_size=1, stride=1, padding=0, bias=False)
    
    
    def forward(self,x):
         # [N, C, H , W]
        b, c, h, w = x.size()
        # [N, C/2, H * W]
        x_phi = self.conv_phi(x).view(b, self.inter_channel, -1)
        # [N, H * W, C/2]
        x_theta = self.conv_theta(x).view(b, self.inter_channel, -1).permute(0, 2, 1).contiguous()
        x_g = self.conv_g(x).view(b, self.inter_channel, -1).permute(0, 2, 1).contiguous()
         # [N, H * W, H * W]
        f = torch.matmul(x_theta, x_phi) 
        # dim = -1 按行进行softmax 
        f_div_C = self.softmax(f,dim = -1)
        
        # [N, H * W, C/2]
        y = torch.matmul(f_div_C,x_g)
        y = y.permute(0,2,1).contiguous()
        # [N, C/2, H, W]
        y = y.view(b,self.inter_channel,h,w)
        # [N, C, H , W]
        mask = self.conv_mask(y)
        out = mask + x

        return out