少走 90% 弯路！《动手学深度学习》编码器-解码器模型代码，精准拆解易踩坑点

from torch import nn

classEncoder(nn.Module):
    """编码器-解码器架构的基本编码器接口"""
    def__init__(self, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        defforward(self, X, *args):
            raise NotImplementedError

classDecoder(nn.Module):
    """编码器-解码器架构的基本解码器接口"""
    def__init__(self, **kwargs):
        super(Decoder, self).__init__(**kwargs)

    definit_state(self, enc_outputs, *args):
        raise NotImplementedError

    defforward(self, X, state):
        raise NotImplementedError

classEncoderDecoder(nn.Module):
    """编码器-解码器架构的基类"""
    def__init__(self, encoder, decoder, **kwargs):
        super(EncoderDecoder, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    defforward(self, enc_X, dec_X, *args):
        enc_outputs = self.encoder(enc_X, *args)
        dec_state = self.decoder.init_state(enc_outputs, *args)
        return self.decoder(dec_X, dec_state)

classSeq2SeqEncoder(Encoder):
    """用于序列到序列学习的循环神经网络编码器"""

    def__init__(self, vocab_size, embed_size, num_hiddens, num_layers,
        dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        # 嵌入层
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers,
            dropout=dropout)

    defforward(self, X, *args):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        X = self.embedding(X)
        # 在循环神经网络模型中，第一个轴对应于时间步
        X = X.permute(1, 0, 2)
        # 如果未提及状态，则默认为0
        output, state = self.rnn(X)
        # output的形状:(num_steps,batch_size,num_hiddens)
        # state的形状:(num_layers,batch_size,num_hiddens)
        return output, state

defrnn(inputs, state, params):
    # inputs的形状：(时间步数量，批量大小，词表大小)
    W_xh, W_hh, b_h, W_hq, b_q = params
    H, = state
    outputs = []
    # X的形状：(批量大小，词表大小)
    for X in inputs:
        H = torch.tanh(torch.mm(X, W_xh) + torch.mm(H, W_hh) + b_h)
        Y = torch.mm(H, W_hq) + b_q
        outputs.append(Y)
return torch.cat(outputs, dim=0), (H,)

    # 测试编码器
    encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
        num_layers=2)
    encoder.eval()
    X = torch.zeros((4, 7), dtype=torch.long)
    output, state = encoder(X)
    print(output.shape, state.shape)

torch.Size([7, 4, 16]) torch.Size([2, 4, 16])

classSeq2SeqDecoder(Decoder):
    """用于序列到序列学习的循环神经网络解码器"""
    def__init__(self, vocab_size, embed_size, num_hiddens, num_layers,
        dropout=0, **kwargs):
        super(Seq2SeqDecoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers,
            dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    definit_state(self, enc_outputs, *args):
        return enc_outputs[1]

    defforward(self, X, state):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        X = self.embedding(X).permute(1, 0, 2)
        # 广播context，使其具有与X相同的num_steps
        context = state[-1].repeat(X.shape[0], 1, 1)
        X_and_context = torch.cat((X, context), 2)
        output, state = self.rnn(X_and_context, state)
        output = self.dense(output).permute(1, 0, 2)
        # output的形状:(batch_size,num_steps,vocab_size)
        # state的形状:(num_layers,batch_size,num_hiddens)
        return output, state

context = state[-1].repeat(X.shape[0], 1, 1)

    # 测试解码器
    decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8, num_hiddens=16,
        num_layers=2)
    decoder.eval()
    state = decoder.init_state(encoder(X))
    output, state = decoder(X, state)
    print(output.shape, state.shape)

torch.Size([4, 7, 10]) torch.Size([2, 4, 16])

defsequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
        device=X.device)[None, :] < valid_len[:, None]
    print(f"mask, {mask}")
    X[~mask] = value
    return X

classMaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""
    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    defforward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction='none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss

    mask = torch.arange((maxlen), dtype=torch.float32,
        device=X.device)[None, :] < valid_len[:, None]

# 假设 maxlen = 5, valid_len = [3, 5, 2]
# torch.arange(5)[None, :] = [[0, 1, 2, 3, 4]]
# valid_len[:, None] = [[3], [5], [2]]

# 广播比较后的结果：
# mask = [[True, True, True, False, False],    # 只有前3个位置有效
#         [True, True, True, True, True],      # 所有5个位置都有效
#         [True, True, False, False, False]]   # 只有前2个位置有效

