使用 PyTorch 学习生成式人工智能——训练 Transformer 进行文本生成

88 阅读36分钟

本章内容包括:

  • 构建一个适合您需求的精简版 GPT-2XL 模型
  • 准备训练 GPT 风格 Transformer 的数据
  • 从零开始训练 GPT 风格的 Transformer
  • 使用训练好的 GPT 模型生成文本

在第 11 章中,我们从零构建了 GPT-2XL 模型,但由于其参数量庞大,无法进行训练。训练一个拥有 15 亿参数的模型需要超级计算资源和大量数据。因此,我们加载了 OpenAI 预训练的权重,并使用该 GPT-2XL 模型生成文本。

然而,学习如何从头训练 Transformer 模型仍然非常重要,原因有三:
首先,尽管本书不直接讲解如何微调预训练模型,但理解训练过程是微调的基础。训练是随机初始化参数,而微调则是加载预训练权重后继续训练。
其次,训练或微调 Transformer 可以让您定制模型以满足特定需求和领域,这能显著提升模型的性能和适用性。
最后,自己训练或微调模型能更好地掌控数据和隐私,这对于敏感应用或处理专有数据尤为关键。总之,掌握 Transformer 的训练与微调技能,对于想要利用语言模型实现特定应用且维护隐私控制的人来说至关重要。

因此,本章将构建一个精简版 GPT 模型,参数约为 500 万。该小型模型遵循 GPT-2XL 架构,主要区别是仅包含 3 个解码块,嵌入维度为 256(而 GPT-2XL 是 48 个解码块和 1600 维嵌入)。通过缩减模型规模,我们可以在普通电脑上训练。

生成文本的风格取决于训练数据。训练文本生成模型时,文本长度和多样性都很关键。训练材料必须足够丰富,以便模型能有效学习和模仿特定写作风格;如果训练材料缺乏多样性,模型可能只是简单复制训练文本;而材料过长则可能带来过高的计算成本。基于此,我们选用欧内斯特·海明威的三部小说作为训练材料:《老人与海》、《永别了,武器》和《丧钟为谁而鸣》。这确保了训练数据既有足够长度和多样性,又不会使训练难以进行。

由于 GPT 模型无法直接处理原始文本,我们首先将文本分词为单词,然后创建词典,将每个唯一词汇映射到不同索引。利用该词典,我们将文本转换成长整数序列,供神经网络输入。

我们将使用长度为 128 的索引序列作为训练输入。和第 8 章及第 10 章类似,输入序列右移一个词作为输出,强制模型基于当前词和前文预测下一个词。

一个关键难题是确定最佳训练轮数(epoch)。目标不是仅仅最小化训练集交叉熵损失,因为那可能导致过拟合,即模型简单复制训练文本。为解决此问题,我们计划训练 40 个 epoch,每 10 个 epoch 保存一次模型,并评估哪个版本能生成连贯且非简单抄袭的文本。或者,也可以像第 2 章一样使用验证集来决定训练终止时机。

训练完成后,我们将使用该 GPT 模型进行自回归文本生成,类似第 11 章的方法。我们将测试不同训练轮次的模型。训练 40 个 epoch 的模型能生成风格连贯的文本,较好地体现海明威的写作风格,但若提示语与训练文本相似,可能部分文本是直接复制。训练 20 个 epoch 的模型也能生成连贯文本,偶有语法错误,但较少复制训练文本。

本章主要目标不是生成最完美连贯的文本,因为这非常困难;而是教您如何从零开始构建 GPT 风格模型,结合实际应用需求量身定制。更重要的是,您将学习如何选择训练文本、分词并转换为索引、准备训练批次,确定训练轮数,以及如何用训练好的模型生成文本并避免直接复制训练材料。

12.1 从零开始构建和训练 GPT

我们的目标是掌握从零开始构建和训练一个针对特定任务的 GPT 模型。这项技能对于将本书的概念应用于实际问题至关重要。

假设你是欧内斯特·海明威作品的狂热粉丝,想训练一个 GPT 模型来生成符合海明威风格的文本。你会如何着手?本节将讨论完成该任务的步骤。

第一步是配置一个适合训练的 GPT 模型。你将创建一个结构类似于第 11 章中构建的 GPT-2 模型,但参数量大幅减少,使训练能在几小时内完成。为此,你需要确定模型的关键超参数,如序列长度、嵌入维度、解码块数量和丢弃率。这些超参数直接影响训练后模型的输出质量和训练速度。

接下来,你将收集几部海明威的原始小说文本并进行清理,确保适合训练。对文本进行分词,并为每个唯一的词汇分配一个整数索引,以便输入模型。准备训练数据时,将文本拆分为一定长度的整数序列作为输入,输入序列向右平移一个词作为输出。这种方法迫使模型基于当前词和所有之前的词预测下一个词。

模型训练完成后,你将用它基于提示生成文本。先将提示文本转换为索引序列,输入训练好的模型。模型基于序列迭代预测最可能的下一个词。随后,你会将模型生成的词序列转回文本。

本节先介绍该任务中 GPT 模型的架构,然后讲述训练步骤。

12.1.1 用于文本生成的 GPT 架构

尽管 GPT-2 有多种规模版本,但它们架构相似。本章构建的 GPT 模型结构与 GPT-2 相同,只是规模小得多,能在无超级计算资源下训练。表 12.1 比较了我们构建的 GPT 模型与 GPT-2 四个版本的差异。

