[人工智能]自注意力机制（self-attention）的理解与pytorch实现

????????attention机制，即注意力机制，为了解决以往的RNN，LSTM等模型对于长距离的上下文分析能力不足的问题。然而，自注意力机制，顾名思义，输出与输入自身有关。对于自注意力机制，最有名的就是在谷歌的transformer模型中所使用。在它的论文attention is all you need中，self-attention被首次重用，以替代NLP领域中的循环神经网络或是CV领域中的卷积神经网络，竟然得到了很不错的效果。下面对于自注意力机制，从初学的角度进行阐述。

? ? ? ? 引用论文中的缩放的点乘注意力机制（论文中还提到了Additive attention，作者采用点乘可以大量减少时间与空间复杂度。）Q,K,V表示三个矩阵。他们是通过一个全连接网络或者核大小为1的卷积层，均由输入得到。Q与K的转置的乘积便是我们需要的output，但为了解决论文中提到的输入维度太大后，点乘导致的数字剧烈增大，而通过softmax函数之后可得到的梯度过小，因此我们还需要乘以dk(key的维度大小)的平方根的倒数，再通过softmax层后与V矩阵相乘。由此我们便得到了注意力得分。

????????

?以下为注意力机制的pytorch代码实现：

class selfattention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels
        self.query = nn.Conv2d(in_channels, in_channels // 8, kernel_size = 1, stride = 1)
        self.key   = nn.Conv2d(in_channels, in_channels // 8, kernel_size = 1, stride = 1)
        self.value = nn.Conv2d(in_channels, in_channels, kernel_size = 1, stride = 1)
        self.gamma = nn.Parameter(torch.zeros(1))  #gamma为一个衰减参数，由torch.zero生成，nn.Parameter的作用是将其转化成为可以训练的参数.
        self.softmax = nn.Softmax(dim = -1)
    def forward(self, input):
        batch_size, channels, height, width = input.shape
        # input: B, C, H, W -> q: B, H * W, C // 8
        q = self.query(input).view(batch_size, -1, height * width).permute(0, 2, 1)
        #input: B, C, H, W -> k: B, C // 8, H * W
        k = self.key(input).view(batch_size, -1, height * width)
        #input: B, C, H, W -> v: B, C, H * W
        v = self.value(input).view(batch_size, -1, height * width)
        #q: B, H * W, C // 8 x k: B, C // 8, H * W -> attn_matrix: B, H * W, H * W
        attn_matrix = torch.bmm(q, k)  #torch.bmm进行tensor矩阵乘法,q与k相乘得到的值为attn_matrix.
        attn_matrix = self.softmax(attn_matrix)#经过一个softmax进行缩放权重大小.
        out = torch.bmm(v, attn_matrix.permute(0, 2, 1))  #tensor.permute将矩阵的指定维进行换位.这里将1于2进行换位。
        out = out.view(*input.shape)

        return self.gamma * out + input