從頭構建和訓練 GPT-2 |實戰

数据科学工厂發表於2024-07-14

引言

該專案將引導您完成構建簡單 GPT-2 模型的所有步驟,並使用 Taylor Swift 和 Ed Sheeran 的一堆歌曲進行訓練。本文的資料集和原始碼將在 Github 上提供。

構建 GPT-2 架構

我們將逐步推進這個專案,不斷最佳化一個基礎的模型框架,並在其基礎上增加新的層次,這些層次都是基於 GPT-2 的原始設計。

我們將按照以下步驟進行:

  • 製作一個定製的分詞工具
  • 開發一個資料載入程式
  • 培養一個基礎的語言處理能力
  • 完成 GPT-2 架構的實現(第二部分)

該專案分為兩個部分,第一個部分介紹語言建模的基礎知識,第二部分直接跳到 GPT-2 實現。我建議您按照本文進行操作並自己構建它,這將使學習 GPT-2 變得更加有趣和有趣。

最終模型:

1. 構建自定義分詞器

語言模型不像我們一樣看到文字。相反,它們將數字序列識別為特定文字的標記。因此,第一步是匯入我們的資料並構建我們自己的角色級別分詞器。

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])

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)

現在,我們有了標記化的張量資料,其中文字中的每個字元都轉換為各自的標記。

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)

2. 構建資料載入器

現在,在構建模型之前,我們必須定義如何將資料輸入模型進行訓練,以及資料的維度和批次大小。

讓我們定義我們的資料載入器如下:

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)

我們現在已經開發了自己的專用資料載入工具,它既可以用於模型的訓練階段,也可以用於評估階段。這個工具包含一個 get_batch 功能,它能夠一次性提供大小為 batch_size 乘以 context_length 的資料批次。

如果你好奇為什麼 x 的範圍是從序列的起始點到結束點,而 y 的範圍則是從 x 的起始點後一位到結束點後一位,這是因為模型的核心任務是預測給定前序序列之後的下一個元素。換句話說,在 y 中會多出一個標記,這樣模型就可以基於 x 中的最後 n 個標記來預測下一個,也就是第 (n+1) 個標記。如果這聽起來有些難以理解,可以參閱下面的圖解說明。

3. 訓練簡單的語言模型

現在,我們即將利用我們剛剛載入的資料,來搭建和訓練一個基礎的語言模型。

在本節中,我們將保持操作的簡潔性,採用一個簡單的二元語法模型,即基於上一個詞來預測下一個詞。如你所見,我們將只利用 Embedding 層,而忽略主解碼模組。

Embedding 層能夠為詞彙表中的每個字元表示出 n = d_model 個獨特的屬性,並且該層會根據字元在詞彙表中的索引來提取這些屬性。

你會驚訝地發現,僅僅依靠 Embedding 層,模型就能表現出色。我們將透過逐步增加更多的層來最佳化模型,所以請耐心等待並繼續關注。

嵌入的維度,也就是 d_model,目前設定為等於詞彙表的大小 vocab_size,這是因為模型的最終輸出需要對應到詞彙表中每個字元的對數機率,以便計算它們各自的機率。在未來,我們會引入一個線性層(Linear 層),它負責將 d_model 的輸出維度轉換為 vocab_size,這樣我們就可以使用自定義的嵌入維度 embedding_dimension

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 layer)和用於生成標記的 Softmax 函式組成。接下來,讓我們觀察一下,當模型接收到一些輸入字元時,它的反應和表現會是怎樣。

現在,我們來到了最後的關鍵步驟——訓練模型,讓它學會識別和理解字元。接下來,我們將配置最佳化器。目前,我們選擇使用一個基礎的 AdamW 最佳化器,設定的學習率為 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

我們取得了相當不錯的損失值。但我們還沒有完全成功。你可以看到,直到訓練的第2000個週期,錯誤率有了顯著的下降,但之後的提升就不明顯了。這是因為模型目前還缺乏足夠的智慧(或者說是層數/神經網路的數量),它僅僅是在比較不同字元的嵌入表示。

現在模型的輸出看起來如下所示:

本文由mdnice多平臺釋出

相關文章