【指南】基于PyTorch从零打造GPT模型

2024年11月11日 由 alex 发表 70 0

今天,我们将暂时离开我们的Vision Transformer系列,转而讨论构建一个生成式预训练变换器(GPT)的基本变体。


更准确地说,我们将构建一个自回归(双词)模型,即每次都会根据所有之前的词元来生成一个新的词元。自回归模型通常会根据已出现的词元顺序地生成新的词元(字符或单词)。例如,在句子“I like to eat”中,<I> <like> <to> <eat>之后的词元可能是<ice-cream>、<cookies>等。


这个不完全可预测或随机的项可以非常松散地[1]与我们在“我喜欢吃”示例中预测的下一个标记相关,我们通过让模型随机选择下一个标记进行预测(即<ice-cream>、<cookies>等)来帮助模型在选择时不那么确定,我们将在本文后面了解这一点。


由于我们正在实现一个非常基本的自回归模型,我们将从零开始,使用一个用于生成威廉·莎士比亚风格文本的数据集。让我们开始吧!


内容

  1. 加载数据——创建数据批量加载器和数据分割。
  2. BigramLanguageModel — 编码语言模型
  3. 训练——训练模型并生成文本。


注意:本文中的代码遵循了Andrej Karparthy 制作的有关 GPT 的视频。他的视频实际上是我对注意力机制的首次实现,在此基础上,我遵循了有关卷积注意力、移位窗口等各种其他架构和论文。


如果你已经看过,那么除了代码的细微变化外,本文中没有太多不同,因此你可以将此内容作为快速修订版。如果你还没有看过……让我们直接开始吧。


加载数据


# Importing torch specific modules
import torch
import torch.nn as nn
from torch.nn import functional as F
# We start by downloading our shakespeare txt file (stored with the name input.txt)
! wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt */
# reading txt file (encode decode)
text = open('input.txt', 'r',).read()
vocab = sorted(list(set(text)))
encode = lambda s: [vocab.index(c) for c in s]
decode = lambda l: [vocab[c] for c in l]


我们将不使用外部分词器,而是实现自定义的lambda函数来对我们的数据进行字符级别的分词。


ids = encode("I like to eat")"I like to eat")
txt = decode(ids)
print(f"ids: {ids}")
print(f"txt: {txt}")
print(f"".join(txt))
# Output:-
ids: [21, 1, 50, 47, 49, 43, 1, 58, 53, 1, 43, 39, 58]
txt: ['I', ' ', 'l', 'i', 'k', 'e', ' ', 't', 'o', ' ', 'e', 'a', 't']
I like to eat


将数据按照90/10的比例进行拆分。


x = int(0.9*len(text)) # text is a big string with our entire data0.9*len(text)) # text is a big string with our entire data
text = torch.tensor(encode(text), dtype=torch.long)
train, val = text[:x], text[x:]


由于我们是在字符级别进行分词,因此生成也将是在字符级别进行。在这种情况下,明智的做法是从语料库中创建包含随机句子的批次,以输入到我们的模型中进行训练。


batch_size = 32 # batch_size - is how many independent sequences will we process in parallel?32 # batch_size - is how many independent sequences will we process in parallel?
block_size = 8 # block_size = is the maximum context length for predictions
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train if split == 'train' elseval  
    ix = torch.randint(len(data) - block_size, (batch_size,)) 
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x.to(device), y.to(device)
xb, yb = get_batch('train')


以下是创建批次的方法……


我们希望从语料库中获取块大小为(8)的随机句子,因此我们生成批次大小(32)的索引(ix),并且对于每个索引,我们取接下来的8个词元ID,并将它们堆叠起来形成批次(ix)中的每个索引对应的序列。然而,我们的目标(y)是通过比x多一个索引来生成的(即i+1, i + block_size + 1),因为我们需要预测序列中的下一个词元。


示例:


ix  = [33]
for i in ix:
    print(train[i:i+18])
    print(train[i+3:i+18+3]) # I've chosen +3 over +1 only for the sake of example
for i in ix:
    print("".join(decode(train[i:i+18])).replace("\n", ""))
    print("".join(decode(train[i+3:i+18+3])).replace("\n", ""))
# Output:-
tensor([39, 52, 63,  1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1])
tensor([ 1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1])
any further, hear 
 further, hear me 


