【代码解读】Transformer: attention-is-all-you-need-pytorch

学习transformer时对GitHub上项目：attention-is-all-you-need-pytorch进行了部分中文注释，主要集中在以下几个文件。

注释后完整代码：attention-is-all-you-need-pytorch，结合这篇文章一起理解。

Models.py

''' Define the Transformer model '''
import torch
import torch.nn as nn
import numpy as np
from transformer.Layers import EncoderLayer, DecoderLayer

__author__ = "Yu-Hsiang Huang"


def get_pad_mask(seq, pad_idx):
    return (seq != pad_idx).unsqueeze(-2)


def get_subsequent_mask(seq):
    ''' For masking out the subsequent info. '''
    sz_b, len_s = seq.size()
    subsequent_mask = (1 - torch.triu(
        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()
    return subsequent_mask


class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()
        # Not a parameter
        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ''' Sinusoid position encoding table '''

        # TODO: make it with torch instead of numpy

        def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        # 设计这个位置函数的用意是位置可以由上一个位置线性表示，参考：https://zhuanlan.zhihu.com/p/308301901#1.4.1
        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)

    def forward(self, x):
        # x.shape = [batch_size, seq_len, embedding_d])
        # pos_table计算最大长度为200的位置编码，这里去前seq_len长度就行，在每个seq上加上[seq_len, embedding_d]的位置编码，相当于每个位置都加[embedding_d]
        # 注意: 这里这是数值上相加，不改变shape
        return x + self.pos_table[:, :x.size(1)].clone().detach()


class Encoder(nn.Module):
    ''' A encoder model with self attention mechanism. '''

    def __init__(
            self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, dropout=0.1, n_position=200, scale_emb=False):

        super().__init__()

        # 训练一个Embedding层，把数据集中（输入数据）n_src_vocab维，压缩为d_word_vec(d_model)维
        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)
        # 位置编码器，一个[d_word_vec, n_position]表
        # n_position表示一个足够长的序列，首先序列会补padding，和batch中最大长度sql_len一样长，n_position-sql_len部分就使用mask置0
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, src_seq, src_mask, return_attns=False):

        enc_slf_attn_list = []
        # -- Forward
        # src_seq是从word到int，最大到n_src_vocab，所以对0~n_src_vocab，进行编码，每个词编码为d_word_vec(d_model)维
        # 比如：开始符号的整数编码为2，那么embedding后的编码为[0, 0, 1, 0, ... ,0]，看做n_src_vocab+1分类
        enc_output = self.src_word_emb(src_seq)
        # 注意此时enc_output的维度为d_word_vec，后面输入enc_layer，而enc_layer要求维度为d_model，所以必须d_word_vec == d_model
        if self.scale_emb:
            enc_output *= self.d_model ** 0.5

        # 加上位置编码
        enc_output = self.dropout(self.position_enc(enc_output))
        enc_output = self.layer_norm(enc_output)
        for enc_layer in self.layer_stack:  # 有n_layers=6层，每层有n_head=8个头
            # 进入Encoder层，输出enc_output作为下层的输入
            enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask)
            enc_slf_attn_list += [enc_slf_attn] if return_attns else []

        if return_attns:
            # 调试或者可视化使用
            return enc_output, enc_slf_attn_list
        return enc_output,


