原始 Transformer（PyTorch）：computer: = :rainbow:

此仓库包含原始 Transformer 论文的 PyTorch 实现（:link: Vaswani 等）。
其目标是让读者能够轻松上手并学习 Transformer 模型。

目录

• 什么是 Transformer？
• 理解 Transformer
• 机器翻译
• 设置
• 使用方法
• 硬件要求

什么是 Transformer

Transformer 最早由 Vaswani 等人在一篇开创性的论文《Attention Is All You Need》中提出（链接）。

你可能或多或少听说过 Transformer。比如GPT-3 和 BERT等知名模型 :unicorn:。其核心思想在于，他们证明了无需使用循环层或卷积层，仅通过简单的架构结合注意力机制就能发挥强大的作用。这带来了更优秀的长距离依赖建模能力，并且该架构本身高度可并行化（:computer::computer::computer:），从而显著提升计算效率！

以下是其简洁而优美的架构示意图：

理解 Transformer

本仓库旨在作为学习资源，帮助大家理解 Transformer 模型，因为原始的 Transformer 已经不再是当前的 SOTA 模型了。

为此，代码中添加了详尽的注释，并且我编写了 playground.py 文件，用于可视化一些难以用文字描述、但通过图像却能直观理解的概念。接下来就让我们开始吧！

位置编码

你能一眼看懂这个公式吗？

我也不行。不过，运行 playground.py 中的 visualize_positional_encodings() 函数后，我们得到如下可视化结果：

根据源语言或目标语言词元的位置，只需“选取这张图中的一行”，将其加到对应的嵌入向量中即可，就这么简单。当然，位置编码也可以通过学习获得，但显然这样实现起来更加优雅！ :nerd_face:

自定义学习率调度

同样地，你能瞬间读懂下面这个公式吗？

不行吧？那我就把它可视化出来：

现在是不是就很容易理解了呢？至于这部分是否对 Transformer 的成功至关重要，我表示怀疑。不过，它确实很酷，也让事情变得更复杂一些。 :nerd_face: （.set_sarcasm(True)）

注：模型维度基本上就是嵌入向量的大小，基准 Transformer 使用的是 512 维，更大的版本则是 1024 维。

标签平滑

初次听到“标签平滑”时，可能会觉得很难理解，但实际上并不复杂。通常情况下，我们会将目标词汇表分布设置为“独热编码”。也就是说，在 3 万个词（或你的词汇表大小是多少）中，只有一个位置的概率为 1，其余均为 0。

而在应用标签平滑时，我们不会直接将概率设为 1，而是例如设为 0.9，然后将剩余的概率均匀分配到其他位置上。 （如上图所示，在一个虚构的 4 个词的词汇表中，这些概率被表示为不同深浅的紫色条带——因此有 4 列。）

注：填充符（Pad token）的概率分布始终为零，因为我们不希望模型预测这些符号！

除了这个仓库之外（这还用说吗？），我还强烈推荐你阅读 Jay Alammar 的精彩博客：这里！

机器翻译

Transformer 最初是在 WMT-14 数据集 上针对 NMT（神经机器翻译）任务进行训练的：

• 英语到德语的翻译任务（BLEU 得分为 28.4）
• 英语到法语的翻译任务（BLEU 得分为 41.8）

目前，我则在规模小得多的 IWSLT 数据集 上，针对英德语言对训练了自己的模型。由于我本人会说这两种语言，调试和实验起来更加方便。

不久之后，我也计划在 WMT-14 数据集上训练模型，请关注 待办事项 部分。

---

总之！让我们来看看这个仓库究竟能为你做些什么吧！没错，它可以进行翻译！

以下是我基于 IWSLT 数据集训练的德英模型的一些简短翻译示例：

输入：Ich bin ein guter Mensch, denke ich.（“金句”：我认为自己是个好人）
输出：['<s>', 'I', 'think', 'I', "'m", 'a', 'good', 'person', '.', '</s>']
或者以人类可读的形式呈现：I think I'm a good person.

这其实相当不错！在我看来，甚至比 Google Translate 的“金句”翻译还要好。

---

当然也有失败的情况，比如：

输入：Hey Alter, wie geht es dir?（哥们儿，最近怎么样？）
输出：['<s>', 'Hey', ',', 'age', 'how', 'are', 'you', '?', '</s>']
或者以人类可读的形式呈现：Hey, age, how are you?

不过这也算不上完全糟糕，原因在于：

• 首先，模型是在 IWSLT 数据集上训练的，该数据集主要包含类似 TED 演讲的对话内容。
• “Alter” 是一种口语化的表达，意为“老兄”或“哥们儿”，但其字面意思确实是“年龄”。

