Training our GPT model#
This notebook builds on the previous one. It is highly recommended to follow the previous notebook before starting this one.
import torch
import torch.nn as nn
from torch.nn import functional as F
Training Parameters#
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
Dataset#
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, Multi-Head, Feed Forward, and TransformerBlock Layers#
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
Model Implementation#
It’s time to implement our own GPT model. In short, it consists of assembling transformer blocks with embedding matrices to convert tokens into embeddings.
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 Training#
It’s now time to train the model. The training duration may vary depending on your computer and the hyperparameters you have set.
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
We can now automatically generate text inspired by Molière!
# 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
As you can see, the text resembles (from afar) a text by Molière. By increasing the number of model parameters, better results can be achieved, but the processing time will be slower, especially if you don’t have a GPU.