BigramLanguageModel(双词语言模型)

双词(Bigram)是一种n-gram,其中n=2,它表示文本中两个连续的词汇单元(如单词或字符)的序列。


对于“i like to eat”这句话:

  • 1-gram(Unigram,单词):["i", "like", "to", "eat"]
  • 2-gram(Bigram,双词):["i like", "like to", "to eat"]
  • 3-gram(Trigram,三词):["i like to", "like to eat"]


由于我们正在执行自回归任务,因此我们需要按照上面代码块中所示的方式,将我们的数据加载为双词格式。


现在,让我们进入文章的核心部分——多头注意力(Multi-Head Attention)。


7


我们看到,GPT借鉴了《Attention is all you need》论文中提出的Transformer架构。然而,它的不同之处在于仅堆叠了解码器部分中的多头注意力(Multi-Head Attention)。


双词语言模型(Bigram Language Model)。


class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, embed_size)
        self.possitional_embedding = nn.Embedding(block_size, embed_size)
        self.linear = nn.Linear(embed_size, vocab_size)
        self.block = nn.Sequential(*[Block(embed_size, num_head) for _ in range(num_layers)])
        self.layer_norm = nn.LayerNorm(embed_size)
    def forward(self, idx, targets=None):
        B, T = idx.shape
        # idx and targets are both (B,T) tensor of integers
        logits = self.token_embedding_table(idx) # (B,T,C)
        ps = self.possitional_embedding(torch.arange(T, device=device))
        x = logits + ps    #(B, T, C)
        logits = self.block(x)     #(B, T, c)
        logits = self.linear(self.layer_norm(logits)) # This suppose to map between embed_size and Vocab_size 
        B, T, C = logits.shape
        logits = logits.view(B*T, C)
        targets = targets.view(B*T)
        loss = F.cross_entropy(logits, targets)
        return logits, loss


在这里,input_idx 是我们之前生成的批次,形状为 (B, T),其中 T 是块大小或词元长度。
前向传播首先为每个词元生成形状为 (B, T, C) 的嵌入。如上图所示,我们需要在词元嵌入中添加位置嵌入。我们为输入(idx)创建嵌入,以便我们可以在固定的嵌入维度中表示词元所持有的信息,但这并不提供关于词元位置(即位置信息)的任何信息,这就是为什么我们必须额外添加位置嵌入来确保模型具有词元位置的上下文信息。
在PyTorch中,nn.Embedding 是一个用于将离散、分类值(如词索引)映射到连续、密集向量的层。该层接受整数索引作为输入,每个索引表示一个唯一的分类项(例如,一个词、词元或其他分类数据)。在内部,nn.Embedding 维护一个形状为 num_embeddings, embedding_dim 的嵌入矩阵,以便为每个词元创建密集表示。由于我们正在实现一个简化版的GPT,我们直接使用 nn.Embeddings 来生成位置嵌入,而不是采用其他标准方法。
接下来,过程就很直接了……我们有我们的块(一堆解码器模块),最后我们生成与输入形状相同但每个词元都包含其前面所有词元信息的新注意力矩阵。
最后,我们应用层归一化(稳定训练的常用方法),然后将其传递给一个线性层,以将嵌入 C 映射到我们的词汇表维度。词汇表维度简单地是我们 input.txt 中所有唯一字符的数量。定义我们预测准确性的一种通用方法是将块(注意力模块)的输出与目标索引进行比较。
输出逻辑应该只是词汇表大小 V 上的概率分布,预测序列中的下一个词元(目标)。因此,使用交叉熵损失来生成一个损失值,以确定我们的输出与目标词元序列的接近程度。


现在我们已经涵盖了双词实现的细节。接下来,让我们看看块是如何使用多头注意力(MHA)来创建注意力矩阵的。


多头注意力

使用多头注意力的原因是我们可以直接将输入 (B, T, embedding_size) 传递给一个注意力块,但一种更快的方法是,我们不是直接生成 Q、K、V 并计算维度为 embedding_size 的注意力权重,而是创建注意力模块的多个部分,分别计算注意力权重,然后在最后将它们连接起来。


