transformer-tutorial-code

🚀 Transformer模型实现教程

📖 概述

本代码库提供了一个基于 PyTorch 的 Transformer 模型完整实现，用于机器翻译任务。该实现遵循 Vaswani 等人在 2017 年论文《Attention is All You Need》中提出的标准 Transformer 架构，包含完整的训练、验证和推理流程。

🔧 核心模块说明

1. MultiHeadAttention（多头注意力机制）

位于 transformer/MultiHeadAttention.py，实现了标准的多头自注意力机制：

• 输入: q, k, v 形状为 (B, T, d_model)，其中 B 是批次大小，T 是序列长度
• 处理流程:
  1. 通过线性层将输入投影到 Q、K、V
  2. 将 d_model 维度拆分为 num_heads 个头，每个头维度为 d_k = d_model // num_heads
  3. 计算缩放点积注意力：scores = (Q @ K^T) / sqrt(d_k)
  4. 应用 mask（如果提供）
  5. Softmax 归一化得到注意力权重
  6. 加权求和得到每个头的输出
  7. 拼接所有头并通过输出投影层
• 输出: (B, T, d_model) 和注意力权重 (B, heads, T, T)

2. PositionalEncoder（位置编码）

位于 transformer/PositionalEncoder.py，实现正弦位置编码：

• 使用公式：PE(pos, 2i) = sin(pos / (10000^(2i/d_model))) 和 PE(pos, 2i+1) = cos(pos / (10000^(2i/d_model)))
• 预计算位置编码矩阵，在训练和推理时直接使用

3. Encoder（编码器）

位于 transformer/Encoder.py：

• 由 N 个 EncoderLayer 堆叠而成
• 每个 EncoderLayer 包含：
  • 多头自注意力层（PreNorm 结构）
  • 残差连接和 Dropout
  • 前馈网络（FFN）
  • 层归一化
• 输入：源序列 token IDs (B, T) 和源序列 mask
• 输出：编码后的特征表示 (B, T, d_model)

4. Decoder（解码器）

位于 transformer/Decoder.py：

• 由 N 个 DecoderLayer 堆叠而成
• 每个 DecoderLayer 包含：
  • 自注意力层：使用因果 mask（causal mask）防止看到未来信息
  • 交叉注意力层：关注编码器的输出
  • 前馈网络
  • 每层都有残差连接和层归一化
• 输入：目标序列 token IDs、编码器输出、源序列 mask 和目标序列 mask
• 输出：解码后的特征表示 (B, T, d_model)

5. Transformer（主模型）

位于 transformer/Transformer.py：

• 整合 Encoder 和 Decoder
• 提供 make_src_mask() 和 make_trg_mask() 方法生成 mask
• 最终输出层将解码器输出映射到目标词汇表大小
• 支持设备管理（CPU/GPU）

6. Modules（辅助模块）

位于 transformer/Modules.py：

• FFN（前馈网络）: 两层全连接网络，中间使用 ReLU 激活函数
• NormLayer（层归一化）: 自定义的层归一化实现，包含可学习的缩放和平移参数

7. BilingualDataset（双语数据集）

位于 dataset.py：

• 处理双语翻译数据（使用 opus_books 数据集）
• 自动进行 tokenization、padding 和 mask 生成
• 返回编码器输入、解码器输入、标签以及相应的 mask

🚀 使用方法

环境配置

使用 conda 环境运行：

conda activate ultralytics-main

配置文件

在 config.py 中可以配置以下参数：

• batch_size: 批次大小（默认 4）
• num_epochs: 训练轮数（默认 20）
• lr: 学习率（默认 1e-4）
• seq_len: 序列最大长度（默认 350）
• d_model: 模型维度（默认 512）
• N: Encoder/Decoder 层数（默认 6）
• heads: 注意力头数（默认 8）
• d_ff: 前馈网络维度（默认 2048）
• dropout: Dropout 比率（默认 0.1）
• lang_src: 源语言（默认 'en'）
• lang_trg: 目标语言（默认 'fr'）

训练模型

运行训练脚本：

python train.py

训练过程包括：

1. 数据加载：

   • 自动从 opus_books 数据集加载双语数据
   • 构建或加载 tokenizer（WordLevel）
   • 将数据分为训练集（90%）和验证集（10%）

