论文解读-《Non-local Neural Networks》

1. 论文介绍

论文题目：Non-local Neural Networks. CVPR 2018.
论文领域：图像识别

2.相关工作

在传统的计算机视觉的non-local means方法的基础下，提出了non-local operation 可以捕获远距离的特征依赖。不同于卷积操作的只能逐层传递局部邻居特征，非局部算子可以通过计算所有位置的加权特征和，直接建立全局连接。
和传统的卷积操作最大的区别是卷积是依赖于邻居周围的特征，而非局部和整个图尝试去建立连接来产生特征。

这个图展示了第一帧的足球的点与后面三帧的有关联的点。点与点之间的关系可以是时间的联系，也可以是空间，也可以是时空的。

3. 论文核心

非局部算子Non-local operation可以定义为：

其中输入x和输出y是相同的尺寸，位置i的输出信号是由位置i的输入信号和所有i和j可能存在的关系的j的信号的乘积，其中f表示i和j的信号连接关系，g是对j信号的一种特征提取。所有的j表示所有的位置都被考虑到了。

上面的公式与全连接层的区别：nl计算不同位置之间的关系的响应，fc使用的是学习权重。（位置i和位置j的关系不是输入数据的函数）

为了简化，把g看做是一个线性嵌入 g(x) = W*x 其中W是一个可学习的权重
关系函数f可以有多个版本

1，高斯函数 Gaussian ，其中使用了点积相似度，同样欧几里得距离也是可行的，但是点积更加实现友好型。

2，嵌套高斯Embedded Gaussian，是高斯函数的拓展版本，分别针对两个点的x进行嵌入，得到另外一个的嵌入空间中的特征

3，点积相似度

5，串联拼接，其中 [**] 内是串联操作，把位置i和j的特征的嵌入空间特征进行串联操作。

非局部块 Non-local block，在非局部算子的基础上，多增加一个residual连接的x，整个计算框架在高级别和子采样下是轻量化的。

4. 典型案例

一个典型的非局部块的例子为下图，非局部算子是嵌入高斯类型。

5. 实验设置

在视频分类问题和目标检测问题都可以发挥优越的性能。

在关于训练曲线上，可以观察到带有非局部块的网络的收敛速度更快。

实验证明，不同的非局部算子之间的效果性能差异较小，说明了非局部框架的有效性，而不是特定实现细节。

实验在resnet里面加入了不同数量的NLB，证明加入越多NLB，能稳定获得更好的效果。

在目标检测上也同样取得了好的效果

6. 代码解释

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: 是否有BatchNorm层
        """

        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

        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)

        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, return_nl_map=False):
        """
        :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_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        f = torch.matmul(theta_x, phi_x)
        f_div_C = F.softmax(f, dim=-1)

        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        if return_nl_map:
            return z, f_div_C
        return z

7. 总结

CV领域和GNN的经典论文，给出了全局性的特征描述。
实验设计很详细，数据分析很到位，做学术的都应该这样。