62序列到序列seq2seq

点击查看代码

import collections
import math
import torch
from torch import nn
from d2l import torch as d2l

# 实现循环神经网络编码器
#@save
class Seq2SeqEncoder(d2l.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)
        # 换成LSTM也可
        # encoder不需要输出层
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers,
                          dropout=dropout)

    def forward(self, X, *args):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        # print('X.device : ', X.device)
        X = self.embedding(X)
        # print('self.embedding(X).shape', X.shape)
        # 在循环神经网络模型中，第一个轴对应于时间步
        # (num_steps,batch_size,embed_size)
        X = X.permute(1, 0, 2)
        # output 最上面的输出
        # state 每一层的输出
        output, state = self.rnn(X)
        # output的形状:(num_steps,batch_size,num_hiddens)
        # state的形状:(num_layers,batch_size,num_hiddens)
        # state[0]的形状:(batch_size,num_hiddens)
        return output, state

# 编码器的实现
print('编码器的实现')
encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
                         num_layers=2)
# eval下，dropout不会生效
encoder.eval()
# (4, 7) --> (batch_size, num_steps)
X = torch.zeros((4, 7), dtype=torch.long)
output, state = encoder(X)
print('output.shape', output.shape)
#output.shape torch.Size([7, 4, 16])
print('state.shape', state.shape)
# state.shape torch.Size([2, 4, 16])
print('state[0].shape', state[0].shape)
# state[0].shape torch.Size([4, 16])
# 解码器
class Seq2SeqDecoder(d2l.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)
        # 假设encoder和decoder隐藏层大小一致
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers,
                          dropout=dropout)
        # 输出层
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, *args):
        # return output, state
        # enc_outputs[1] -> state
        return enc_outputs[1]

    def forward(self, X, state):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        # 在循环神经网络模型中，第一个轴对应于时间步
        X = self.embedding(X).permute(1, 0, 2)
        # 广播context，使其具有与X相同的num_steps
        # state[-1] rnn隐藏状态最后一刻最后一层输出
        print('X.shape : ', X.shape)
        print('state[-1].shape', state[-1].shape)
        context = state[-1].repeat(X.shape[0], 1, 1)
        X_and_context = torch.cat((X, context), 2)
        print('context.shape : ', context.shape)
        print('X_and_context.shape : ', X_and_context.shape)
        """
        X.shape :  torch.Size([7, 4, 8])
        state[-1].shape torch.Size([4, 16])
        context.shape :  torch.Size([7, 4, 16])
        X_and_context.shape :  torch.Size([7, 4, 24])
        """
        # 𝐬𝑡′=𝑔(𝑦𝑡′−1,𝐜,𝐬𝑡′−1)
        output, state = self.rnn(X_and_context, state)
        # 将batch_size调到前面
        output = self.dense(output).permute(1, 0, 2)
        # output的形状:(batch_size,num_steps,vocab_size)
        # state的形状:(num_layers,batch_size,num_hiddens)
        # state[0]的形状:(batch_size,num_hiddens)
        return output, state


# 实例化解码器
print('实例化解码器')
# X = torch.zeros((4, 7), dtype=torch.long)
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', output.shape)
"""
(batch_size, num_steps, vocab_size)
output.shape torch.Size([4, 7, 10])
"""
print('state.shape', state.shape)
"""
(num_layers, batch_size, num_hiddens)
state.shape torch.Size([2, 4, 16])
"""

# 损失函数
# 零值化屏蔽不相关的项
# 将句子中填充项赋0
print('损失函数 零值化屏蔽不相关的项')
#@save
# value 填充项
# mask NLP处理变长的一个常见操作
def sequence_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('mask', mask)
    """
    mask 
    tensor([[ True, False, False],
            [ True,  True, False]])
    """
    # print(~mask)
    """
    tensor([[False,  True,  True],
            [False, False,  True]])
    """
    X[~mask] = value
    return X

X = torch.tensor([[1, 2, 3], [4, 5, 6]])
print('sequence_mask(X, torch.tensor([1, 2]))',
      sequence_mask(X, torch.tensor([1, 2])))
"""
sequence_mask(X, torch.tensor([1, 2])) 
tensor([[1, 0, 0],
        [4, 5, 0]])
"""
# 使用此函数屏蔽最后几个轴上的所有项
print('使用此函数屏蔽最后几个轴上的所有项')
X = torch.ones(2, 3, 4)
print(sequence_mask(X, torch.tensor([1, 2]), value=-1))

# 扩展softmax交叉熵损失函数来遮蔽不相关的预测
#@save
class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""
    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        # 有效保留，其余变为0
        weights = sequence_mask(weights, valid_len)
        self.reduction='none'
        # 在[none, mean, sum]中选，string型。
        # none表示不降维，返回和target相同形状；
        # mean表示对一个batch的损失求均值；
        # sum表示对一个batch的损失求和。
        unweighted_loss = super().forward(
            # pytorch要将预测维度放中间
            pred.permute(0, 2, 1), label)
        # 每个句子取平均，每个样本取loss
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss

# 代码健全性检查
print('代码健全性检查')
loss = MaskedSoftmaxCELoss()
# (3, 4, 10)
# (batch_size,num_steps,vocab_size)
print(loss(torch.ones(3, 4, 10), torch.ones((3, 4), dtype=torch.long),
     torch.tensor([4, 2, 0])))
"""
tensor([2.3026, 1.1513, 0.0000])
"""