deftrain_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):
    """训练序列到序列模型"""
    defxavier_init_weights(m):
        iftype(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        iftype(m) == nn.GRU:
            for param in m._flat_weights_names:
                if"weight"in param:
                    nn.init.xavier_uniform_(m._parameters[param])

    net.apply(xavier_init_weights)
    net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    loss = MaskedSoftmaxCELoss()
    net.train()
    animator = Animator(xlabel='epoch', ylabel='loss',
        xlim=[10, num_epochs])

    for epoch inrange(num_epochs):
        print(f'Epoch {epoch + 1}/{num_epochs}')
        timer = d2l.Timer()
        metric = d2l.Accumulator(2) # 训练损失总和，词元数量
        for batch in data_iter:
            optimizer.zero_grad()
            X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]
            bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],
                device=device).reshape(-1, 1)
            dec_input = torch.cat([bos, Y[:, :-1]], 1) # 强制教学
            Y_hat, _ = net(X, dec_input, X_valid_len)
            l = loss(Y_hat, Y, Y_valid_len)
            l.sum().backward() # 损失函数的标量进行“反向传播”
            rnn.grad_clipping(net, 1)
            num_tokens = Y_valid_len.sum()
            optimizer.step()
            with torch.no_grad():
                metric.add(l.sum(), num_tokens)
        if (epoch + 1) % 10 == 0:
            animator.add(epoch + 1, (metric[0] / metric[1],))
    print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} '
        f'tokens/sec on {str(device)}')
    return metric[0] / metric[1], metric[1] / timer.stop()

    embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
    batch_size, num_steps = 64, 10
    lr, num_epochs, device = 0.005, 300, d2l.try_gpu()
    train_iter, src_vocab, tgt_vocab = dataset.load_data_nmt(batch_size, num_steps)
    encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hiddens, num_layers,
        dropout)
    decoder = Seq2SeqDecoder(len(tgt_vocab), embed_size, num_hiddens, num_layers,
        dropout)
    net = enc_dec.EncoderDecoder(encoder, decoder)
    train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

loss 0.019, 11063.6 tokens/sec on cpu

defbleu(pred_seq, label_seq, k):#@save
    """计算BLEU"""
    pred_tokens, label_tokens = pred_seq.split(' '), label_seq.split(' ')
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    score = math.exp(min(0, 1 - len_label / len_pred))
    for n inrange(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        for i inrange(len_label - n + 1):
            label_subs[' '.join(label_tokens[i: i + n])] += 1
        for i inrange(len_pred - n + 1):
            if label_subs[' '.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[' '.join(pred_tokens[i: i + n])] -= 1
        score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    return score

defpredict_seq2seq(net, src_sentence, src_vocab, tgt_vocab, num_steps,
    device, save_attention_weights=False):
    """序列到序列模型的预测"""

    # 在预测时将net设置为评估模式
    net.eval()

    src_tokens = src_vocab[src_sentence.lower().split(' ')] + [
        src_vocab['<eos>']]
    enc_valid_len = torch.tensor([len(src_tokens)], device=device)
    src_tokens = dataset.truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])

    # 添加批量轴
    enc_X = torch.unsqueeze(
        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)
    enc_outputs = net.encoder(enc_X, enc_valid_len)
    dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)

    # 添加批量轴
    dec_X = torch.unsqueeze(torch.tensor(
    [tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)
    output_seq, attention_weight_seq = [], []

    for _ inrange(num_steps):
        Y, dec_state = net.decoder(dec_X, dec_state)
        # 我们使用具有预测最高可能性的词元，作为解码器在下一时间步的输入
        dec_X = Y.argmax(dim=2)
        pred = dec_X.squeeze(dim=0).type(torch.int32).item()
        # 保存注意力权重（稍后讨论）
        if save_attention_weights:
            attention_weight_seq.append(net.decoder.attention_weights)
        # 一旦序列结束词元被预测，输出序列的生成就完成了
        if pred == tgt_vocab['<eos>']:
            break
        output_seq.append(pred)
    return' '.join(tgt_vocab.to_tokens(output_seq)), attention_weight_seq

    engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
    fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
    for eng, fra inzip(engs, fras):
        translation, attention_weight_seq = predict_seq2seq(
            net, eng, src_vocab, tgt_vocab, num_steps, d2l.try_gpu())
        print(f'{eng} => {translation}, bleu {bleu(translation, fra, k=2):.3f}')

go . => va !, bleu 1.000
i lost . => j'ai perdu <unk> ., bleu 0.658
he's calm . => il est paresseux emporté !, bleu 0.447
i'm home . => je suis chez du c'est c'est c'est c'est c'est c'est, bleu 0.376