基于 GPT2 实现生成小说任务 – 数据处理

前面使用 GPT2 实现生成对联任务，这两天想用 GPT2 实现小说生成。这两个任务看似都是文本生成任务，但还是不同的。对联任务生成的内容很短小，生成小说则内容很长。从实现过程来看，虽然 GPT 可以生成长文本的内容，但是发现越到后面生成的内容其实越像为了生成而生成了，内容之间不是很连贯，这也可能由于我使用的数据集较小。

另外，也在这个过程中碰到一些问题。比如：transformers 库中对 GPT2LMHeadModel 模型在输入时，如果输入 labels 则模型会帮助计算损失，我们直接使用损失反向传播即可。以前使用时，并没有直接使用这个带 head 的 Model，而是自己写 GPT2Model + nn.Linear，所以习惯了自己构造好输入数据和标签输入传入，例如：

完整: abcdef
输入: abcde
标签: bcdef

也就是说，自己会提前构建好数据送入模型。此时，如果使用 GPT2LMHeadModel 的话，程序也不会报错，但是计算的损失就大错特错了，为什么？这是因为该模型内部对输入的数据又做了一个偏移，如下所示：

完整: abcdef
输入: abcde
标签: bcdef

# 模型计算损失时对输入和标签的改动
输入: abcd
标签: cdef

原来，我们输入 a 预测 b，现在变成让模型输入 a 去预测 c 了，显然驴唇不对马嘴了，这导致训练时，损失持续降低，但是预测时惨不忍睹。下面是 GPT2LMHeadModel 损失部分的计算代码：

loss = None
if labels is not None:
    # Shift so that tokens < n predict n
    # 类似于 data[:-1]
    shift_logits = lm_logits[..., :-1, :].contiguous()
    # 类似于 data[1:]
    shift_labels = labels[..., 1:].contiguous()
    # Flatten the tokens
    loss_fct = CrossEntropyLoss()
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

构建输入数据时，我也思考了很多，进行了简单的尝试：

按小说句子送入模型训练，这就使得模型学习第一句和第二句是有些割裂的。比如模型生成了第一句，再生成第二句时，没有学习到对第一句的依赖。
如果按照滑动窗口的方式进行采样，如果窗口太小，则模型学习不到对前面内容更长的依赖。所以，最终就设置了滑动窗口大小为 500，句子长度为 1000，反正 GPT2 的输入支持 1024 长度。此时，能够生成较为可以的长文本内容了。这里面的参数设置仅仅是初步尝试，并不是最优的。

上面，只是做了一些简单的尝试。在具体训练时，由于我的显存只有 6G，模型以及优化器等相关的参数加载到显存之后，可用显存就所剩无几，只能每次送入一个样本，这样使得训练速度也较慢。所以，训练时使用了累加梯度的方式加快训练。

为了能够更加便捷的编写训练代码，使用了 transformers 的 datasets 库、以及 ignite 高阶的、方便的训练工具，最终以 Clear ML 来可视化训练过程，训练过程也仅仅可视化了损失变化，以及学习率的变化。

学习率调整策略使用的是从开始到训练结束，学习率呈现一个线性的下降，使用的是 ignite 里的 PiecewiseLinear 学习率调整策略，简单的截了下图：

模型的存储时，根据损失值作为指标，保存整个训练过程中，损失值最低的 2 个模型，另外存储时，除了存储 model，也一并存储了 optimizer、trainer 对象，如果后面需要对模型继续进行训练，可能会用到这些对象。

下面的是对数据的预处理代码，处理之后由训练文件直接读取，所有的数据有 30 万行左右，这里为了便于快速训练，只使用其中 1 万行数据进行预处理。

语料词汇总量: 485791 字
训练语料数量: 10000 行
训练词表大小: 3563 字
训练数据信息: ['text'] 972 个

预处理代码：

import re
import numpy as np
from transformers import BertTokenizer
import math
import pickle
from datasets import Dataset

def build_vocab():

    lines = []
    words = []
    # None 表示使用所有数据
    max_line = 10000
    with open('data/novel.txt') as file:
        for line in file:
            # 去除空行
            line = line.strip()
            if not len(line):
                continue

            # 去除 "第12章" 这些行
            result = re.findall(r'第\d+章 ', line)
            if len(result) > 0:
                continue

            if max_line is not None and len(lines) == max_line:
                break

            lines.append(line)
            words.extend(line)

    unique_words = np.unique(words).tolist()
    unique_words.insert(0, '[SEP]')
    unique_words.insert(0, '[UNK]')
    unique_words.insert(0, '[PAD]')


    # 存储训练词表
    with open('temp/vocab.txt', 'w') as file:
        for word in unique_words:
            file.write(word + '\n')

    # 生成分词器
    tokenizer = BertTokenizer(vocab_file='temp/vocab.txt')
    tokenizer.save_pretrained('model')

    # 可视化训练语料
    with open('temp/train_corpus.txt', 'w') as file:
        file.write(''.join(lines))

    print('语料词汇总量:', len(words), '字')
    print('训练语料数量:', len(lines), '行')
    print('训练词表大小:', len(unique_words), '字')


# 构建训练输入语料
def build_inputs():

    window = 500
    length = 1000
    # 读取语料
    novel = open('temp/train_corpus.txt').read()
    number = math.ceil(len(novel) / window)

    data_inputs = []
    for start in range(number):
        start = start * window
        end = start + length
        text = novel[start: end]
        data_inputs.append(text)

    dataset = Dataset.from_dict({'text': data_inputs})
    dataset.save_to_disk('temp/train_data')

    print('训练数据信息:', list(dataset.features.keys()), dataset.num_rows, '个')


if __name__ == '__main__':
    build_vocab()
    build_inputs()