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