模型嵌入维度解码层数注意力头数序列长度词汇表大小参数数量
GPT-2S76812121024502571.24 亿
GPT-2M102424161024502573.5 亿
GPT-2L128036201024502577.74 亿
GPT-2XL1600482510245025715.58 亿
我们的 GPT2563412810600512 万

本章构建的 GPT 模型包含 3 个解码层,嵌入维度为 256(即词嵌入后每个词由 256 维向量表示)。如第 11 章所述,GPT 模型采用不同于 2017 年论文《Attention Is All You Need》的位置编码方法,使用嵌入层学习序列中不同位置的编码,因此序列中每个位置也用一个 256 维向量表示。

在计算因果自注意力时,我们使用 4 个并行注意力头,分别捕捉序列中词语意义的不同方面。每个头的维度为 256/4=64,与 GPT-2 模型类似。例如 GPT-2XL 中每个注意力头维度是 1600/25=64。

我们的 GPT 模型最大序列长度为 128,远小于 GPT-2 模型的 1024。这种缩减确保模型参数量可控,但即使如此,模型仍能学习序列中词与词之间的关系并生成连贯文本。

虽然 GPT-2 词汇表大小为 50257,我们的模型词汇表较小,仅有 10600。词汇表大小主要由训练数据决定,若使用更多文本训练,词汇表规模可能会增大。

图 12.1 展示了本章将创建的仅解码器 Transformer 架构,类似于第 11 章中见到的 GPT-2,但规模更小。总参数约为 512 万,而 GPT-2XL 为 15.58 亿。图中还显示了训练各阶段的数据规模。

image.png

图 12.1 仅解码器 Transformer 的架构,设计用于生成文本。文本来源于三部海明威小说,先被分词,然后转换成索引。我们将 128 个索引排列成一个序列,每个批次包含 32 个这样的序列。输入首先经过词嵌入和位置编码,输入嵌入是这两部分的加和。接着,输入嵌入依次通过三个解码层。随后,输出经过层归一化并传入线性层,输出大小为 10,600,对应词汇表中唯一标记的数量。

我们构建的 GPT 模型的输入是输入嵌入,如图 12.1 底部所示。下一小节将详细介绍这些嵌入的计算方法。简而言之,它们是输入序列的词嵌入与位置编码之和。

输入嵌入随后依次通过三个解码层。与第 11 章构建的 GPT-2XL 模型类似,每个解码层包含两个子层:因果自注意力层和前馈网络。此外,每个子层都应用层归一化和残差连接。之后,输出经过层归一化和线性层。我们 GPT 模型的输出数量对应词汇表中唯一标记的数量,即 10,600。模型的输出是下一个标记的 logits,后续会对其应用 softmax 函数,得到词汇表上的概率分布。该模型旨在基于当前标记及其之前所有标记来预测下一个标记。

12.1.2 GPT 模型生成文本的训练过程

既然我们已经了解了如何构建用于文本生成的 GPT 模型,接下来让我们探讨训练该模型所涉及的步骤。在深入代码实现之前,我们先对训练流程进行一个概览。

生成文本的风格受训练文本的影响。由于我们的目标是训练模型生成具有欧内斯特·海明威风格的文本,因此我们将选用他的三部小说作为训练数据:《老人与海》、《永别了,武器》和《丧钟为谁而鸣》。如果只选择一部小说,训练数据的多样性会不足,导致模型记忆小说中的段落,生成的文本与训练数据完全相同。反之,若选用过多小说,会使唯一标记的数量增多,从而增加在短时间内有效训练模型的难度。因此,我们通过选择三部小说并将它们合并作为训练数据,实现了多样性和训练效率之间的平衡。

图 12.2 展示了训练 GPT 模型以生成文本所涉及的步骤。

image.png

和前面三章一样,训练过程的第一步是将文本转换为数值形式,以便我们能够将训练数据输入模型。具体来说,我们首先使用词级分词(word-level tokenization)将三部小说的文本拆分成词元,就像第8章中所做的那样。在这种情况下,每个词元是一个完整的单词或标点符号(例如冒号、括号或逗号)。词级分词实现简单,而且我们可以控制唯一词元的数量。分词后,我们为每个词元分配唯一的索引(即整数),将训练文本转换成整数序列(见图12.2中的步骤1)。

接下来,我们通过将这段整数序列划分为等长的序列来转换为训练数据(图12.2中的步骤2)。每个序列最多包含128个索引。选择128是为了在捕捉句子中词元的长距离依赖关系的同时,使模型规模保持可控。不过,128并非神奇数字:如果改成100或150,结果也会类似。这些序列作为模型的特征(即x变量)。和前几章一样,我们将输入序列向右平移一个词元,并将其用作训练数据中的输出(即y变量;图12.2中的步骤3)。

输入和输出的成对序列作为训练数据(x, y)。以句子“the old man and the sea”为例,我们使用表示“the old man and the”的索引作为输入x,然后将输入序列右移一个词元,使用“old man and the sea”的索引作为输出y。在第一个时间步,模型用“the”预测“old”;第二个时间步用“the old”预测“man”,依此类推。