class Decoder(nn.Module):
    ''' A decoder model with self attention mechanism. '''

    def __init__(
            self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):

        super().__init__()
        # 训练一个Embedding层，把数据集中（标注数据）n_trg_vocab维，压缩为d_word_vec(d_model)维
        self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=pad_idx)
        # 位置编码器，一个[d_word_vec, n_position]表
        # n_position表示一个足够长的序列，首先序列会补padding，和batch中最大长度sql_len一样长，n_position-sql_len部分就使用mask置0
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):

        dec_slf_attn_list, dec_enc_attn_list = [], []

        # -- Forward
        # 训练时，标注数据和encoder输出 作为decoder的输入，而在测试时，上一时刻的decoder的最终输出和encoder输出 作为下一时刻的decoder输入
        dec_output = self.trg_word_emb(trg_seq)
        if self.scale_emb:
            dec_output *= self.d_model ** 0.5
        dec_output = self.dropout(self.position_enc(dec_output))
        dec_output = self.layer_norm(dec_output)

        for dec_layer in self.layer_stack:
            # 进入Decoder层，enc_output用于cross—attention，输出dec_output作为下层的输入
            dec_output, dec_slf_attn, dec_enc_attn = dec_layer(dec_output, enc_output, slf_attn_mask=trg_mask, dec_enc_attn_mask=src_mask)
            dec_slf_attn_list += [dec_slf_attn] if return_attns else []
            dec_enc_attn_list += [dec_enc_attn] if return_attns else []

        if return_attns:
            # 调试或者可视化使用
            return dec_output, dec_slf_attn_list, dec_enc_attn_list
        return dec_output,


class Transformer(nn.Module):
    ''' A sequence to sequence model with attention mechanism. '''

    def __init__(
            self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx,
            d_word_vec=512, d_model=512, d_inner=2048,
            n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
            trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True,
            scale_emb_or_prj='prj'):

        super().__init__()

        self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx

        # In section 3.4 of paper "Attention Is All You Need", there is such detail:
        # "In our model, we share the same weight matrix between the two
        # embedding layers and the pre-softmax linear transformation...
        # In the embedding layers, we multiply those weights by \sqrt{d_model}".
        #
        # Options here:
        #   'emb': multiply \sqrt{d_model} to embedding output
        #   'prj': multiply (\sqrt{d_model} ^ -1) to linear projection output
        #   'none': no multiplication

        assert scale_emb_or_prj in ['emb', 'prj', 'none']
        scale_emb = (scale_emb_or_prj == 'emb') if trg_emb_prj_weight_sharing else False
        self.scale_prj = (scale_emb_or_prj == 'prj') if trg_emb_prj_weight_sharing else False
        self.d_model = d_model

        self.encoder = Encoder(
            n_src_vocab=n_src_vocab,  # 数据集中词的个数，比如10000个，那么就是[0~9999]中不同词，转成one-hot分类编码一个词就有10000维，显然可以使用embedding进行压缩
            n_position=n_position,  # 由于每个batch序列长度不一样，所以给个最大长度，我认为可以求数据集中序列最大长度作为n_position
            d_word_vec=d_word_vec,  # 经过embedding压缩后的维度，且作为后续EncoderLayer层的输入，所以要==d_model
            d_model=d_model,  # d_word_vec == d_model
            d_inner=d_inner,  # 前反馈层中间fc层维度，比如: d_model x d_inner x d_model，经过这3层，对一个EncoderLayer做最后处理
            n_layers=n_layers,  # 有多少EncoderLayer
            n_head=n_head,  # 多头注意力，每头注意力侧重点不一样，可以看到后面对多头进行拆分单头注意力分别计算，
            d_k=d_k,  # 注意query和key的维度一定要一样，用于匹配attn，所以在attention时d_k=d_q
            d_v=d_v,  # value维度，默认和k和v一样
            pad_idx=src_pad_idx,  # Embedding Layer 用，猜测padding需要特殊处理，所以需要先知道其编码(索引)
            dropout=dropout,
            scale_emb=scale_emb)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab,  # 目标词个数
            n_position=n_position,  # 和encoder一样
            d_word_vec=d_word_vec,
            d_model=d_model,
            d_inner=d_inner,
            n_layers=n_layers,
            n_head=n_head,
            d_k=d_k,
            d_v=d_v,
            pad_idx=trg_pad_idx,
            dropout=dropout,
            scale_emb=scale_emb)

        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

        assert d_model == d_word_vec, \
            'To facilitate the residual connections, \
             the dimensions of all module outputs shall be the same.'

        if trg_emb_prj_weight_sharing:
            # Share the weight between target word embedding & last dense layer
            self.trg_word_prj.weight = self.decoder.trg_word_emb.weight

        if emb_src_trg_weight_sharing:
            self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight

    def forward(self, src_seq, trg_seq):

        src_mask = get_pad_mask(src_seq, self.src_pad_idx)  # mask作用为将padding的attn置0
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)  # 当解码第一个字的时候，第一个字只能与第一个字计算相关性，当解出第二个字的时候，只能计算出第二个字与第一个字和第二个字的相关性
        # 注意shape：src_mask[batch_size, 1, src_seq_len] trg_mask[batch_size, trg_seq_len, trg_seq_len]
        enc_output, *_ = self.encoder(src_seq, src_mask)  # enc_output用于decoder中的cross-attention
        dec_output, *_ = self.decoder(trg_seq, trg_mask, enc_output, src_mask)

        # dec_output的输出表示了seq上每个词的embedding向量（维度d_model），通过一个fc，转为词分类概率（维度n_trg_vocab）
        seq_logit = self.trg_word_prj(dec_output)
        # softmax放在外边的loss函数中了 :)

        if self.scale_prj:
            seq_logit *= self.d_model ** -0.5

        return seq_logit.view(-1, seq_logit.size(2))  # 合并batch