# 训练
#@save
def train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):
    """训练序列到序列模型"""
    # 初始化参数
    def xavier_init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        if type(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 = d2l.Animator(xlabel='epoch', ylabel='loss',
    #                  xlim=[10, num_epochs])
    for epoch in range(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开始标识
            bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],
                          device=device).reshape(-1, 1)
            dec_input = torch.cat([bos, Y[:, :-1]], 1)
            # X_valid_len 注意力使用到
            Y_hat, _ = net(X, dec_input, X_valid_len)
            l = loss(Y_hat, Y, Y_valid_len)
            # https://zhuanlan.zhihu.com/p/427853673
            # 一个向量是不进行backward操作的，而sum（）后，由于梯度为1，
            # 所以对结果不产生影响。反向传播算法一定要是一个标量才能进行计算。
            l.sum().backward()
            # 梯度裁剪
            d2l.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('metric[0] / metric[1]:', metric[0] / metric[1])
    print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} '
        f'tokens/sec on {str(device)}')

# 创建和训练一个循环神经网络“编码器－解码器”模型
print('创建和训练一个循环神经网络“编码器－解码器”模型')
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, torch.device('cpu')
                         # d2l.try_gpu()

train_iter, src_vocab, tgt_vocab = d2l.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 = d2l.EncoderDecoder(encoder, decoder)
# train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

# 预测
print('-------------------------')
print('predict')
#@save
def predict_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 = d2l.truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])
    # 添加批量轴
    # torch.unsqueeze()这个函数主要是对数据维度进行扩充。
    print('src_tokens.shape : ',
          torch.tensor(src_tokens, dtype=torch.long, device=device).shape)
    """src_tokens.shape :  torch.Size([10])"""
    enc_X = torch.unsqueeze(
        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)
    print('enc_X.shape : ', enc_X.shape)
    """enc_X.shape :  torch.Size([1, 10])"""
    enc_outputs = net.encoder(enc_X, enc_valid_len)
    print('enc_outputs : ', type(enc_outputs[0]), type(enc_outputs[1]))
    """enc_outputs :  <class 'torch.Tensor'> <class 'torch.Tensor'>"""
    dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)
    print('dec_state.shape : ', dec_state.shape)
    """dec_state.shape :  torch.Size([2, 1, 32])"""
    # 添加批量轴
    print('tgt_vocab.shape : ',
          torch.tensor(
              [tgt_vocab['<bos>']], dtype=torch.long, device=device).shape)
    """tgt_vocab.shape :  torch.Size([1])"""
    dec_X = torch.unsqueeze(torch.tensor(
        [tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)
    print('dec_X.shape : ', dec_X.shape)
    """dec_X.shape :  torch.Size([1, 1])"""
    output_seq, attention_weight_seq = [], []
    for _ in range(num_steps):
        Y, dec_state = net.decoder(dec_X, dec_state)
        print('Y.shape : ', Y.shape)
        """Y.shape :  torch.Size([1, 1, 194])"""
        # 我们使用具有预测最高可能性的词元，作为解码器在下一时间步的输入
        dec_X = Y.argmax(dim=2)
        pred = dec_X.squeeze(dim=0).type(torch.int32).item()
        print('pred : ', pred)
        """pred :  111"""
        # 保存注意力权重（稍后讨论）
        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

# 预测序列的评估
# BLEU
# exp(min(0,1−lenlabel / lenpred))∏𝑛=1𝑘 𝑝1/2𝑛 𝑛,
def bleu(pred_seq, label_seq, k):  #@save
    """计算BLEU"""
    # len_label 表示标签序列中的词元数
    # len_label 表示预测序列中的词元数
    # 𝑘 是用于匹配的最长的𝑛元语法
    # 𝑝𝑛 表示 𝑛 元语法的精确度
    # 它是两个数量的比值： 第一个是预测序列与标签序列中匹配的 𝑛 元语法的数量，
    #                  第二个是预测序列中 𝑛 元语法的数量的比率。
    #                  给定标签序列 𝐴 、 𝐵 、 𝐶 、 𝐷 、 𝐸 、 𝐹  和
    #                  预测序列 𝐴 、 𝐵 、 𝐵 、 𝐶 、 𝐷 ，
    #                  𝑝1=4/5 、 𝑝2=3/4 、 𝑝3=1/3 和 𝑝4=0 。
    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_label))
    for n in range(1, k + 1):
        # 使用普通的字典时，用法一般是dict={},添加元素的只需要dict[element] =value
        # 调用的时候也是如此，dict[element] = xxx,但前提是element字典里，如果不在字典里就会报错
        # defaultdict的作用是在于，当字典里的key不存在但被查找时，返回的不是keyError而是一个默认值
        # dict =defaultdict( factory_function)
        # 这个factory_function可以是list、set、str等等，
        # 作用是当key不存在时，返回的是工厂函数的默认值，
        # 比如list对应[ ]，str对应的是空字符串，set对应set( )，int对应0，
        num_matches, label_subs = 0, collections.defaultdict(int)
        for i in range(len_label - n + 1):
            print(' '.join(label_tokens[i: i + n]))
            label_subs[' '.join(label_tokens[i: i + n])] += 1
        for i in range(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

# 翻译
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 in zip(engs, fras):
    translation, attention_weight_seq = predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device)
    # n-grams  n = 2
    print(f'{eng} => {translation}, bleu {bleu(translation, fra, k=2):.3f}')
    break