逐行解析多头注意力机制

多头注意力机制是NLP算法岗常考的代码题，本篇文章将逐行梳理多头注意力机制的代码。

全部代码

import math
import torch
import torch.nn as nn


class MutilHeadAttention(nn.Module):
	def __init__(self,nums_head,d_model):
		self.nums_head = nums_head
		self.d_model = d_model
		assert d_model % nums_head == 0
		self.head_size = d_model // nums_head
		self.linear_q = nn.Linear(d_model, d_model)
		self.linear_k = nn.Linear(d_model, d_model)
		self.linear_v = nn.Linear(d_model, d_model)
		self.linear_o = nn.Linear(d_model, d_model)
	def split_head(self, x):
		batch_size, seq_length, hidden_size = x.size()
		return x.view(batch_size, seq_length, self.nums_head, self.head_size).transpose(1,2)
	def forward(self,q,k,v,mask=None):
		query = self.linear(q)
		key = self.linear(k)
		value = self.linear(v)
		query = self.split_head(query)
		key = self.split_head(key)
		value =  self.split_head(value)
		score = torch.matmul(query, key.transpose(-2,-1)) / math.sqrt(self.head_size))
		if mask is not None:
			score = score + score.masked_fill(mask==0, -float("inf))
		weight = torch.softmax(socre, dim=-1)
		att_output = torch.matmul(value, weight)
		batch_size, _ , seq_length, hidden_size = att_output.size()
		att_output = att_output.transpose(1,2).view(batch_size,seq_length,self.d_model)
		att_output = self.linear_o(att_output)
		return att_output
		

细节解析

设置断言，保证原始维度可以整除注意力头数。

assert self.d_model % self.nums_head == 0

将x的维度进行转化，并且将seq_length和nums_head的维度交换(允许对每个序列位置和每个注意力头独立地处理数据。这样，模型就可以并行地关注输入序列的不同部分，从而提高处理复杂输入序列的能力)。

def split_head(self, x):
    batch_size, seq_length, hidden_size = x.size()
    return x.view(batch_size, seq_length, self.nums_head, self.split).transpose(1, 2)

mask == 0：这部分代码会生成一个与mask形状相同的布尔张量（Boolean Tensor），其中mask中值为0的位置在布尔张量中对应为True，其他位置为False。
scores.masked_fill(mask == 0, -1e9)：masked_fill函数接受一个布尔张量和一个标量值作为输入。它会遍历布尔张量，并将原始张量（这里是scores）中对应为True的位置填充为指定的标量值（这里是-1e9）。因此，如果mask中的某个位置是0，那么scores中对应位置的值将被替换为-float(‘inf’)。
scores += …：最后，这个替换后的张量（其中mask为0的位置被替换为-float(‘inf’)）被加回到原始的scores张量上。然而，由于masked_fill实际上返回了一个新张量（即原始张量的一个副本，但其中某些位置被修改了），这里的加法操作实际上是将原始scores张量中的对应位置更新为-1e9加上它们原来的值。但是，由于-float(‘inf’)是一个非常大的负数，并且softmax函数在计算时会将输入值转换为概率分布，所以任何接近-float(‘inf’)的值在softmax中都会接近于0。因此，这个操作实际上是在告诉模型在计算注意力权重时忽略mask为0的位置。

if mask is not None:
   score = score + score.masked_fill(mask == 0, -float('inf'))