生成式聊天机器人 -- 基于Pytorch + Global Attention + 双向 GRU 实现的SeqToSeq模型 -- 下

生成式聊天机器人 – 基于Pytorch + Global Attention + 双向 GRU 实现的SeqToSeq模型 – 上

训练

Masked 损失

encoder 和 decoder 的 forward 函数实现之后，我们就需要计算loss。seq2seq有两个RNN，Encoder RNN是没有直接定义损失函数的，它是通过影响Decoder从而影响最终的输出以及loss。

Decoder输出一个序列，前面我们介绍的是Decoder在预测时的过程，它的长度是不固定的，只有遇到EOS才结束。给定一个问答句对，我们可以把问题输入Encoder，然后用Decoder得到一个输出序列，但是这个输出序列和”真实”的答案长度并不相同。而且即使长度相同并且语义相似，也很难直接知道预测的答案和真实的答案是否类似。

那么我们怎么计算loss呢？

比如输入是”What is your name?”，训练数据中的答案是”I am LiLi”。假设模型有两种预测：”I am fine”和”My name is LiLi”。从语义上显然第二种答案更好，但是如果字面上比较的话可能第一种更好。

但是让机器知道”I am LiLi”和”My name is LiLi”的语义很接近这是非常困难的，所以实际上我们通常还是通过字面上进行比较。我们会限制Decoder的输出，使得Decoder的输出长度和”真实”答案一样，然后逐个时刻比较。Decoder输出的是每个词的概率分布，因此可以使用交叉熵损失函数。但是这里还有一个问题，因为是一个batch的数据里有一些是padding的，因此这些位置的预测是没有必要计算loss的，因此我们需要使用前面的mask矩阵把对应位置的loss去掉，我们可以通过下面的函数来实现计算Masked的loss。

# 带掩码的负对数似然损失函数
# 一次求出一个Time step一批词的预测损失 
# inp(batch_size,7826),词汇表总共7826个词
# target(batch_size,),给出当前批次中各个批当前Time step正确输出词在词汇表中的索引ID
# mask(batch_size,),给出当前批次中各个批当前Time step对应的词是否为padding填充词
def maskNLLLoss(inp, target, mask):
    # 计算实际的词的个数，因为padding是0，非padding是1，因此sum就可以得到词的个数
    nTotal = mask.sum()
    # 计算交叉熵损失,返回结果维度: (batch_size,) 
    crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)).squeeze(1))
    # 通过mask矩阵筛选出非padding词的预测损失,然后求和算平均损失
    # 此处计算得到的平均损失是某个Time Step输入的一批词的平均损失
    loss = crossEntropy.masked_select(mask).mean()
    loss = loss.to(device)
    return loss, nTotal.item()

这段代码有两个函数比较难理解，下面来重点讲解一下，首先是gather函数:

为了实现交叉熵这里使用了gather函数，这是一种比较底层的实现方法，更简便的方法应该使用CrossEntropyLoss或者NLLLoss，其中CrossEntropy等价与LogSoftmax+NLLLoss。

交叉熵的定义为：$H(p,q)=-\sum p(x)logq(x)$。其中p和q是两个随机变量的概率分布，这里是离散的随机变量，如果是连续的需要把求和变成积分。在我们这里p是真实的分布，也就是one-hot的，而q是模型预测的softmax的输出。因为p是one-hot的，所以只需要计算真实分类对应的那个值。

比如假设一个5分类的问题，当前正确分类是$2$(下标从$0$~$4$)，而模型的预测是$(0.1,0.1,0.4,0.2,0.2)$，则$H=-log(0.4)$。用交叉熵作为分类的$Loss$是比较合理的，正确的分类是2，那么模型在下标为2的地方预测的概率$q_2$越大，则$-log{q}_2$越小，也就是$loss$越小。

假设inp是：

0.3 0.2 0.4 0.1
0.2 0.1 0.4 0.3

