牙牙学语：使用 PyTorch 从零开始构建和训练 GPT-2

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]^ 的架构设计增加模型层。

以下是本模型的构建指南：

1. 构建个性化分词器（Building a custom Tokenizer）
2. 构建数据加载器（Building a Data Loader）
3. 从简单的模型入手，为后续复杂模型的搭建奠定基础（Train a simple language model）
4. 实现 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）。

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）内每个字符都有专属的索引位置，将这些索引与输入文本中的相应字符一一对应起来。

# 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:

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。

到目前为止，代码如下:

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）：

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:

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。至于如何进一步有效优化模型效果，我们留待后续章节再做讨论。

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...