别再死记硬背Transformer了！用Python+PyTorch手写一个简易版，5分钟搞懂注意力机制

用Python+PyTorch手写Transformer：5行代码拆解注意力机制

当你第一次听说Transformer时，可能被那些高大上的术语吓到——自注意力、多头机制、位置编码...但真相是，它的核心思想简单得令人发指。今天我们不谈论文不背公式，直接打开Colab，用不到50行代码实现一个可运行的Transformer编码器层。你会发现，那些看似复杂的矩阵运算，不过是几个张量拼接的游戏。

1. 准备工作：理解注意力机制的三个关键张量

在开始写代码前，我们需要明确三个核心概念：

• Query（Q）：当前正在处理的"问题"，就像你在搜索引擎输入的关键词
• Key（K）：所有可能匹配的"答案索引"，相当于网页的元数据
• Value（V）：实际的"答案内容"，也就是网页本身

它们的数学关系可以用这个公式表示：

Attention(Q, K, V) = softmax(QK^T/√d_k)V

先安装必要的库：

pip install torch matplotlib

2. 从零实现缩放点积注意力

让我们用PyTorch实现最基础的注意力计算。创建一个新的Python文件，导入必要的库：

import torch import torch.nn as nn import torch.nn.functional as F import math

接着定义缩放点积注意力函数：

def scaled_dot_product_attention(q, k, v, mask=None): # 计算Q和K的点积 matmul_qk = torch.matmul(q, k.transpose(-2, -1)) # 缩放因子 d_k = q.size()[-1] scaled_attention_logits = matmul_qk / math.sqrt(d_k) # 可选：应用mask（解码器使用） if mask is not None: scaled_attention_logits += (mask * -1e9) # softmax归一化 attention_weights = F.softmax(scaled_attention_logits, dim=-1) # 加权求和 output = torch.matmul(attention_weights, v) return output, attention_weights

测试我们的实现：

# 假设输入序列长度为4，嵌入维度为8 seq_len, d_model = 4, 8 q = torch.rand(1, seq_len, d_model) # 随机生成Query k = torch.rand(1, seq_len, d_model) # 随机生成Key v = torch.rand(1, seq_len, d_model) # 随机生成Value output, attn_weights = scaled_dot_product_attention(q, k, v) print("注意力权重形状:", attn_weights.shape) print("输出形状:", output.shape)

你会看到类似这样的输出：

注意力权重形状: torch.Size([1, 4, 4]) 输出形状: torch.Size([1, 4, 8])

3. 实现多头注意力机制

单头注意力就像只用一只眼睛看世界，而多头机制让我们能从多个角度理解输入。下面是实现方法：

class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() self.num_heads = num_heads self.d_model = d_model assert d_model % num_heads == 0 self.depth = d_model // num_heads # 定义线性变换层 self.wq = nn.Linear(d_model, d_model) self.wk = nn.Linear(d_model, d_model) self.wv = nn.Linear(d_model, d_model) self.dense = nn.Linear(d_model, d_model) def split_heads(self, x, batch_size): x = x.view(batch_size, -1, self.num_heads, self.depth) return x.transpose(1, 2) def forward(self, q, k, v, mask=None): batch_size = q.size(0) # 线性变换 q = self.wq(q) k = self.wk(k) v = self.wv(v) # 分割多头 q = self.split_heads(q, batch_size) k = self.split_heads(k, batch_size) v = self.split_heads(v, batch_size) # 计算缩放点积注意力 scaled_attention, attention_weights = scaled_dot_product_attention( q, k, v, mask) # 合并多头 scaled_attention = scaled_attention.transpose(1, 2) concat_attention = scaled_attention.reshape( batch_size, -1, self.d_model) # 最终线性变换 output = self.dense(concat_attention) return output, attention_weights

测试多头注意力：

d_model, num_heads = 512, 8 mha = MultiHeadAttention(d_model, num_heads) # 模拟输入 (batch_size=1, seq_len=10, d_model=512) x = torch.rand(1, 10, 512) output, attn = mha(x, x, x) print("多头注意力输出形状:", output.shape) print("注意力权重形状:", attn.shape)

4. 构建完整的Transformer编码器层

现在我们把多头注意力包装成一个完整的编码器层：

