Entraînons notre modèle GPT¶
Ce notebook s'appuie sur le notebook précédent. Il est fortement conseillé de faire le notebook précédent avant celui-ci.
In [1]:
Copied!
import torch
import torch.nn as nn
from torch.nn import functional as F
import torch
import torch.nn as nn
from torch.nn import functional as F
Hyperparamètres¶
In [2]:
Copied!
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
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¶
In [3]:
Copied!
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
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
Couches Head, Multi-Head, Feed Forward et TransformerBlock¶
In [4]:
Copied!
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
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
Implémentation de notre modèle¶
Il est temps d'implémenter notre propre GPT. Dans l'idée, il s'agit d'un agencement de block transformer avec également des matrices d'embeddings pour convertir les tokens en embeddings.
In [5]:
Copied!
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
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
Entrainement du modèle¶
Il est maintenant temps d'entraîner le modèle. L'entraînement peut être assez long selon votre ordinateur et les hyperparamètres que vous avez défini.
In [6]:
Copied!
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()
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
On peut maintenant générer du "Molière" automatiquement !
In [9]:
Copied!
# 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()))
# 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
Comme vous le voyez, ça ressemble (de loin) à un texte de Molière. En augmentant le nombre de paramètres du modèle, on peut obtenir beaucoup mieux mais le temps de traitement sera très lent (surtout si vous n'avez pas de GPU).