class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.head_size = head_size
        self.key = nn.Linear(embed_size, head_size, bias=False)
        self.query = nn.Linear(embed_size, head_size, bias=False)
        self.value = nn.Linear(embed_size, head_size, bias=False)
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        v = self.value(x)
        wei = q@k.transpose(2, 1)/self.head_size**0.5


8


根据上面的逻辑,我们创建了一个单一的注意力头。现在让我们来深入理解一下。


在添加位置嵌入后,我们有一个维度为(批量大小,词元长度,嵌入维度)的输入。在这里,输入中的每个词元都由一个嵌入维度(64)表示。但是,没有任何词元包含关于其前面所有词元的信息。


为了创建包含这种信息的丰富嵌入,我们使用注意力机制,通过生成键(Key)、查询(Query)和值(Value)向量来实现。


Head 类中的注意力机制旨在帮助模型在生成输出时关注输入序列的不同部分,这在语言建模等任务中特别有用。


键、查询和值投影来源于查询序列中每个词元上下文相关信息的概念。每个词元都由一个向量(x)表示,通过将其线性转换为单独的键、查询和值向量,我们可以计算出序列中哪些词元应该相互关注。


当查询(q)与键(k)进行点积运算时,结果(wei)告诉我们每个词元与其他词元之间的“相关性”或“注意力”得分。较高的得分意味着在某个上下文中,一个词元对另一个词元更具相关性或“重要性”。缩放因子 1 / sqrt(head_size) 用于防止这些得分过大,这可能会使 softmax 分布过于尖锐,从而更难优化。


class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.head_size = head_size
        self.key = nn.Linear(embed_size, head_size, bias=False)
        self.query = nn.Linear(embed_size, head_size, bias=False)
        self.value = nn.Linear(embed_size, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        v = self.value(x)
        wei = q@k.transpose(2, 1)/self.head_size**0.5
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=2)    # (B , block_size, block_size)
        wei = self.dropout(wei)
        out = wei@v
        return out


因果掩码(tril)的应用是为了确保每个词元只能“看到”自己以及之前的词元。这对于自回归任务(如文本生成)至关重要,因为在预测下一个词元时,模型不应该提前查看未来的词元。通过将不相关位置设置为负无穷大(通过masked_fill实现),这些位置在softmax之后将变为零,因此它们不会对最终的注意力计算产生贡献。这样做是为了防止模型通过查看未来词元来“作弊”,因为我们的目标是预测未来的词元。最后,我们对权重和值向量进行点积运算,并返回输出结果。


你可以查看这个示例输出,以更好地理解所发生的变换:


q = torch.randint(10, (1, 3, 3))10, (1, 3, 3))
v = torch.randint(10, (1, 3, 3))
print("Query:\n",q)
print("Value:\n",v)
wei = q@v.transpose(2, 1)/3**0.5
print("weights:\n", wei)
tril = torch.tril(torch.ones(3, 3))
print("Triangular Metrics:\n",tril)
wei = wei.masked_fill(tril == 0, float('-inf'))
print("Masked Weights\n", wei)
print("Softmax ( e^-inf = 0 )\n", F.softmax(wei, dim=2))
# Output:-
Query:
 tensor([[[2, 8, 8],
         [4, 2, 4],
         [1, 2, 9]]])
Value:
 tensor([[[9, 5, 7],
         [3, 1, 4],
         [6, 2, 9]]])
weights:
 tensor([[[65.8179, 26.5581, 57.7350],
         [42.7239, 17.3205, 36.9504],
         [47.3427, 23.6714, 52.5389]]])
Triangular Metrics:
 tensor([[1., 0., 0.],
        [1., 1., 0.],
        [1., 1., 1.]])
Masked Weights
 tensor([[[65.8179,    -inf,    -inf],
         [42.7239, 17.3205,    -inf],
         [47.3427, 23.6714, 52.5389]]])
Softmax ( e^-inf = 0 )
 tensor([[[1.0000e+00, 0.0000e+00, 0.0000e+00],
         [1.0000e+00, 9.2777e-12, 0.0000e+00],
         [5.5073e-03, 2.8880e-13, 9.9449e-01]]])


由于我们使用的是多头注意力机制,因此我们将按照以下方式来实现它:


class MultiHeadAttention(nn.Module):
    def __init__(self, head_size, num_head):
        super().__init__()
        self.sa_head = nn.ModuleList([Head(head_size) for _ in range(num_head)])
        self.dropout = nn.Dropout(dropout)
        self.proj = nn.Linear(embed_size, embed_size)
    def forward(self, x):
        x = torch.cat([head(x) for head in self.sa_head], dim= -1)
        x = self.dropout(self.proj(x))
        return x


在这里,我们将输入x(维度为B, T, E)传递给不同的注意力头,每个注意力头返回一个最终向量,其大小为(B, T, head_size),其中head_size = E(64)/ num_heads(4)= 16。因为我们在一个范围为num_heads(4)的for循环中执行此操作,所以我们会将结果拼接回其原始大小(B, T, 4*16)。


多头注意力机制在嵌入维度更大的情况下,通常被认为速度更快且效率更高。


拼接之后,我们将最终输出传递给线性投影层。这样做的目的是使最终向量中的嵌入能够进一步交流它们在注意力权重计算过程中彼此学习到的信息。之后,它会被传递给一个丢弃(dropout)层,并返回结果。


综上所述,标准的解码器块按照上面图1所示的方式实现。


9


class FeedForward(nn.Module):
    def __init__(self, embed_size):
        super().__init__()
    
        self.ff = nn.Sequential(
              nn.Linear(embed_size, 4*embed_size),
              nn.ReLU(),
              nn.Linear(4*embed_size, embed_size),
              nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.ff(x)
    
class Block(nn.Module):
    def __init__(self, embed_size, num_head):
        super().__init__()
        head_size = embed_size // num_head
        self.multihead = MultiHeadAttention(head_size, num_head)
        self.ff = FeedForward(embed_size)
        self.ll1 = nn.LayerNorm(embed_size)
        self.ll2 = nn.LayerNorm(embed_size)
    def forward(self, x):
        x = x + self.multihead(self.ll1(x))
        x = x + self.ff(self.ll2(x))
        return x


如前所述计算头大小,输入首先通过层归一化(Layer Normalization),然后传递给我们的多头注意力网络,接着再通过另一个层归一化,最后传递给一个前馈网络(Feed-Forward Network)。


回到Bigram Model的话题:


class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, embed_size)
        self.possitional_embedding = nn.Embedding(block_size, embed_size)
        self.linear = nn.Linear(embed_size, vocab_size)
        self.block = nn.Sequential(*[Block(embed_size, num_head) for _ in range(num_layers)])
        self.layer_norm = nn.LayerNorm(embed_size)
    def forward(self, idx, targets=None):
        B, T = idx.shape
        # idx and targets are both (B,T) tensor of integers
        logits = self.token_embedding_table(idx) # (B,T,C)
        ps = self.possitional_embedding(torch.arange(T, device=device))
        x = logits + ps    #(B, T, C)
        logits = self.block(x)     #(B, T, c)
        logits = self.linear(self.layer_norm(logits)) # This suppose to map between head_size and Vocab_size 
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        return logits, loss
    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # crop idx to the last block_size tokens
            crop_idx= idx[:, -block_size:].to(device)
            # get the predictions
            logits, loss = self(crop_idx)
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C) from (B, T, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            # We sample one index from the filtered distribution
            idx_next = torch.multinomial(probs, num_samples=1).to(device)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx


在这里,我们看到Block是在Sequential层下通过层的数量范围来调用的。


生成词元…


首先,我们将单维度索引张量(即我们的词元索引)以及我们想要生成的新词元的最大数量一起传递给generate函数。由于我们的模型是为块大小8构建的,因此我们一次只能传递8个词元,所以我们裁剪idx中的最后8个词元(如果idx长度小于块大小,则选择所有词元)。


我们将裁剪后的idx传递给BigramLanguageModel。由于我们生成的logits包含了目标词元的概率分布,而我们只关心最后一个词元,因为目标(y)的最后一个词元是序列(x)中的下一个词元(这在批处理加载器部分有解释)。


我们现在得到的logits形状为(B, C),其中C是词汇表大小,它表示最后一个词元索引在整个词汇表上的概率分布。接下来,我们只需对其应用softmax函数,将这个向量转换为概率向量(即元素之和为1)。


现在,还记得我们在文章开头提到的不可完全预测或随机项,以及我们如何让模型随机选择序列中的下一个词元吗?为此,我们使用torch.multinomial,这是一种从给定概率分布中抽取样本的统计策略。在这里,它根据指定的概率随机抽样索引。


然后,我们终于得到了预测的下一个索引,将其与之前的索引拼接,并继续for循环,基于之前的索引不断生成下一个索引,直到达到最大词元数。


训练

幸运的是,训练部分非常简单明了。


m = BigramLanguageModel(65).to(device)65).to(device)
optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)
# training the model, cause I won't give up without a fight 
batch_size = 32
for epoch in range(5000):
    # Printing the Training and Validation Loss
    if epoch00==0: 
        m.eval()
        Loss= 0.0
        Val_Loss = 0.0
        for k in range(200):
            x, y = get_batch(True)
            
            val_ , val_loss = m(x, y)
            x1, y1 = get_batch(False)
            _, train_loss = m(x1, y1)            
            Loss += train_loss.item()
            Val_Loss += val_loss.item()
        avg_loss = Val_Loss/(k+1)
        avg_train_loss = Loss/(k+1)
        m.train()
        
        print("Epoch: {} \n The validation loss is:{}    The Loss is:{}".format(epoch, avg_loss, avg_train_loss))
    # Forward
    data, target = get_batch(False)
    logits, loss = m(data, target)
    #Backward
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()