Modules.py

import torch
import torch.nn as nn
import torch.nn.functional as F

__author__ = "Yu-Hsiang Huang"

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):

        # query和key进行match，这里match是直接相乘，其意义是query在每维key中的匹配情况
        # 除以sqrt{d}是为了防止维度过高内积值太大
        # 单头：q[seq_len, single_d] * k[single_d， seq_len]，其代表有seq_len个query[single_d]
        # 每个query和seq_len个key[single_d]计算match（也是做weight sum），所以1个query和seq_len个key计算得到attn[k_seq_len]
        # 而一共有seq_len个query，所以有[q_seq_len, k_seq_len]个attn

        #        q_s_d               k_l              k_l
        #    |—— —— —— ——|        |——|——|——|      |——|——|——|  // 单个query得到的attn[k_seq_len]
        # q_l|—— —— —— ——| x k_s_d|  |  |  | = q_l|——|——|——|
        #    |—— —— —— ——|        |  |  |  |      |——|——|——|
        #                         |——|——|——|
        # 解释为什么要把多头拆分：观察上面s_d是单头的维度，此时得到的基于单头的match(q_s_d, k_s_d)做weight sum得到一个数，所以每个头得到的atte是不一样的，意义就是学习不同类型的注意力
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            # padding 部分设为极小，经过softmax后为0，表示忽略
            # 注意区分维度： encoder[1, src_seq_len]  decoder[trg_seq_len, trg_seq_len]
            # encoder时每个query（理解为从左往右对应的query）其对应的是完整的一个序列，所以[1, src_seq_len]的mask对应到每个query上，无差别
            # decoder时每个query（理解为从左往右对应的query）其对应的不是完整的一个序列，所以需要trg_seq_len个[1, trg_seq_len]的不同mask对应到每个query上
            # 补充一下decoder时每个query，使用mask的下三角形就能实现，如第一个query只有前1个词，第二个query有前2个词...，当然训练时可以这样批量放在矩阵操作，测试时则需要按照时间来逐个输出
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))

        # 应用attention
        # attn[q_seq_len, k_seq_len] * v[v_seq_len, single_d] 代表每个query计算得到的attn[k_seq_len]和v_seq_len个v[single_d]计算输出[single_d](每个维度都应用weight sum过程)
        # 所以seq_len个query得到的输出为[q_seq_len, single_d]

        #        k_l            v_s_d               v_s_d
        #    |——|——|——|      |—— —— —— ——|      |—— —— —— ——| // 再在每个维度上进行，得到v[single_d]
        # q_l|——|——|——| x v_l|—— —— —— ——| = q_l|—— —— —— ——|
        #    |——|——|——|      |—— —— —— ——|      |—— —— —— ——|
        #                                         \ 应用attn，然后weight sum过程

        # 注意q_s_d和k_s_d的维度一定要一样，用于匹配attn，应用attn时，是在v的每个维度上都应用attn，所以v_s_d不需要等于k_s_d或q_s_d
        output = torch.matmul(attn, v)

        # 输出attn，仅用于调试或可视化
        return output, attn

