Here's how you can build and train GPT-2 from scratch using PyTorch
原文链接:differ.blog/p/here-s-ho...
原文作者:Amit Kharel
译者:菜小鸟魔王
如若你已经厌倦了使用 ChatGPT,并对自己动手构建属于自己的语言模型充满好奇,那么恭喜你,你来对地方了!今天,我们将从零开始,动手构建 GPT-2(OpenAI 推出的一款较为强大的语言模型),这款模型能通过预测文本序列中的下一个单词(the next word in a sequence),生成接近类似人类创造的文本内容。
想要更透彻地理解 GPT-2 的理论框架和设计结构,我强烈推荐阅读 Jay Alammar 的《The Illustrated GPT-2》^[1]^。这篇文章对 GPT-2 的核心原理进行了非常形象直观的介绍。在讲解过程中,我将借助该文中的部分插图,帮助大家更好地理解关键概念。
我已竭力简化整个讲解过程,以期各位读者都能够理解本文内容。
Resources
接下来,本文将手把手教你构建一个简易版的 GPT-2 模型,并利用泰勒·斯威夫特(Taylor Swift)和艾德·希兰(Ed Sheeran)的一些歌曲进行训练。让我们拭目以待,看看最终构建出的模型会给我们带来哪些惊喜!:-)
本教程涉及的数据集和源代码,都会发布于 GitHub^[2]^。
Building GPT-2 Architecture
我们将按部就班地构建本项目,从一个最基础的模型(bare-bone model)起步,逐步完善,并参照 GPT-2^[3]^ 的架构设计增加模型层。
以下是本模型的构建指南:
- 构建个性化分词器(Building a custom Tokenizer)
- 构建数据加载器(Building a Data Loader)
- 从简单的模型入手,为后续复杂模型的搭建奠定基础(Train a simple language model)
- 实现 GPT-2 架构(第二部分)^[4]^
该项目被拆分为两部分,第一部分主要介绍构建语言模型的入门知识,而第二部分则直接切入构建 GPT-2 的实战过程。强烈建议各位读者边浏览本文边构建模型,亲自动手实践。
Final Model:
Final Model output:
Your summer has a matter likely you trying I wish you would call Oh-oh, I'll be a lot of everyoneI just walked You're sorry"Your standing in love out, And something would wait forever bring 'Don't you think about the storyIf you're perfectly I want your beautiful You had sneak for you make me This ain't think that it wanted you this enough for lonely thing It's a duchess and I did nothin' home was no head Oh, but you left me Was all the less pair of the applause Honey, he owns me now But've looks for us?" If I see you'll be alright You understand, a out of theWait for me I can't call Everything Oh, no words don't read about me You should've been so You're doing what you so tired, If you, you got perfect fall
如果这首歌让你有一点点心动,那就别犹豫,让我们一起动手开始吧!
构建个性化分词器(Building a custom Tokenizer)
语言模型处理文本的方式与我们大相径庭 ------ 它们不会直接"阅读"文字,而是将文本处理为文本词元(tokens)(译者注:机器可读的数字形式)。因此,我们的首要任务便是导入所需数据,并着手构建专属的字符级分词器(character level Tokenizer)。
py
data_dir = "data.txt"
text = open(data_dir, 'r').read() # load all the data as simple string
# Get all unique characters in the text as vocabulary
chars = list(set(text))
vocab_size = len(chars)
Example:
上方的输出结果是:初始化阶段,文本数据中识别出的所有不同字符。字符级别的分词(Character tokenization)其奥秘就在于,其词汇表(vocabulary)内每个字符都有专属的索引位置,将这些索引与输入文本中的相应字符一一对应起来。
py
# build the character level tokenizer
chr_to_idx = {c:i for i, c in enumerate(chars)}
idx_to_chr = {i:c for i, c in enumerate(chars)}
def encode(input_text: str) -> list[int]:
return [chr_to_idx[t] for t in input_text]
def decode(input_tokens: list[int]) -> str:
return "".join([idx_to_chr[i] for i in input_tokens])
Example:
将文本数据转换为 tokens:
- 安装相关 Python 库:
pip install torch
- Code:
py
import torch
# use cpu or gpu based on your system
device = "cpu"
if torch.cuda.is_available():
device = "cuda"
# convert our text data into tokenized tensor
data = torch.tensor(encode(text), dtyppe=torch.long, device=device)
现在我们已经处理好了文本数据,将其转变为了张量形式;在这个过程中,文本里的每一个字符都被精准地转化为了对应的 tokens。
到目前为止,代码如下:
py
import torch
data_dir = "data.txt"
text = open(data_dir, 'r').read() # load all the data as simple string
# Get all unique characters in the text as vocabulary
chars = list(set(text))
vocab_size = len(chars)
# build the character level tokenizer
chr_to_idx = {c:i for i, c in enumerate(chars)}
idx_to_chr = {i:c for i, c in enumerate(chars)}
def encode(input_text: str) -> list[int]:
return [chr_to_idx[t] for t in input_text]
def decode(input_tokens: list[int]) -> str:
return "".join([idx_to_chr[i] for i in input_tokens])
# convert our text data into tokenized tensor
data = torch.tensor(encode(text), dtyppe=torch.long, device=device)
在着手构建模型前,有一项关键任务摆在我们面前:我们需要明确数据是如何被送入模型进行训练的,同时,也要搞清楚数据的 dimensions(译者注:数据集的结构和形状) 和 batch size (译者注:把数据集分几次"喂"给模型)。
按照如下方式定义数据加载器(data loader):
py
train_batch_size = 16 # training batch size
eval_batch_size = 8 # evaluation batch size
context_length = 256 # number of tokens processed in a single batch
train_split = 0.8 # percentage of data to use from total data for training
# split data into trian and eval
n_data = len(data)
train_data = data[:int(n_data * train_split)]
eval_data = data[int(n_data * train_split):]
class DataLoader:
def __init__(self, tokens, batch_size, context_length) -> None:
self.tokens = tokens
self.batch_size = batch_size
self.context_length = context_length
self.current_position = 0
def get_batch(self) -> torch.tensor:
b, c = self.batch_size, self.context_length
start_pos = self.current_position
end_pos = self.current_position + b * c + 1
# if the batch exceeds total length, get the data till last token
# and take remaining from starting token to avoid always excluding some data
add_data = -1 # n, if length exceeds and we need `n` additional tokens from start
if end_pos > len(self.tokens):
add_data = end_pos - len(self.tokens) - 1
end_pos = len(self.tokens) - 1
d = self.tokens[start_pos:end_pos]
if add_data != -1:
d = torch.cat([d, self.tokens[:add_data]])
x = (d[:-1]).view(b, c) # inputs
y = (d[1:]).view(b, c) # targets
self.current_position += b * c # set the next position
return x, y
train_loader = DataLoader(train_data, train_batch_size, context_length)
eval_loader = DataLoader(eval_data, eval_batch_size, context_length)
Example:
至此,我们已经成功打造了一款可用于训练和评估阶段的数据加载器。这款加载器内含一个 get_batch
函数,能够生成大小为 batch_size * context_length
的数据集子集,以供模型训练之需。
或许你会感到好奇,为何 x
覆盖的区间是从数据序列的起始点到终点,而y
的起始点是x
的第二个元素,y
的终点则是x
序列末尾元素的下一个元素呢?这是因为模型的核心任务是基于过往的文本序列信息,预测其后续文本内容。正因如此,y
中特意预留了一个 token 位置,目的就是为了在掌握了 x
序列中前 n
个 tokens 的基础上,预测出第 (n+1)
个 token 的具体内容。倘若这一概念初听之下有些晦涩难懂,不妨参考下方的图解,相信会对你理解这部分内容有所帮助。
Figure 2: GPT-2 Input & Output flow from "The Illustrated GPT-2" by Jay Alammar.
构建数据加载器(Building a Data Loader)
此刻,万事俱备,只欠东风 ------ 我们可以利用刚刚加载的数据,构建并训练一个简易语言模型咯✌️。
在本环节,我们将秉持着"大道至简"这一理念,构建一个简单的 Bi-Gram 模型。就是基于上一个 token ,预测出下一个 token 。你或许会看到,我们这次只使用了嵌入层(Embedding layer),而有意忽略了更为复杂的模型解码部分的核心组件(main decoder block)。
诶!说起嵌入层(Embedding layer),那可堪称模型中的"魔法师🧙♂️",能够捕捉并量化我们词汇库中每个字符的n = d_model
种独特特征。具体而言,该层巧妙地运用 token 的索引,或者说,字符在词汇表中的位置编码,来提取出与之对应的内在属性。
你可能会惊叹,即便仅凭"小小"嵌入层,模型就能表现的如此出色。a piece of cake 我们后续还会循序渐进,通过不断叠加新的模型层,日益完善这一模型。所以,请系好安全带,坐稳(嘀嘀)🚗!
从简单的模型入手,为后续复杂模型的搭建奠定基础(Train a simple language model)
Initialization:
这一环节主要用于确定嵌入的大小(size of embeddings):d_model = vocab_size
当前情况下,嵌入维度(embedding dimension)d_model
的大小为vocab_size
,这是基于这样一个需求:模型的最终输出需要对应到词汇表(vocab)中每个字符的 logits (译者注:模型原始的、未经过任何激活函数变换的输出值),从而计算出各个字符出现的概率。
后续的模型升级方案已规划完毕 ------ 计划先加入一个线性层(Linear layer which),在个模型层将扮演"桥梁"的角色,将d_model
再度映射至vocab_size
,这样一来,我们就能够自由设置embedding_dimension
,以满足更加精细化的模型需求。
Model:
py
import torch.nn as nn
import torch.nn.functional as F
class GPT(nn.Module):
def __init__(self, vocab_size, d_model):
super().__init__()
self.wte = nn.Embedding(vocab_size, d_model) # word token embeddings
def forward(self, inputs, targets = None):
logits = self.wte(inputs) # dim -> batch_size, sequence_length, d_model
loss = None
if targets != None:
batch_size, sequence_length, d_model = logits.shape
# to calculate loss for all token embeddings in a batch
# kind of a requirement for cross_entropy
logits = logits.view(batch_size * sequence_length, d_model)
targets = targets.view(batch_size * sequence_length)
loss = F.cross_entropy(logits, targets)
return logits, loss
def generate(self, inputs, max_new_tokens):
# this will store the model outputs along with the initial input sequence
# make a copy so that it doesn't interfare with model
for _ in range(max_new_tokens):
# we only pass targets on training to calculate loss
logits, _ = self(inputs)
# for all the batches, get the embeds for last predicted sequence
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=1)
# get the probable token based on the input probs
idx_next = torch.multinomial(probs, num_samples=1)
inputs = torch.cat([inputs, idx_next], dim=1)
# as the inputs has all model outputs + initial inputs, we can use it as final output
return inputs
m = GPT(vocab_size=vocab_size, d_model=d_model).to(device)
现在,我们只使用单个嵌入层(Embedding)和 softmax 函数就成功构建了这一模型,可专门用于生成 tokens。现在,我们来观察一下,当该模型接收到输入的文本内容时,会有怎样的行为表现。
😄确实非常有趣啊!不过,还有最后一步路要走。
接下来,我们要训练模型,让它掌握一些 characters (译者注:模型需要识别和处理的文字或符号)的相关知识。先配置优化器:现阶段,我们暂时采用 AdamW 这种简单优化器(optimizer),设定的学习率(learning rate)是 0.001。至于如何进一步有效优化模型效果,我们留待后续章节再做讨论。
py
lr = 1e-3
optim = torch.optim.AdamW(m.parameters(), lr=lr)
Below is a very simple training loop.
epochs = 5000
eval_steps = 1000 # perform evaluation in every n steps
for ep in range(epochs):
xb, yb = train_loader.get_batch()
logits, loss = m(xb, yb)
optim.zero_grad(set_to_none=True)
loss.backward()
optim.step()
if ep % eval_steps == 0 or ep == epochs-1:
m.eval()
with torch.no_grad():
xvb, yvb = eval_loader.get_batch()
_, e_loss = m(xvb, yvb)
print(f"Epoch: {ep}tlr: {lr}ttrain_loss: {loss}teval_loss: {e_loss}")
m.train() # back to training mode
Let's run:
我们确实取得了相当不错的结果,损失值(loss result)表现良好。但是似乎没有达到预期目标,从图中可以看出,在前 2000 次迭代中显著下降,之后的幅度就不那么明显了。这主要是因为模型当前的"智力水平"有限(具体来说,神经网络的层数不足),它所做的仅仅是对比不同字符的嵌入表征。
目前的输出结果如下所示:
😮哇哦!虽说目前的模型效果不尽如人意,但肯定比没有经过任何训练的初始版本有所改进(这一点毋庸置疑)。模型已经开始学会歌曲的结构布局,包括歌词行文的规律,真令人刮目相看。
References
Automatic Arabic Poem Generation with GPT-2 --- Scientific Figure on ResearchGate. Available from: https://www.researchgate.net/figure/GPT-2-architecture-Heilbron-et-al-2019_fig1_358654229
Alammar, J (2018). The Illustrated GPT-2 [Blog post]. Retrieved from jalammar.github.io/illustrated...
文中链接
[1]jalammar.github.io/illustrated... [2]github.com/ajeetkharel... [3]cdn.openai.com/better-lang... [4]medium.com/@mramitkhar...