自注意力机制(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矩阵相乘。由此我们便得到了注意力得分。

自注意力机制(self-attention)的理解与pytorch实现_第1张图片

        自注意力机制(self-attention)的理解与pytorch实现_第2张图片

 以下为注意力机制的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

你可能感兴趣的