也就是batch=2，而分类数(词典大小)是4，inp是模型预测的分类概率。 而target = [2,3] ，表示第一个样本的正确分类是第三个类别(概率是0.4），第二个样本的正确分类是第四个类别(概率是0.3)。因此我们需要计算的是 $-log(0.4)-log(0.3)$。怎么不用for循环求出来呢？我们可以使用torch.gather函数首先把0.4和0.3选出来：

inp = torch.tensor([[0.3, 0.2, 0.4, 0.1], [0.2, 0.1, 0.4, 0.3]])
target = torch.tensor([2, 3])
selected = torch.gather(inp, 1, target.view(-1, 1))
print(selected)

输出：

tensor([[ 0.4000],
        [ 0.3000]])

关于masked_select函数，我们来看一个例子:

>>> x = torch.randn(3, 4)
>>> x
tensor([[ 0.3552, -2.3825, -0.8297,  0.3477],
        [-1.2035,  1.2252,  0.5002,  0.6248],
        [ 0.1307, -2.0608,  0.1244,  2.0139]])
>>> mask = x.ge(0.5)
>>> mask
tensor([[ 0,  0,  0,  0],
        [ 0,  1,  1,  1],
        [ 0,  0,  0,  1]], dtype=torch.uint8)
>>> torch.masked_select(x, mask)
tensor([ 1.2252,  0.5002,  0.6248,  2.0139])

它要求mask和被mask的tensor的shape是一样的，然后从crossEntropy选出mask值为1的那些值。输出的维度会减1。

单次训练过程

函数train实现一个batch数据的训练。前面我们提到过，在训练的时候我们会限制Decoder的输出，使得Decoder的输出长度和”真实”答案一样长。但是我们在训练的时候如果让Decoder自行输出，那么收敛可能会比较慢，因为Decoder在t时刻的输入来自t-1时刻的输出。

如果前面预测错了，那么后面很可能都会错下去。另外一种方法叫做teacher forcing，它不管模型在t-1时刻做什么预测都把t-1时刻的正确答案作为t时刻的输入。但是如果只用teacher forcing也有问题，因为真实的Decoder是没有老师来帮它纠正错误的。所以比较好的方法是追加一个teacher_forcing_ratio参数随机的来确定本次训练是否teacher forcing。

另外使用到的一个技巧是梯度裁剪(gradient clipping) 。这个技巧通常是为了防止梯度爆炸(exploding gradient)，它把参数限制在一个范围之内，从而可以避免梯度的梯度过大或者出现NaN等问题。注意：虽然它的名字叫梯度裁剪，但实际它是对模型的参数进行裁剪，它把整个参数看成一个向量，如果这个向量的模大于max_norm，那么就把这个向量除以一个值使得模等于max_norm，因此也等价于把这个向量投影到半径为max_norm的球上。它的效果如下图所示。

单次具体训练过程为:

1. 把整个batch的输入传入encoder
2. 把decoder的输入设置为特殊的，初始隐状态设置为encoder最后时刻的隐状态
3. decoder每次处理一个时刻的forward计算
4. 如果是teacher forcing，把上个时刻的"正确的"词作为当前输入，否则用上一个时刻的输出作为当前时刻的输入
5. 计算loss
6. 反向计算梯度
7. 对梯度进行裁剪
8. 更新模型(包括encoder和decoder)参数

注意，PyTorch的RNN模块(RNN, LSTM, GRU)也可以当成普通的非循环的网络来使用。在Encoder部分，我们是直接把所有时刻的数据都传入RNN，让它一次计算出所有的结果，但是在Decoder的时候(非teacher forcing)后一个时刻的输入来自前一个时刻的输出，因此无法一次计算。

# input_var维度:(max_length,batch_size)
# target_var维度: (max_length,batch_size)
# lengths维度: (batch_size)
# mask维度: (max_length,batch_size)
def train(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder, embedding,
          encoder_optimizer, decoder_optimizer, batch_size, clip, max_length=MAX_LENGTH):

    # 梯度清空
    encoder_optimizer.zero_grad()
    decoder_optimizer.zero_grad()

    # 设置device，从而支持GPU，当然如果没有GPU也能工作。
    input_variable = input_variable.to(device)
    lengths = lengths.to(device)
    target_variable = target_variable.to(device)
    mask = mask.to(device)

    # 初始化变量
    loss = 0
    print_losses = []
    n_totals = 0

    # encoder的Forward计算 --- 输入数据维度: (max_len,batch_size) , (batch_size)
    encoder_outputs, encoder_hidden = encoder(input_variable, lengths)

    # Decoder的初始输入是SOS，我们需要构造(1, batch)的输入，表示第一个时刻batch个输入。
    decoder_input = torch.LongTensor([[SOS_token for _ in range(batch_size)]])
    decoder_input = decoder_input.to(device)

    # 注意：Encoder是双向的，而Decoder是单向的，因此从下往上取n_layers个
    decoder_hidden = encoder_hidden[:decoder.n_layers]

    # 确定是否teacher forcing
    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False

    # 一次处理一个时刻
    if use_teacher_forcing:
        for t in range(max_target_len):
            decoder_output, decoder_hidden = decoder(
                decoder_input, decoder_hidden, encoder_outputs
            )
            # Teacher forcing: 下一个时刻的输入是当前正确答案
            decoder_input = target_variable[t].view(1, -1)
            # 计算累计的loss
            mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])
            loss += mask_loss
            print_losses.append(mask_loss.item() * nTotal)
            n_totals += nTotal
    else:
        for t in range(max_target_len):
            decoder_output, decoder_hidden = decoder(
                decoder_input, decoder_hidden, encoder_outputs
            )
            # 不是teacher forcing: 下一个时刻的输入是当前模型预测概率最高的值
            _, topi = decoder_output.topk(1)
            decoder_input = torch.LongTensor([[topi[i][0] for i in range(batch_size)]])
            decoder_input = decoder_input.to(device)
            # 计算累计的loss
            mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])
            loss += mask_loss
            print_losses.append(mask_loss.item() * nTotal)
            n_totals += nTotal

    # 反向计算
    loss.backward()

    # 对encoder和decoder进行梯度裁剪
    _ = torch.nn.utils.clip_grad_norm_(encoder.parameters(), clip)
    _ = torch.nn.utils.clip_grad_norm_(decoder.parameters(), clip)

    # 更新参数
    encoder_optimizer.step()
    decoder_optimizer.step()

    return sum(print_losses) / n_totals