类似的逻辑也适用于英德翻译模型。

设置

我们已经讨论了 Transformer 是什么以及它能为你做什么（以及其他用途）。
现在就让我们把它跑起来吧！请按照以下步骤操作：

1. git clone https://github.com/gordicaleksa/pytorch-original-transformer
2. 打开 Anaconda 控制台，进入项目目录：cd path_to_repo
3. 在项目目录下运行 conda env create（这将创建一个新的 conda 环境）。
4. 运行 activate pytorch-transformer（以便从控制台运行脚本，或在 IDE 中设置解释器）。

就这样！执行 environment.yml 文件后，应该可以开箱即用，自动处理所有依赖关系。
由于需要自动下载 SpaCy 的英语和德语统计模型，这个过程可能需要一些时间。

---

虽然 PyTorch 的 pip 包会自带某个版本的 CUDA/cuDNN，但强烈建议你事先在系统级别安装 CUDA，主要是为了确保 GPU 驱动程序的兼容性。 此外，我推荐使用 Miniconda 安装程序来在你的系统上部署 conda。 请参考 这篇设置指南 中的第 1 和第 2 步， 并为你的系统选择最新版本的 Miniconda 和 CUDA/cuDNN。

使用方法

选项 1：Jupyter Notebook

只需在 Anaconda 控制台中运行 jupyter notebook，它就会在您的默认浏览器中打开会话。
打开 The Annotated Transformer ++.ipynb，您就可以开始玩了！

---

注意： 如果您遇到 DLL 加载失败，在导入 win32api 时：指定的模块未找到 的错误
只需执行 pip uninstall pywin32，然后要么 pip install pywin32，要么 conda install pywin32 应该可以解决这个问题！

选项 2：使用您选择的 IDE

您只需要将您在 设置 部分创建的 Python 环境链接到您的 IDE 即可。

训练

要运行训练，请启动 training_script.py。有几个参数您需要指定：

• --batch_size - 这个参数很重要，应设置为不会导致 CUDA 内存不足的最大值。
• --dataset_name - 在 IWSLT 和 WMT14 之间选择（在添加多 GPU 支持之前，不建议使用 WMT14 见待办事项）。
• --language_direction - 在 E2G 和 G2E 之间选择。

因此，一个示例运行（从控制台）看起来如下：
python training_script.py --batch_size 1500 --dataset_name IWSLT --language_direction G2E

代码中有详细的注释，因此您应该能够理解训练过程是如何进行的。

该脚本会：

• 将检查点模型以 .pth 格式保存到 models/checkpoints/ 目录下。
• 将最终模型以 .pth 格式保存到 models/binaries/ 目录下。
• 下载 IWSLT/WMT-14 数据集（首次运行时），并将其放置在 data/ 目录下。
• 将 tensorboard 数据 保存到 runs/ 目录下，您只需在 Anaconda 中运行 tensorboard --logdir=runs 即可。
• 定期将一些训练元数据输出到控制台。

注意：torch text 中的数据加载速度较慢，因此我实现了一个自定义包装器，增加了缓存机制，使加载速度提高了约 30 倍！（首次运行时仍然会比较慢）

推理（翻译）

第二部分是关于使用这些模型进行翻译，并观察它们的表现！
要进行翻译，请启动 translation_script.py。有几个参数您需要设置：

• --source_sentence - 根据您指定的模型，这应该是英语或德语句子。
• --model_name - 预训练模型名称之一：iwslt_e2g、iwslt_g2e，或者您自己的模型(*)。
• --dataset_name - 与模型保持一致，如果模型是在 IWSLT 数据集上训练的，则设置为 IWSLT。
• --language_direction - 与模型保持一致，如果模型是用于英译德的，则设置为 E2G。

(*) 注意：在您训练完模型后，它会被保存到 models/binaries 目录下。请查看其文件名，并通过 --model_name 参数指定，以便用于翻译任务。如果您指定了一些预训练模型，它们将在您首次运行翻译脚本时 自动下载。

我也会在此处提供 IWSLT 预训练模型的链接：英译德 和 德译英。

就是这样！您还可以可视化注意力机制，详情请参阅 可视化注意力 部分。

评估 NMT 模型

我在训练过程中跟踪了三条曲线：

• 训练损失（KL 散度，批次均值）
• 验证损失（KL 散度，批次均值）
• BLEU-4 分数

BLEU 是一种基于 n-gram 的指标，用于定量评估机器翻译模型的质量。
我使用了由优秀的 nltk Python 模块提供的 BLEU-4 指标。

