别再死磕RNN了！用Python从零实现一个简易Transformer（附完整代码）

用Python从零构建微型Transformer实战编码与核心原理剖析在自然语言处理领域Transformer架构已经彻底改变了序列建模的游戏规则。但对于许多开发者来说论文中的数学公式和框架图总是隔着一层朦胧的面纱。本文将带您用PyTorch亲手搭建一个简化版Transformer通过代码透视自注意力、位置编码等核心机制。不同于理论讲解我们将聚焦于可运行的实现——最终完成的迷你模型虽只有2层编码器却能完整展示Transformer的精髓。1. 环境准备与基础组件在开始构建Transformer之前我们需要准备好开发环境并实现几个基础组件。这些模块将作为整个模型的构建基石。首先安装必要的Python库pip install torch numpy matplotlib1.1 词嵌入层实现词嵌入是将离散的单词索引转换为连续向量表示的关键步骤。以下是使用PyTorch实现的可训练嵌入层import torch import torch.nn as nn class EmbeddingLayer(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.d_model d_model def forward(self, x): # x shape: (batch_size, seq_len) return self.embedding(x) * (self.d_model ** 0.5) # 缩放嵌入向量这个简单的嵌入层会将每个单词映射到d_model维的向量空间对嵌入结果进行缩放保持数值稳定性自动处理批量输入和变长序列1.2 位置编码解析与实现Transformer需要显式地编码位置信息因为自注意力机制本身不具备序列顺序的概念。以下是正弦位置编码的实现import math def positional_encoding(seq_len, d_model): position torch.arange(seq_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe torch.zeros(seq_len, d_model) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe位置编码的特点使用不同频率的正弦和余弦函数能够表示绝对和相对位置信息可以扩展到训练时未见过的序列长度2. 自注意力机制深度实现自注意力是Transformer最核心的创新它允许模型在处理每个单词时关注序列中的所有其他单词。2.1 缩放点积注意力def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attention torch.softmax(scores, dim-1) return torch.matmul(attention, V)关键参数说明参数维度作用Q(..., seq_len, d_k)查询向量K(..., seq_len, d_k)键向量V(..., seq_len, d_v)值向量mask(..., seq_len, seq_len)防止关注填充token2.2 多头注意力实现多头注意力并行运行多个注意力机制使模型能够同时关注不同位置的子空间class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() 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 split_heads(self, x): batch_size x.size(0) return x.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) def forward(self, Q, K, V, maskNone): Q self.split_heads(self.W_q(Q)) K self.split_heads(self.W_k(K)) V self.split_heads(self.W_v(V)) attention scaled_dot_product_attention(Q, K, V, mask) attention attention.transpose(1, 2).contiguous() attention attention.view(attention.size(0), -1, self.d_model) return self.W_o(attention)多头注意力的优势并行计算多个注意力头每个头学习不同的关注模式最后将各头的输出拼接并线性变换3. 构建Transformer编码器层现在我们将前面实现的组件组合起来构建完整的Transformer编码器层。3.1 前馈网络实现每个编码器层包含一个全连接前馈网络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) self.dropout nn.Dropout(0.1) def forward(self, x): return self.linear2(self.dropout(torch.relu(self.linear1(x))))3.2 编码器层完整实现class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff): super().__init__() self.self_attn MultiHeadAttention(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(0.1) def forward(self, x, mask): attn_output self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) ffn_output self.ffn(x) x self.norm2(x self.dropout(ffn_output)) return x编码器层的关键特性残差连接保留原始输入信息层归一化稳定训练过程Dropout防止过拟合4. 组装完整Transformer模型4.1 模型架构整合现在我们将所有组件整合成一个完整的Transformer模型class MiniTransformer(nn.Module): def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers): super().__init__() self.embedding EmbeddingLayer(vocab_size, d_model) self.pos_encoding positional_encoding(1000, d_model) # 假设最大序列长度为1000 self.encoder_layers nn.ModuleList( [EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)] ) self.linear nn.Linear(d_model, vocab_size) def forward(self, x, maskNone): seq_len x.size(1) x self.embedding(x) x self.pos_encoding[:seq_len, :].to(x.device) for layer in self.encoder_layers: x layer(x, mask) return self.linear(x)4.2 模型训练与验证让我们在IMDb电影评论数据集上测试这个迷你Transformerfrom torch.utils.data import DataLoader from torchtext.datasets import IMDB from torchtext.data.utils import get_tokenizer from torchtext.vocab import build_vocab_from_iterator # 数据预处理 tokenizer get_tokenizer(basic_english) train_iter IMDB(splittrain) def yield_tokens(data_iter): for _, text in data_iter: yield tokenizer(text) vocab build_vocab_from_iterator(yield_tokens(train_iter), specials[unk, pad]) vocab.set_default_index(vocab[unk]) # 模型配置 model MiniTransformer( vocab_sizelen(vocab), d_model128, num_heads4, d_ff256, num_layers2 ).to(device) # 训练循环 optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion nn.CrossEntropyLoss() for epoch in range(10): for batch in train_loader: inputs, targets batch.text.to(device), batch.label.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs.view(-1, len(vocab)), targets.view(-1)) loss.backward() optimizer.step()训练过程中的常见问题与解决方案梯度爆炸使用梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)过拟合增加dropout率或使用早停法训练缓慢减小模型尺寸或使用混合精度训练这个简化实现虽然性能不及完整Transformer但完整保留了核心机制。在实际项目中可以考虑以下优化方向增加解码器部分实现序列到序列任务实现更高效的位置编码方案添加学习率调度器实现批处理掩码和填充处理