Transformer的代码实现 day02（Multi-Head Self-Attention）

Multi-Head Self-Attention的理论部分

• Self-Attention是指Q、K、V同源，同时对输入x乘以三个矩阵$W^Q$、$W^K$、$W^V$得到Q、K、V，之后Q与K相乘得到概率，概率和V相乘得到最后的注意力值Z，流程图如下：
  
• Multi-Head Self-Attention是将输入x拆分为8块，分别得到它们的Q、K、V，并进行自注意力计算，得到各自的注意力值Z后，拼接起来，经一个线性变换得到和输入x相同维度的Z‘作为最后的输出，流程图如下：
  

Multi-Head Self-Attention代码实现

• Self-Attention的代码实现：

def self_attention(query, key, value, dropout=None, mask=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    # mask的操作在QK之后，softmax之前
    if mask is not None:
        mask.cuda()
        scores = scores.masked_fill(mask == 0, -1e9)
    # self_atten为概率值，还没有乘V
    self_attn = F.softmax(scores, dim=-1)
    if dropout is not None:
        self_attn = dropout(self_attn)
    return torch.matmul(self_attn, value), self_attn

• Multi-Head Attention的代码实现：

class MultiHeadAttention(nn.Module):

    def __init__(self):
        super(MultiHeadAttention, self).__init__()


    def forward(self,  head, d_model, query, key, value, dropout=0.1,mask=None):
        """
		# 注意：head、d_model、query、key、value都可以放在init()中
        :param head: 头数，默认 8
        :param d_model: 输入的维度 512
        :param query: Q
        :param key: K
        :param value: V
        :param dropout:
        :param mask:
        :return:
        """
        assert (d_model % head == 0)
        self.d_k = d_model // head
        self.head = head
        self.d_model = d_model

        self.linear_query = nn.Linear(d_model, d_model)
        self.linear_key = nn.Linear(d_model, d_model)
        self.linear_value = nn.Linear(d_model, d_model)

        # 自注意力机制的 QKV 同源，线性变换
		# 为了保证多头自注意力的最后输出和输入维度相同，所以还要加一层线性变换
        self.linear_out = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(p=dropout)
        self.attn = None



        # if mask is not None:
        #     # 多头注意力机制的线性变换层是4维，是把query[batch, frame_num, d_model]变成[batch, -1, head, d_k]
        #     # 再1，2维交换变成[batch, head, -1, d_k], 所以mask要在第一维添加一维，与后面的self attention计算维度一样
        #     mask = mask.unsqueeze(1)

        n_batch = query.size(0)

        # 多头需要对这个 X 切分成多头

        # query==key==value
        # [b,1,512]
        # [b,8,1,64]

        # [b,32,512]
        # [b,8,32,64]
		# view方法是用来改变张量的形状的。它返回一个新的张量，其数据与原始张量共享，但具有不同的形状。
		# -1 为占位符，表示自动计算维数大小，以保持总元素数不变
        query = self.linear_query(query).view(n_batch, -1, self.head, self.d_k).transpose(1, 2)  # [b, 8, 32, 64]
        key = self.linear_key(key).view(n_batch, -1, self.head, self.d_k).transpose(1, 2)  # [b, 8, 32, 64]
        value = self.linear_value(value).view(n_batch, -1, self.head, self.d_k).transpose(1, 2)  # [b, 8, 32, 64]

        x, self.attn = self_attention(query, key, value, dropout=self.dropout, mask=mask)
        # [b,8,32,64]
        # [b,32,512]
        # 变为三维， 或者说是concat head
        # contiguous()函数用来确保transpose交换维度后的张量在内存中是连续存储的。因为view()函数需要张量在内存中是连续存储
        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.head * self.d_k)

        return self.linear_out(x)

参考文献

1. 03 Transformer 中的多头注意力（Multi-Head Attention）Pytorch代码实现