基于PyTorch的Transformer模型Python实现详解

2026年1月8日 互联网

基于PyTorch的Transformer模型Python实现详解

Transformer模型凭借其高效的并行计算能力和对长序列的优秀处理能力，已成为自然语言处理领域的核心架构。本文将以PyTorch框架为基础，通过Python代码详细解析Transformer模型的关键组件实现，并提供完整的模型构建与训练流程，帮助开发者快速掌握这一技术的核心实现。

一、Transformer模型核心架构解析

Transformer模型的核心由编码器（Encoder）和解码器（Decoder）组成，两者均采用堆叠的多层结构。每个编码器层包含两个子层：多头自注意力机制（Multi-Head Self-Attention）和前馈神经网络（Feed Forward Network），而解码器层在此基础上增加了编码器-解码器注意力（Encoder-Decoder Attention）模块。

1.1 自注意力机制实现

自注意力机制是Transformer的核心创新，其通过计算输入序列中各位置与其他位置的关联权重，实现动态的上下文感知。以下是使用PyTorch实现缩放点积注意力（Scaled Dot-Product Attention）的代码示例：

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  5. class ScaledDotProductAttention(nn.Module):
  6.     def __init__(self, d_model):
  7.         super().__init__()
  8.         self.sqrt_d_k = torch.sqrt(torch.tensor(d_model, dtype=torch.float32))
  10.     def forward(self, Q, K, V, mask=None):
  11.         # Q, K, V的形状均为[batch_size, seq_len, d_model]
  12.         attn_scores = torch.bmm(Q, K.transpose(1, 2)) / self.sqrt_d_k
  13.         if mask is not None:
  14.             attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
  15.         attn_weights = F.softmax(attn_scores, dim=-1)
  16.         return torch.bmm(attn_weights, V)

1.2 多头注意力机制实现

多头注意力通过将输入投影到多个子空间并行计算注意力，增强模型对不同特征的捕捉能力。以下是多头注意力层的完整实现：

  1. class MultiHeadAttention(nn.Module):
  2.     def __init__(self, d_model, num_heads):
  3.         super().__init__()
  4.         self.d_model = d_model
  5.         self.num_heads = num_heads
  6.         self.d_head = d_model // num_heads
  8.         # 线性投影层
  9.         self.W_q = nn.Linear(d_model, d_model)
  10.         self.W_k = nn.Linear(d_model, d_model)
  11.         self.W_v = nn.Linear(d_model, d_model)
  12.         self.W_o = nn.Linear(d_model, d_model)
  14.     def split_heads(self, x):
  15.         # x形状:[batch_size, seq_len, d_model]
  16.         batch_size, seq_len, _ = x.size()
  17.         return x.view(batch_size, seq_len, self.num_heads, self.d_head)\
  18.                 .transpose(1, 2)  # [batch_size, num_heads, seq_len, d_head]
  20.     def forward(self, Q, K, V, mask=None):
  21.         # 线性投影
  22.         Q = self.W_q(Q)
  23.         K = self.W_k(K)
  24.         V = self.W_v(V)
  26.         # 分割多头
  27.         Q = self.split_heads(Q)
  28.         K = self.split_heads(K)
  29.         V = self.split_heads(V)
  31.         # 计算注意力
  32.         attn_output = ScaledDotProductAttention(self.d_head)(Q, K, V, mask)
  34.         # 合并多头
  35.         attn_output = attn_output.transpose(1, 2)\
  36.                      .contiguous()\
  37.                      .view(Q.size(0), -1, self.d_model)
  39.         # 输出投影
  40.         return self.W_o(attn_output)

二、Transformer编码器层实现

编码器层由多头自注意力、残差连接、层归一化和前馈网络组成。以下是完整的编码器层实现：

  1. class EncoderLayer(nn.Module):
  2.     def __init__(self, d_model, num_heads, d_ff):
  3.         super().__init__()
  4.         self.self_attn = MultiHeadAttention(d_model, num_heads)
  5.         self.ffn = nn.Sequential(
  6.             nn.Linear(d_model, d_ff),
  7.             nn.ReLU(),
  8.             nn.Linear(d_ff, d_model)
  9.         )
  10.         self.norm1 = nn.LayerNorm(d_model)
  11.         self.norm2 = nn.LayerNorm(d_model)
  12.         self.dropout = nn.Dropout(0.1)
  14.     def forward(self, x, mask=None):
  15.         # 自注意力子层
  16.         attn_output = self.self_attn(x, x, x, mask)
  17.         x = x + self.dropout(attn_output)
  18.         x = self.norm1(x)
  20.         # 前馈子层
  21.         ffn_output = self.ffn(x)
  22.         x = x + self.dropout(ffn_output)
  23.         x = self.norm2(x)
  25.         return x

