Transformer架构实战：从零开始手把手实现一个简易版（Python代码示例）

Transformer架构实战从零开始手把手实现一个简易版Python代码示例在人工智能领域Transformer架构已经彻底改变了自然语言处理的游戏规则。不同于传统的循环神经网络RNNTransformer通过自注意力机制实现了并行化处理大幅提升了模型训练和推理的效率。本文将带你从零开始用Python实现一个简化版的Transformer模型深入理解其核心组件的工作原理。1. 环境准备与基础概念在开始编码之前我们需要确保开发环境配置正确。推荐使用Python 3.8和PyTorch 1.10这些版本提供了良好的兼容性和性能优化。pip install torch numpy matplotlibTransformer的核心是自注意力机制它允许模型在处理序列数据时动态地为不同位置的元素分配不同的权重。这种机制模拟了人类阅读时的注意力分配方式——我们会更关注句子中重要的词语而忽略不太相关的部分。自注意力机制与传统RNN的关键区别RNN必须按顺序处理序列难以并行化Transformer可以同时处理整个序列充分利用现代GPU的并行计算能力2. 实现位置编码由于Transformer没有循环结构它需要一种特殊的方式来保留序列中词语的位置信息。这就是位置编码Positional Encoding的作用。import torch import math def positional_encoding(max_len, d_model): 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) return pe # 示例生成长度为10维度为512的位置编码 pe positional_encoding(10, 512) print(pe.shape) # 输出: torch.Size([10, 512])位置编码的关键特性每个位置有唯一的编码编码值在-1到1之间编码可以扩展到任意长度的序列提示位置编码不需要训练它是固定的数学函数生成的。这种设计允许模型处理比训练时更长的序列。3. 实现自注意力机制自注意力是Transformer最核心的组件它由三个主要部分组成查询Query、键Key和值Value。import torch.nn as nn class SelfAttention(nn.Module): def __init__(self, embed_size, heads): super(SelfAttention, self).__init__() self.embed_size embed_size self.heads heads self.head_dim embed_size // heads assert (self.head_dim * heads embed_size), Embed size needs to be divisible by heads self.values nn.Linear(self.head_dim, self.head_dim, biasFalse) self.keys nn.Linear(self.head_dim, self.head_dim, biasFalse) self.queries nn.Linear(self.head_dim, self.head_dim, biasFalse) self.fc_out nn.Linear(heads * self.head_dim, embed_size) def forward(self, values, keys, query, mask): N query.shape[0] value_len, key_len, query_len values.shape[1], keys.shape[1], query.shape[1] # 分割嵌入维度到多个头 values values.reshape(N, value_len, self.heads, self.head_dim) keys keys.reshape(N, key_len, self.heads, self.head_dim) queries query.reshape(N, query_len, self.heads, self.head_dim) values self.values(values) keys self.keys(keys) queries self.queries(queries) # 计算注意力分数 energy torch.einsum(nqhd,nkhd-nhqk, [queries, keys]) if mask is not None: energy energy.masked_fill(mask 0, float(-1e20)) attention torch.softmax(energy / (self.embed_size ** (1/2)), dim3) # 应用注意力权重到值上 out torch.einsum(nhql,nlhd-nqhd, [attention, values]).reshape( N, query_len, self.heads * self.head_dim ) out self.fc_out(out) return out自注意力计算过程详解线性变换将输入分别转换为Q、K、V分数计算计算Q和K的点积然后缩放Softmax归一化得到注意力权重加权求和用注意力权重对V进行加权4. 构建Transformer块现在我们可以将自注意力机制与其他组件组合起来构建完整的Transformer块。class TransformerBlock(nn.Module): def __init__(self, embed_size, heads, dropout, forward_expansion): super(TransformerBlock, self).__init__() self.attention SelfAttention(embed_size, heads) self.norm1 nn.LayerNorm(embed_size) self.norm2 nn.LayerNorm(embed_size) self.feed_forward nn.Sequential( nn.Linear(embed_size, forward_expansion * embed_size), nn.ReLU(), nn.Linear(forward_expansion * embed_size, embed_size) ) self.dropout nn.Dropout(dropout) def forward(self, value, key, query, mask): attention self.attention(value, key, query, mask) # 残差连接和层归一化 x self.dropout(self.norm1(attention query)) forward self.feed_forward(x) out self.dropout(self.norm2(forward x)) return outTransformer块的关键组件多头自注意力捕获序列中不同位置的关系前馈网络增加模型的非线性表达能力层归一化稳定训练过程残差连接缓解梯度消失问题5. 构建完整Transformer模型现在我们可以将所有组件组合起来构建完整的Transformer模型架构。class Transformer(nn.Module): def __init__( self, src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, embed_size512, num_layers6, forward_expansion4, heads8, dropout0, devicecpu, max_length100 ): super(Transformer, self).__init__() self.encoder_embedding nn.Embedding(src_vocab_size, embed_size) self.decoder_embedding nn.Embedding(trg_vocab_size, embed_size) self.positional_encoding positional_encoding(max_length, embed_size) self.encoder_layers nn.ModuleList( [ TransformerBlock( embed_size, heads, dropoutdropout, forward_expansionforward_expansion ) for _ in range(num_layers) ] ) self.decoder_layers nn.ModuleList( [ TransformerBlock( embed_size, heads, dropoutdropout, forward_expansionforward_expansion ) for _ in range(num_layers) ] ) self.fc_out nn.Linear(embed_size, trg_vocab_size) self.dropout nn.Dropout(dropout) self.src_pad_idx src_pad_idx self.trg_pad_idx trg_pad_idx self.device device def make_src_mask(self, src): src_mask (src ! self.src_pad_idx).unsqueeze(1).unsqueeze(2) return src_mask.to(self.device) def make_trg_mask(self, trg): N, trg_len trg.shape trg_mask torch.tril(torch.ones((trg_len, trg_len))).expand( N, 1, trg_len, trg_len ) return trg_mask.to(self.device) def forward(self, src, trg): src_mask self.make_src_mask(src) trg_mask self.make_trg_mask(trg) src_embedded self.dropout( (self.encoder_embedding(src) self.positional_encoding[:src.shape[1], :]) ) trg_embedded self.dropout( (self.decoder_embedding(trg) self.positional_encoding[:trg.shape[1], :]) ) enc_out src_embedded for layer in self.encoder_layers: enc_out layer(enc_out, enc_out, enc_out, src_mask) dec_out trg_embedded for layer in self.decoder_layers: dec_out layer(enc_out, enc_out, dec_out, trg_mask) out self.fc_out(dec_out) return out6. 模型训练与评估有了完整的Transformer模型后我们需要设置训练流程和评估指标。# 示例训练代码 device torch.device(cuda if torch.cuda.is_available() else cpu) # 假设我们有一些示例数据 src_vocab_size 5000 trg_vocab_size 5000 src_pad_idx 0 trg_pad_idx 0 model Transformer( src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, devicedevice ).to(device) optimizer torch.optim.Adam(model.parameters(), lr3e-4) criterion nn.CrossEntropyLoss(ignore_indextrg_pad_idx) def train(model, iterator, optimizer, criterion, clip): model.train() epoch_loss 0 for i, batch in enumerate(iterator): src batch.src.to(device) trg batch.trg.to(device) optimizer.zero_grad() output model(src, trg[:, :-1]) output_dim output.shape[-1] output output.contiguous().view(-1, output_dim) trg trg[:, 1:].contiguous().view(-1) loss criterion(output, trg) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), clip) optimizer.step() epoch_loss loss.item() return epoch_loss / len(iterator)训练Transformer模型时需要注意的几个关键点学习率调度使用warmup策略逐步提高学习率梯度裁剪防止梯度爆炸批处理充分利用GPU并行能力正则化适当使用dropout防止过拟合注意在实际应用中你可能需要使用更大的数据集和更长时间的训练才能获得良好的性能。这个简化版主要用于教学目的帮助你理解Transformer的核心原理。7. 模型优化技巧要让Transformer模型发挥最佳性能可以考虑以下优化技巧学习率调度class CustomSchedule(torch.optim.lr_scheduler._LRScheduler): def __init__(self, optimizer, d_model, warmup_steps4000): self.d_model d_model self.warmup_steps warmup_steps super().__init__(optimizer) def get_lr(self): step self._step_count 1 arg1 step ** (-0.5) arg2 step * (self.warmup_steps ** (-1.5)) return [base_lr * (self.d_model ** (-0.5) * min(step ** (-0.5), step * self.warmup_steps ** (-1.5))) for base_lr in self.base_lrs]标签平滑减轻模型对预测结果的过度自信criterion nn.CrossEntropyLoss( ignore_indextrg_pad_idx, label_smoothing0.1 )混合精度训练减少内存占用并加速训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): output model(src, trg[:, :-1]) loss criterion(output, trg[:, 1:]) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()8. 实际应用中的注意事项在实际项目中应用Transformer模型时有几个关键因素需要考虑数据预处理文本清洗和标准化适当的tokenization策略子词分割如BPE处理罕见词模型大小选择小数据集4-6层512维度8个头中等数据集6-12层768维度12个头大数据集12-24层1024维度16个头硬件考虑GPU内存限制批处理大小多GPU训练策略混合精度训练推理优化模型量化减小部署体积缓存注意力计算结果束搜索(beam search)参数调优# 示例推理代码 def translate_sentence(sentence, model, device, max_length50): model.eval() # 简单的tokenization tokens sentence.lower().split() tokens [sos] tokens [eos] # 转换为索引 src_indexes [src_field.vocab.stoi[token] for token in tokens] src_tensor torch.LongTensor(src_indexes).unsqueeze(0).to(device) # 创建目标序列 trg_indexes [trg_field.vocab.stoi[sos]] for i in range(max_length): trg_tensor torch.LongTensor(trg_indexes).unsqueeze(0).to(device) with torch.no_grad(): output model(src_tensor, trg_tensor) pred_token output.argmax(2)[:,-1].item() trg_indexes.append(pred_token) if pred_token trg_field.vocab.stoi[eos]: break trg_tokens [trg_field.vocab.itos[i] for i in trg_indexes] return trg_tokens[1:]