别再死记硬背Transformer了！用PyTorch手把手实现一个简易翻译模型（附完整代码）

用PyTorch从零构建Transformer翻译模型代码驱动的深度学习实践如果你已经读过Transformer的论文或看过相关教程却依然对如何实现这个革命性架构感到迷茫那么这篇文章正是为你准备的。我们将避开繁琐的理论推导直接进入代码层面通过构建一个英中翻译模型来掌握Transformer的核心实现技巧。1. 环境准备与数据预处理在开始构建模型之前我们需要准备好开发环境并处理翻译任务所需的数据集。这个阶段虽然看似简单却直接影响后续模型训练的效果。1.1 安装必要的Python库首先确保你的Python环境建议3.8中已安装以下关键库pip install torch torchtext spacy sentencepiece然后下载中英文语言模型用于分词python -m spacy download en_core_web_sm python -m spacy download zh_core_web_sm1.2 构建双语数据集我们将使用IWSLT2017英中翻译数据集它包含约20万条平行句对。以下是数据加载和预处理的完整代码from torchtext.datasets import IWSLT2017 from torchtext.data.utils import get_tokenizer from torchtext.vocab import build_vocab_from_iterator # 初始化分词器 en_tokenizer get_tokenizer(spacy, languageen_core_web_sm) zh_tokenizer get_tokenizer(spacy, languagezh_core_web_sm) def yield_tokens(data_iter, tokenizer, language): for data in data_iter: yield tokenizer(data[language]) # 构建词汇表 train_iter IWSLT2017(splittrain) en_vocab build_vocab_from_iterator(yield_tokens(train_iter, en_tokenizer, en), specials[unk, pad, bos, eos]) zh_vocab build_vocab_from_iterator(yield_tokens(train_iter, zh_tokenizer, zh), specials[unk, pad, bos, eos]) # 设置默认未知词索引 en_vocab.set_default_index(en_vocab[unk]) zh_vocab.set_default_index(zh_vocab[unk])注意实际应用中应考虑限制词汇表大小如保留前30000个高频词以控制模型规模2. Transformer核心组件实现现在我们来构建Transformer的核心模块。与原始论文不同我们会做一些简化以提升代码可读性同时保持架构的关键特性。2.1 位置编码序列顺序的数字化表达Transformer抛弃了RNN的循环结构因此需要显式地注入位置信息。以下是改进后的位置编码实现import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int 5000): super().__init__() position torch.arange(max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe torch.zeros(max_len, d_model) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe) def forward(self, x: torch.Tensor) - torch.Tensor: 参数: x: 形状为 [batch_size, seq_len, embedding_dim] 的张量 返回: 添加位置编码后的张量 return x self.pe[:x.size(1)]可视化位置编码可以帮助我们理解其工作原理import matplotlib.pyplot as plt plt.figure(figsize(12, 6)) pe PositionalEncoding(256) y pe(torch.zeros(1, 100, 256)) plt.plot(y[0, :, 4:8].data.numpy()) plt.legend([dim %d % p for p in [4,5,6,7]]) plt.title(Positional Encoding Visualization) plt.show()2.2 多头注意力机制的实现多头注意力是Transformer最具创新性的部分下面是其PyTorch实现class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, dropout: float 0.1): super().__init__() assert d_model % num_heads 0, d_model必须能被num_heads整除 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) self.dropout nn.Dropout(dropout) def scaled_dot_product_attention(self, q, k, v, maskNone): attn_scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: attn_scores attn_scores.masked_fill(mask 0, -1e9) attn_probs torch.softmax(attn_scores, dim-1) attn_probs self.dropout(attn_probs) output torch.matmul(attn_probs, v) return output, attn_probs def forward(self, q, k, v, maskNone): batch_size q.size(0) # 线性变换并分头 q self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 计算注意力 attn_output, attn_probs self.scaled_dot_product_attention(q, k, v, mask) # 合并多头 attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 最终线性变换 output self.w_o(attn_output) return output, attn_probs3. 编码器与解码器架构有了核心组件后我们可以构建完整的编码器和解码器结构。3.1 编码器层的实现每个编码器层包含一个自注意力机制和一个前馈网络class EncoderLayer(nn.Module): def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float 0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 自注意力子层 attn_output, _ self.self_attn(x, x, x, mask) x x self.dropout(attn_output) x self.norm1(x) # 前馈网络子层 ff_output self.feed_forward(x) x x self.dropout(ff_output) x self.norm2(x) return x3.2 解码器层的实现解码器层比编码器更复杂包含三种注意力机制class DecoderLayer(nn.Module): def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float 0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.cross_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, encoder_output, src_maskNone, tgt_maskNone): # 自注意力子层目标序列的自注意力 attn_output, _ self.self_attn(x, x, x, tgt_mask) x x self.dropout(attn_output) x self.norm1(x) # 交叉注意力子层查询来自解码器键值来自编码器 cross_output, _ self.cross_attn(x, encoder_output, encoder_output, src_mask) x x self.dropout(cross_output) x self.norm2(x) # 前馈网络子层 ff_output self.feed_forward(x) x x self.dropout(ff_output) x self.norm3(x) return x4. 完整模型组装与训练现在我们将所有组件组合成完整的Transformer模型并实现训练流程。4.1 模型组装class Transformer(nn.Module): def __init__(self, src_vocab_size: int, tgt_vocab_size: int, d_model: int 512, num_heads: int 8, num_layers: int 6, d_ff: int 2048, dropout: float 0.1, max_seq_len: int 100): super().__init__() # 词嵌入层 self.src_embed nn.Embedding(src_vocab_size, d_model) self.tgt_embed nn.Embedding(tgt_vocab_size, d_model) # 位置编码 self.pos_encoding PositionalEncoding(d_model, max_seq_len) # 编码器 self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) # 解码器 self.decoder_layers nn.ModuleList([ DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) # 输出层 self.linear nn.Linear(d_model, tgt_vocab_size) self.dropout nn.Dropout(dropout) def encode(self, src, src_mask): src_embedded self.dropout(self.pos_encoding(self.src_embed(src))) for layer in self.encoder_layers: src_embedded layer(src_embedded, src_mask) return src_embedded def decode(self, tgt, encoder_output, src_mask, tgt_mask): tgt_embedded self.dropout(self.pos_encoding(self.tgt_embed(tgt))) for layer in self.decoder_layers: tgt_embedded layer(tgt_embedded, encoder_output, src_mask, tgt_mask) return tgt_embedded def forward(self, src, tgt, src_maskNone, tgt_maskNone): encoder_output self.encode(src, src_mask) decoder_output self.decode(tgt, encoder_output, src_mask, tgt_mask) output self.linear(decoder_output) return output4.2 训练流程实现以下是简化的训练循环包含学习率调度和梯度裁剪def train_model(model, train_loader, val_loader, epochs10, lr0.0001): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss(ignore_index1) # 忽略padding索引 optimizer torch.optim.Adam(model.parameters(), lrlr, betas(0.9, 0.98), eps1e-9) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size1, gamma0.95) for epoch in range(epochs): model.train() total_loss 0 for batch in train_loader: src, tgt batch.src.to(device), batch.tgt.to(device) # 创建掩码 src_mask (src ! 1).unsqueeze(1).unsqueeze(2) # padding索引为1 tgt_mask (tgt ! 1).unsqueeze(1).unsqueeze(2) seq_len tgt.size(1) nopeak_mask torch.triu(torch.ones(1, seq_len, seq_len) 1).transpose(1, 2) nopeak_mask nopeak_mask.float().to(device) tgt_mask tgt_mask nopeak_mask optimizer.zero_grad() # 前向传播 output model(src, tgt[:, :-1], src_mask, tgt_mask[:, :-1, :-1]) # 计算损失 loss criterion(output.reshape(-1, output.size(-1)), tgt[:, 1:].reshape(-1)) # 反向传播 loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() scheduler.step() avg_loss total_loss / len(train_loader) print(fEpoch {epoch1}, Loss: {avg_loss:.4f}) # 验证 model.eval() val_loss evaluate(model, val_loader, criterion, device) print(fValidation Loss: {val_loss:.4f}) return model4.3 推理与翻译示例训练完成后我们可以用以下函数进行翻译def translate(model, sentence, src_vocab, tgt_vocab, max_len50): model.eval() device next(model.parameters()).device # 预处理输入句子 tokens [src_vocab[bos]] [src_vocab[token] for token in en_tokenizer(sentence)] [src_vocab[eos]] src torch.LongTensor(tokens).unsqueeze(0).to(device) src_mask (src ! 1).unsqueeze(1).unsqueeze(2) # 初始化目标序列 tgt torch.LongTensor([[tgt_vocab[bos]]]).to(device) for i in range(max_len): tgt_mask (tgt ! 1).unsqueeze(1).unsqueeze(2) seq_len tgt.size(1) nopeak_mask torch.triu(torch.ones(1, seq_len, seq_len) 1).transpose(1, 2) nopeak_mask nopeak_mask.float().to(device) tgt_mask tgt_mask nopeak_mask output model(src, tgt, src_mask, tgt_mask) next_token output.argmax(-1)[:, -1].unsqueeze(1) tgt torch.cat([tgt, next_token], dim1) if next_token.item() tgt_vocab[eos]: break # 将索引转换为单词 translated [tgt_vocab.lookup_token(idx) for idx in tgt.squeeze().tolist()] return .join(translated[1:-1]) # 去掉bos和eos在实际项目中我发现在解码阶段使用束搜索(beam search)比贪婪解码能获得更流畅的翻译结果。此外添加标签平滑(label smoothing)也能有效缓解模型过度自信的问题。