PyTorch 2.0 实现 Transformer 编码器：从零构建 6 层模型并可视化注意力

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, mask=None): # 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, weights

2.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, mask=None): 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, mask=None): 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 x

3. 完整 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_mask=None): 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, cmap='viridis') plt.colorbar() plt.title(f"Layer {layer_idx+1} Head {head_idx+1} 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_maps

4.3 交互式可视化

def interactive_attention_visualization(attention_maps, tokens): import ipywidgets as widgets from IPython.display import display layer_slider = widgets.IntSlider( value=0, min=0, max=len(attention_maps)-1, description='Layer:' ) head_slider = widgets.IntSlider( value=0, min=0, max=attention_maps[0].shape[1]-1, description='Head:' ) def update_plot(layer, head): fig, ax = plt.subplots(figsize=(12, 10)) im = ax.imshow(attention_maps[layer][0, head], cmap='viridis') ax.set_xticks(range(len(tokens))) ax.set_yticks(range(len(tokens))) ax.set_xticklabels(tokens, rotation=90) ax.set_yticklabels(tokens) plt.colorbar(im) plt.show() widgets.interactive(update_plot, layer=layer_slider, head=head_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(split='train') 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 0

5.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(dim=1), 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(dim=1), label.float()) total_loss += loss.item() pred = (output.mean(dim=1) > 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 训练中的常见问题及解决方案：

1. 梯度消失/爆炸

   • 使用 Layer Normalization
   • 适当调整学习率
   • 梯度裁剪

2. 过拟合

   • 增加 Dropout 比例
   • 使用标签平滑
   • 早停策略

3. 训练不稳定

   • 使用学习率预热
   • 检查初始化方法
   • 验证注意力权重分布

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 model

7.2 模型压缩

减小模型尺寸的技术对比：

7.3 多模态扩展

将 Transformer 应用于视觉任务：

class VisionTransformerEncoder(nn.Module): def __init__(self, config, image_size=224, patch_size=16): super().__init__() self.patch_embedding = nn.Conv2d(3, config.d_model, kernel_size=patch_size, stride=patch_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), dim=1) x = x + self.position_embedding return self.transformer(x)

通过本文的实现，您不仅掌握了 Transformer 编码器的核心原理，还获得了可立即应用于实际项目的完整代码。理解注意力机制的可视化结果将帮助您调试模型并解释其决策过程。