在这里,我们进行了5000个周期的训练,这在4 GB VRAM的Nvidia RTX 3050上大约需要2分钟。


我们开始前向传播,通过get_batch()获取批次数据,并将其传递给BigramLanguageModel。然后设置optimizer.zero_grad(),执行loss.backward(),并执行optimizer.step()。我们使用了AdamW优化器,这对于我们的需求来说已经足够了。


我们为以下句子创建了一个张量:


ids = torch.tensor(encode("i like to eat food"), dtype=torch.long).unsqueeze(0)


ids的形状是(1, 18)(批次大小,词元)。我们起始只有18个词元(代表我们词汇表中的索引),并将生成额外的2000个字符。为了明确上下文,我们的词汇表是在之前的数据加载部分实现的,即它是input.txt中所有唯一字符的集合,具体实现为vocab = sorted(list(set(text)))。


print("".join(decode(m.generate(torch.zeros([1,1], dtype=torch.long) , max_new_tokens=2000)[0].tolist())))


Output:-
i like to eat food, noBANIO:
Here and
I shake married entreature by colantied at to women oword this swamind-betweet you are
As Grave eare his sun.
LUCENTIO:
Go what a doubled mistressed well.
Taildoes, not to memble, the peashat you;--are master, in thou comsand of the for slawake to bound and to of off this couchio;
Petruchio?
Fece poor this cockepopen neve so it do old loaps islied I'comment and curh
and blate sure poccient you the miad e'er a to partink,
Unory speitied where buzzarr'd formorns,
Pitedame,
Beach, and whom I firit.
ANDO:
O the virtuus a parros that that is acleast, not for suck could mighreature well; thy,
I'll toence after counteent,
Signior to paptista?
Shile you cappier?
BIANCA:
Where womand betire asleck him snall conglithing.
PROSPERO:
I, as expase caspierfed success,
This all no be trutes from the good the island mognied buzent; tensting in this garmortwant;
Do be marriage.
TRANIO:
'Tis, jointer.
GRUCHIO:
Soubt sI'll show I freek born.
PROSPETRUCHIO:
The vant mine; it it 


我知道……这看起来可能没什么意义。但我们必须认识到,语言模型的训练并不仅仅基于莎士比亚的著作这样的数据集。它需要大量的GPU计算能力、更先进的分词技术以及真正庞大的数据集。


由于我们是在一个非常小的模型上进行训练的,所以我们的表现其实并不差。输出仍然有意义,并且模型已经学会了使用实际的英语单词,而不是生成随机的胡言乱语。


你可以尝试用不同的数据集、词元大小、批次大小、层数等参数来实验这个模型。


这篇文章的主要目的是详细向你解释如何从零开始构建自己的语言模型,并在你的数据集上进行训练!



文章来源:https://medium.com/gitconnected/lets-build-our-own-gpt-model-from-scratch-with-pytorch-236a65a1fb54
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消