代码详解——Transformer

整体架构

源码地址（pytorch）： https://github.com/jadore801120/attention-is-all-you-need-pytorch
论文地址：Attention is All You Need
✨✨✨强烈建议先去看《详解注意力机制和Transformer》 理解Transformer的机制后再去理解本篇中的代码。

项目的整体架构如下，其中Transformer 包下的文件是用于主要构建Transfomer模型的代码，包外的其他文件是Transfomer用于完成特定翻译任务的预处理文件和训练测试代码。

本文重点讲解红框内的代码(构建Transformer的核心代码), 即实现了下图所示的Transformer的架构。
因为Transformer经常被用到其他的任务中，所以这部分的核心代码也常被移植到其他的项目代码中。

Modules.py

Models.py文件主要就是定义了一个缩放点积注意力 （下图红框中的部分）

缩放点积的计算公式如下：
$$\mathrm{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d}}\right)\mathbf{V}\in {\mathbb{R}}^{n\times v}.$$

ScaledDotProductAttention

# 缩放点积注意力
class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):
        # q:  [sz_b,n_head,len,d_q]
        # k:  [sz_b,n_head,len,d_k] ->  transpose 后：[sz_b,n_head,d_k,len]
        # v:  [sz_b,n_head,len,d_v]
        # 一般来说，d_q=d_k=d_v
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3)) # score= qk^T/tempreture
        # attn: [sz_b,n_head,len,len]
        if mask is not None: # 判断是否有mask
            attn = attn.masked_fill(mask == 0, -1e9) # Mask
        attn = self.dropout(F.softmax(attn, dim=-1)) # a=softmax(Score) 然后 dropout
        output = torch.matmul(attn, v) # z=a*v
        # output: [sz_b,n_head,len,d_v]
        return output, attn

相关参数的含义：

• q,k, v 分别表示的query,key,value, 对应下图中的Q K V；它们的大小均是[sz_b,n_head,len,d_x] (d_x代表d_q、d_v、d_k)
  • sz_b 表示batch size
  • n_head 表示多头注意力的head 数量
  • len 表示单词的个数，如下图就是2。
  • d_x 表示特征的个数，如下图是64。
• temperature 就是的$\sqrt{d_{m}odel}$, $d_{m}odel$表示的是特征个数的，用作是归一化。如下图中就是$\sqrt{64}=8$
• mask 表示是否传入mask,在Transformer中有两种mask，分别是padding mask和sequence mask
  

相关代码解读：

• attn = torch.matmul(q / self.temperature, k.transpose(2, 3)) : 就是计算注意力得分，并归一化$$score=\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d}}$$
  k.transpose(2, 3) 表示在k的后两个维度(len_q,d_q)进行转置。
  根据矩阵乘法的原理，得到的attn的大小为[sz_b,n_head,len_q,len_k]， 如下图就是[sz_b,n_head,2,2]
  

• attn = attn.masked_fill(mask == 0, -1e9)然后判断是否传入的mask, 如果有mask (mask参数值不为None)，则把mask为0的位置，将对应位置的attn的值设为无穷小的负数$-e^9$
  为什么要设置为无穷小呢？如下图展示了softmax函数，当x为无穷小时，softmax的输出趋近于0，attn的值就为0，就相当于是被mask掉了。
  

• attn = self.dropout(F.softmax(attn, dim=-1)) 就是对刚才得到的注意力得分attn在d_q维度上进行softmax操作,把attn转换成一个值分布在[0,1]之间的 α概率分布矩阵
  
  然后softmax后使用dropout操作防止过拟合。

• output = torch.matmul(attn, v) 最终得到的输出就把上述的attn和value相乘。最终的输出大小为[sz_b,n_head,len_q,d_v], 如下图就是[sz_b,n_head,2,64] 。可以发现得到的输出和输入的K,Q,V的大小相同。
  

SubLayers.py

MultiHeadAttention

MultiHeadAttention定义了一个多头注意力和 Add&Norm。（下图中的红框部分）
可以实现如下三种多头注意力:
1）Multi-Head Self-Attention: K、Q、V的来源相同
2）Masked Multi-Head Self-Attention ： 传入sequence mask 的mask参数，且K、Q、V的来源相同
3）Multi-Head Cross-Attention ： K、V和Q的来源不同

# 多头注意力
class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()

        self.n_head = n_head # head数量
        self.d_k = d_k # key 的维度
        self.d_v = d_v # v 的维度

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) # [sz_b,len_q,d_model]->[sz_b,len_q,n*d_k]
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5) # 缩放点积注意力

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)


    def forward(self, q, k, v, mask=None):
        # 原始输入 q/k/v:[sz_d,len,d_model]
        # sz_b: batch_size 
        # len: 单词的个数 (一般来说：len=len_q=len_k=len_v)
        # d_model：单词嵌入的维度 (一般来说：d_model=d_k=d_v)
        # n_head : head的个数
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b