自然语言处理|Attention is all you need pytorch实现源码解析03 - 模型的训练（2）- transformer模型的代码实现以及结构论文解析以及实现|机器学习

我们继续分析著名的attention is all you need 论文的pytorch实现的源码解析。
由于项目很大，所以我们会分开几讲来进行讲解。
上一讲连接在此：
Attention is all you need pytorch实现源码解析01 - 数据预处理、词表的构建 - https://blog.csdn.net/weixin_42744102/article/details/87006081
Attention is all you need pytorch实现源码解析02 - 模型的训练（1）- 模型的训练代码 - https://blog.csdn.net/weixin_42744102/article/details/87076089
先上github源码：https://github.com/Eathoublu/attention-is-all-you-need-pytorch
【自然语言处理|Attention is all you need pytorch实现源码解析03 - 模型的训练（2）- transformer模型的代码实现以及结构】项目结构：

-transfomer
—__init__.py
—Beam.py
—Constants.py
—Layers.py
—Models.py
—Module.py
—Optim.py
—SubLayers.py
—Translator.py
datasets.py
preprocess.py
train.py
translate.py

话不多说，我们现在就看一下代码以及解析，代码顺序是逐级深入的，transformer这个里面又包括了encoder模型以及decoder模型，encoder模型又是由n_layers个encoderlayer平行组成的，每个encoderlayer层又有一个MultiHeadAttention子层和一个PositionwiseFeedForward子层。下面请看源码，源码贯穿了tenplates文件夹下的Models.py，Layers.py，sublayers.py，请阅读从1到14的注释。