训练过程中,你将遍历训练数据。在前向传播时,将输入序列x传入GPT模型(步骤4)。模型根据当前参数做出预测(步骤5)。通过将预测的下一个词元与步骤3中的输出进行比较,计算交叉熵损失(步骤6)。也就是说,将模型预测结果与真实标签进行比较。最后,调整GPT模型中的参数,使下一次迭代时模型预测更接近真实输出,最小化交叉熵损失(步骤7)。注意,模型本质上是在解决一个多分类问题:预测下一个词元是词汇表中所有唯一词元中的哪一个。

你将多次重复步骤3到步骤7。每次迭代后,模型参数都会调整,以提升下一个词元的预测准确性。我们将重复这一过程40个epoch,并每10个epoch保存一次训练好的模型。正如你稍后会看到的,训练时间过长会导致过拟合,模型记忆训练文本的段落,生成的文本就和原小说中的内容一模一样。我们将事后测试哪个版本的模型既能生成连贯文本,又不会简单复制训练数据。

12.2 处理海明威小说文本的分词

既然你已经了解了GPT模型的架构和训练过程,让我们从第一步开始:对海明威小说文本进行分词和索引。

首先,我们将对文本数据进行处理,为训练做准备。我们会像第8章那样,将文本拆分为独立的词元。由于深度神经网络无法直接处理原始文本,我们需要创建一个字典,为每个词元分配一个唯一的索引,将其映射为整数。之后,我们会将这些索引整理成训练数据批次,这对后续训练GPT模型非常重要。

我们采用词级分词,因为它简单,能方便地将文本划分成单词,而不必像子词分词那样涉及复杂的语言结构理解。此外,词级分词产生的唯一词元数量比子词分词要少,从而减少GPT模型的参数数量。

12.2.1 文本分词

为了训练GPT模型,我们将使用欧内斯特·海明威的三部小说的原始文本文件:《老人与海》、《永别了,武器》和《丧钟为谁而鸣》。这些文本文件来自Faded Page网站:www.fadedpage.com。我已将文本清理,去除了不属于原书的开头和结尾段落。准备你自己的训练文本时,务必删除所有无关信息,如供应商信息、格式和版权声明,确保模型专注于学习文本的写作风格。我还删除了章节间与正文无关的文字。你可以从本书GitHub仓库下载这三份文件 OldManAndSea.txt、FarewellToArms.txt 和 ToWhomTheBellTolls.txt,放入电脑的 /files/ 文件夹。