Layers.py

''' Define the Layers '''
import torch.nn as nn
import torch
from transformer.SubLayers import MultiHeadAttention, PositionwiseFeedForward


__author__ = "Yu-Hsiang Huang"


class EncoderLayer(nn.Module):
    ''' Compose with two layers '''

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

        # 由于q和k的维度默一定要一样，这里区分k和v的维度
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)

        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, enc_input, slf_attn_mask=None):
        # Encoder只有self-attention，即q=v=k，在Encoder中只有slf_attn，输入enc_input_i来自前一个输出enc_output_i-1
        enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)

        # 上面说到当前输出enc_output_i会作为下一层的输入enc_input_i+1，但是不能直接传过去，加一个前馈神经网络层pos_ffn
        enc_output = self.pos_ffn(enc_output)
        return enc_output, enc_slf_attn


class DecoderLayer(nn.Module):
    ''' Compose with three layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(DecoderLayer, self).__init__()
        # 由于q和k的维度默一定要一样，这里区分k和v的维度
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)

        # cross-attention 用于提取src上对应当前时刻输出的注意力
        self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)

        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, dec_input, enc_output, slf_attn_mask=None, dec_enc_attn_mask=None):
        # self-attention，即q=v=k，输入dec_input_i来自前一个输出dec_output_i-1
        dec_output, dec_slf_attn = self.slf_attn(dec_input, dec_input, dec_input, mask=slf_attn_mask)
        # 输出dec_output_i，作为cross-attention的query

        # cross-attention，即q!=v=k，其中query来自self-attention的输出（由于输出是一个一个输出的，所以在decoder上的self-attention不能够从src中提取注意力），
        # 而注意力是基于全局的，此时应用和query匹配的是具有全局完整信息的encoder上的key，然后应用到encoder上的value，用于提取src上对应当前时刻输出的注意力
        dec_output, dec_enc_attn = self.enc_attn(dec_output, enc_output, enc_output, mask=dec_enc_attn_mask)

        # 和encoder同理不能直接传过去，加一个前馈神经网络层pos_ffn
        dec_output = self.pos_ffn(dec_output)
        return dec_output, dec_slf_attn, dec_enc_attn

SubLayers.py

''' Define the sublayers in encoder/decoder layer '''
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from transformer.Modules import ScaledDotProductAttention

__author__ = "Yu-Hsiang Huang"

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
        self.d_k = d_k
        self.d_v = d_v

        # 一个可学习的矩阵[embedding_d, n_head * single_d]
        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        # q和k计算attn时是做weight sum match，所以d_k=d_q
        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):

        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        # 残差是没有进行attention之前的值
        residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        # 在普通attention机制中，q来自外部，v=k来自自身，attention = math(q, k)
        # 而在self-attention中，q=v=k都来自自身，那么，attention = math(q, k) 的结果最相似的，所以我们需要产生q != k，所以使用一个nn.Linear，一个可学习的矩阵[embedding_d, n_head * single_d]
        # 比如embedding后一个词表示为512维，q[seq_len， 512] * w[512, 8 * 64]，所以得到的query为8个(多头)64维向量，同理k和v
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        # n_head 移到前面
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            # 多头拆分了，从batch, seq_len, d_model, 到batch, seq_len, n_head, d_single_model(single_d)
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        # 为什么多头要拆开？看attention里解释
        q, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)  # 后面没有在涉及attn计算，合并多头便于计算
        q = self.dropout(self.fc(q))
        q += residual

        # channel方向做归一化
        q = self.layer_norm(q)

        # 输出attn，仅用于调试或可视化
        return q, attn


class PositionwiseFeedForward(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid) # position-wise
        self.w_2 = nn.Linear(d_hid, d_in) # position-wise
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        residual = x

        x = self.w_2(F.relu(self.w_1(x)))
        x = self.dropout(x)
        x += residual

        x = self.layer_norm(x)

        return x