从零开始用PyTorch实现一个简易版ChatGPT，仅需500行代码

简介

你是否好奇ChatGPT背后的技术原理？本教程将带你从零开始，使用PyTorch实现一个简易的GPT模型。虽然只有500行代码，但包含了Transformer的核心机制。

环境准备

pip install torch numpy tqdm

1. 数据预处理

我们使用一个小型文本数据集（例如莎士比亚作品）进行训练。

import torch
from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, text, seq_len):
        # 字符级编码
        chars = sorted(list(set(text)))
        self.vocab_size = len(chars)
        self.char_to_idx = {ch:i for i,ch in enumerate(chars)}
        self.idx_to_char = {i:ch for i,ch in enumerate(chars)}
        self.data = [self.char_to_idx[ch] for ch in text]
        self.seq_len = seq_len

    def __len__(self):
        return len(self.data) - self.seq_len

    def __getitem__(self, idx):
        x = self.data[idx:idx+self.seq_len]
        y = self.data[idx+1:idx+self.seq_len+1]
        return torch.tensor(x), torch.tensor(y)

2. 定义GPT模型

我们的模型包含嵌入层、多个Transformer解码器块和输出层。

import torch.nn as nn

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model, n_head, n_layers, seq_len):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(seq_len, d_model)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_head) for _ in range(n_layers)
        ])
        self.ln = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        B, T = x.shape
        tok_emb = self.token_embed(x)
        pos = torch.arange(T, device=x.device).unsqueeze(0)
        pos_emb = self.pos_embed(pos)
        x = tok_emb + pos_emb
        for block in self.blocks:
            x = block(x)
        x = self.ln(x)
        logits = self.head(x)
        return logits

其中TransformerBlock包含多头自注意力和前馈网络：

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4*d_model),
            nn.GELU(),
            nn.Linear(4*d_model, d_model),
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # 自注意力 + 残差
        attn_out, _ = self.attn(x, x, x, need_weights=False)
        x = self.ln1(x + attn_out)
        # 前馈 + 残差
        ffn_out = self.ffn(x)
        x = self.ln2(x + ffn_out)
        return x

3. 训练循环

def train(model, dataloader, epochs, lr):
    optimizer = torch.optim.AdamW(model.parameters(), lr)
    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(epochs):
        for x, y in dataloader:
            logits = model(x)
            loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

4. 文本生成

def generate(model, start_str, max_new_tokens, device):
    model.eval()
    with torch.no_grad():
        input_ids = [char_to_idx[ch] for ch in start_str]
        input_tensor = torch.tensor([input_ids]).to(device)
        for _ in range(max_new_tokens):
            logits = model(input_tensor)
            # 只取最后一个时间步
            logits_last = logits[:, -1, :]
            probs = torch.softmax(logits_last, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            input_tensor = torch.cat([input_tensor, next_id], dim=1)
        output_ids = input_tensor.squeeze().tolist()
        output_str = ''.join([idx_to_char[i] for i in output_ids])
    return output_str

5. 完整代码与运行

完整代码已上传至GitHub（链接略），运行后即可生成类似莎士比亚风格的文本。注意：由于模型较小，生成内容可能不够连贯，但原理与ChatGPT相同。

总结

通过本教程，你已掌握GPT的核心实现。进一步学习可尝试：

• 使用更大的数据集（如维基百科）
• 增加模型层数和维度
• 加入注意力掩码实现因果注意力

欢迎在评论区交流！