迭代训练过程

最后是把前面的代码组合起来进行训练。函数trainIters用于进行n_iterations次minibatch的训练。

值得注意的是我们定期会保存模型，我们会保存一个tar包，包括encoder和decoder的state_dicts(参数),优化器(optimizers)的state_dicts, loss和迭代次数。这样保存模型的好处是从中恢复后我们既可以进行预测也可以进行训练(因为有优化器的参数和迭代的次数)。

def trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer, decoder_optimizer,
               embedding, encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size,
               print_every, save_every, clip, corpus_name, loadFilename):
    # 随机选择n_iteration个batch的数据(pair)
    # 一次性生成n_iteration轮迭代需要的全部数据 (n_iteration,batch)
    # batch 表示当前轮迭代需要的大小为batch_size的数据
    training_batches = [batch2TrainData(voc, [random.choice(pairs) for _ in range(batch_size)])
                        for _ in range(n_iteration)]

    # 初始化
    print('Initializing ...')
    start_iteration = 1
    print_loss = 0
    if loadFilename:
        start_iteration = checkpoint['iteration'] + 1

    # 训练
    print("Training...")
    for iteration in range(start_iteration, n_iteration + 1):
        training_batch = training_batches[iteration - 1]
        # input_var维度:(max_length,batch_size)
        # target_var维度: (max_length,batch_size)
        # lengths维度: (batch_size)
        # mask维度: (max_length,batch_size)
        input_variable, lengths, target_variable, mask, max_target_len = training_batch

        # 训练一个batch的数据
        loss = train(input_variable, lengths, target_variable, mask, max_target_len, encoder,
                     decoder, embedding, encoder_optimizer, decoder_optimizer, batch_size, clip)
        print_loss += loss

        # 进度
        if iteration % print_every == 0:
            print_loss_avg = print_loss / print_every
            print("Iteration: {}; Percent complete: {:.1f}%; Average loss: {:.4f}"
                  .format(iteration, iteration / n_iteration * 100, print_loss_avg))
            print_loss = 0

        # 保存checkpoint
        if (iteration % save_every == 0):
            directory = os.path.join(save_dir, model_name, corpus_name, '{}-{}_{}'
                                     .format(encoder_n_layers, decoder_n_layers, hidden_size))
            if not os.path.exists(directory):
                os.makedirs(directory)
            torch.save({
                'iteration': iteration,
                'en': encoder.state_dict(),
                'de': decoder.state_dict(),
                'en_opt': encoder_optimizer.state_dict(),
                'de_opt': decoder_optimizer.state_dict(),
                'loss': loss,
                'voc_dict': voc.__dict__,
                'embedding': embedding.state_dict()
            }, os.path.join(directory, '{}_{}.tar'.format(iteration, 'checkpoint')))

测试

模型训练完成之后，我们需要测试它的效果。最简单直接的方法就是和chatbot来聊天。因此我们需要用Decoder来生成一个响应。

贪心解码(Greedy decoding)算法

