PyTorch实战：语言模型

结合前面介绍的PyTorch基础，我们将进行一系列的实战操练，首先从NLP中最基本的语言模型入手。传统的语言模型是一个使用句子的前文若干个单词预测后文的下一个单词，很明显这是一个序列问题，应该使用循环神经网络来建模。我们先导入相关的包，并选择合适的算力设备：

import torch,pickle
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
from torch.nn.utils import clip_grad_norm_

device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')

预处理

我们下载 PTB数据集，解压后使用 data 文件夹下的 ptb.train.txt 作为构建语言模型的训练数据。首先需要对数据进行预处理，包括分词、单词编号等操作，我们可以直接调用前面章节《PyTorch数据预处理》中的 build_vocab 方法，需要注意的是，PTB 的训练数据里已经将低频词替换为 \ 了，所以我们不需要显式地在词表中增加 \，同时由于语料是英文的，可以简单地将 tokenizer 实现为空格分隔。

def preprocessing():
    word2idx, idx2word = build_vocab("./data/NNLM/ptb.train.txt", space_tokenizer, 50000, 0)
    fout = open('./data/NNLM/vocab.pkl', 'wb')
    pickle.dump((word2idx, idx2word), fout)
    fout.close()
    print("vocab size is: {}".format(str(len(word2idx))))

需要将 build_vocab 返回单词与编号的索引保存下来，同时输出并记录下单词的数量，作为后续模型的超参数传入。

数据加载

数据的读取与加载我们使用 PyTorch 提供的 Dataset 与 DataLoader 来实现：

class MyDataset(Dataset):
    def __init__(self, vocab_file, input_file):
        word2idx, _ = pickle.load(open(vocab_file, 'rb'))
        data = []
        with open(input_file, encoding='utf-8') as fin:
            for sline in fin:
                word_list = sline.strip().split(' ')
                data.append([word2idx[w] for w in word_list])
        # 语言模型的target可以根据data来生成，在train的过程中去做就行
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

我们继承 Dataset 实现自己的数据读取，由于语言模型的目标 label 完全可以通过输入 data 移位来获取，所以我们只需要 Dataset 返回一个 data 就可以了，在类的构造函数中实现从文件读取训练数据并转换为单词 ID 序列。将 Dataset 实例化后与 DataLoader 配合实现数据的批量加载：

def create_loader(vocab_file, input_file, batch_size):
    dataset = MyDataset(vocab_file, input_file)

    def collate_fn(batch_data, pad=0):
        # 以批量数据中的最长序列为基准补齐
        max_len = max([len(seq) for seq in batch_data])  
        data = [seq + [pad] * (max_len - len(seq)) for seq in batch_data]
        data = torch.LongTensor(data).to(device)
        return data

    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, 
                      collate_fn=collate_fn, drop_last=True)

我们通过实现一个 collate_fn 函数来对序列数据进行补齐，最后我们要舍弃掉训练数据按 batch_size 划分时最后一个 batch 中不足一个 batch_size 的那些数据，因为我们后面使用的模型（LSTM）需要多次初始化隐藏层（其大小与 batch_size 相关，所以要保持每个 batch 的数据量相同），当然也可以通过把最后一个 batch 做一个批层面的补齐来保留这部分数据。

定义网络

我们定义一个使用循环神经网络实现的语言模型：它包含一个 embedding 层，用于将每个单词编号转换成一个连续空间里的向量；然后是 LSTM 层，用来对序列的每一步进行预测，为了与其他网络层的输入要求保持一致，我们在定义 LSTM 层的时候设置了 batch_first=True，以便让 batch_size 位于第0轴的位置；由于语言模型是用来预测下一个单词的，所以 LSTM 的输出需要接一个全连接层，以实现对预测位置上的单词的还原（分类）。

class RNNLM(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, num_layers=1, dropout_p=0.5):
        super(RNNLM, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(dropout_p)

        self.hidden_size = hidden_size
        self.num_layers = num_layers

    def init_hidden(self, batch_size):
        states = (torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device),
                  torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device))
        return states

    def forward(self, x, h):
        x = self.embed(x)
        out, h = self.lstm(x, h)
        out = self.fc(self.dropout(out))
        return out, h

在使用 LSTM 的时候，需要额外实现一个 init_hidden( ) 的函数，因为训练过程中 loss 反向传播的时候，PyTorch 会试图把 hidden state 也反向传播，但是在新的一轮 batch 的时候 hidden state 已经被内存释放了，所以需要每个 batch 都要重新初始化隐层状态。该网络在前向传播时需要注意的是，输入的 x 形状是 (batch_size, seq_len, embed_size)，输出的 out 的形状是 (batch_size, seq_len, vocab_size)，输入和输出都是一个完整的序列，我们可以只从宏观上数据整体的形状进行把控，而 LSTM 内部会实现序列的按步展开与逐步运算。

开始训练

启动训练之前，我们需要设置好相关的超参数，我们不妨把这些参数预先定义到一个配置类里：

class Config(object):
    batch_size = 64         # 数据批量的大小
    epochs = 30             # 训练的迭代次数

    vocab_size = 10000      # 词汇表的大小（预处理时统计而得）
    embed_size = 64         # 词向量的维度
    hidden_size = 128       # 隐藏层的维度
    num_layers = 1          # LSTM的层数
    learning_rate = 0.002   # 学习率

然后，我们开始正式的训练过程，包括初始化超参数、加载数据、初始化网络、定义损失、选择优化器、迭代训练（每一轮迭代都要完成一次批量遍历，每一个批量的数据训练后都要进行一次参数更新）、记录损失变化：

