用PyTorch从零实现Tiny Transformer：手把手教你构建简化版注意力模型

用PyTorch从零实现Tiny Transformer：手把手教你构建简化版注意力模型

在深度学习领域，Transformer架构已经彻底改变了序列建模的范式。不同于传统RNN的串行处理方式，Transformer通过自注意力机制实现了并行化计算，大幅提升了长序列处理的效率。本文将带您用PyTorch构建一个精简但功能完整的Tiny Transformer，通过代码级实现深入理解这一革命性架构的核心原理。

1. 理解Transformer的核心组件

1.1 自注意力机制的本质

自注意力机制的核心思想是让序列中的每个元素都能"关注"序列中的所有其他元素，并通过动态权重计算来决定信息交互的重要性。这种机制解决了传统RNN难以捕捉长距离依赖的问题。

def scaled_dot_product_attention(Q, K, V, mask=None): """ Q: Query矩阵 [batch_size, seq_len, d_k] K: Key矩阵 [batch_size, seq_len, d_k] V: Value矩阵 [batch_size, seq_len, d_v] """ d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attention = torch.softmax(scores, dim=-1) return torch.matmul(attention, V)

关键点说明：

• Q(查询)、K(键)、V(值)矩阵共同决定了注意力的分布
• 缩放因子1/√d_k防止点积结果过大导致softmax梯度消失
• 可选的mask机制用于处理变长序列或防止信息泄露

1.2 多头注意力的并行计算优势

多头注意力通过将输入投影到多个子空间，使模型能够同时关注不同位置的不同特征表示：

class MultiHeadAttention(nn.Module): def __init__(self, d_model=512, num_heads=8): super().__init__() assert d_model % num_heads == 0 self.d_k = d_model // num_heads self.num_heads = 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) def split_heads(self, x): batch_size = x.size(0) return x.view(batch_size, -1, self.num_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)) attn_output = scaled_dot_product_attention(Q, K, V, mask) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(attn_output.size(0), -1, self.num_heads * self.d_k) return self.W_o(attn_output)

提示：实际应用中，num_heads通常选择4-16之间的值，d_model建议设置为64的倍数以便均匀分割

2. 构建Transformer的基础模块

2.1 位置编码：弥补无递归结构的缺陷

由于Transformer没有循环结构，需要显式地注入位置信息：

class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=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): return x + self.pe[:x.size(1)]

为什么使用正弦/余弦函数？

• 可以学习到相对位置关系
• 对长序列有良好的外推性
• 计算高效且易于优化

2.2 前馈网络的非线性变换

每个注意力层后都接一个两层的全连接网络：

class FeedForward(nn.Module): def __init__(self, d_model, d_ff=2048, dropout=0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(d_ff, d_model) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))

3. 组装编码器与解码器

3.1 编码器层的堆叠结构

class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): attn_output = self.self_attn(x, x, x, mask) x = self.norm1(x + self.dropout(attn_output)) ffn_output = self.ffn(x) return self.norm2(x + self.dropout(ffn_output))

3.2 解码器层的特殊设计

解码器除了自注意力外，还增加了对编码器输出的交叉注意力：

class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.cross_attn = MultiHeadAttention(d_model, num_heads) self.ffn = FeedForward(d_model, d_ff, dropout) 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, enc_output, src_mask=None, tgt_mask=None): # 自注意力（带掩码） attn1 = self.self_attn(x, x, x, tgt_mask) x = self.norm1(x + self.dropout(attn1)) # 交叉注意力（编码器输出作为K,V） attn2 = self.cross_attn(x, enc_output, enc_output, src_mask) x = self.norm2(x + self.dropout(attn2)) # 前馈网络 ffn_output = self.ffn(x) return self.norm3(x + self.dropout(ffn_output))

4. 完整Tiny Transformer实现

4.1 模型组装与初始化

class TinyTransformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, num_encoder_layers=6, num_decoder_layers=6, d_ff=2048, dropout=0.1): super().__init__() self.encoder_embedding = nn.Embedding(src_vocab_size, d_model) self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model) self.pos_encoding = PositionalEncoding(d_model) self.encoder_layers = nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_encoder_layers) ]) self.decoder_layers = nn.ModuleList([ DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_decoder_layers) ]) self.fc_out = nn.Linear(d_model, tgt_vocab_size) self.dropout = nn.Dropout(dropout) def encode(self, src, src_mask): src = self.dropout(self.pos_encoding(self.encoder_embedding(src))) for layer in self.encoder_layers: src = layer(src, src_mask) return src def decode(self, tgt, enc_output, src_mask, tgt_mask): tgt = self.dropout(self.pos_encoding(self.decoder_embedding(tgt))) for layer in self.decoder_layers: tgt = layer(tgt, enc_output, src_mask, tgt_mask) return tgt def forward(self, src, tgt, src_mask=None, tgt_mask=None): enc_output = self.encode(src, src_mask) dec_output = self.decode(tgt, enc_output, src_mask, tgt_mask) return self.fc_out(dec_output)

4.2 训练技巧与优化策略

学习率调度：使用Transformer特有的warmup策略

class TransformerOptimizer: def __init__(self, optimizer, d_model, warmup_steps=4000): self.optimizer = optimizer self.d_model = d_model self.warmup_steps = warmup_steps self.current_step = 0 def step(self): self.current_step += 1 lr = self.d_model ** -0.5 * min( self.current_step ** -0.5, self.current_step * self.warmup_steps ** -1.5 ) for param_group in self.optimizer.param_groups: param_group['lr'] = lr self.optimizer.step()

批处理与掩码生成：

def create_padding_mask(seq): return (seq != 0).unsqueeze(1).unsqueeze(2) # [batch_size, 1, 1, seq_len] def create_look_ahead_mask(size): return torch.triu(torch.ones(size, size), diagonal=1).bool()

5. 实战应用与性能调优

5.1 模型初始化技巧

def initialize_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0, std=m.embedding_dim ** -0.5) model.apply(initialize_weights)

5.2 常见问题排查指南

5.3 性能优化技巧

• 混合精度训练：减少显存占用，加速计算

scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(src, tgt) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

• 梯度累积：模拟更大batch size

for i, (src, tgt) in enumerate(train_loader): loss = model(src, tgt) loss = loss / accumulation_steps loss.backward() if (i+1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()