PyTorch 2.0 实现 Transformer 编码器从零构建 6 层模型并可视化注意力在自然语言处理领域Transformer 架构已经成为事实上的标准。本文将带您从零开始使用 PyTorch 2.0 实现一个完整的 6 层 Transformer 编码器并深入探索其核心组件——注意力机制的可视化。1. 环境准备与基础配置首先确保您已安装 PyTorch 2.0 或更高版本。我们将使用以下关键组件import torch import torch.nn as nn import torch.nn.functional as F import math import matplotlib.pyplot as plt import numpy as np为保持实验可复现性设置随机种子torch.manual_seed(42) np.random.seed(42)定义模型的基本参数class Config: def __init__(self): self.d_model 512 # 嵌入维度 self.n_heads 8 # 注意力头数量 self.d_ff 2048 # 前馈网络隐藏层维度 self.n_layers 6 # 编码器层数 self.dropout 0.1 # dropout概率 self.max_len 100 # 最大序列长度2. 核心组件实现2.1 缩放点积注意力这是 Transformer 中最基础的计算单元class ScaledDotProductAttention(nn.Module): def __init__(self, config): super().__init__() self.d_k config.d_model // config.n_heads self.scale math.sqrt(self.d_k) def forward(self, Q, K, V, maskNone): # Q, K, V形状: [batch_size, n_heads, seq_len, d_k] scores torch.matmul(Q, K.transpose(-2, -1)) / self.scale if mask is not None: scores scores.masked_fill(mask 0, -1e9) weights F.softmax(scores, dim-1) output torch.matmul(weights, V) return output, weights2.2 多头注意力机制多头注意力允许模型同时关注不同表示子空间的信息class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() self.d_model config.d_model self.n_heads config.n_heads self.d_k config.d_model // config.n_heads self.W_Q nn.Linear(config.d_model, config.d_model) self.W_K nn.Linear(config.d_model, config.d_model) self.W_V nn.Linear(config.d_model, config.d_model) self.W_O nn.Linear(config.d_model, config.d_model) self.attention ScaledDotProductAttention(config) self.dropout nn.Dropout(config.dropout) def split_heads(self, x): batch_size x.size(0) return x.view(batch_size, -1, self.n_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)) x, self.weights self.attention(q, k, v, mask) x x.transpose(1, 2).contiguous().view(Q.size(0), -1, self.d_model) return self.W_O(x)2.3 位置编码由于 Transformer 没有循环结构我们需要显式地注入位置信息class PositionalEncoding(nn.Module): def __init__(self, config): super().__init__() self.dropout nn.Dropout(config.dropout) position torch.arange(config.max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, config.d_model, 2) * (-math.log(10000.0) / config.d_model)) pe torch.zeros(config.max_len, config.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): x x self.pe[:x.size(1)] return self.dropout(x)2.4 前馈网络每个编码器层中的全连接前馈网络class FeedForward(nn.Module): def __init__(self, config): super().__init__() self.linear1 nn.Linear(config.d_model, config.d_ff) self.linear2 nn.Linear(config.d_ff, config.d_model) self.dropout nn.Dropout(config.dropout) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))2.5 编码器层完整的编码器层实现class EncoderLayer(nn.Module): def __init__(self, config): super().__init__() self.self_attn MultiHeadAttention(config) self.ffn FeedForward(config) self.norm1 nn.LayerNorm(config.d_model) self.norm2 nn.LayerNorm(config.d_model) self.dropout1 nn.Dropout(config.dropout) self.dropout2 nn.Dropout(config.dropout) def forward(self, x, maskNone): attn_output self.self_attn(x, x, x, mask) x x self.dropout1(attn_output) x self.norm1(x) ffn_output self.ffn(x) x x self.dropout2(ffn_output) x self.norm2(x) return x3. 完整 Transformer 编码器现在我们可以组装完整的 6 层编码器class TransformerEncoder(nn.Module): def __init__(self, config): super().__init__() self.embedding nn.Embedding(config.vocab_size, config.d_model) self.pos_encoding PositionalEncoding(config) self.layers nn.ModuleList([EncoderLayer(config) for _ in range(config.n_layers)]) self.norm nn.LayerNorm(config.d_model) def forward(self, src, src_maskNone): x self.embedding(src) x self.pos_encoding(x) for layer in self.layers: x layer(x, src_mask) return self.norm(x)4. 注意力可视化技术理解模型如何分配注意力是解释 Transformer 行为的关键。我们将实现三种可视化方法4.1 热力图绘制def plot_attention_heatmap(attention_weights, layer_idx, head_idx): plt.figure(figsize(10, 8)) plt.imshow(attention_weights, cmapviridis) plt.colorbar() plt.title(fLayer {layer_idx1} Head {head_idx1} Attention Weights) plt.xlabel(Key Position) plt.ylabel(Query Position) plt.show()4.2 注意力模式分析def analyze_attention_patterns(model, sample_input): model.eval() with torch.no_grad(): output model(sample_input) # 收集各层的注意力权重 attention_maps [] for layer in model.layers: attention_maps.append(layer.self_attn.weights.cpu().numpy()) return attention_maps4.3 交互式可视化def interactive_attention_visualization(attention_maps, tokens): import ipywidgets as widgets from IPython.display import display layer_slider widgets.IntSlider( value0, min0, maxlen(attention_maps)-1, descriptionLayer: ) head_slider widgets.IntSlider( value0, min0, maxattention_maps[0].shape[1]-1, descriptionHead: ) def update_plot(layer, head): fig, ax plt.subplots(figsize(12, 10)) im ax.imshow(attention_maps[layer][0, head], cmapviridis) ax.set_xticks(range(len(tokens))) ax.set_yticks(range(len(tokens))) ax.set_xticklabels(tokens, rotation90) ax.set_yticklabels(tokens) plt.colorbar(im) plt.show() widgets.interactive(update_plot, layerlayer_slider, headhead_slider)5. 模型训练与验证5.1 数据准备我们使用一个简单的文本分类任务作为示例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) def yield_tokens(data_iter): for _, text in data_iter: yield tokenizer(text) train_iter IMDB(splittrain) vocab build_vocab_from_iterator(yield_tokens(train_iter), specials[unk, pad]) vocab.set_default_index(vocab[unk]) text_pipeline lambda x: vocab(tokenizer(x)) label_pipeline lambda x: 1 if x pos else 05.2 训练循环def train(model, dataloader, criterion, optimizer, device): model.train() total_loss 0 for batch in dataloader: text, label batch.text.to(device), batch.label.to(device) optimizer.zero_grad() output model(text) loss criterion(output.mean(dim1), label.float()) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader)5.3 验证与测试def evaluate(model, dataloader, criterion, device): model.eval() total_loss 0 correct 0 with torch.no_grad(): for batch in dataloader: text, label batch.text.to(device), batch.label.to(device) output model(text) loss criterion(output.mean(dim1), label.float()) total_loss loss.item() pred (output.mean(dim1) 0.5).long() correct (pred label).sum().item() return total_loss / len(dataloader), correct / len(dataloader.dataset)6. 实际应用与优化技巧6.1 性能优化PyTorch 2.0 引入了多项优化技术# 启用混合精度训练 scaler torch.cuda.amp.GradScaler() # 使用torch.compile加速 model torch.compile(model) # 激活Flash Attention (需要兼容的GPU) torch.backends.cuda.enable_flash_sdp(True)6.2 注意力模式分析不同层的注意力头通常会学习不同的模式层数典型注意力模式功能描述低层局部注意力关注相邻词和短语结构中层句法注意力捕捉主谓一致等句法关系高层语义注意力关注跨句子的语义关联6.3 常见问题排查Transformer 训练中的常见问题及解决方案梯度消失/爆炸使用 Layer Normalization适当调整学习率梯度裁剪过拟合增加 Dropout 比例使用标签平滑早停策略训练不稳定使用学习率预热检查初始化方法验证注意力权重分布7. 扩展应用与进阶技巧7.1 迁移学习预训练 Transformer 的微调策略def fine_tune_pretrained(config, pretrained_path): # 加载预训练权重 pretrained_dict torch.load(pretrained_path) model TransformerEncoder(config) model_dict model.state_dict() # 过滤不匹配的键 pretrained_dict {k: v for k, v in pretrained_dict.items() if k in model_dict and v.size() model_dict[k].size()} model_dict.update(pretrained_dict) model.load_state_dict(model_dict) # 冻结部分层 for name, param in model.named_parameters(): if embedding in name or pos_encoding in name: param.requires_grad False return model7.2 模型压缩减小模型尺寸的技术对比技术压缩率精度损失实现难度量化4x低简单剪枝2-10x中中等知识蒸馏2-4x低复杂7.3 多模态扩展将 Transformer 应用于视觉任务class VisionTransformerEncoder(nn.Module): def __init__(self, config, image_size224, patch_size16): super().__init__() self.patch_embedding nn.Conv2d(3, config.d_model, kernel_sizepatch_size, stridepatch_size) num_patches (image_size // patch_size) ** 2 self.position_embedding nn.Parameter( torch.randn(1, num_patches 1, config.d_model)) self.cls_token nn.Parameter(torch.randn(1, 1, config.d_model)) self.transformer TransformerEncoder(config) def forward(self, x): B x.shape[0] x self.patch_embedding(x).flatten(2).transpose(1, 2) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat((cls_tokens, x), dim1) x x self.position_embedding return self.transformer(x)通过本文的实现您不仅掌握了 Transformer 编码器的核心原理还获得了可立即应用于实际项目的完整代码。理解注意力机制的可视化结果将帮助您调试模型并解释其决策过程。