def train():
    config = Config()

    train_loader = create_loader('./data/NNLM/vocab.pkl', './data/NNLM/ptb.train.txt', config.batch_size)

    net = RNNLM(config.vocab_size, config.embed_size, config.hidden_size, config.num_layers, dropout_p=0.5)
    net = net.to(device)
    # 计算损失的时候忽略补齐序列所padding的部分
    criterion = nn.CrossEntropyLoss(ignore_index=0)
    optimizer = torch.optim.Adam(net.parameters(), lr=config.learning_rate)

    training_loss = []
    training_accuracy = []
    for epoch in range(config.epochs):
        losses = 0
        accuracy = 0
        batch_num = 0
        for batch_sentence in train_loader:
            # target与src相差一个位置
            src, tgt = batch_sentence[:], batch_sentence[:, 1:]
            # 每个batch都要初始化一次隐层状态
            states = net.init_hidden(config.batch_size)
            outputs, states = net(src, states)
            # 计算损失时，预测结果在seq_len轴上的最后一个结果要去掉，调整形状为二阶张量
            # 其中第1轴的维度为分类类别即词表大小：(batch*seq_len, vocab_size)
            y_pred = outputs[:,:-1,:].reshape(-1, config.vocab_size)
            # 计算损失时，真实结果也要调整形状为一阶张量
            y_true = tgt.reshape(-1)
            loss = criterion(y_pred, y_true)
            optimizer.zero_grad()
            loss.backward()
            clip_grad_norm_(net.parameters(), 0.5)
            optimizer.step()
            losses += loss.item()
            accuracy += compute_accuracy(y_pred, y_true)
            batch_num += 1
            print("Batch index = %d, Batch loss = %f." % (batch_num, loss.item()))
        training_loss.append(losses/batch_num)
        training_accuracy.append(accuracy/batch_num)
        print('Epoch = %d, Train Loss = %f, Train Accuracy = %f.'
              % (epoch, training_loss[epoch], training_accuracy[epoch]))
    torch.save(net, './model/NNLM_epoch%d_loss_%f.pth' %(epoch, training_loss[-1]))

训练过程中需要注意的是，一定要确保数据和网络都放到相同的设备（CPU或GPU）上进行运算；另外，我们没有对序列增加特殊的句首句尾标记 \ 和 \，所以序列的预测结果中最后一个位置上的值要丢弃（并不存在），其实增加句首句尾标记是有用的，大家可以自行尝试加入。我们除了观察损失的变化，也计算了训练时的准确率进行输出：

def compute_accuracy(y_pred, y_true, mask_index=0):
    _, y_pred_indices = y_pred.max(dim=1)

    correct_indices = torch.eq(y_pred_indices, y_true).float()
    valid_indices = torch.ne(y_true, mask_index).float()

    n_correct = (correct_indices * valid_indices).sum().item()
    n_valid = valid_indices.sum().item()

    return n_correct / n_valid * 100

使用模型

使用语言模型的一个案例就是进行文本生成，给定一个输入，根据语言模型的预测逐步输出下文：

from torch.autograd import Variable
import torch.nn.functional as F

def generate(input_word, vocab_file, model_file, word_len=10, temperature=1.0):
    config = Config()
    # 加载模型
    model = torch.load(model_file, map_location=device)
    # 加载词典
    word2idx, idx2word = pickle.load(open(vocab_file, 'rb'))
    # 使用eval()使 Dropout 层冻结
    model.eval()
    # 使用Variable定义隐藏层，不进行梯度回传，batch_size=1
    hidden = (Variable(torch.zeros(config.num_layers, 1, config.hidden_size)).to(device),
              Variable(torch.zeros(config.num_layers, 1, config.hidden_size)).to(device))
    # 输入是二阶张量，batch_size=1, seq_len=1
    input_tensor = torch.tensor([[word2idx[input_word]]])
    input = input_tensor.to(device)
    print(input_word, end=' ')
    for i in range(word_len): 
        output, hidden = model(input, hidden)
        # output的形状是 (batch_size=1, seq_len=1, vocab_size=10000)
        # 压缩维度把第0、1轴删除了，进而通过softmax转换成概率
        word_weights = F.softmax(output.squeeze()/temperature)
        # 按照概率大小进行采样
        word_idx = torch.multinomial(word_weights, 1)[0]
        input.data.fill_(word_idx)
        word = idx2word[word_idx.item()]
        print(word, end=' ')

上面的代码我们引入了两个小技巧：一个是在进行softmax计算时加入了温度参数，可以参考论文；还有一个就是在生成序列时并不是通过 argmax 直接取概率最大的那个词，而是按概率大小进行采样，当然也可以引入 beam search 等算法。

可视化

tensorboradX 只是一个辅助工具，可以把 PyTorch 中的参数以 tensorboard 可以呈现的格式保存，我们需要启动 tensorboard 来可视化这些训练过程的记录信息。由于 tensorboard 是 tensorflow 的工具，我们可以单独创建一个 tensorflow 的虚拟环境，并切换到这个环境下启动 tensorboard，通过执行命令 tensorboard —logdir=’./log’ 启动 tensorboard，并在浏览器上通过 http://localhost:6006 查看训练过程记录指标的可视化图形。如果我们是在服务器上训练模型，而想在自己的 PC 上打开可是化界面，则可以增加一个参数 —host=0.0.0.0 来启动。

模块来记录训练过程中各种指标的变化，该模块可以方便可视化训练结束后我们会发现，在 log 目录下会有多个已 “日期_时间”为名称的子目录，每一个目录对应一次训练过程的记录信息。我们可以超参数的调优往往都是通过多次实验，在指标的变化曲线上确定的。