当前结果：模型均训练了 20 个 epoch（DE 代表德语，即 German in German :nerd_face:）：

模型	BLEU 分数	数据集
基准 Transformer（EN-DE）	27.8	IWSLT 验证集
基准 Transformer（DE-EN）	33.2	IWSLT 验证集
基准 Transformer（EN-DE）	x	WMT-14 验证集
基准 Transformer（DE-EN）	x	WMT-14 验证集

这些结果是使用贪婪解码得到的，因此是一个偏保守的估计。我计划很快加入束搜索解码 见待办事项。

重要提示： 初始化对 Transformer 的性能影响很大！我最初认为其他实现中使用的 Xavier 初始化只是一种任意的启发式方法，PyTorch 的默认初始化就足够了——但我错了：

您可以看到这里有 3 次运行，下面的 2 次使用了 PyTorch 的默认初始化（其中一次使用 KL 散度损失的 mean，另一次使用 batchmean），而上面的一次则使用了 Xavier 均匀 初始化！

---

想法：您也可以定期为一组参考源句生成翻译结果。
这样可以为您提供关于 Transformer 行为的定性见解，尽管我没有这样做。类似的做法也常用于难以定量评估模型的情况，例如 GANs 和 NST 领域。

使用 TensorBoard 跟踪

上述图表是我使用 Azure ML 运行时截取的一部分，但在本地运行时，我会使用 TensorBoard。

只需在 Anaconda 控制台中运行 tensorboard --logdir=runs，即可在训练过程中跟踪各项指标。

可视化注意力

您可以使用 translation_script.py，并将 --visualize_attention 设置为 True，以进一步了解您的模型在源句和目标句中“关注”的内容。

以下是针对输入句子 Ich bin ein guter Mensch, denke ich. 所获得的注意力图。

这些属于编码器的第 6 层。您可以看到所有的 8 个多头注意力头。

而这张图则属于解码器第 6 层的自注意力 MHA（多头注意力）模块。
您会注意到一个有趣的 三角形模式，这是由于目标词无法“向前看”造成的！

第三种类型的 MHA 模块是源端注意力模块，其图示与您之前看到的编码器注意力图相似。
您可以根据自己的节奏随意探索这些注意力图！

注意：显然，这个模型存在一些偏差问题，但在这里我们暂不深入分析。

硬件要求

如果你想在 WMT-14 数据集上训练 Transformer 模型，确实需要一台性能不错的硬件。

作者们使用了：

• 在 8 张 P100 GPU 上花费 12 小时 来训练基准模型，而训练大型模型则用了 3.5 天。

如果我的计算没错的话，这相当于基准模型大约进行了 ~19 个 epoch（共 10 万个 step，每个 step 约有 2.5 万个 token，而 WMT-14 总共有约 1.3 亿个源端/目标端 token），而大型模型则是其三倍，即 30 万个 step。

另一方面，在 IWSLT 数据集上训练模型则要容易得多。我在自己的 RTX 2080 机器上（8 GB 显存）的训练时间为：

• 每个 epoch 约 13.2 分钟（1500 个 token 的 batch 大小）； 而在 Azure ML 的 K80s 上（24 GB 显存）则约为：
• 每个 epoch 约 34 分钟（1500 个 token 的 batch 大小）。

我本可以把 K80s 的 batch 大小提高到 3500+ 个 token，但遇到了一些 CUDA 内存不足的问题。

待办事项：

最后还有一些待办事项，我希望很快就能添加上去：

• 多 GPU/多节点训练支持（这样你就可以在 WMT-14 数据集上训练 19 个 epoch 的模型）；
• 束搜索解码（事实证明实现起来并不那么容易！）；
• BPE 和共享的源目标词汇表（我现在使用的是 SpaCy）。

这个仓库已经具备了所需的一切功能，这些只是额外的补充。我已经测试过从环境搭建到自动下载模型等所有内容。

视频学习资料

如果你在理解代码时遇到困难，我在这段视频中对这篇论文做了深入的概述：点击观看

我还有一些其他视频可以帮助你更好地理解 Transformer 模型：

• 我理解 NLP/Transformer 的方法；
• 另一篇关于该论文的概述（层次稍高一点）；
• 关于该项目开发过程的案例分析。

致谢

在开发这个项目的过程中，以下资源对我很有帮助：

• The Annotated Transformer；
• PyTorch 官方实现。

我在设计模型时从 The Annotated Transformer 中获得了一些灵感，但它的内容比较难懂，而且还存在一些 bug。它主要是为研究人员编写的。希望这个仓库也能让普通大众更容易理解 Transformer 模型！ :nerd_face:

引用

如果你觉得这段代码有用，请引用如下：

@misc{Gordić2020PyTorchOriginalTransformer,
  author = {Gordić, Aleksa},
  title = {pytorch-original-transformer},
  year = {2020},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/gordicaleksa/pytorch-original-transformer}},
}

与我联系

如果你想在生活中多接触一些与人工智能相关的内容 :nerd_face:，不妨考虑：

• 订阅我的 YouTube 频道 The AI Epiphany :bell:；
• 在 LinkedIn 和 Twitter 上关注我 :bulb:；
• 在 Medium 上关注我 :books: :heart:。

许可证

pytorch-original-transformer 快速上手指南

本指南基于 pytorch-original-transformer 项目，旨在帮助开发者快速理解并运行原始 Transformer 架构（Vaswani et al. "Attention Is All You Need"）的 PyTorch 实现。该项目代码注释详细，非常适合学习 Transformer 内部机制及进行机器翻译实验。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 环境: 推荐安装 Miniconda 或 Anaconda 以管理依赖。
• GPU 支持 (可选但推荐):
  • 虽然项目包含 PyTorch 自带的 CUDA 版本，但强烈建议预先安装系统级的 CUDA 和 cuDNN 以确保 GPU 驱动兼容性和最佳性能。
  • 训练 Transformer 模型对显存有一定要求，建议使用支持 CUDA 的 NVIDIA 显卡。
• 网络环境: 首次运行时会自动下载 SpaCy 统计模型（英文/德文）及数据集，请确保网络畅通。

安装步骤

本项目通过 Conda 环境文件自动处理所有依赖关系。请按以下步骤操作：

1. 克隆仓库

   git clone https://github.com/gordicaleksa/pytorch-original-transformer
   cd pytorch-original-transformer
   

2. 创建 Conda 环境 在项目根目录下运行以下命令，这将根据 environment.yml 创建一个名为 pytorch-transformer 的新环境（首次运行可能需要几分钟下载依赖）：

   conda env create
   

3. 激活环境

   conda activate pytorch-transformer
   

   注：如果您在 IDE（如 PyCharm 或 VS Code）中开发，请将解释器设置为该环境下的 Python。

   Windows 用户注意：如果在运行 Jupyter Notebook 时遇到 DLL load failed while importing win32api 错误，请执行以下命令修复：

   pip uninstall pywin32
   pip install pywin32
   

基本使用

安装完成后，您可以通过以下两种方式使用本项目：交互式学习（Jupyter Notebook）或命令行训练/推理。

方式一：交互式学习（推荐初学者）

项目提供了一个详细的注解版 Notebook，可视化了位置编码、学习率调度等核心概念。

1. 启动 Jupyter Notebook：

   jupyter notebook
   

2. 在浏览器中打开 The Annotated Transformer ++.ipynb 文件。
3. 按顺序运行单元格，即可观察数据流动和模型内部状态。

方式二：命令行训练与推理

1. 训练模型 (Training)

使用 training_script.py 开始训练。以下示例展示了如何在 IWSLT 数据集上训练一个 德语到英语 (G2E) 的翻译模型：

python training_script.py --batch_size 1500 --dataset_name IWSLT --language_direction G2E

关键参数说明：

• --batch_size: 根据显存大小调整，避免 CUDA OOM 错误。
• --dataset_name: 选择 IWSLT (较小，适合调试) 或 WMT14 (较大，需多卡支持)。
• --language_direction: E2G (英译德) 或 G2E (德译英)。

训练过程中，脚本会自动：

• 下载数据集至 data/ 目录。
• 保存检查点至 models/checkpoints/。
• 生成 TensorBoard 日志至 runs/ (查看命令：tensorboard --logdir=runs)。

2. 模型推理/翻译 (Inference)

训练完成后（或使用预训练模型），使用 translation_script.py 进行翻译测试。

使用预训练模型示例：

python translation_script.py --source_sentence "Ich bin ein guter Mensch, denke ich." --model_name iwslt_g2e --dataset_name IWSLT --language_direction G2E

使用自定义训练模型： 如果您训练了自己的模型，它会被保存在 models/binaries/ 目录下。将 --model_name 替换为该文件名（不含 .pth 后缀）即可：

python translation_script.py --source_sentence "Your input sentence" --model_name your_trained_model_name --dataset_name IWSLT --language_direction G2E

输出示例： 程序将输出分词后的结果及人类可读的翻译句子，例如： I think I'm a good person.

---

提示：项目默认使用 Xavier Uniform 初始化，这对 Transformer 的性能至关重要，请勿随意更改为默认初始化方式。