2. 模型初始化：

   • 根据配置创建 Transformer 模型
   • 支持从检查点恢复训练（设置 config['preload']）

3. 训练循环：

   • 使用 Adam 优化器
   • CrossEntropyLoss 损失函数（忽略 PAD token，使用 label smoothing）
   • 每个 epoch 结束后自动保存模型权重到 weights/ 目录
   • 使用 TensorBoard 记录训练损失

4. 验证：

   • 每个 epoch 结束后在验证集上运行验证
   • 使用贪心解码（greedy decoding）生成翻译结果
   • 打印源文本、目标文本和预测文本的对比

推理（Inference）

使用 inference.ipynb Jupyter Notebook 进行推理：

1. 加载模型：

   from train import get_model, get_ds
   from config import get_config, get_weights_file_path
   
   config = get_config()
   train_dataloader, val_dataloader, tokenizer_src, tokenizer_trg = get_ds(config)
   model = get_model(config, len(tokenizer_src.get_vocab()), len(tokenizer_trg.get_vocab())).to(device)
   
   # 加载训练好的权重
   model_filename = get_weights_file_path(config, "29")  # 加载第 29 个 epoch 的权重
   state = torch.load(model_filename)
   model.load_state_dict(state['model_state_dict'])

2. 运行验证：

   from train import run_validation
   
   run_validation(
       model, 
       val_dataloader, 
       tokenizer_src, 
       tokenizer_trg, 
       config['seq_len'], 
       device, 
       lambda msg: print(msg),
       0,
       None,
       num_examples=2
   )

推理流程说明

推理使用**贪心解码（Greedy Decoding）**策略：

1. 编码阶段：

   • 将源序列输入编码器，得到编码器输出
   • 编码器输出在整个解码过程中保持不变（只需计算一次）

2. 解码阶段（自回归生成）：

   • 初始化解码器输入为 [SOS] token
   • 循环生成每个 token：
     • 使用当前解码器输入和编码器输出计算解码器输出
     • 通过输出层得到词汇表上的概率分布
     • 选择概率最大的 token（贪心策略）
     • 将新 token 添加到解码器输入
     • 如果生成 [EOS] token 或达到最大长度，停止生成
   • 返回生成的 token 序列

📊 Mask 机制说明

1. 源序列 Mask（Padding Mask）

• 作用：屏蔽填充 token（PAD），防止模型关注无效位置
• 形状：(B, 1, 1, T)，其中 1 表示有效位置，0 表示填充位置
• 应用位置：编码器的自注意力层和解码器的交叉注意力层

2. 目标序列 Mask（Causal Mask）

• 作用：防止解码器在生成当前 token 时看到未来的 token（自回归约束）
• 形状：(B, 1, T, T)，下三角矩阵，上三角为 0（屏蔽未来信息）
• 应用位置：解码器的自注意力层

🔍 关键设计细节

1. PreNorm vs PostNorm：本实现使用 PreNorm 结构（先归一化再计算注意力/前馈网络）

2. 位置编码：使用正弦位置编码，预计算并注册为 buffer（不参与梯度更新）

3. 设备管理：模型自动检测 CUDA 可用性，支持 CPU 和 GPU 训练

4. Mask 处理：支持从数据集传入 mask 或由模型自动生成

5. Tokenizer：使用 WordLevel tokenizer，支持特殊 token（[UNK], [PAD], [SOS], [EOS]）

📝 注意事项

• 确保 d_model 能被 heads 整除
• 序列长度不能超过 max_seq_len（在 PositionalEncoder 中预定义）
• 训练时使用 label_smoothing=0.1 提高模型泛化能力
• 模型权重保存在 weights/ 目录，文件名格式为 tmodel_{epoch:02d}.pth
• TensorBoard 日志保存在 runs/tmodel/ 目录

🎯 扩展建议

• 实现 Beam Search 解码（当前使用贪心解码）
• 添加 BLEU 分数评估
• 支持更多语言对
• 实现学习率调度器
• 添加梯度裁剪防止梯度爆炸

希望这个实现能帮助你深入理解 Transformer 的工作原理，并能够将其应用到自己的项目中！