class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, dff, rate=0.1): super(EncoderLayer, self).__init__() self.mha = MultiHeadAttention(d_model, num_heads) self.ffn = nn.Sequential( nn.Linear(d_model, dff), nn.ReLU(), nn.Linear(dff, d_model) ) self.layernorm1 = nn.LayerNorm(d_model) self.layernorm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(rate) self.dropout2 = nn.Dropout(rate) def forward(self, x, mask=None): # 多头注意力 attn_output, _ = self.mha(x, x, x, mask) attn_output = self.dropout1(attn_output) out1 = self.layernorm1(x + attn_output) # 前馈网络 ffn_output = self.ffn(out1) ffn_output = self.dropout2(ffn_output) out2 = self.layernorm2(out1 + ffn_output) return out2

使用示例：

encoder_layer = EncoderLayer(d_model=512, num_heads=8, dff=2048) sample_input = torch.rand(1, 10, 512) # (batch_size, seq_len, d_model) sample_output = encoder_layer(sample_input) print("编码器层输出形状:", sample_output.shape) # 应该保持输入形状

5. 可视化注意力权重

理解Transformer最好的方式就是看它"看"到了什么。让我们可视化注意力权重：

import matplotlib.pyplot as plt def plot_attention_weights(attention, sentence, layer_name): fig = plt.figure(figsize=(16, 8)) attention = attention.squeeze(0).cpu().detach().numpy() ax = fig.add_subplot(1, 1, 1) ax.matshow(attention, cmap='viridis') fontdict = {'fontsize': 14} ax.set_xticks(range(len(sentence))) ax.set_yticks(range(len(sentence))) ax.set_xticklabels(sentence, fontdict=fontdict, rotation=90) ax.set_yticklabels(sentence, fontdict=fontdict) ax.set_xlabel('Key') ax.set_ylabel('Query') ax.set_title(f'Attention weights at {layer_name}') plt.tight_layout() plt.show()

测试可视化：

# 模拟输入句子 sentence = ["The", "animal", "didn't", "cross", "the", "street", "because", "it", "was", "too", "tired"] # 随机生成注意力权重 (模拟单头注意力) attention = torch.rand(1, len(sentence), len(sentence)) plot_attention_weights(attention, sentence, "layer1_head1")

6. 常见问题与调试技巧

在实现过程中，你可能会遇到以下问题：

维度不匹配错误

• 检查所有张量的形状，特别是矩阵乘法前后的维度
• 记住矩阵乘法规则：(m,n) @ (n,p) → (m,p)
• 使用print(tensor.shape)调试

梯度消失/爆炸

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

多头注意力的常见错误

• 确保d_model能被num_heads整除
• 分割多头后记得转置维度
• 合并多头时要恢复原始形状

性能优化技巧

# 使用PyTorch的优化操作 @torch.jit.script # 启用JIT编译 def optimized_attention(q, k, v): return F.scaled_dot_product_attention(q, k, v)

7. 扩展应用：自定义注意力变体

基础实现完成后，可以尝试这些有趣的变体：

局部注意力

def local_attention(q, k, v, window_size=3): # 只计算局部窗口内的注意力 seq_len = q.size(-2) mask = torch.ones_like(q @ k.transpose(-2, -1)) for i in range(seq_len): start = max(0, i - window_size // 2) end = min(seq_len, i + window_size // 2 + 1) mask[..., i, :start] = -1e9 mask[..., i, end:] = -1e9 return scaled_dot_product_attention(q, k, v, mask)

稀疏注意力

def sparse_attention(q, k, v, stride=2): # 每隔stride个token计算一次注意力 seq_len = q.size(-2) mask = torch.zeros_like(q @ k.transpose(-2, -1)) for i in range(0, seq_len, stride): mask[..., i, :] = 1 mask = (1.0 - mask) * -1e9 return scaled_dot_product_attention(q, k, v, mask)

线性注意力（降低复杂度）

def linear_attention(q, k, v): # 近似计算，复杂度O(n)而非O(n^2) kv = torch.einsum('nshd,nshm->nhmd', k, v) z = 1 / (torch.einsum('nlhd,nhd->nlh', q, k.sum(dim=1)) + 1e-6) return torch.einsum('nlhd,nhmd,nlh->nlhm', q, kv, z)