class Transformer(nn.Module):# 1 - transformer模型继承自pytorch的nn.model（这是一个接口类，留下一个forward方法必须由子类实现），init传入的参数有词表的大小、句子序列的最大长度，剩余为可选参数，其中attention的头数为8 ''' A sequence to sequence model with attention mechanism. '''def __init__( self, n_src_vocab, n_tgt_vocab, len_max_seq, 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, tgt_emb_prj_weight_sharing=True, emb_src_tgt_weight_sharing=True):super().__init__()self.encoder = Encoder(# 2 - 初始化了encoder模型，下面会讲到，是用于组成encoder-decoder的组件 n_src_vocab=n_src_vocab, len_max_seq=len_max_seq, 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, dropout=dropout)self.decoder = Decoder(# 3 - 初始化了decoder模型，下面会讲到，是用于组成encoder-decoder的组件 n_tgt_vocab=n_tgt_vocab, len_max_seq=len_max_seq, 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, dropout=dropout)# 4 - 下面是一些参数的设置，这里就不多讲self.tgt_word_prj = nn.Linear(d_model, n_tgt_vocab, bias=False) nn.init.xavier_normal_(self.tgt_word_prj.weight)assert d_model == d_word_vec, \ 'To facilitate the residual connections, \ the dimensions of all module outputs shall be the same.'if tgt_emb_prj_weight_sharing: # Share the weight matrix between target word embedding & the final logit dense layer self.tgt_word_prj.weight = self.decoder.tgt_word_emb.weight self.x_logit_scale = (d_model ** -0.5) else: self.x_logit_scale = 1.if emb_src_tgt_weight_sharing: # Share the weight matrix between source & target word embeddings assert n_src_vocab == n_tgt_vocab, \ "To share word embedding table, the vocabulary size of src/tgt shall be the same." self.encoder.src_word_emb.weight = self.decoder.tgt_word_emb.weightdef forward(self, src_seq, src_pos, tgt_seq, tgt_pos):# 5 - forward方法的实现，写着forward的就是在训练的时候调用的。tgt_seq, tgt_pos = tgt_seq[:, :-1], tgt_pos[:, :-1]# 6 - 这里将传入的target的序列和位置的最后一位丢掉（这里有点不太明白，可能是因为最后一位是符号）enc_output, *_ = self.encoder(src_seq, src_pos)# 7 - 将训练集的data传入encoder模型，得到encoder的output dec_output, *_ = self.decoder(tgt_seq, tgt_pos, src_seq, enc_output)# 8 - 将训练集的target、data以及encoder得到的output传入decoder，得到decoder的output seq_logit = self.tgt_word_prj(dec_output) * self.x_logit_scale# 9 - 将decoder的输出结果进行一个线性变化，并在return的时候返回，好，后面我们具体讲一讲encoder以及decoder层。return seq_logit.view(-1, seq_logit.size(2))class Encoder(nn.Module): # 10 - 这是encoder模型，我们往下看： ''' A encoder model with self attention mechanism. '''def __init__(# 11 - 初始化，继承了部分父类（transformer类）的属性 self, n_src_vocab, len_max_seq, d_word_vec, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):super().__init__()n_position = len_max_seq + 1# 12 - 这个+1应该是为了防止溢出self.src_word_emb = nn.Embedding(# 13 - 先说一下nn.embedding是什么：Pytorch官网的解释是：一个保存了固定字典和大小的简单查找表。这个模块常用来保存词嵌入和用下标检索它们。模块的输入是一个下标的列表，输出是对应的词嵌入。 n_src_vocab, d_word_vec, padding_idx=Constants.PAD)self.position_enc = nn.Embedding.from_pretrained( get_sinusoid_encoding_table(n_position, d_word_vec, padding_idx=0), freeze=True)self.layer_stack = nn.ModuleList([ EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)]) # 14 - 这里复制n个一摸一样的EncoderLayer层，接下来，我们看Encoderlayer层。def forward(self, src_seq, src_pos, return_attns=False):enc_slf_attn_list = []# -- Prepare masks slf_attn_mask = get_attn_key_pad_mask(seq_k=src_seq, seq_q=src_seq) non_pad_mask = get_non_pad_mask(src_seq)# -- Forward enc_output = self.src_word_emb(src_seq) + self.position_enc(src_pos)for enc_layer in self.layer_stack: enc_output, enc_slf_attn = enc_layer( enc_output, non_pad_mask=non_pad_mask, slf_attn_mask=slf_attn_mask) if return_attns: enc_slf_attn_list += [enc_slf_attn]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_tgt_vocab, len_max_seq, d_word_vec, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):super().__init__() n_position = len_max_seq + 1self.tgt_word_emb = nn.Embedding( n_tgt_vocab, d_word_vec, padding_idx=Constants.PAD)self.position_enc = nn.Embedding.from_pretrained( get_sinusoid_encoding_table(n_position, d_word_vec, padding_idx=0), freeze=True)self.layer_stack = nn.ModuleList([ DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)])def forward(self, tgt_seq, tgt_pos, src_seq, enc_output, return_attns=False):dec_slf_attn_list, dec_enc_attn_list = [], []# -- Prepare masks non_pad_mask = get_non_pad_mask(tgt_seq)slf_attn_mask_subseq = get_subsequent_mask(tgt_seq) slf_attn_mask_keypad = get_attn_key_pad_mask(seq_k=tgt_seq, seq_q=tgt_seq) slf_attn_mask = (slf_attn_mask_keypad + slf_attn_mask_subseq).gt(0)dec_enc_attn_mask = get_attn_key_pad_mask(seq_k=src_seq, seq_q=tgt_seq)# -- Forward dec_output = self.tgt_word_emb(tgt_seq) + self.position_enc(tgt_pos)for dec_layer in self.layer_stack: dec_output, dec_slf_attn, dec_enc_attn = dec_layer( dec_output, enc_output, non_pad_mask=non_pad_mask, slf_attn_mask=slf_attn_mask, dec_enc_attn_mask=dec_enc_attn_mask)if return_attns: dec_slf_attn_list += [dec_slf_attn] dec_enc_attn_list += [dec_enc_attn]if return_attns: return dec_output, dec_slf_attn_list, dec_enc_attn_list return dec_output,class EncoderLayer(nn.Module):# 10 - 这是单层的encoderLayer ''' Compose with two layers '''def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1): super(EncoderLayer, self).__init__() self.slf_attn = MultiHeadAttention(# 11 - 实例化多头注意力模型 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, non_pad_mask=None, slf_attn_mask=None): # 12 - 使用多头注意力模型进行训练 enc_output, enc_slf_attn = self.slf_attn( enc_input, enc_input, enc_input, mask=slf_attn_mask) enc_output *= non_pad_maskenc_output = self.pos_ffn(enc_output)# 13 - 将注意力的output传到前馈模型中得到outut enc_output *= non_pad_maskreturn enc_output, enc_slf_attnclass MultiHeadAttention(nn.Module):# 14 - 这就是encoderlayer里面实例化的多头注意力子层 ''' 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_vself.w_qs = nn.Linear(d_model, n_head * d_k) self.w_ks = nn.Linear(d_model, n_head * d_k) self.w_vs = nn.Linear(d_model, n_head * d_v) nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5)) self.layer_norm = nn.LayerNorm(d_model)self.fc = nn.Linear(n_head * d_v, d_model) nn.init.xavier_normal_(self.fc.weight)self.dropout = nn.Dropout(dropout)def forward(self, q, k, v, mask=None):d_k, d_v, n_head = self.d_k, self.d_v, self.n_headsz_b, len_q, _ = q.size() sz_b, len_k, _ = k.size() sz_b, len_v, _ = v.size()residual = qq = 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)q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dvmask = mask.repeat(n_head, 1, 1) # (n*b) x .. x .. output, attn = self.attention(q, k, v, mask=mask)output = output.view(n_head, sz_b, len_q, d_v) output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv)output = self.dropout(self.fc(output)) output = self.layer_norm(output + residual)return output, attn