從頭構建gpt2 基于Transformer
VX關注{曉理紫|小李子},獲取技術推送信息,如感興趣,請轉發給有需要的同學,謝謝支持!!
如果你感覺對你有所幫助,請關注我。
源碼獲取 VX關注曉理紫并回復“chatgpt-0”
- 頭文件以及超參數
import torch
import torch.nn as nn
from torch.nn import functional as F
# 加入為了擴大網絡進行修改 head ,注意力、前向網絡添加了dropout和設置蹭數目
#超參數
batch_size = 64
block_size = 34 #塊大小 現在有34個上下文字符來預測
max_iters = 5000
eval_interval = 500
learning_rate=3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384 #嵌入維度
n_head = 6 #有6個頭,每個頭有284/6維
n_layer = 6 # 6層
dropout = 0.2torch.manual_seed(1337)
- 數據處理
with open('input.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]
decode = lambda l: ''.join([itos[i] for i in l])data = torch.tensor(encode(text),dtype=torch.long)n = int(0.9*len(data))
train_data = data[:n]
val_data = data[n:]def get_batch(split):data = train_data if split=="train" else val_dataix = torch.randint(len(data)-batch_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])x,y = x.to(device),y.to(device)return x,y
- 估計損失
@torch.no_grad()
def estimate_loos(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)logits,loss = model(x,y)losses[k] = loss.mean()out[split] = losses.mean()model.train()return out
- 單頭注意力
class Head(nn.Module):"""one head of self-attention"""def __init__(self, head_size):super(Head,self).__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)))self.dropout = nn.Dropout(dropout)def forward(self,x):B,T,C = x.shapek = self.key(x) #(B,T,C)q = self.query(x) #(B,T,C)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)out = wei@vreturn out
- 多頭注意力
class MultiHeadAttention(nn.Module):"""multiple heads of self-attention in parallel"""def __init__(self, num_heads,head_size) -> None:super(MultiHeadAttention,self).__init__()self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])self.proj = nn.Linear(n_embd,n_embd) #投影,為了方便使用慘差跳連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):"""a simple linear layer followed by a non-linearity"""def __init__(self,n_embd):super().__init__()self.net = nn.Sequential(nn.Linear(n_embd,4*n_embd), #從512變成2048nn.ReLU(),nn.Linear(4*n_embd,n_embd), #從2048變成512nn.Dropout(dropout), #Dropout 是可以在慘差鏈接之前加的東西)def forward(self,x):out = self.net(x)return out
- 塊
class Block(nn.Module):"""Transformer block:communication followed by computation"""def __init__(self, n_embd,n_head) -> None:#n_embd 需要嵌入維度中的嵌入數量#n_head 頭部數量super().__init__()head_size = n_embd//n_headself.sa = MultiHeadAttention(n_head,head_size) #通過多頭注意力進行計算self.ffwd = FeedFoward(n_embd) # 對注意力計算的結果進行提要完成self.ln1 = nn.LayerNorm(n_embd) #層規范 對于優化深層網絡很重要 論文Layer Normalizationself.ln2 = nn.LayerNorm(n_embd) #層規范def forward(self,x):# 通過使用殘差網絡的跳連進行x = x + self.sa(self.ln1(x)) x = x + self.ffwd(self.ln2(x))return x
- 整個語言模型
class BigramLangeNodel(nn.Module):def __init__(self):super(BigramLangeNodel,self).__init__()self.token_embedding_table = nn.Embedding(vocab_size,n_embd) #令牌嵌入表,對標記的身份進行編碼self.position_embedding_table = nn.Embedding(block_size,n_embd) #位置嵌入表,對標記的位置進行編碼。從0到block_size大小減一的每個位置將獲得自己的嵌入向量self.blocks = nn.Sequential(*[Block(n_embd,n_head=n_head) for _ in range(n_layer)]) #通過n_layer設置構建的曾數self.ln_f = nn.LayerNorm(n_embd)self.lm_head = nn.Linear(n_embd,vocab_size) #進行令牌嵌入到logits的轉換,這是語言頭def forward(self,idx,targets=None):B,T = idx.shapetok_emb= self.token_embedding_table(idx) #(B,T,C) C是嵌入大小 根據idx內的令牌的身份進行編碼pos_emb = self.position_embedding_table(torch.arange(T,device=device)) #(T,C) 從0到T減一的整數都嵌入到表中x = tok_emb+pos_emb #(B,T,C) 標記的身份嵌入與位置嵌入相加。x保存了身份以及身份出現的位置# x = self.sa_head(x) #(B,T,C)x = self.blocks(x)x = self.ln_f(x)logits = self.lm_head(x) #(B,T,vocab_size)if targets is None:loss = Noneelse:B,T,C = logits.shapelogits = logits.view(B*T,C)targets = targets.view(B*T)loss = F.cross_entropy(logits,targets)return logits,lossdef generate(self,idx,max_new_tokens):for _ in range(max_new_tokens):idx_cond = idx[:,-block_size] logits,loss= self(idx_cond)logits = logits[:,-1,:]# becomes (B,C)probs = F.softmax(logits,dim=-1)idx_next = torch.multinomial(probs,num_samples=1)idx = torch.cat((idx,idx_next),dim=1) #(B,T+1)return idx
- 訓練
model = BigramLangeNodel()
m = model.to(device)optimizer = torch.optim.AdamW(model.parameters(),lr = learning_rate)for iter in range(max_iters):if iter % eval_interval==0:losses = estimate_loos(model)print(f"step {iter}:train loss {losses['train']:.4f},val loss{losses['val']:.4f}")xb,yb = get_batch('train')logits,loss = model(xb,yb)optimizer.zero_grad(set_to_none=True)loss.backward()optimizer.step()context = torch.zeros((1,1),dtype=torch.long,device=device)
print(decode(m.generate(context,max_new_tokens=500)[0].tolist()))
- 訓練損失
step 0:train loss 4.3975,val loss4.3983
step 500:train loss 1.8497,val loss1.9600
step 1000:train loss 1.6500,val loss1.8210
step 1500:train loss 1.5530,val loss1.7376
step 2000:train loss 1.5034,val loss1.6891
step 2500:train loss 1.4665,val loss1.6638
step 3000:train loss 1.4431,val loss1.6457
step 3500:train loss 1.4156,val loss1.6209
step 4000:train loss 1.3958,val loss1.6025
step 4500:train loss 1.3855,val loss1.5988
簡單實現自注意力
VX 關注{曉理紫|小李子},獲取技術推送信息,如感興趣,請轉發給有需要的同學,謝謝支持!!
如果你感覺對你有所幫助,請關注我。
源碼獲取 VX關注曉理紫并回復“chatgpt-0”