在《老人与海》的文本文件中,开头和结尾的双引号都用直引号(")表示,而其他两本小说中则不是。因此,我们加载《老人与海》的文本,将直引号替换成开引号或闭引号,以区分它们。这也方便后续生成文本的格式化:去除开引号后的空格和闭引号前的空格。下面代码实现了这一操作:

with open("files/OldManAndSea.txt","r", encoding='utf-8-sig') as f:
    text=f.read()
text=list(text)                                             # ①
for i in range(len(text)):
    if text[i]=='"':
        if text[i+1]==' ' or text[i+1]=='\n':
            text[i]='”'                                     # ②
        if text[i+1]!=' ' and text[i+1]!='\n':
            text[i]='“'                                     # ③
    if text[i]=="'":
        if text[i-1]!=' ' and text[i-1]!='\n':
            text[i]="’"                                     # ④
text="".join(text)                                          # ⑤

① 读取文本并拆分为单个字符
② 如果直双引号后面跟空格或换行,改为闭引号
③ 否则改为开引号
④ 直单引号改为撇号(闭单引号)
⑤ 将字符重新合并成字符串

接着,我们加载另外两部小说文本,将三部小说合并成一个文件:

with open("files/ToWhomTheBellTolls.txt","r", encoding='utf-8-sig') as f:
    text1=f.read()
  
with open("files/FarewellToArms.txt","r", encoding='utf-8-sig') as f:
    text2=f.read()
  
text=text+" "+text1+" "+text2

with open("files/ThreeNovels.txt","w", encoding='utf-8-sig') as f:
    f.write(text)
print(text[:250])

① 读取第三部小说文本
② 读取第二部小说文本
③ 合并三部小说文本
④ 保存合并后的文本

输出为合并文本的前250个字符。

我们使用空格作为分词符进行分词。正如上面输出所示,句号(.)、连字符(-)和撇号(’)等标点直接贴在前面的单词后,无空格。因此,需要在所有标点符号前后添加空格。同时,将换行符(\n)替换为空格,防止换行符被当作词元。还将所有单词转换为小写,这样“The”和“the”会被识别为同一个词元,有助于减少唯一词元数量,提高训练效率。以下代码实现了上述文本清理:

text=text.lower().replace("\n", " ")                         # ①

chars=set(text.lower())
punctuations=[i for i in chars if i.isalpha()==False and i.isdigit()==False]  # ②
print(punctuations)

for x in punctuations:
    text=text.replace(f"{x}", f" {x} ")                      # ③
text_tokenized=text.split()

unique_tokens=set(text_tokenized)
print(len(unique_tokens))                                    # ④

① 换行符替换为空格
② 找出所有标点符号
③ 在标点符号前后插入空格
④ 统计唯一词元数量

执行后输出所有标点符号列表和唯一词元数量:

[')', '.', '&', ':', '(', ';', '-', '!', '"', ' ', "'", '"', '?', ',', "'"]
10599

总共标点如上所示,词元数量为10,599,远小于GPT-2的50,257个词元,大大降低了模型规模和训练时间。

我们还添加一个“UNK”词元,代表未知词元,防止遇到词表外的词时程序崩溃。例如,提示词中出现“technology”,如果该词未包含在词典word_to_int里,程序将报错。加入“UNK”后,可以将未知词转换为“UNK”索引输入模型。训练自己的GPT模型时应始终包含“UNK”词元,因为不可能囊括所有词元。具体实现如下:

from collections import Counter   

word_counts=Counter(text_tokenized)    
words=sorted(word_counts, key=word_counts.get, reverse=True)     
words.append("UNK")                                            # ①
text_length=len(text_tokenized)
ntokens=len(words)                                             # ②
print(f"the text contains {text_length} words")
print(f"there are {ntokens} unique tokens")  
word_to_int={v:k for k,v in enumerate(words)}                  # ③
int_to_word={v:k for k,v in word_to_int.items()}               # ④
print({k:v for k,v in word_to_int.items() if k in words[:10]})
print({k:v for k,v in int_to_word.items() if v in words[:10]})

① 将“UNK”添加到唯一词元列表
② 统计词表大小ntokens,作为模型超参数
③ 词元映射到索引
④ 索引映射回词元

输出:

the text contains 698207 words
there are 10600 unique tokens
{'.': 0, 'the': 1, ',': 2, '"': 3, '"': 4, 'and': 5, 'i': 6, 'to': 7, 'he': 8, 'it': 9}
{0: '.', 1: 'the', 2: ',', 3: '"', 4: '"', 5: 'and', 6: 'i', 7: 'to', 8: 'he', 9: 'it'}

三部小说文本共计698,207个词元,加入“UNK”后词表大小为10,600。字典word_to_int为每个唯一词元分配索引,比如频率最高的句号“.”对应索引0,“the”对应1。字典int_to_word将索引映射回词元,例如索引3对应开引号(“),索引4对应闭引号(”)。

我们打印文本中前20个词元及对应索引:

print(text_tokenized[0:20])
wordidx=[word_to_int[w] for w in text_tokenized]
print([word_to_int[w] for w in text_tokenized[0:20]])

输出:

['he', 'was', 'an', 'old', 'man', 'who', 'fished', 'alone', 'in', 'a', 'skiff', 'in', 'the', 'gulf', 'stream', 'and', 'he', 'had', 'gone', 'eighty']
[8, 16, 98, 110, 67, 85, 6052, 314, 14, 11, 1039, 14, 1, 3193, 507, 5, 8,25, 223, 3125] 

接下来,我们将把索引划分为等长序列,用作训练数据。

12.2.2 创建训练批次

我们将使用长度为128的词元序列作为模型的输入。然后将该序列向右平移一个词元,作为模型的输出。

具体来说,我们创建用于训练的(x, y)对。每个x是一个包含128个索引的序列。选择128是为了在训练速度和模型捕捉长距离依赖能力之间取得平衡。设置过大会降低训练速度,设置过小则难以有效捕捉长距离依赖。

得到序列x后,我们将序列窗口向右滑动一个词元,作为目标y。在序列生成中,将序列右移一个词元作为输出是训练语言模型(包括GPT)时的常用技巧。我们在第8至10章中也用过。下面的代码块创建训练数据:

import torch

seq_len=128                                                 # ①
xys=[]
for n in range(0, len(wordidx)-seq_len-1):
    x = wordidx[n:n+seq_len]                                # ②
    y = wordidx[n+1:n+seq_len+1]                            # ③
    xys.append((torch.tensor(x),(torch.tensor(y))))         # ④

① 设置序列长度为128个索引
② 输入序列x包含训练文本中连续的128个索引
③ 将x右移一个位置,作为输出y
④ 将(x, y)对加入训练数据列表

我们创建了一个列表xys,存储(x, y)对作为训练数据。与之前章节类似,我们将训练数据组织成批次以稳定训练,批次大小设为32:

from torch.utils.data import DataLoader

torch.manual_seed(42)
batch_size=32
loader = DataLoader(xys, batch_size=batch_size, shuffle=True)

x,y=next(iter(loader))
print(x)
print(y)
print(x.shape,y.shape)

打印输出示例如下:

tensor([[   3,  129,    9,  ...,   11,  251,   10],
        [   5,   41,   32,  ...,  995,   52,   23],
        [   6,   25,   11,  ...,   15,    0,   24],
        ...,
        [1254,    0,    4,  ...,   15,    0,    3],
        [  17,    8, 1388,  ...,    0,    8,   16],
        [  55,   20,  156,  ...,   74,   76,   12]])
tensor([[ 129,    9,   23,  ...,  251,   10,    1],
        [  41,   32,   34,  ...,   52,   23,    1],
        [  25,   11,   59,  ...,    0,   24,   25],
        ...,
        [   0,    4,    3,  ...,    0,    3,   93],
        [   8, 1388,    1,  ...,    8,   16, 1437],
        [  20,  156,  970,  ...,   76,   12,   29]])
torch.Size([32, 128]) torch.Size([32, 128])

每个x和y的形状为(32, 128),表示每个训练批次中有32对序列,每个序列含128个索引。

当索引通过nn.Embedding()层时,PyTorch会在嵌入矩阵中查找对应行,返回该索引的嵌入向量,从而避免生成非常大的独热向量。因此,当x通过词嵌入层时,相当于x先被转换成形状为(32, 128, 256)的独热张量。同理,当x通过位置编码层(由nn.Embedding()实现)时,相当于x先被转换成形状为(32, 128, 128)的独热张量。

12.3 构建一个用于生成文本的 GPT 模型

现在我们已经准备好训练数据,接下来将从零开始创建一个 GPT 模型用于生成文本。我们构建的模型架构与第11章中构建的 GPT-2XL 模型类似,但我们只使用3个解码层,而非48个。嵌入维度和词汇表大小也都大幅缩小,正如本章前面所解释的。因此,我们的 GPT 模型参数数量远少于 GPT-2XL。

我们将按照第11章的步骤进行,同时重点说明我们的 GPT 模型与 GPT-2XL 的不同之处以及做出这些修改的原因。

12.3.1 模型超参数

解码器块中的前馈网络采用高斯误差线性单元(GELU)激活函数。GELU 在深度学习任务中,尤其是自然语言处理领域被证明能够提升模型性能,这已成为 GPT 模型中的标准做法。因此,我们定义如下 GELU 类,正如第11章所做:

import torch
from torch import nn
import math

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

class GELU(nn.Module):
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * \
                       (x + 0.044715 * torch.pow(x, 3.0))))