三、位置编码与嵌入层实现

Transformer通过位置编码注入序列顺序信息，以下是正弦位置编码的实现：

  1. class PositionalEncoding(nn.Module):
  2.     def __init__(self, d_model, max_len=5000):
  3.         super().__init__()
  4.         position = torch.arange(max_len).unsqueeze(1)
  5.         div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
  6.                            (-math.log(10000.0) / d_model))
  7.         pe = torch.zeros(max_len, d_model)
  8.         pe[:, 0::2] = torch.sin(position * div_term)
  9.         pe[:, 1::2] = torch.cos(position * div_term)
  10.         self.register_buffer('pe', pe.unsqueeze(0))
  12.     def forward(self, x):
  13.         # x形状:[batch_size, seq_len, d_model]
  14.         return x + self.pe[:, :x.size(1)]

四、完整Transformer模型实现

结合上述组件，完整的Transformer编码器模型实现如下：

  1. class TransformerEncoder(nn.Module):
  2.     def __init__(self, vocab_size, d_model, num_heads, 
  3.                  num_layers, d_ff, max_len, dropout=0.1):
  4.         super().__init__()
  5.         self.embedding = nn.Embedding(vocab_size, d_model)
  6.         self.pos_encoding = PositionalEncoding(d_model, max_len)
  7.         self.layers = nn.ModuleList([
  8.             EncoderLayer(d_model, num_heads, d_ff) 
  9.             for _ in range(num_layers)
  10.         ])
  11.         self.dropout = nn.Dropout(dropout)
  12.         self.scale = torch.sqrt(torch.tensor(d_model, dtype=torch.float32))
  14.     def forward(self, src, mask=None):
  15.         # src形状:[batch_size, seq_len]
  16.         src = self.embedding(src) * self.scale
  17.         src = self.pos_encoding(src)
  18.         src = self.dropout(src)
  20.         for layer in self.layers:
  21.             src = layer(src, mask)
  23.         return src

五、模型训练与优化实践

5.1 训练流程设计

完整的训练流程包含数据预处理、模型初始化、损失计算和优化器配置：

  1. def train_model(model, train_loader, criterion, optimizer, device):
  2.     model.train()
  3.     total_loss = 0
  4.     for batch in train_loader:
  5.         src, tgt = batch
  6.         src, tgt = src.to(device), tgt.to(device)
  8.         optimizer.zero_grad()
  9.         output = model(src)
  10.         loss = criterion(output.view(-1, output.size(-1)), tgt.view(-1))
  12.         loss.backward()
  13.         optimizer.step()
  14.         total_loss += loss.item()
  16.     return total_loss / len(train_loader)

5.2 性能优化技巧

  1. 混合精度训练：使用torch.cuda.amp实现自动混合精度，减少显存占用并加速训练
  2. 梯度累积：对于大batch场景，可通过多次前向传播累积梯度后再更新参数
  3. 学习率调度：采用torch.optim.lr_scheduler实现动态学习率调整
  4. 分布式训练：使用torch.nn.parallel.DistributedDataParallel实现多GPU并行训练

六、实际应用中的注意事项

  1. 序列长度处理：对于变长序列，需通过填充（Padding）和掩码（Mask）机制处理
  2. 模型压缩：可通过知识蒸馏、量化等技术将大模型压缩为轻量级版本
  3. 部署优化：使用ONNX格式导出模型，配合TensorRT等推理引擎提升部署效率
  4. 超参数调优：重点关注d_model、num_heads和num_layers的组合效果

七、进阶实现方向

  1. 预训练模型集成：接入预训练的Transformer权重（如BERT、GPT）
  2. 多模态扩展：修改输入嵌入层以支持图像、音频等多模态数据
  3. 稀疏注意力：采用局部敏感哈希（LSH）等技术降低注意力计算复杂度
  4. 自适应计算：实现动态调整计算深度的机制，提升长序列处理效率

通过上述实现，开发者可以快速构建基于PyTorch的Transformer模型，并根据实际需求进行扩展和优化。在实际应用中，建议结合具体任务场景进行模型结构的调整和超参数的优化，以获得最佳性能表现。

最新文章