多头注意力机制是NLP算法岗常考的代码题，本篇文章将逐行梳理多头注意力机制的代码。
全部代码
import math
import torch
import torch.nn as nnclass MultiHeadAttention(nn.Module):def __init__(self, d_model, nums_head):super(MultiHeadAttention, self).__init__()self.d_model = d_modelself.nums_head = nums_headassert self.d_model % self.nums_head == 0self.split = d_model // nums_headself.q_linear = nn.linear(self.d_model, self.d_model)self.k_linear = nn.linear(self.d_model, self.d_model)self.v_linear = nn.linear(self.d_model, self.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.split).transpose(1, 2)def forward(self, q, k, v, mask=None):query = self.q_linear(q)key = self.k_linear(k)value = self.v_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.split)if mask is not None:score = score + score.masked_fill(mask == 0, -float('inf'))weight = torch.softmax(score, dim=-1)att_output = torch.matmul(weight, value)batch_size, _, seq_length, hidden_size = att_output.size()att_output = att_output.transpose(1, 2).view(batch_size, seq_length, self.d_model)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'))