在第11章中,即使在文本生成阶段,我们也没有使用 GPU,因为模型太大,普通 GPU 容易内存不足,无法加载该模型。

然而,本章的模型明显更小,我们将把模型移动到 GPU 以加快训练速度,并在 GPU 上使用模型进行文本生成。

我们使用 Config() 类来包含模型中使用的所有超参数:

class Config():
    def __init__(self):
        self.n_layer = 3
        self.n_head = 4
        self.n_embd = 256
        self.vocab_size = ntokens
        self.block_size = 128
        self.embd_pdrop = 0.1
        self.resid_pdrop = 0.1
        self.attn_pdrop = 0.1

config = Config()

Config() 类中的属性用作 GPT 模型的超参数。我们将 n_layer 属性设为3,表示 GPT 模型有3个解码层。n_head 设为4,意味着在计算因果自注意力时,将查询 Q、键 K、值 V 向量分成4个并行头。n_embd 设为256,表示嵌入维度是256:每个词元用一个256维的向量表示。vocab_size 由词汇表中唯一词元的数量决定,正如上一节所述,我们的训练文本中有10,600个唯一词元。block_size 设为128,表示输入序列最长包含128个词元。我们将 dropout 比例设为0.1,和第11章保持一致。

12.3.2 因果自注意力机制建模

因果自注意力的定义与第11章相同:

import torch.nn.functional as F
class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        self.register_buffer("bias", torch.tril(torch.ones(
                   config.block_size, config.block_size))
             .view(1, 1, config.block_size, config.block_size))
        self.n_head = config.n_head
        self.n_embd = config.n_embd

    def forward(self, x):
        B, T, C = x.size() 
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
        hs = C // self.n_head
        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
        v = v.view(B, T, self.n_head, hs).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v 
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

在计算因果自注意力时,输入嵌入通过三个神经网络得到查询 Q、键 K、值 V。然后将它们各自拆分为四个并行头,在每个头内部计算掩码自注意力。接着,将四个注意力向量合并回一个单一的注意力向量,作为 CausalSelfAttention() 类的输出。

12.3.3 构建 GPT 模型

我们将前馈网络与因果自注意力子层结合,构成解码块。前馈网络为模型注入非线性,没有它,Transformer 只是线性操作的堆叠,难以捕捉复杂的数据关系。此外,前馈网络独立且一致地处理每个位置,转换自注意力机制识别的特征,从而增强模型表达能力。解码块定义如下:

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd),
            c_proj = nn.Linear(4 * config.n_embd, config.n_embd),
            act    = GELU(),
            dropout = nn.Dropout(config.resid_pdrop),
        ))
        m = self.mlp
        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x))))

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlpf(self.ln_2(x))
        return x

我们 GPT 模型中的每个解码块包含两个子层:因果自注意力子层和前馈网络。每个子层都应用了层归一化和残差连接,以提升稳定性和性能。接着,我们将三个解码层堆叠组成 GPT 模型的主体。

构建 GPT 模型的代码如下:

class Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.block_size = config.block_size
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.embd_pdrop),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0,
                  std=0.02 / math.sqrt(2 * config.n_layer))
    
    def forward(self, idx, targets=None):
        b, t = idx.size()
        pos = torch.arange(0, t, dtype=torch.long).unsqueeze(0).to(device)  # ①
        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        return logits

① 将位置编码移动到支持 CUDA 的 GPU(若可用),确保模型所有输入都在同一设备上,否则会报错。

模型的输入是一系列对应词汇表中词元的索引序列。输入先经过词嵌入和位置编码,两者相加形成输入嵌入。然后,输入嵌入依次通过三个解码块。之后,输出经过层归一化,并附加一个线性头,输出维度是 10,600,即词汇表中唯一词元的数量。模型输出的是对应 10,600 个词元的 logits,后续生成文本时,我们会对 logits 应用 softmax 函数,获得词汇表中词元的概率分布。模型设计目的是基于当前词元及之前所有词元预测下一个词元。

接下来,我们实例化之前定义的 Model() 类,创建 GPT 模型:

model = Model(config)
model.to(device)
num = sum(p.numel() for p in model.transformer.parameters())
print("number of parameters: %.2fM" % (num / 1e6,))
print(model)

输出如下:

number of parameters: 5.12M
Model(
  (transformer): ModuleDict(
    (wte): Embedding(10600, 256)
    (wpe): Embedding(128, 256)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-2): 3 x Block(
        (ln_1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (attn): CausalSelfAttention(
          (c_attn): Linear(in_features=256, out_features=768, bias=True)
          (c_proj): Linear(in_features=256, out_features=256, bias=True)
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (mlp): ModuleDict(
          (c_fc): Linear(in_features=256, out_features=1024, bias=True)
          (c_proj): Linear(in_features=1024, out_features=256, bias=True)
          (act): GELU()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=256, out_features=10600, bias=False)
)

我们的 GPT 模型拥有 512 万参数。模型结构与 GPT-2XL 相似。如果对比第11章中的模型输出,可以看到主要差异在超参数上,比如嵌入维度、解码层数、词汇大小等。

12.4 训练 GPT 模型以生成文本

本节中,你将使用本章前面准备好的训练数据批次,训练刚刚构建的 GPT 模型。一个相关的问题是我们应该训练多少个 epoch。训练 epoch 太少可能导致生成文本不连贯,而训练太多又可能造成模型过拟合,生成的文本会与训练语料中的内容完全一致。

因此,我们计划训练模型 40 个 epoch。每隔 10 个 epoch 保存一次模型,并评估哪一版本的模型能生成连贯且没有简单复制训练文本片段的文本。另一种方法是创建验证集,当模型在验证集上的性能收敛时停止训练,类似第2章的做法。

12.4.1 训练 GPT 模型

我们仍然使用 Adam 优化器。由于 GPT 模型本质上是多类别分类问题,我们采用交叉熵损失函数:

lr = 0.0001
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_func = nn.CrossEntropyLoss()

下面展示训练 GPT 模型 40 个 epoch 的代码示例:

model.train()
for i in range(1, 41):
    tloss = 0.
    for idx, (x, y) in enumerate(loader):                   # ①
        x, y = x.to(device), y.to(device)
        output = model(x)
        loss = loss_func(output.view(-1, output.size(-1)), y.view(-1))  # ②
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1)     # ③
        optimizer.step()                                     # ④
        tloss += loss.item()
    print(f'epoch {i} loss {tloss/(idx+1)}')
    if i % 10 == 0:
        torch.save(model.state_dict(), f'files/GPTe{i}.pth')  # ⑤

① 遍历所有训练数据批次
② 计算模型预测值与实际输出的损失
③ 将梯度范数裁剪至最大为1,防止梯度爆炸
④ 调整模型参数以最小化损失
⑤ 每训练10个epoch保存一次模型

训练时,我们将批次中的所有输入序列 x 输入模型获得预测结果,然后与对应的输出序列 y 比较,计算交叉熵损失。接着调整模型参数以减少损失。注意,我们进行了梯度范数裁剪(最大为1),以防止梯度爆炸问题。

梯度范数裁剪

梯度范数裁剪是神经网络训练中防止梯度爆炸的技术。当损失函数对模型参数的梯度过大时,会导致训练不稳定、模型表现差。通过范数裁剪,将超过阈值的梯度按比例缩小,保证梯度不过大,从而稳定训练并加快收敛。

如果你有支持 CUDA 的 GPU,训练大约需要几个小时;如果用 CPU,可能需要一整天。训练完成后,你的电脑中会保存四个模型文件:GPTe10.pth、GPTe20.pth、GPTe30.pth 和 GPTe40.pth。你也可以从我的网站下载这些训练好的模型:gattonweb.uky.edu/faculty/liu…

12.4.2 生成文本的函数

现在我们有多个训练版本的模型,可以开始文本生成,并比较各版本表现。我们判断哪个版本表现最佳,并用它来生成文本。

和 GPT-2XL 类似,文本生成从向模型输入一段索引序列(表示词元)开始。模型预测下一个词元的索引,并将其添加到输入序列末尾,形成新的序列。新的序列继续作为输入,循环预测直到生成指定数量的新词元。

为此,我们定义了 sample() 函数。该函数以当前文本的索引序列作为输入,迭代预测并添加新的索引,直到达到 max_new_tokens 限定的生成长度。实现如下:

def sample(idx, weights, max_new_tokens, temperature=1.0, top_k=None):
    model.eval()
    model.load_state_dict(torch.load(weights, map_location=device))   # ①
    original_length = len(idx[0])
    for _ in range(max_new_tokens):                                    # ②
        if idx.size(1) <= config.block_size:
            idx_cond = idx
        else:
            idx_cond = idx[:, -config.block_size:]
        logits = model(idx_cond.to(device))                            # ③
        logits = logits[:, -1, :] / temperature
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = -float('Inf')
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, idx_next.cpu()), dim=1)                  # ④
    return idx[:, original_length:]                                    # ⑤

① 加载训练好的模型权重
② 生成固定数量的新索引
③ 利用模型进行预测
④ 将预测的新索引拼接到序列末尾
⑤ 返回仅包含新生成索引的序列

