当前位置：首页 > news >正文

别再只调API了！手把手带你用PyTorch从零复现GPT-1的Transformer Decoder结构

news 2026/6/11 2:43:01

从零构建GPT-1的Transformer Decoder：PyTorch实战指南

在当今AI领域，预训练语言模型已成为自然语言处理的核心工具。许多开发者习惯于直接调用现成的API，却对模型内部的精妙设计知之甚少。本文将带你用PyTorch从零开始实现GPT-1的核心组件——Transformer Decoder结构，通过动手实践深入理解这一革命性架构的工作原理。

1. 环境准备与基础配置

在开始编码前，我们需要搭建合适的开发环境。推荐使用Python 3.8+和PyTorch 1.12+版本，这些版本在稳定性和功能支持上都有良好表现。

conda create -n gpt python=3.8 conda activate gpt pip install torch torchvision torchaudio

GPT-1的模型配置如下表所示：

参数名称	配置值
层数	12
隐藏层维度	768
注意力头数	12
Feed Forward维度	3072
激活函数	GeLU
位置编码	可学习

提示：在实际实验中，为节省计算资源，可以按比例缩小这些参数进行原型验证。

2. 核心组件实现

2.1 可学习位置编码

与传统Transformer不同，GPT-1采用了可学习的位置编码。这种设计让模型能够自适应地学习最适合任务的位置表示。

class LearnablePositionalEncoding(nn.Module): def __init__(self, d_model, max_len=512): super().__init__() self.position_emb = nn.Parameter(torch.zeros(max_len, d_model)) def forward(self, x): # x shape: (batch_size, seq_len, d_model) seq_len = x.size(1) positions = self.position_emb[:seq_len, :] return x + positions.unsqueeze(0)

2.2 Masked多头注意力机制

这是GPT-1的核心组件之一，负责捕捉序列中的上下文关系，同时确保自回归特性。

class MaskedMultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): batch_size = x.size(0) # 线性变换并分头 Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 应用mask if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # 计算注意力权重 attn_weights = F.softmax(scores, dim=-1) # 应用注意力权重 context = torch.matmul(attn_weights, V) # 合并多头 context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 输出线性变换 output = self.W_o(context) return output

2.3 前馈网络

GPT-1的前馈网络由两个线性变换和一个GeLU激活函数组成。

class FeedForward(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) def forward(self, x): return self.linear2(F.gelu(self.linear1(x)))

3. 构建Decoder Block

将上述组件组合起来，形成完整的Decoder Block：

class DecoderBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MaskedMultiHeadAttention(d_model, num_heads) self.ffn = FeedForward(d_model, d_ff) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): # 自注意力子层 attn_output = self.self_attn(x, mask) x = x + self.dropout(attn_output) x = self.norm1(x) # 前馈网络子层 ffn_output = self.ffn(x) x = x + self.dropout(ffn_output) x = self.norm2(x) return x

4. 完整GPT-1模型组装

现在我们可以将所有Decoder Block堆叠起来，构建完整的GPT-1模型：

class GPT1(nn.Module): def __init__(self, vocab_size, d_model=768, num_layers=12, num_heads=12, d_ff=3072, max_len=512, dropout=0.1): super().__init__() self.token_emb = nn.Embedding(vocab_size, d_model) self.pos_emb = LearnablePositionalEncoding(d_model, max_len) self.dropout = nn.Dropout(dropout) self.layers = nn.ModuleList([ DecoderBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.norm = nn.LayerNorm(d_model) self.lm_head = nn.Linear(d_model, vocab_size, bias=False) def forward(self, x, mask=None): # 嵌入层 token_embeddings = self.token_emb(x) pos_embeddings = self.pos_emb(token_embeddings) x = self.dropout(pos_embeddings) # 通过所有Decoder层 for layer in self.layers: x = layer(x, mask) # 最终归一化 x = self.norm(x) # 语言模型头 logits = self.lm_head(x) return logits

5. 模型训练与推理

5.1 数据准备

我们需要准备适合语言模型训练的文本数据。以下是一个简单的数据加载器实现：

class TextDataset(Dataset): def __init__(self, texts, tokenizer, seq_len=512): self.tokenizer = tokenizer self.seq_len = seq_len self.data = self.process_texts(texts) def process_texts(self, texts): token_ids = [] for text in texts: tokens = self.tokenizer.encode(text) token_ids.extend(tokens) return token_ids def __len__(self): return len(self.data) // self.seq_len def __getitem__(self, idx): start = idx * self.seq_len end = start + self.seq_len segment = self.data[start:end] x = torch.tensor(segment[:-1], dtype=torch.long) y = torch.tensor(segment[1:], dtype=torch.long) return x, y

5.2 训练循环

实现一个基本的训练循环：

def train(model, dataloader, optimizer, device, epochs=10): model.train() criterion = nn.CrossEntropyLoss() for epoch in range(epochs): total_loss = 0 for batch_idx, (x, y) in enumerate(dataloader): x, y = x.to(device), y.to(device) # 创建mask mask = torch.triu(torch.ones(x.size(1), x.size(1)), diagonal=1).bool() mask = mask.to(device) optimizer.zero_grad() outputs = model(x, mask) loss = criterion(outputs.view(-1, outputs.size(-1)), y.view(-1)) loss.backward() optimizer.step() total_loss += loss.item() if batch_idx % 100 == 0: print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}') print(f'Epoch {epoch}, Avg Loss: {total_loss/len(dataloader):.4f}')

5.3 文本生成

实现一个简单的自回归文本生成函数：

def generate_text(model, tokenizer, prompt, max_len=50, temperature=1.0): model.eval() tokens = tokenizer.encode(prompt) input_tensor = torch.tensor([tokens], dtype=torch.long).to(device) for _ in range(max_len): with torch.no_grad(): # 创建mask mask = torch.triu(torch.ones(input_tensor.size(1), input_tensor.size(1)), diagonal=1).bool() mask = mask.to(device) outputs = model(input_tensor, mask) next_token_logits = outputs[:, -1, :] / temperature next_token = torch.multinomial(F.softmax(next_token_logits, dim=-1), num_samples=1) input_tensor = torch.cat([input_tensor, next_token], dim=1) generated = input_tensor[0].cpu().numpy() return tokenizer.decode(generated)