非局部神经网络.

作者提出了一个Non-Local block，实现了对卷积网络提取特征的自注意力机制。

对该网络的解释如下，$g$函数用$1 \times 1$卷积实现$W_g$（本质是线性映射）；$f$函数用Embedded Gaussian函数，即用卷积拟合映射$\phi$和$\theta$，再通过Softmax函数计算权重。注意到此处的$g$、$\phi$和$\theta$相当于NLP中自注意力机制的value、key和query。

该方法在视频分类任务上效果很好，比$3D$卷积效率更高、性能更好，但该方法计算量偏大，建议不要在底层网络使用，可以适当在高层网络中使用。

自注意力机制在计算输入位置$i$的特征$y_i$时，考虑所有位置$j$的加权：

\[y_i = \sum_{j}^{} \frac{e^{f(x_i)^Tg(x_j)}}{\sum_j e^{f(x_i)^Tg(x_j)}} h(x_j)\]

自注意力机制的实现步骤如下：

  1. $f(x)$、$g(x)$和$h(x)$通过三个$1\times 1$卷积层实现，$f(x)$和$g(x)$改变了通道数(缩小为$C/8$)，$h(x)$维持通道数不变；
  2. 将空间尺寸合并为$H\times W$，将$f(x)$的输出转置后和$g(x)$的输出进行矩阵相乘，经过softmax归一化得到尺寸为$[H\times W,H\times W]$的注意力图；
  3. 将注意力图与$h(x)$的输出进行矩阵相乘，得到尺寸为$[H\times W,C]$的特征图，经过$1\times 1$卷积层并把输出尺寸调整为$[H,W,C]$；
  4. 最终输出的特征可以通过标量缩放$\gamma$和残差连接构造：$y = γy + x$。

自注意力机制的实现可参考：

class SelfAttention(nn.Module):
    def __init__(self, in_channels, k=8):
        super(SelfAttention, self).__init__()
        self.inter_channels = int(in_channels//k)
        self.f = nn.Conv2d(in_channels, self.inter_channels, 1)
        self.g = nn.Conv2d(in_channels, self.inter_channels, 1)
        self.h = nn.Conv2d(in_channels, in_channels, 1)
        self.o = nn.Conv2d(in_channels, in_channels, 1)
        self.gamma = torch.zeros(1).requires_grad_(True)

    def forward(self, x):
        b, c, h, w = x.shape
        fx = self.f(x).view(b, self.inter_channels, -1) # [b, c', hw]
        fx = fx.permute(0, 2, 1) # [b, hw, c']
        gx = self.g(x).view(b, self.inter_channels, -1) # [b, c', hw]
        attn = torch.matmul(fx, gx) # [b, hw, hw]
        attn = F.softmax(attn, dim=2) # 按行归一化

        hx = self.h(x).view(b, c, -1) # [b, c, hw]
        hx = hx.permute(0, 2, 1) # [b, hw, c]
        y = torch.matmul(attn, hx) # [b, hw, c]
        y = y.permute(0, 2, 1).contiguous() # [b, c, hw]
        y = y.view(b, c, h, w)
        y = self.o(y)

        return self.gamma*y + x