sample() 函数的参数 weights 是已保存的模型权重文件路径。与第11章的 sample() 函数不同,这里只返回新生成的索引,不包含原始输入的索引。这样可以保证当提示中含未知词元时,最终输出仍包含原始提示,而不会被替换为 “UNK”。


接着定义 generate() 函数,实现基于提示词生成文本的功能:

UNK = word_to_int["UNK"]

def generate(prompt, weights, max_new_tokens, temperature=1.0, top_k=None):
    assert len(prompt) > 0, "prompt must contain at least one token"   # ①
    text = prompt.lower().replace("\n", " ")
    for x in punctuations:
        text = text.replace(f"{x}", f" {x} ")
    text_tokenized = text.split()
    idx = [word_to_int.get(w, UNK) for w in text_tokenized]             # ②
    idx = torch.LongTensor(idx).unsqueeze(0)
    idx = sample(idx, weights, max_new_tokens, temperature=temperature, top_k=top_k)  # ③
    tokens = [int_to_word[i] for i in idx.squeeze().numpy()]             # ④
    text = " ".join(tokens)
    for x in '''").:;!?,-''''':
        text = text.replace(f" {x}", f"{x}")
    for x in '''"(-''''':
        text = text.replace(f"{x} ", f"{x}")
    return prompt + " " + text

① 确保提示词不为空,否则报错
② 将提示词转为索引序列
③ 调用 sample() 函数生成新索引
④ 将生成的索引序列转换回文本

generate() 函数允许你指定要使用的训练模型版本(权重文件)。例如,你可以指定 'files/GPTe10.pth' 作为参数 weights 的值。函数先将提示词转换成索引序列,传入模型预测下一个索引,生成指定数量的新索引后,再将完整索引序列转换为文本输出。

12.4.3 使用不同训练版本的模型进行文本生成

接下来,我们将使用训练得到的不同版本模型来生成文本进行实验。

我们可以使用未知词元 “UNK” 作为提示词,进行无条件文本生成。在我们的场景下,这种方式特别有用,因为我们想检验生成的文本是否直接复制了训练文本。虽然一个与训练文本截然不同的唯一提示词不太可能生成训练文本中的片段,但无条件生成的文本更可能来源于训练文本。

首先,我们使用训练了20个epoch后的模型进行无条件文本生成:

prompt = "UNK"
for i in range(10):
    torch.manual_seed(i)
    print(generate(prompt, 'files/GPTe20.pth', max_new_tokens=20)[4:])

输出示例:

way." "kümmel," i said. "it's the way to talk about it
--------------------------------------------------
," robert jordan said. "but do not realize how far he is ruined." "pero
--------------------------------------------------
in the fog, robert jordan thought. and then, without looking at last, so 
good, he 
--------------------------------------------------
pot of yellow rice and fish and the boy loved him. "no," the boy said.
--------------------------------------------------
the line now. it's wonderful." "he's crazy about the brave."
--------------------------------------------------
candle to us. "and if the maria kisses thee again i will commence kissing 
thee myself. it 
--------------------------------------------------
?" "do you have to for the moment." robert jordan got up and walked away in
--------------------------------------------------
. a uniform for my father, he thought. i'll say them later. just then he
--------------------------------------------------
and more practical to read and relax in the evening; of all the things he 
had enjoyed the next 
--------------------------------------------------
in bed and rolled himself a cigarette. when he gave them a log to a second 
grenade. " 
--------------------------------------------------

我们将提示词设为 “UNK”,调用 generate() 函数进行无条件生成,每次生成20个新词元,重复10次。通过 manual_seed() 方法固定随机种子,以保证结果可复现。如你所见,这10段短文都符合语法,且风格类似海明威小说。例如第一段中的 “kummel” 是《永别了,武器》中频繁出现的一种利口酒。同时,上述10段文本均未直接复制训练语料。

接着,我们使用训练了40个epoch的模型进行无条件文本生成,观察结果:

prompt = "UNK"
for i in range(10):
    torch.manual_seed(i)
    print(generate(prompt, 'files/GPTe40.pth', max_new_tokens=20)[4:])

输出示例:

way." "kümmel, and i will enjoy the killing. they must have brought me a spit
--------------------------------------------------
," robert jordan said. "but do not tell me that he saw anything." "not
--------------------------------------------------
in the first time he had bit the ear like that and held onto it, his neck 
and jaws
--------------------------------------------------
pot of yellow rice with fish. it was cold now in the head and he could not 
see the
--------------------------------------------------
the line of his mouth. he thought." "the laughing hurt him." "i can
--------------------------------------------------
candle made? that was the worst day of my life until one other day." "don'
--------------------------------------------------
?" "do you have to for the moment." robert jordan took the glasses and 
opened the
--------------------------------------------------
. that's what they don't marry." i reached for her hand. "don
--------------------------------------------------
and more grenades. that was the last for next year. it crossed the river 
away from the front
--------------------------------------------------
in a revolutionary army," robert jordan said. "that's really nonsense. it's
--------------------------------------------------

这10段文本依旧语法正确,风格也像海明威小说。然而仔细观察,第八段的很大一部分是直接复制自小说《永别了,武器》的片段:“they don't marry." i reached for her hand. "don”。你可以在之前保存的 ThreeNovels.txt 文件中搜索验证。

练习 12.1
使用训练了10个epoch的模型,无条件生成50个新词元的文本,设置随机种子为42,温度和 top-K 采样保持默认。检查生成文本是否语法正确,是否存在直接复制训练文本的情况。

