注意力机制论文:Non-Local neural networks及其Pytorch实现

Non-Local neural networks
PDF: https://arxiv.org/pdf/1711.07971.pdf
PyTorch代码: https://github.com/shanglianlm0525/PyTorch-Networks

Non-Local Neural Network和Non-Local Means非局部均值去噪滤波有点相似。普通的滤波都是3×3的卷积核，然后在整个图片上进行移动，处理的是3×3局部的信息。Non-Local Means操作则是结合了一个比较大的搜索范围，并进行加权。

1 概述

• non-local operations通过计算任意两个位置之间的交互直接捕捉远程依赖，而不用局限于相邻点，其相当于构造了一个和特征图谱尺寸一样大的卷积核, 从而可以维持更多信息。
• non-local可以作为一个组件，和其它网络结构结合，用于其他视觉任务中。
• Non-local在视频分类上效果可观

2 Non-local operation

Non-local 操作可以表示为

其中
g函数是一个线性转换

f函数用于计算i和j相似度的函数, 文中列举中四种具体实现
Gaussian:

Embedded Gaussian:

Dot product:

Concatenation:

汇总起来就是

3 Non-local block

3-1 抽象图

3-2 细节图

4 Ablations

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

import torch
import torch.nn as nn
import torchvision


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, c, -1)
        # [N, H * W, C/2]
        x_theta = self.conv_theta(x).view(b, c, -1).permute(0, 2, 1).contiguous()
        x_g = self.conv_g(x).view(b, c, -1).permute(0, 2, 1).contiguous()
        # [N, H * W, H * W]
        mul_theta_phi = torch.matmul(x_theta, x_phi)
        mul_theta_phi = self.softmax(mul_theta_phi)
        # [N, H * W, C/2]
        mul_theta_phi_g = torch.matmul(mul_theta_phi, x_g)
        # [N, C/2, H, W]
        mul_theta_phi_g = mul_theta_phi_g.permute(0,2,1).contiguous().view(b,self.inter_channel, h, w)
        # [N, C, H , W]
        mask = self.conv_mask(mul_theta_phi_g)
        out = mask + x
        return out


if __name__=='__main__':
    model = NonLocalBlock(channel=16)
    print(model)

    input = torch.randn(1, 16, 64, 64)
    out = model(input)
    print(out.shape)