训练我们的 GPT 模型#
本笔记本基于前一笔记本的内容。强烈建议在开始本笔记本之前，先完成前一笔记本的学习。
import torch
import torch.nn as nn
from torch.nn import functional as F
训练参数#
batch_size = 64
block_size = 128 # Longueur du contexte
max_iters = 5000 # Nombre d'itérations d'entraînement
eval_interval = 500 # Intervalle pour l'évaluation sur les données de validation
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 128 # Dimension de la couche d'attention
n_head = 4 # Nombre de head d'attention (pour la multi-head attention)
n_layer = 3 # Nombre de couches d'attention
dropout = 0.2
数据集#
with open('moliere.txt', 'r', encoding='utf-8') as f:
text = f.read()
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encore : prend un string et output une liste d'entiers
decode = lambda l: ''.join([itos[i] for i in l]) # decode: prend une liste d'entiers et output un string
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data)) # 90% pour le train et 10% pour la validation
train_data = data[:n]
val_data = data[n:]
def get_batch(split):
# On genere un batch de données (sur train ou val)
data = train_data if split == 'train' else val_data
#On génére batch_size indice de début de séquence pris au hasard dans le dataset
ix = torch.randint(len(data) - block_size, (batch_size,))
# On stocke dans notre tenseur torch
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])
x, y = x.to(device), y.to(device) # On met les sur le GPU si on en a un
return x, y
@torch.no_grad()
# Fonction pour estimer le loss plus précisement
def estimate_loss(model):
out = {}
model.eval()
for split in ['train', 'val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
_, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
Head 层、多头注意力层、前馈层与 Transformer 块#
class Head(nn.Module):
""" Couche de self-attention unique """
def __init__(self, head_size):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
# Ajout de dropout pour la regularization
self.dropout = nn.Dropout(dropout)
def forward(self, x):
B,T,C = x.shape
k = self.key(x) # (B,T,C)
q = self.query(x) # (B,T,C)
# Le * C**-0.5 correspond à la normalisation par la racine de head_size
wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
wei = self.dropout(wei)
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
class MultiHeadAttention(nn.Module):
""" Plusieurs couches de self attention en parallèle"""
def __init__(self, num_heads, head_size):
super().__init__()
# Création de num_head couches head de taille head_size
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
# Couche pour Linear (voir schema) après concatenation
self.proj = nn.Linear(n_embd, n_embd)
# Dropout si besoin
self.dropout = nn.Dropout(dropout)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
class FeedFoward(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
# 4*n_embd comme dans le papier
nn.Linear(n_embd, 4 * n_embd),
nn.GeLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class TransformerBlock(nn.Module):
""" Block transformer"""
def __init__(self, n_embd, n_head):
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size)
self.ffwd = FeedFoward(n_embd)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
x = x + self.sa(self.ln1(x)) # x+ car c'est une connexion résiduelle
x = x + self.ffwd(self.ln2(x))
return x
模型实现#
现在是时候实现我们自己的 GPT 模型了。总体来说，该模型由多个 Transformer 块组成，并使用 嵌入矩阵 将 token 转换为 嵌入向量。
class GPT(nn.Module):
def __init__(self):
super().__init__()
# Chaque token recupere son embedding à partir d'une look-up table (entrainable)
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
# Agencement de n_layer TransformerBlock de taille n_embed avec n_head heads.
self.blocks = nn.Sequential(*[TransformerBlock(n_embd, n_head=n_head) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd) # final layer norm
self.lm_head = nn.Linear(n_embd, vocab_size)
# Une manière optimile d'initialiser les poids (à regarder plus en détails si ça vous interesse)
self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx et targets sont des tenseurs d'entier de taille (B,T)
tok_emb = self.token_embedding_table(idx) # (B,T,C)
# Position embedding pour ajouter une information spatiale sur les éléments
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,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
# Fonction pour générer du texte : très proche de celle utilisée dans le modèle bigramme du notebook précédent
def generate(self, idx, max_new_tokens):
# idx contient des entiers dans un tenseur de taille (B, T), c'est le contexte actuel
for _ in range(max_new_tokens):
# On limite le contexte à block_size (maximum que le réseau peut prendre)
idx_cond = idx[:, -block_size:]
# On calcule la prédictions du prochain token
logits, _ = self(idx_cond)
logits = logits[:, -1, :] # becomes (B, C)
probs = F.softmax(logits, dim=-1) # (B, C)
# On sample depuis les probabilités obtenues avec le softmax
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# On ajouter l'élément sample à notre texte
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
模型训练#
现在可以开始训练模型了。训练时间会根据您的电脑性能以及设定的超参数而有所不同。
model = GPT()
m = model.to(device)
# print the number of parameters in the model
print("Nombre de paramètres du modèle : ",sum(p.numel() for p in m.parameters()))
# Optimizer AdamW
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
for iter in range(max_iters):
# On evalue le loss sur train et validation de temps en temps (tous les eval_interval itérations)
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss(model)
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
# On recupère un batch aléatoire du dataset
xb, yb = get_batch('train')
# Calcul du loss
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
Nombre de paramètres du modèle : 632149
/home/aquilae/anaconda3/envs/dev/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
step 0: train loss 4.4823, val loss 4.4809
step 500: train loss 1.9838, val loss 2.0682
step 1000: train loss 1.6177, val loss 1.7505
step 1500: train loss 1.4865, val loss 1.6591
step 2000: train loss 1.4180, val loss 1.6133
step 2500: train loss 1.3637, val loss 1.5702
step 3000: train loss 1.3256, val loss 1.5459
step 3500: train loss 1.3022, val loss 1.5306
step 4000: train loss 1.2784, val loss 1.5076
step 4500: train loss 1.2622, val loss 1.4943
step 4999: train loss 1.2467, val loss 1.4835
现在我们可以自动生成受莫里哀风格启发的文本了！
# Generation d'un texte par le modèle
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))
FMEROT.
Il a rien qu'il mieux a plus que me le sortir engage.
DONE ELVIRE.
Ah! ot, et vous bons aime aussi seule abuer?
SGANARELLE.
Vous plaissez. Vous savez non pas!
LE MONCTAGNE.
Vous êtes, cet autres que changus!
DON JUAN.
Ne sont point tout ce qui ne vous désie.
VALÈRE.
Une son méclai que expère
Que je suis sous-moi qu'en le consie,
Que je tiens peurté aun la régle, que je suis en me ducens mois.
Au comoqu'on disant fait ce le vite avisenve
Qui voir de cettte réjoustoisse dis, et
如您所见，生成的文本在某种程度上类似于莫里哀的风格。通过增加模型的参数数量，可以获得更好的效果，但处理时间也会更长，尤其是在没有 GPU 的情况下。