另一种做法是使用训练语料中不存在的唯一提示词来生成文本。例如,使用“the old man saw the shark near the”作为提示词,请 generate() 函数为其生成20个新词元,重复10次:

prompt = "the old man saw the shark near the"
for i in range(10):
    torch.manual_seed(i)
    print(generate(prompt, 'files/GPTe40.pth', max_new_tokens=20))
    print("-" * 50)

输出示例:

the old man saw the shark near the old man's head with his tail out and the old man hit him squarely in the center of
--------------------------------------------------
the old man saw the shark near the boat with one hand. he had no feeling of
the morning but he started to pull on it gently
--------------------------------------------------
the old man saw the shark near the old man's head. then he went back to 
another man in and leaned over and dipped the
--------------------------------------------------
the old man saw the shark near the fish now, and the old man was asleep in 
the water as he rowed he was out of the
--------------------------------------------------
the old man saw the shark near the boat. it was a nice-boat. he saw the old
 man's head and he started
--------------------------------------------------
the old man saw the shark near the boat to see him clearly and he was 
afraid that he was higher out of the water and the old
--------------------------------------------------
the old man saw the shark near the old man's head and then, with his tail 
lashing and his jaws clicking, the shark plowed
--------------------------------------------------
the old man saw the shark near the line with his tail which was not sweet 
smelling it. the old man knew that the fish was coming
--------------------------------------------------
the old man saw the shark near the fish with his jaws hooked and the old 
man stabbed him in his left eye. the shark still hung
--------------------------------------------------
the old man saw the shark near the fish and he started to shake his head 
again. the old man was asleep in the stern and he
--------------------------------------------------

生成文本语法正确,连贯,风格和海明威的《老人与海》相似。由于使用了训练了40个epoch的模型,生成文本更有可能直接复制训练数据中的内容。但使用唯一提示词则能降低复制的概率。

通过设置温度参数和使用 top-K 采样,我们还能进一步控制生成文本的多样性。以下示例使用训练了20个epoch的模型,温度0.9,top-K取50,提示词同样为“the old man saw the shark near the”:

prompt = "the old man saw the shark near the"
for i in range(10):
    torch.manual_seed(i)
    print(generate(prompt, 'files/GPTe20.pth', max_new_tokens=20,
                   temperature=0.9, top_k=50))
    print("-" * 50)

输出示例:

The old man saw the shark near the boat. then he swung the great fish that 
was more comfortable in the sun. the old man could
--------------------------------------------------
the old man saw the shark near the boat with one hand. he wore his overcoat
 and carried the submachine gun muzzle down, carrying it in
--------------------------------------------------
the old man saw the shark near the boat with its long dip sharply and the 
old man stabbed him in the morning. he could not see
--------------------------------------------------
the old man saw the shark near the fish that was now heavy and long and 
grave he had taken no part in. he was still under
--------------------------------------------------
the old man saw the shark near the boat. it was a nice little light. then 
he rowed out and the old man was asleep over
--------------------------------------------------
the old man saw the shark near the boat to come. "old man's shack and i'll 
fill the water with him in
--------------------------------------------------
the old man saw the shark near the boat and then rose with his lines close 
him over the stern. "no," the old man
--------------------------------------------------
the old man saw the shark near the line with his tail go under. he was 
cutting away onto the bow and his face was just a
--------------------------------------------------
the old man saw the shark near the fish with his tail that he swung him in.
 the shark's head was out of water and
--------------------------------------------------
the old man saw the shark near the boat and he started to cry. he could 
almost have them come down and whipped him in again.
--------------------------------------------------

由于使用的是训练了20个epoch的模型,输出文本连贯度较低,偶尔有语法错误,例如第三段中的 “with its long dip sharply” 语法不正确,但生成文本直接复制训练文本的风险较小。

练习 12.2

使用训练了40个epoch的模型,以“the old man saw the shark near the”为提示词,无条件生成50个新词元文本。设置随机种子为42,温度为0.95,top_k为100。检查生成的文本是否语法正确,是否存在直接复制训练文本的情况。

本章中,你学习了如何从头构建和训练一个 GPT 风格的 Transformer 模型。你创建了一个仅有512万参数的简化版 GPT-2 模型,使用了海明威的三部小说作为训练语料,成功训练了模型,并生成了风格连贯、符合海明威写作风格的文本。

总结

  • GPT模型生成文本的风格在很大程度上受训练数据的影响。为了实现有效的文本生成,训练材料中应在文本长度和多样性之间保持平衡。训练数据集应足够庞大,以便模型能够准确学习并模仿特定的写作风格。然而,如果数据集缺乏多样性,模型可能会直接复制训练文本中的片段。相反,过长的训练数据集又会导致训练所需的计算资源过大。

  • 在GPT模型中选择合适的超参数对成功训练模型和生成文本至关重要。超参数设置过大可能导致参数数量过多,从而延长训练时间并导致模型过拟合。超参数设置过小则可能影响模型的学习能力,难以有效捕捉训练数据中的写作风格,导致生成文本缺乏连贯性。

  • 适当的训练轮数对于文本生成也十分关键。训练轮数过少可能导致生成文本缺乏连贯性,而训练轮数过多则可能造成模型过拟合,生成与训练文本完全相同的内容。