From cab8cc7de48a18368f3c8e9cdd515ff5da894ff9 Mon Sep 17 00:00:00 2001 From: Petr Masopust Date: Sat, 9 Nov 2019 19:12:24 +0100 Subject: [PATCH 1/3] pytorch transformer implementation --- train.py | 4 -- transformer/Layers.py | 49 ------------------ transformer/Models.py | 104 +++++++++++++-------------------------- transformer/Modules.py | 28 ----------- transformer/SubLayers.py | 81 ------------------------------ transformer/__init__.py | 6 +-- translate.py | 2 - 7 files changed, 35 insertions(+), 239 deletions(-) delete mode 100644 transformer/Layers.py delete mode 100644 transformer/Modules.py delete mode 100644 transformer/SubLayers.py diff --git a/train.py b/train.py index c96a862..306cf75 100644 --- a/train.py +++ b/train.py @@ -248,8 +248,6 @@ def main(): parser.add_argument('--d_model', type=int, default=512) parser.add_argument('--d_inner_hid', type=int, default=2048) - parser.add_argument('--d_k', type=int, default=64) - parser.add_argument('--d_v', type=int, default=64) parser.add_argument('--n_head', type=int, default=8) parser.add_argument('--n_layers', type=int, default=6) @@ -291,8 +289,6 @@ def main(): opt.max_token_seq_len, tgt_emb_prj_weight_sharing=opt.proj_share_weight, emb_src_tgt_weight_sharing=False, - d_k=opt.d_k, - d_v=opt.d_v, d_model=opt.d_model, d_word_vec=opt.d_word_vec, d_inner=opt.d_inner_hid, diff --git a/transformer/Layers.py b/transformer/Layers.py deleted file mode 100644 index f1b45be..0000000 --- a/transformer/Layers.py +++ /dev/null @@ -1,49 +0,0 @@ -''' Define the Layers ''' -import torch.nn as nn -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__() - 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, non_pad_mask=None, slf_attn_mask=None): - enc_output, enc_slf_attn = self.slf_attn( - enc_input, enc_input, enc_input, mask=slf_attn_mask) - enc_output *= non_pad_mask - - enc_output = self.pos_ffn(enc_output) - enc_output *= non_pad_mask - - 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__() - self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout) - 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, non_pad_mask=None, slf_attn_mask=None, dec_enc_attn_mask=None): - dec_output, dec_slf_attn = self.slf_attn( - dec_input, dec_input, dec_input, mask=slf_attn_mask) - dec_output *= non_pad_mask - - dec_output, dec_enc_attn = self.enc_attn( - dec_output, enc_output, enc_output, mask=dec_enc_attn_mask) - dec_output *= non_pad_mask - - dec_output = self.pos_ffn(dec_output) - dec_output *= non_pad_mask - - return dec_output, dec_slf_attn, dec_enc_attn diff --git a/transformer/Models.py b/transformer/Models.py index d323e78..580ac96 100644 --- a/transformer/Models.py +++ b/transformer/Models.py @@ -3,13 +3,17 @@ import torch.nn as nn import numpy as np import transformer.Constants as Constants -from transformer.Layers import EncoderLayer, DecoderLayer +from torch.nn.modules.transformer import TransformerEncoderLayer, TransformerDecoderLayer, \ + TransformerEncoder, TransformerDecoder +from torch.nn.modules.normalization import LayerNorm __author__ = "Yu-Hsiang Huang" + def get_non_pad_mask(seq): assert seq.dim() == 2 - return seq.ne(Constants.PAD).type(torch.float).unsqueeze(-1) + return seq.eq(Constants.PAD) + def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None): ''' Sinusoid position encoding table ''' @@ -31,25 +35,6 @@ def get_posi_angle_vec(position): return torch.FloatTensor(sinusoid_table) -def get_attn_key_pad_mask(seq_k, seq_q): - ''' For masking out the padding part of key sequence. ''' - - # Expand to fit the shape of key query attention matrix. - len_q = seq_q.size(1) - padding_mask = seq_k.eq(Constants.PAD) - padding_mask = padding_mask.unsqueeze(1).expand(-1, len_q, -1) # b x lq x lk - - return padding_mask - -def get_subsequent_mask(seq): - ''' For masking out the subsequent info. ''' - - sz_b, len_s = seq.size() - subsequent_mask = torch.triu( - torch.ones((len_s, len_s), device=seq.device, dtype=torch.uint8), diagonal=1) - subsequent_mask = subsequent_mask.unsqueeze(0).expand(sz_b, -1, -1) # b x ls x ls - - return subsequent_mask class Encoder(nn.Module): ''' A encoder model with self attention mechanism. ''' @@ -57,9 +42,8 @@ class Encoder(nn.Module): def __init__( self, n_src_vocab, len_max_seq, d_word_vec, - n_layers, n_head, d_k, d_v, + n_layers, n_head, d_model, d_inner, dropout=0.1): - super().__init__() n_position = len_max_seq + 1 @@ -71,32 +55,26 @@ def __init__( 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)]) - - def forward(self, src_seq, src_pos, return_attns=False): + self.encoder = TransformerEncoder( + TransformerEncoderLayer(d_model, n_head, dim_feedforward=d_inner, dropout=dropout), + n_layers, norm=LayerNorm(d_model)) + def forward(self, src_seq, src_pos): 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] + enc_output = enc_output.permute(1, 0, 2) + enc_output = self.encoder(enc_output, + src_key_padding_mask=non_pad_mask) + + # print(enc_output) + return enc_output.permute(1, 0, 2), - if return_attns: - return enc_output, enc_slf_attn_list - return enc_output, class Decoder(nn.Module): ''' A decoder model with self attention mechanism. ''' @@ -104,9 +82,8 @@ class Decoder(nn.Module): def __init__( self, n_tgt_vocab, len_max_seq, d_word_vec, - n_layers, n_head, d_k, d_v, + n_layers, n_head, d_model, d_inner, dropout=0.1): - super().__init__() n_position = len_max_seq + 1 @@ -117,40 +94,27 @@ def __init__( 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): + self.decoder = TransformerDecoder( + TransformerDecoderLayer(d_model, n_head, dim_feedforward=d_inner, dropout=dropout), + n_layers, norm=LayerNorm(d_model)) + def forward(self, tgt_seq, tgt_pos, src_seq, enc_output): 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) + dec_enc_attn_mask = get_non_pad_mask(src_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) + dec_output = self.decoder(dec_output.permute(1, 0, 2), + enc_output.permute(1, 0, 2), + tgt_key_padding_mask=non_pad_mask, + memory_key_padding_mask=dec_enc_attn_mask) - if return_attns: - dec_slf_attn_list += [dec_slf_attn] - dec_enc_attn_list += [dec_enc_attn] + return dec_output.permute(1, 0, 2), - 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. ''' @@ -159,7 +123,7 @@ 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, + n_layers=6, n_head=8, dropout=0.1, tgt_emb_prj_weight_sharing=True, emb_src_tgt_weight_sharing=True): @@ -168,21 +132,21 @@ def __init__( self.encoder = Encoder( 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, + n_layers=n_layers, n_head=n_head, dropout=dropout) self.decoder = 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, + n_layers=n_layers, n_head=n_head, dropout=dropout) 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.' + '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 @@ -194,7 +158,7 @@ def __init__( 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." + "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.weight def forward(self, src_seq, src_pos, tgt_seq, tgt_pos): diff --git a/transformer/Modules.py b/transformer/Modules.py deleted file mode 100644 index c711f44..0000000 --- a/transformer/Modules.py +++ /dev/null @@ -1,28 +0,0 @@ -import torch -import torch.nn as nn -import numpy as np - -__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) - self.softmax = nn.Softmax(dim=2) - - def forward(self, q, k, v, mask=None): - - attn = torch.bmm(q, k.transpose(1, 2)) - attn = attn / self.temperature - - if mask is not None: - attn = attn.masked_fill(mask, -np.inf) - - attn = self.softmax(attn) - attn = self.dropout(attn) - output = torch.bmm(attn, v) - - return output, attn diff --git a/transformer/SubLayers.py b/transformer/SubLayers.py deleted file mode 100644 index 0298a19..0000000 --- a/transformer/SubLayers.py +++ /dev/null @@ -1,81 +0,0 @@ -''' 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 - - self.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_head - - sz_b, len_q, _ = q.size() - sz_b, len_k, _ = k.size() - sz_b, len_v, _ = v.size() - - residual = q - - 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) - - 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 dv - - mask = 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 - -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.Conv1d(d_in, d_hid, 1) # position-wise - self.w_2 = nn.Conv1d(d_hid, d_in, 1) # position-wise - self.layer_norm = nn.LayerNorm(d_in) - self.dropout = nn.Dropout(dropout) - - def forward(self, x): - residual = x - output = x.transpose(1, 2) - output = self.w_2(F.relu(self.w_1(output))) - output = output.transpose(1, 2) - output = self.dropout(output) - output = self.layer_norm(output + residual) - return output diff --git a/transformer/__init__.py b/transformer/__init__.py index 88675cd..e14f549 100644 --- a/transformer/__init__.py +++ b/transformer/__init__.py @@ -1,7 +1,4 @@ import transformer.Constants -import transformer.Modules -import transformer.Layers -import transformer.SubLayers import transformer.Models import transformer.Translator import transformer.Optim @@ -9,6 +6,5 @@ import transformer.Beam __all__ = [ - transformer.Constants, transformer.Modules, transformer.Layers, - transformer.SubLayers, transformer.Models, transformer.Optim, + transformer.Constants, transformer.Models, transformer.Optim, transformer.Translator, transformer.Beam, transformer.BatchWrapper] diff --git a/translate.py b/translate.py index d74e3f9..e614302 100644 --- a/translate.py +++ b/translate.py @@ -25,8 +25,6 @@ def load_model(opt): model_opt.max_token_seq_len, tgt_emb_prj_weight_sharing=model_opt.proj_share_weight, emb_src_tgt_weight_sharing=False, - d_k=model_opt.d_k, - d_v=model_opt.d_v, d_model=model_opt.d_model, d_word_vec=model_opt.d_word_vec, d_inner=model_opt.d_inner_hid, From 790e62eb128cf84505cfde2e6fd3e8302aa15c8d Mon Sep 17 00:00:00 2001 From: Petr Masopust Date: Sat, 9 Nov 2019 19:32:34 +0100 Subject: [PATCH 2/3] Remove unnecessary permutes --- transformer/Models.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer/Models.py b/transformer/Models.py index 580ac96..2fc13de 100644 --- a/transformer/Models.py +++ b/transformer/Models.py @@ -73,7 +73,7 @@ def forward(self, src_seq, src_pos): src_key_padding_mask=non_pad_mask) # print(enc_output) - return enc_output.permute(1, 0, 2), + return enc_output, class Decoder(nn.Module): @@ -108,8 +108,9 @@ def forward(self, tgt_seq, tgt_pos, src_seq, enc_output): # -- Forward dec_output = self.tgt_word_emb(tgt_seq) + self.position_enc(tgt_pos) - dec_output = self.decoder(dec_output.permute(1, 0, 2), - enc_output.permute(1, 0, 2), + dec_output = dec_output.permute(1, 0, 2) + dec_output = self.decoder(dec_output, + enc_output, tgt_key_padding_mask=non_pad_mask, memory_key_padding_mask=dec_enc_attn_mask) From 9767190d49494c38c152555fb1758732747ab908 Mon Sep 17 00:00:00 2001 From: Petr Masopust Date: Mon, 11 Nov 2019 11:19:19 +0100 Subject: [PATCH 3/3] New positional encoding implementation --- transformer/BatchWrapper.py | 3 ++ transformer/Models.py | 79 +++++++++++++++++++++++-------------- 2 files changed, 52 insertions(+), 30 deletions(-) diff --git a/transformer/BatchWrapper.py b/transformer/BatchWrapper.py index e5e7e6d..b31562a 100644 --- a/transformer/BatchWrapper.py +++ b/transformer/BatchWrapper.py @@ -9,6 +9,9 @@ def __init__(self, dl, fields, device=None): self.fields = fields self.device = device + if len(self.fields) > 1: + assert Constants.PAD == self.fields[1][1].vocab.stoi[Constants.PAD_WORD] + self.xpad = self.fields[0][1].vocab.stoi[Constants.PAD_WORD] self.ypad = Constants.PAD diff --git a/transformer/Models.py b/transformer/Models.py index 2fc13de..85d0366 100644 --- a/transformer/Models.py +++ b/transformer/Models.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn import numpy as np +import math import transformer.Constants as Constants from torch.nn.modules.transformer import TransformerEncoderLayer, TransformerDecoderLayer, \ TransformerEncoder, TransformerDecoder @@ -15,25 +16,48 @@ def get_non_pad_mask(seq): return seq.eq(Constants.PAD) -def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None): - ''' Sinusoid position encoding table ''' - - def cal_angle(position, hid_idx): - return position / np.power(10000, 2 * (hid_idx // 2) / d_hid) - - def get_posi_angle_vec(position): - return [cal_angle(position, hid_j) for hid_j in range(d_hid)] - - sinusoid_table = np.array([get_posi_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 - - if padding_idx is not None: - # zero vector for padding dimension - sinusoid_table[padding_idx] = 0. - - return torch.FloatTensor(sinusoid_table) +class PositionalEncoding(nn.Module): + r"""Inject some information about the relative or absolute position of the tokens + in the sequence. The positional encodings have the same dimension as + the embeddings, so that the two can be summed. Here, we use sine and cosine + functions of different frequencies. + .. math:: + \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model)) + \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model)) + \text{where pos is the word position and i is the embed idx) + Args: + d_model: the embed dim (required). + dropout: the dropout value (default=0.1). + max_len: the max. length of the incoming sequence (default=5000). + Examples: + >>> pos_encoder = PositionalEncoding(d_model) + """ + + def __init__(self, d_model, dropout=0.1, max_len=5000): + super(PositionalEncoding, self).__init__() + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0).transpose(0, 1) + self.register_buffer('pe', pe) + + def forward(self, x): + r"""Inputs of forward function + Args: + x: the sequence fed to the positional encoder model (required). + Shape: + x: [sequence length, batch size, embed dim] + output: [sequence length, batch size, embed dim] + Examples: + >>> output = pos_encoder(x) + """ + + x = x + self.pe[:x.size(0), :] + return self.dropout(x) class Encoder(nn.Module): @@ -46,14 +70,10 @@ def __init__( d_model, d_inner, dropout=0.1): super().__init__() - n_position = len_max_seq + 1 - self.src_word_emb = nn.Embedding( 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.position_enc = PositionalEncoding(d_model) self.encoder = TransformerEncoder( TransformerEncoderLayer(d_model, n_head, dim_feedforward=d_inner, dropout=dropout), @@ -66,7 +86,8 @@ def forward(self, src_seq, src_pos): non_pad_mask = get_non_pad_mask(src_seq) # -- Forward - enc_output = self.src_word_emb(src_seq) + self.position_enc(src_pos) + enc_output = self.src_word_emb(src_seq) + enc_output = self.position_enc(enc_output) enc_output = enc_output.permute(1, 0, 2) enc_output = self.encoder(enc_output, @@ -85,14 +106,11 @@ def __init__( n_layers, n_head, d_model, d_inner, dropout=0.1): super().__init__() - n_position = len_max_seq + 1 self.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.position_enc = PositionalEncoding(d_model) self.decoder = TransformerDecoder( TransformerDecoderLayer(d_model, n_head, dim_feedforward=d_inner, dropout=dropout), @@ -106,7 +124,8 @@ def forward(self, tgt_seq, tgt_pos, src_seq, enc_output): dec_enc_attn_mask = get_non_pad_mask(src_seq) # -- Forward - dec_output = self.tgt_word_emb(tgt_seq) + self.position_enc(tgt_pos) + dec_output = self.tgt_word_emb(tgt_seq) + dec_output = self.position_enc(dec_output) dec_output = dec_output.permute(1, 0, 2) dec_output = self.decoder(dec_output,