最简单的解码算法是贪心算法，也就是每次都选择概率最高的那个词，然后把这个词作为下一个时刻的输入，直到遇到EOS结束解码或者达到一个最大长度。但是贪心算法不一定能得到最优解，因为某个答案可能开始的几个词的概率并不太高，但是后来概率会很大。因此除了贪心算法，我们通常也可以使用Beam-Search算法，也就是每个时刻保留概率最高的Top K个结果，然后下一个时刻尝试把这K个结果输入(当然需要能恢复RNN的状态)，然后再从中选择概率最高的K个。

为了实现贪心解码算法，我们定义一个GreedySearchDecoder类。这个类的forwar的方法需要传入一个输入序列(input_seq)，其shape是(input_seq length, 1)， 输入长度input_length和最大输出长度max_length。过程如下：

1. 把输入传给Encoder，得到所有时刻的输出和最后一个时刻的隐状态。
2. 把Encoder最后时刻的隐状态作为Decoder的初始状态。
3. Decoder的第一时刻输入初始化为SOS。
4. 定义保存解码结果的tensor。
5. 循环直到最大解码长度。
   1. 把当前输入传入Decoder。
   2. 得到概率最大的词以及概率。
   3. 把这个词和概率保存下来。
   4. 把当前输出的词作为下一个时刻的输入。
   5. 返回所有的词和概率。

import torch
from torch import nn
from Trainer import device, decoder, encoder
from utils import SOS_token, normalizeString, indexesFromSentence, MAX_LENGTH, voc

class GreedySearchDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super(GreedySearchDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
    
    # input_seq维度: (max_len,batch_size)
    def forward(self, input_seq, input_length, max_length):
        # Encoder的Forward计算
        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)
        # 把Encoder最后时刻的隐状态作为Decoder的初始值
        decoder_hidden = encoder_hidden[:decoder.n_layers]
        # 因为我们的函数都是要求(time,batch)，因此即使只有一个数据，也要做出二维的。
        # Decoder的初始输入是SOS
        decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token
        # 用于保存解码结果的tensor
        all_tokens = torch.zeros([0], device=device, dtype=torch.long)
        all_scores = torch.zeros([0], device=device)
        # 循环，这里只使用长度限制，后面处理的时候把EOS去掉了。
        for _ in range(max_length):
            # Decoder forward一步
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden,
								encoder_outputs)
            # decoder_outputs是(batch=1, vob_size)
            # 使用max返回概率最大的词和得分
            decoder_scores, decoder_input = torch.max(decoder_output, dim=1)
            # 把解码结果保存到all_tokens和all_scores里
            all_tokens = torch.cat((all_tokens, decoder_input), dim=0)
            all_scores = torch.cat((all_scores, decoder_scores), dim=0)
            # decoder_input是当前时刻输出的词的ID，这是个一维的向量，因为max会减少一维。
            # 但是decoder要求有一个batch维度，因此用unsqueeze增加batch维度。
            decoder_input = torch.unsqueeze(decoder_input, 0)
        # 返回所有的词和得分。
        return all_tokens, all_scores

实现对话函数

解码方法完成后，我们写一个函数来测试从终端输入一个句子然后来看看chatbot的回复。我们需要用前面的函数来把句子分词，然后变成ID传入解码器，得到输出的ID后再转换成文字。我们会实现一个evaluate函数，由它来完成这些工作。

我们需要把一个句子变成输入需要的格式——shape为(max_length , batch)，即使只有一个输入也需要增加一个batch维度。我们首先把句子分词，然后变成ID的序列，然后转置成合适的格式。此外我们还需要创建一个名为lengths的tensor，虽然只有一个，来表示输入的实际长度。接着我们构造类GreedySearchDecoder的实例searcher，然后用searcher来进行解码得到输出的ID，最后我们把这些ID变成词并且去掉EOS之后的内容。

另外一个evaluateInput函数作为chatbot的用户接口，当运行它的时候，它会首先提示用户输入一个句子，然后使用evaluate来生成回复。然后继续对话直到用户输入”q”或者”quit”。如果用户输入的词不在词典里，我们会输出错误信息(当然还有一种办法是忽略这些词)然后提示用户重新输入。

