class TransformerEncoderLayer(nn.Module):
    """编码器中的一个层 """
    def __init__(self, d_model, heads, d_ff, dropout, max_relative_positions=0):
        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=dropout,
            max_relative_positions=max_relative_positions)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, inputs, mask):
        # 1. LayerNorm 2. self-attention 3. dropout 4. Feed-Forward
        input_norm = self.layer_norm(inputs)
        context, _ = self.self_attn(input_norm, input_norm, input_norm, mask=mask, type="self")
        out = self.dropout(context) + inputs
        return self.feed_forward(out)

class TransformerEncoder(EncoderBase):
    def __init__(self, num_layers, d_model, heads, d_ff, dropout, embeddings, max_relative_positions):
        ...
        self.embeddings = embeddings
        self.transformer = nn.ModuleList([TransformerEncoderLayer(...) for i in range(num_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, src, lengths=None):
        ...
        # 单词的编码
        emb = self.embeddings(src)

        out = emb.transpose(0, 1).contiguous()
        words = src[:, :, 0].transpose(0, 1)
        w_batch, w_len = words.size()
        padding_idx = self.embeddings.word_padding_idx
        mask = words.data.eq(padding_idx).unsqueeze(1)  # [B, 1, T]
        # 遍历多层网络
        for layer in self.transformer:
            out = layer(out, mask)
        out = self.layer_norm(out)
        return emb, out.transpose(0, 1).contiguous(), lengths

class TransformerDecoderLayer(nn.Module):
    '''解码器中一个单层'''
    def __init__(self, d_model, heads, d_ff, dropout,
                 self_attn_type="scaled-dot", max_relative_positions=0):
        
        # 定义self-attention类型
        if self_attn_type == "scaled-dot":
            self.self_attn = MultiHeadedAttention(heads, d_model, dropout, max_relative_positions)
        elif self_attn_type == "average":
            self.self_attn = AverageAttention(d_model, dropout=dropout)
        # 定义对encoder所有输出单词的context attention
        self.context_attn = MultiHeadedAttention(heads, d_model, dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-6)
        self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-6)
        self.drop = nn.Dropout(dropout)

    def forward(self, inputs, memory_bank, src_pad_mask, tgt_pad_mask, layer_cache=None, step=None):
        # 对还未生成的单词做masking
        dec_mask = None
        if step is None:
            tgt_len = tgt_pad_mask.size(-1)
            future_mask = torch.ones(...)
            future_mask = future_mask.triu_(1).view(1, tgt_len, tgt_len)
            dec_mask = torch.gt(tgt_pad_mask + future_mask, 0)
        
        input_norm = self.layer_norm_1(inputs)
        if isinstance(self.self_attn, MultiHeadedAttention):
            query, attn = self.self_attn(input_norm, input_norm, input_norm, ..., type="self")
        elif isinstance(self.self_attn, AverageAttention):
            query, attn = self.self_attn(input_norm, mask=dec_mask, layer_cache=layer_cache, step=step)
        # 对encoder输出做context attention
        query = self.drop(query) + inputs
        query_norm = self.layer_norm_2(query)
        mid, attn = self.context_attn(memory_bank, memory_bank, query_norm, ..., type="context")
        output = self.feed_forward(self.drop(mid) + query)
        return output, attn
        
class TransformerDecoder(DecoderBase):
    def __init__(self, num_layers, d_model, heads, d_ff, attn_type,
                 copy_attn, self_attn_type, dropout, embeddings,
                 max_relative_positions):
        self.embeddings = embeddings

        # Decoder State
        self.state = {}
        self.transformer_layers = nn.ModuleList([TransformerDecoderLayer(...) for i in range(num_layers)])

        self._copy = copy_attn
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, tgt, memory_bank, step=None, **kwargs):
        ...
        src = self.state["src"]
        src_words = src[:, :, 0].transpose(0, 1)
        tgt_words = tgt[:, :, 0].transpose(0, 1)
        src_batch, src_len = src_words.size()
        tgt_batch, tgt_len = tgt_words.size()

        emb = self.embeddings(tgt, step=step) # len x batch x embedding_dim

        output = emb.transpose(0, 1).contiguous()
        src_memory_bank = memory_bank.transpose(0, 1).contiguous()
        
        # Masking
        pad_idx = self.embeddings.word_padding_idx
        src_pad_mask = src_words.data.eq(pad_idx).unsqueeze(1)  # [B, 1, T_src]
        tgt_pad_mask = tgt_words.data.eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]

        for i, layer in enumerate(self.transformer_layers):
            ...
            # 多层网络
            output, attn = layer(output, src_memory_bank, src_pad_mask, tgt_pad_mask, layer_cache, step)
        output = self.layer_norm(output)
        dec_outs = output.transpose(0, 1).contiguous()
        attn = attn.transpose(0, 1).contiguous()
        attns = {"std": attn}
        return dec_outs, attns

class NMTModel(nn.Module):
    def __init__(self, encoder, decoder):
        super(NMTModel, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, src, tgt, lengths, bptt=False):
        tgt = tgt[:-1]  # 从输入中去掉最后一个单词
        # 编码
        enc_state, memory_bank, lengths = self.encoder(src, lengths)
        if bptt is False:
            self.decoder.init_state(src, memory_bank, enc_state)
        # 解码
        dec_out, attns = self.decoder(tgt, memory_bank, memory_lengths=lengths)
        return dec_out, attns

class Trainer(object):
    # 定义训练的类
    def __init__(self, model, train_loss, valid_loss, optim,
                 trunc_size=0, shard_size=32,
                 norm_method="sents", grad_accum_count=1, n_gpu=1, gpu_rank=1,
                 gpu_verbose_level=0, report_manager=None, model_saver=None,
                 average_decay=0, average_every=1):
        ...
        # 把模型设置成训练状态
        self.model.train()

    def train(self, train_iter, train_steps, save_checkpoint_steps=5000, valid_iter=None, valid_steps=10000):
        ...
        # 开始遍历每个mini-batch
        for i, (batches, normalization) in enumerate(
                self._accum_batches(train_iter)):
            step = self.optim.training_step
            ...
            # 累积多个mini-batch算出来的gradient
            self._gradient_accumulation(
                batches, normalization, total_stats,
                report_stats)
            # 把多个gradient求平均值并更新模型参数
            if self.average_decay > 0 and i % self.average_every == 0:
                self._update_average(step)
            ...
            # 做validation，可用于判断是否需要终止训练
            if valid_iter is not None and step % valid_steps == 0:
                valid_stats = self.validate(valid_iter, moving_average=self.moving_average)
                ...
                
            # 保存model
            if (self.model_saver is not None
                and (save_checkpoint_steps != 0
                     and step % save_checkpoint_steps == 0)):
                self.model_saver.save(step, moving_average=self.moving_average)

            if train_steps > 0 and step >= train_steps:
                break
        return total_stats