def evaluate(encoder, decoder, searcher, voc, sentence, max_length=MAX_LENGTH):
    # 把输入的一个batch句子变成id
    indexes_batch = [indexesFromSentence(voc, sentence)]
    # 创建lengths tensor
    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])
    # 转置 -- (max_len,batch_size)
    input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)
    # 放到合适的设备上(比如GPU)
    input_batch = input_batch.to(device)
    lengths = lengths.to(device)
    # 用searcher解码
    tokens, scores = searcher(input_batch, lengths, max_length)
    # ID变成词。
    decoded_words = [voc.index2word[token.item()] for token in tokens]
    return decoded_words

def evaluateInput(encoder, decoder, searcher, voc):
    input_sentence = ''
    while(1):
        try:
            # 得到用户终端的输入
            input_sentence = input()
            # 是否退出
            if input_sentence == 'q' or input_sentence == 'quit': break
            # 句子归一化
            input_sentence = normalizeString(input_sentence)
            # 生成响应Evaluate sentence
            output_words = evaluate(encoder, decoder, searcher, voc, input_sentence)
            # 去掉EOS后面的内容
            words = []
            for word in output_words:
                if word == 'EOS':
                    break
                elif word != 'PAD':
                    words.append(word)
            print('Bot:', ' '.join(words))

        except KeyError:
            print("Error: Encountered unknown word.")

训练和测试模型

不论是我们是训练模型还是测试对话，我们都需要初始化encoder和decoder模型参数。在下面的代码，我们从头开始训练模型或者从某个checkpoint加载模型。读者可以尝试不同的超参数配置来进行调优。

USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda" if USE_CUDA else "cpu")

model_name = 'cb_model'
attn_model = 'dot'
# attn_model = 'general'
# attn_model = 'concat'
hidden_size = 500
encoder_n_layers = 2
decoder_n_layers = 2
dropout = 0.1
batch_size = 64

# 从哪个checkpoint恢复，如果是None，那么从头开始训练。
#loadFilename = "data/save/cb_model/cornell movie-dialogs corpus/2-2_500/4000_checkpoint.tar"
loadFilename = None
checkpoint_iter = 4000

# 如果loadFilename不空，则从中加载模型
if loadFilename:
    # 如果训练和加载是一条机器，那么直接加载
    checkpoint = torch.load(loadFilename)
    # 否则比如checkpoint是在GPU上得到的，但是我们现在又用CPU来训练或者测试，那么注释掉下面的代码
    # checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))
    encoder_sd = checkpoint['en']
    decoder_sd = checkpoint['de']
    encoder_optimizer_sd = checkpoint['en_opt']
    decoder_optimizer_sd = checkpoint['de_opt']
    embedding_sd = checkpoint['embedding']
    voc.__dict__ = checkpoint['voc_dict']

print('Building encoder and decoder ...')
# 初始化word embedding
embedding = nn.Embedding(voc.num_words, hidden_size)
if loadFilename:
    embedding.load_state_dict(embedding_sd)
# 初始化encoder和decoder模型
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = GlobalAttnDecoderRNN(attn_model, embedding, hidden_size, voc.num_words,
                               decoder_n_layers, dropout)
if loadFilename:
    encoder.load_state_dict(encoder_sd)
    decoder.load_state_dict(decoder_sd)
# 使用合适的设备
encoder = encoder.to(device)
decoder = decoder.to(device)
print('Models built and ready to go!')

训练前，我们需要设置一些训练的超参数。初始化优化器，最后调用函数trainIters进行训练。

# 配置训练的超参数和优化器
clip = 50.0
teacher_forcing_ratio = 1.0
learning_rate = 0.0001
decoder_learning_ratio = 5.0
n_iteration = 4000
print_every = 1
save_every = 500

# 设置进入训练模式，从而开启dropout
encoder.train()
decoder.train()

# 初始化优化器
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=learning_rate * decoder_learning_ratio)
if loadFilename:
    encoder_optimizer.load_state_dict(encoder_optimizer_sd)
    decoder_optimizer.load_state_dict(decoder_optimizer_sd)

# 开始训练
print("Starting Training!")
trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer, decoder_optimizer,
           embedding, encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size,
           print_every, save_every, clip, corpus_name, loadFilename)

训练完毕后，我们使用下面的代码进行测试。

# 进入eval模式，从而去掉dropout。 
encoder.eval()
decoder.eval()

# 构造searcher对象 
searcher = GreedySearchDecoder(encoder, decoder)

# 测试
evaluateInput(encoder, decoder, searcher, voc)

下面是测试的一些例子：

> hello
Bot: hello .
> what's your name?
Bot: jacob .
> I am sorry.
Bot: you re not .
> where are you from?
Bot: southern .
> q

完整代码

码云仓库链接