Transformer代码理解（注释版本）

 原文链接http://nlp.seas.harvard.edu/annotated-transformer/

# !pip install -r requirements.txt
# # Uncomment for colab
# #
# !pip install -q torchdata==0.3.0 torchtext==0.12 spacy==3.2 altair GPUtil
# !python -m spacy download de_core_news_sm
# !python -m spacy download en_core_web_sm

import os
from os.path import exists
import torch
import torch.nn as nn
from torch.nn.functional import log_softmax, pad
import math
import copy
import time
from torch.optim.lr_scheduler import LambdaLR
import pandas as pd
import altair as alt
from torchtext.data.functional import to_map_style_dataset
from torch.utils.data import DataLoader
from torchtext.vocab import build_vocab_from_iterator
import torchtext.datasets as datasets
import spacy
import GPUtil
import warnings
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP


# Set to False to skip notebook execution (e.g. for debugging)
warnings.filterwarnings("ignore")
RUN_EXAMPLES = True

# Some convenience helper functions used throughout the notebook


def is_interactive_notebook():
    return __name__ == "__main__"


def show_example(fn, args=[]):
    if __name__ == "__main__" and RUN_EXAMPLES:
        return fn(*args)


def execute_example(fn, args=[]):
    if __name__ == "__main__" and RUN_EXAMPLES:
        fn(*args)


class DummyOptimizer(torch.optim.Optimizer):
    def __init__(self):
        self.param_groups = [{"lr": 0}]
        None

    def step(self):
        None

    def zero_grad(self, set_to_none=False):
        None


class DummyScheduler:
    def step(self):
        None

Background 

减少顺序计算的目标也构成了Extended Neural GPU、ByteNet和ConvS2S的基础，所有这些模型都使用卷积神经网络作为基本构建块，并行计算所有输入和输出位置的隐藏表示。在这些模型中，将两个任意输入或输出位置的信号关联起来的所需操作数量随着位置之间的距离而增长，ConvS2S线性增长，ByteNet对数增长。这使得学习远距离位置之间的依赖关系更加困难。在Transformer中，这减少到常数个操作，尽管由于平均注意力加权位置导致有效分辨率降低，我们通过多头注意力来抵消这种影响。

自注意力（有时称为内部注意力）是一种将单个序列的不同位置相关联的机制，以计算该序列的表示。自注意力已成功用于各种任务，包括阅读理解、抽象总结、文本蕴含和学习任务无关的句子表示。端到端记忆网络基于循环注意力机制而不是序列对齐的循环，已证明在简单语言问答和语言建模任务上表现良好。

然而，据我们所知，Transformer是第一个完全依靠自注意力来计算其输入和输出表示的转导模型，没有使用序列对齐的RNN或卷积。

 Model Architecture

编码器encoder将输入x=(x1，x2...xn)映射到序列z=(z1,z2...zn),然后解码器产生输出序列(y1,y2...yn)

class EncoderDecoder(nn.Module):
    """
    #这是一个标准的编码器-解码器架构，并且是其他许多模型的基础
    A standard Encoder-Decoder architecture. Base for this and many
    other models.
    """

    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder#编码器
        self.decoder = decoder#解码器
        self.src_embed = src_embed#源语言序列嵌入层
        self.tgt_embed = tgt_embed#目标语言序列嵌入层
        self.generator = generator#生成器

    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
    
    #首先将输入的源语言序列通过源语言嵌入层进行嵌入，然后通过编码器进行编码
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    #首先将目标语言序列通过目标语言嵌入层进行嵌入，然后调用解码器对嵌入后的目标语言序列以及编码结果进行解码 memory是编码器输出结果 
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

class Generator(nn.Module):
    #将模型输出的隐藏表示转化为词汇表上的概率分布。
    "Define standard linear + softmax generation step."

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        #proj是一个线性变换层（nn.Linear），它将输入的维度为d_model的隐藏表示映射到词汇表大小为vocab的空间上
        self.proj = nn.Linear(d_model, vocab)
       
    def forward(self, x):
        #接收一个张量x作为输入，将x通过线性变换self.proj得到输出，然后通过log_softmax函数进行softmax操作并返回结果。在这里，dim=-1表示沿着最后一个维度进行softmax操作。
        return log_softmax(self.proj(x), dim=-1)

 Encoder and Decoder Stacks

Encoder

def clones(module, N):
    "Produce N identical layers."
    #创建了一个包含 N 个 module 的深拷贝的 nn.ModuleList 对象
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

每两个子层（多头注意力层和归一化层）添加一个残差连接，然后加层归一化： 

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."

    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        #创建了两个可学习的参数 a_2 和 b_2，它们分别用于缩放和平移归一化后的特征
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        #很小的数值eps，用于防止分母为0
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)#计算均值
        std = x.std(-1, keepdim=True)#计算标准差
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2#层归一化

编码器： 

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        #调用clones函数，复制n个相同的模块，存储在self.layers中
        self.layers = clones(layer, N)
        #对经过所有层处理后的输出进行归一化处理
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        #接收输入张量x和掩码mask，然后将输入依次通过堆叠的N个层。最后，对经过所有层处理后的输出进行归一化处理
        #x是源语言序列嵌入向量，mask是源语言序列遮罩
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

 子层（自注意力层和前馈神经网络层）的输出：LayerNorm(x+dropout(Sublayer(x)))

注：按照代码应该是x+dropout(Sublayer(LayerNorm(x)))，此处存疑

为了方便残差连接，模型中的所有子层以及嵌入层都产生维度为d_model = 512的输出

编码器/解码器中的一层的子层容器：

#一个子层容器，sublayer是自注意力层或者前馈神经网络层
class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """

    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        #接收输入张量x和一个子层sublayer作为输入，首先对x进行层归一化处理，然后通过sublayer进行前向计算并应用Dropout操作，最后将原始输入x与处理后的结果相加，实现了残差连接
        #self.norm(x)对上个子层的输出层归一化，sublayer()经过这个子层，然后随机失活，然后+x进行残差连接
        return x + self.dropout(sublayer(self.norm(x)))

 编码器中的一层：

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"

    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn#自注意力层
        self.feed_forward = feed_forward#前馈神经网络层
        self.sublayer = clones(SublayerConnection(size, dropout), 2)#创建了两个SublayerConnection对象，即两个子层容器
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))#实现自注意力层
        return self.sublayer[1](x, self.feed_forward)#返回前馈神经网络层的输出

 Decoder

 解码器：

class Decoder(nn.Module):
    "Generic N layer decoder with masking."

    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        #x是目标语言序列嵌入向量 memory是编码器输出结果  src_mask是源语言序列遮罩 tgt_mask是目标语言序列遮罩
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

解码器中的一层： 

class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"

    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn#自注意力层
        self.src_attn = src_attn#源注意力层
        self.feed_forward = feed_forward#前馈神经网络层
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        "Follow Figure 1 (right) for connections."
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

遮罩矩阵： 

def subsequent_mask(size):
    "Mask out subsequent positions."
    #创建一个遮盖矩阵，用于在注意力机制中屏蔽后续位置的信息
    attn_shape = (1, size, size)
    #利用PyTorch的triu函数创建了一个上三角矩阵，diagonal参数指定了对角线的偏移量，将主对角线以下的元素都置为0，而主对角线及以上的元素保持不变，最后将得到的矩阵转换为torch.uint8类型
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )
    return subsequent_mask == 0 #返回的矩阵由等于0的元素构成

 Attention

查询query和键key点积，然后除以根号dk，然后使用softmax函数计算出权重，然后与值value点积得到输出

 查询query和键key维度为dk，值value维度为dv

多头注意力计算：

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    #mask参数用于进行遮盖操作，如果指定了mask，使用masked_fill方法将mask为0的位置的分数替换为一个非常小的负数（-1e9），从而实现遮罩操作。
    #dropout参数用于进行Dropout操作，如果指定了dropout，则在计算softmax后进行Dropout操作。
    d_k = query.size(-1)
    #transpose(-2, -1)表示将数组或矩阵的倒数第二维和倒数第一维进行交换，实现转置
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

         两种最常用的注意力函数是加性注意力和点积（乘法）注意力。点积注意力除了缩放因子根号dk与我们的算法完全相同，加性注意力使用具有单个隐藏层的前馈网络来计算兼容性函数。虽然这两种机制在理论上的复杂性相似，但点积注意力在实际中更快、更节省空间，因为它可以使用高度优化的矩阵乘法代码来实现。

        虽然对于小的dk值，这两种机制的表现相似，但加性注意力在不进行缩放的情况下优于点积注意力，对于大的dk值。我们怀疑对于大的dk值，点积的幅度会增大（q,k服从（0，1）分布，q*k服从（0，dk）分布），将softmax函数推入其梯度极小的区域，为了抵消这种影响，我们对点积进行缩放，除以根号dk。

 维度信息：

d_model=512,h=8,dk=dv=d_model/h=64

Q,K,V:(d_model,d_model)  注：Q,K,V不应该是(L,d_model)吗，此处存疑

WQ:(d_model,dk) Q':(d_model,dk)

WK:(d_model,dk) K':(d_model,dk)

WV:(d_model,dv) V':(d_model,dv)

headi:(d_model,dv)

concat（head1,...,headh）:(d_model,dv*h)

WO:(dv*h,d_model) 

由于每个头的维度减小，总计算成本与具有完整维度的单头注意力机制相似。

自注意力层：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0#确保 d_model 能够被 h 整除，否则会抛出异常
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)#四个线性变换层
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        #判断是否有 mask 参数，如果存在，则对 mask 进行维度调整以便后续计算使用
        #unsqueeze(1)的作用是在mask张量的第一个维度（索引从0开始）前插入一个新的维度，使得原本的形状从[a, b, c]变为[a, 1, b, c]
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)#获取输入 query 的 batch 大小，获取第一个维度

        # 1) Do all the linear projections in batch from d_model => h x d_k
        #列表推导式，使用前三个线性层lin(x)将query, key, value分别进行线性变换，类似于乘三个权重矩阵
        #.view()将输出张量重新塑形为四维张量，nbatches批大小，-1推断自输入尺寸，self.h多头注意力的头数，self.d_k每个头的维度，.transpose(1, 2)将张量的第二维和第三维互换
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]
        #输出维度(nbatches, self.h, seq_len, self.d_k) 例如：(16, 8, 10, 64)

        # 2) Apply attention on all the projected vectors in batch.
        #调用attention 的函数实现多头注意力机制的计算，x是计算结果，self.attn是相似度（权重）矩阵
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )

        # 3) "Concat" using a view and apply a final linear.
        #将注意力向量 x 拼接起来，.contiguous()确保张量在内存中是连续的，.view()将张量重新塑形（拼接）为三维
        #输出维度(nbatches, seq_len, self.h*self.d_k) 例如：(16, 10, 8*64)
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(nbatches, -1, self.h * self.d_k)
        )
        del query
        del key
        del value
        return self.linears[-1](x)#将拼接后向量应用最后一个线性变换并返回输出

查询来自于解码器前面的层，键和值来自于编码器的输出。这意味着解码器中的每个位置都可以关注输入序列中的所有位置。

在自注意力层中，所有键、值和查询都来自同一位置，在本例中是编码器中上一层的输出。这意味着编码器中的每个位置都可以关注编码器中上一层的所有位置。

我们需要防止解码器中的左向信息流来保留自回归属性。这意味着解码器中的某个位置不能关注解码器中后面的位置，在标量点积注意力中实现，方法是将与非法连接对应的所有值从softmax输入中mask掉。

通俗解释：

举个例子，假设我们要翻译一句英文句子“I love you”。在解码器中，我们需要生成每句话的每个单词。对于第一个单词“I”，它可以关注输入序列中的所有单词，包括“I”、“love”、“you”。这意味着解码器可以根据输入序列中的所有信息来生成第一个单词“I”。

我们将输入序列“I love you”转换为一个向量序列。对于每个位置，我们都可以计算它与其他位置的注意力分数。这意味着编码器可以根据输入序列中每个单词之间的关系来生成输出。

自回归属性是指解码器中的每个位置只能依赖于它之前的位置来生成输出。解码器中的某个位置如果关注解码器中后面的位置。这可能会导致模型生成不一致的输出

Position-wise Feed-Forward Networks

前馈神经网络是全连接的，包括两个线性变换，并在它们之间加入ReLU激活。

线性变换每层参数不同,另一种描述方式是使用两个内核大小为1的卷积。

输入输出维度d_model=512,内层维度d_ff=2048

前馈神经网络：

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        #映射到维度d_ff，然后将映射后的向量应用ReLU 激活函数使其成为非线性，然后将激活后的向量通过Dropout层，然后通过另一个线性变换w_2映射回原来的维度空间
        return self.w_2(self.dropout(self.w_1(x).relu()))

 Embeddings and Softmax

 我们还使用通常学到的线性变换和 softmax 函数将解码器输出转换为预测下一个token的概率。在我们的模型中，我们在两个嵌入层和预 softmax 线性变换之间共享相同的权重矩阵。

在嵌入层中，我们将这些权重乘以 根号d_model，以确保损失函数的梯度能够正确地流过网络。

嵌入表：

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        #通过nn.Embedding创建了一个嵌入表（look-up table）（LUT），其大小为vocab x d_model，vocab参数指定词汇表的大小，d_model参数指定嵌入向量的维度
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        #self.lut(x) 操作将 x 中的每个标记 ID 映射到其对应的嵌入向量，* math.sqrt(self.d_model) 操作缩放生成的向量。
        return self.lut(x) * math.sqrt(self.d_model)

Positional Encoding 

        由于我们的模型中不包含循环和卷积，为了使模型利用序列的顺序，我们必须注入一些关于标记在序列中的相对或绝对位置的信息。为此，我们在编码器和解码器堆栈的底部将“位置编码”添加到输入嵌入中。位置编码具有与嵌入相同的维度，这样两者就可以相加。位置编码有很多选择，可以学习也可以固定。

在本工作中，我们使用了不同频率的正弦和余弦函数：

pos表示位置，i表示维度。也就是说，位置编码的每个维度都对应一个正弦波。波长形成一个从2π到10000×2π的几何级数。我们选择这个函数是因为我们假设它可以让模型通过相对位置来轻松学习注意力，因为对于任何固定的偏移量k，PEpos+k 可以表示为PEpos 的线性函数。

此外，我们在编码器和解码器堆栈中的嵌入和位置编码的和上应用dropout。对于基本模型，我们使用率为Pdrop =0.1。

正弦函数确保位置编码在序列中是平滑连续的。

为嵌入添加位置信息（位置编码）：

class PositionalEncoding(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        #初始化一个大小为（max_len，d_model）的零张量 pe
        pe = torch.zeros(max_len, d_model)
        #创建一个大小为（max_len，1）的张量 position，其中包含从0到max_len-1的所有位置。
        position = torch.arange(0, max_len).unsqueeze(1)
        #创建一个大小为d_model的张量 div_term，其中包含正弦位置编码的缩放因子。缩放因子从 0 到 1 以指数方式增加，随着嵌入维度的增加
        #div_term 中的每个元素可以表示为：div_term[i] = exp(-2i * log(10000) / d_model)，torch.arange(0, d_model, 1)中d_model表示生成d_model个元素，非到d_model为止
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)#赋给 pe 的偶数索引元素
        pe[:, 1::2] = torch.cos(position * div_term)#赋给 pe 的奇数索引元素
        pe = pe.unsqueeze(0)#将 pe 的形状改为（1，max_len，d_model）
        self.register_buffer("pe", pe)#将 pe 注册为一个缓冲区，使其可以从 self.pe 属性访问

    def forward(self, x):
        #将指定序列长度的所有位置编码添加到输入嵌入中，并将 requires_grad 标志设置为 False，防止位置编码在反向传播过程中更新。pe 缓冲区存储了所有可能序列长度的预计算位置编码
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

在位置编码中，将基于位置添加正弦波。每个维度的波的频率和偏移量都不同。

def example_positional():
    pe = PositionalEncoding(20, 0)
    y = pe.forward(torch.zeros(1, 100, 20))

    data = pd.concat(
        [
            pd.DataFrame(
                {
                    "embedding": y[0, :, dim],
                    "dimension": dim,
                    "position": list(range(100)),
                }
            )
            for dim in [4, 5, 6, 7]
        ]
    )

    return (
        alt.Chart(data)
        .mark_line()
        .properties(width=800)
        .encode(x="position", y="embedding", color="dimension:N")
        .interactive()
    )


show_example(example_positional)

 

我们还尝试使用学到的位置嵌入(cite)，并发现两个版本产生了几乎相同的结果。我们选择了正弦版本，因为它允许模型推断训练过程中遇到的序列长度比训练中遇到的更长。 

 Full Model

整个模型：

def make_model(
    src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):
    "Helper: Construct a model from hyperparameters."
    c = copy.deepcopy#用于浅拷贝对象
    attn = MultiHeadedAttention(h, d_model)#自注意力层
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)#前馈神经网络层
    position = PositionalEncoding(d_model, dropout)#位置编码器
    model = EncoderDecoder(
        #创建了一个 Encoder 模型，其中包含 N 个编码器层。每个编码器层由一个多头注意力机制、一个前馈网络和一个 dropout 层组成。
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        #创建了一个 Decoder 模型，其中包含 N 个解码器层。每个解码器层由一个多头注意力机制、两个多头注意力机制和一个前馈网络组成。
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        #创建了一个 Embedding 层，用于将源语言和目标语言的词语编码为向量。该层首先使用一个嵌入矩阵将词语编码为嵌入向量，然后使用位置编码器添加位置信息。
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        #创建了一个 Generator 层，用于将解码器输出转换为目标语言的词语。
        Generator(d_model, tgt_vocab),
    )

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
       #将 Transformer 模型的所有矩阵参数初始化为 Xavier 均匀分布，Xavier 均匀初始化方法可以使参数在学习过程中更加稳定。
    for p in model.parameters():
        if p.dim() > 1:#如果大于 1，则表示参数是一个矩阵
            nn.init.xavier_uniform_(p)
    return model

 Inference

生成模型的预测。我们试图使用我们的Transformer记住输入。正如您将看到的那样，由于模型尚未经过训练，因此输出是随机生成的。

测试 Transformer 模型的推理能力：

def inference_test():
    #构建一个 Transformer 模型，并将其设置为评估模式
    test_model = make_model(11, 11, 2)
    test_model.eval()
    src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])#测试序列
    src_mask = torch.ones(1, 1, 10)

    memory = test_model.encode(src, src_mask)#使用编码器将测试序列编码为memory
    ys = torch.zeros(1, 1).type_as(src)#预测序列初始化

    for i in range(9):
        #使用解码器解码memory，并生成下一个词语的概率分布（prob）
        out = test_model.decode(
            memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
        )
        prob = test_model.generator(out[:, -1])
        #从概率分布中选择概率最高的词语（next_word
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]
        #将预测的词语添加到预测序列（ys）中
        ys = torch.cat(
            [ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
        )

    print("Example Untrained Model Prediction:", ys)


def run_tests():
    for _ in range(10):
        inference_test()


show_example(run_tests)

输出： 

 Part 2: Model Training

首先，我们定义一个batch对象，用于保存源句和目标句以便训练，并构建掩码。

Batches and Masking

class Batch:
    """Object for holding a batch of data with mask during training."""

    def __init__(self, src, tgt=None, pad=2):  # 2 = <blank> 此处2用来表示填充词！
        #src: 源语言的句子。src_mask: 源语言句子的掩码，用于忽略填充词。
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2) #维度(batch_size, 1, seq_len)
        if tgt is not None:
            #tgt: 目标语言的句子。tgt_mask: 目标语言句子的掩码，用于忽略填充词和未来的词语。
            self.tgt = tgt[:, :-1]
            self.tgt_y = tgt[:, 1:]
            self.tgt_mask = self.make_std_mask(self.tgt, pad)
            self.ntokens = (self.tgt_y != pad).data.sum()#计算目标语言句子中非填充词的数量，并将其赋值给 self.ntokens 属性
    
    #静态方法 make_std_mask(), 该方法用于创建目标语言句子的掩码
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        #subsequent_mask() 函数用于创建一个下三角矩阵，表示未来的词语被忽略。.type_as(tgt_mask.data)将subsequent_mask张量转换为与tgt_mask张量相同的类型。
        tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
            tgt_mask.data
        )
        return tgt_mask

接下来，我们创建一个通用的训练和评分函数来跟踪损失。我们传入一个通用的损失计算函数，该函数也处理参数更新。 

Training Loop

class TrainState:
    """Track number of steps, examples, and tokens processed"""

    step: int = 0  # Steps in the current epoch当前 epoch 中的步数
    accum_step: int = 0  # Number of gradient accumulation steps梯度累积的步数
    samples: int = 0  # total # of examples used已使用总样本数
    tokens: int = 0  # total # of tokens processed已处理总词元数

def run_epoch(
    data_iter,#数据迭代器
    model,#模型
    loss_compute,#损失计算函数
    optimizer,#优化器
    scheduler,#学习率调度器
    mode="train",#训练模式，可以是 train、train+log 或 valid
    accum_iter=1,#梯度累积的步数
    train_state=TrainState(),#训练状态对象
):
    """Train a single epoch"""
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    n_accum = 0
    for i, batch in enumerate(data_iter):
        #使用模型对 batch 进行预测 一个batch里面有一个源语言序列和一个目标语言序列
        out = model.forward(
            batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
        )
        #使用 loss_compute 函数计算损失和梯度
        loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
        # loss_node = loss_node / accum_iter
        #如果 mode 为 train 或 train+log，则将梯度反向传播并更新模型参数
        if mode == "train" or mode == "train+log":
            loss_node.backward()
            train_state.step += 1
            train_state.samples += batch.src.shape[0]
            train_state.tokens += batch.ntokens
            if i % accum_iter == 0:
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                n_accum += 1
                train_state.accum_step += 1
            scheduler.step()

        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens

        #每隔 40 步打印一次训练进度
        if i % 40 == 1 and (mode == "train" or mode == "train+log"):
            lr = optimizer.param_groups[0]["lr"]
            elapsed = time.time() - start
            print(
                (
                    "Epoch Step: %6d | Accumulation Step: %3d | Loss: %6.2f "
                    + "| Tokens / Sec: %7.1f | Learning Rate: %6.1e"
                )
                % (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr)
            )
            start = time.time()
            tokens = 0
        del loss
        del loss_node
    return total_loss / total_tokens, train_state

 Training Data and Batching

 我们使用标准的2014年WMT英语-德语数据集进行训练，该数据集包含约450万句对。句子使用字节对编码进行编码，该编码具有约37000个令牌的共享源目标词汇。对于英语-法语，我们使用了更大的2014年WMT英语-法语数据集，该数据集包含3600万句子，并将令牌拆分为32000个词块词汇。

将句子对按照大致的序列长度一起进行批处理。每个训练批次包含一组句子对，包含约25000个源令牌和25000个目标令牌。

Hardware and Schedule

我们在一台配备8个NVIDIA P100 GPU的机器上训练了我们的模型。对于我们使用本文中描述的超参数的基本模型，每个训练步骤大约需要0.4秒。我们总共训练了10万步或12小时的基本模型。对于我们的大型模型，步长时间为1.0秒。大型模型训练了30万步（3.5天）。

Optimizer

我们使用了Adam优化器，其中：α=0.9、β1=0.9、β2=0.98、ϵ=10−9。在训练过程中，我们根据以下公式调整学习率：

这与在第一个 warmup_steps（预热步数）训练步骤中学习率将线性增加，此后按步骤数的倒数的平方根成比例降低。我们使用 warmup_steps=4000。

注意：这一部分非常重要。需要使用该模型设置进行训练。

学习率调整公式：

def rate(step, model_size, factor, warmup):
    """
    we have to default the step to 1 for LambdaLR function
    to avoid zero raising to negative power.
    """
    if step == 0:
        step = 1
    return factor * (
        model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
    )

三种调整学习率策略对比： 

def example_learning_schedule():
    #模型大小 学习率因子 预热步数
    opts = [
        [512, 1, 4000],  # example 1
        [512, 1, 8000],  # example 2
        [256, 1, 4000],  # example 3
    ]
    #创建虚拟模型
    dummy_model = torch.nn.Linear(1, 1)
    learning_rates = []

    # we have 3 examples in opts list.
    for idx, example in enumerate(opts):
        # run 20000 epoch for each example
        #使用 Adam 优化器创建一个优化器 optimizer，并将其绑定到虚拟模型 dummy_model
        optimizer = torch.optim.Adam(
            dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
        )
        #创建一个学习率调度器 lr_scheduler，并将其绑定到优化器 optimizer
        lr_scheduler = LambdaLR(
            optimizer=optimizer, lr_lambda=lambda step: rate(step, *example)
        )
        tmp = []
        # take 20K dummy training steps, save the learning rate at each step
        for step in range(20000):
            tmp.append(optimizer.param_groups[0]["lr"])
            optimizer.step()#训练
            lr_scheduler.step()#更新学习率
        learning_rates.append(tmp)

    learning_rates = torch.tensor(learning_rates)

    # Enable altair to handle more than 5000 rows
    alt.data_transformers.disable_max_rows()

    opts_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "Learning Rate": learning_rates[warmup_idx, :],
                    "model_size:warmup": ["512:4000", "512:8000", "256:4000"][
                        warmup_idx
                    ],
                    "step": range(20000),
                }
            )
            for warmup_idx in [0, 1, 2]
        ]
    )

    return (
        alt.Chart(opts_data)
        .mark_line()
        .properties(width=600)
        .encode(x="step", y="Learning Rate", color="model_size:warmup:N")
        .interactive()
    )


example_learning_schedule()

Regularization

在训练过程中，我们采用了值为ϵ=0.1的标签平滑（Label Smoothing）。这会增加困惑度，因为模型学会了更加不确定，但提高了准确率和BLEU分数。（在训练中）我们利用KL div loss实现标签平滑。我们没有使用one-hot目标分布，而是创建了一个分布，它具有正确单词的置信度，其余的平滑质量在整个词汇表中分布。

（标签平滑（label smoothing）是一种用于正则化神经网络输出的技术。它可以通过在训练过程中向目标标签添加噪声来防止模型过拟合。）

class LabelSmoothing(nn.Module):
    "Implement label smoothing."

    def __init__(self, size, padding_idx, smoothing=0.0):
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(reduction="sum")#用于计算交叉熵损失的损失函数
        self.padding_idx = padding_idx#填充词的索引
        self.confidence = 1.0 - smoothing#模型置信度
        self.smoothing = smoothing#标签平滑的程度
        self.size = size#标签的数量
        self.true_dist = None#真实的分布

    #计算标签平滑后的损失 x: 模型的输出。target: 目标标签
    def forward(self, x, target):
        assert x.size(1) == self.size
        ##创建一个真实的分布 true_dist，其中包含每个标签的真实概率
        true_dist = x.data.clone()
        #将 true_dist 填充为标签平滑后的真实分布。标签平滑后的真实分布是一个均匀分布，其中每个标签的概率为 smoothing / (size - 2)。 将 true_dist 的所有元素填充为 smoothing / (size - 2)
        true_dist.fill_(self.smoothing / (self.size - 2))
        #将 target.data 复制到 true_dist 的 target.data.unsqueeze(1) 位置，并将概率设置为 self.confidence
        #也就是说true_dist中，target标签被设置成confidence，其余仍然保持self.confidence/(size-2)
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        #将 true_dist 中填充词的概率设置为 0。
        true_dist[:, self.padding_idx] = 0
        #创建一个掩码 mask，用于标记填充词的位置。
        mask = torch.nonzero(target.data == self.padding_idx)
        #如果 mask 的维度大于 0，则表示 target.data 中存在填充词。
        if mask.dim() > 0:、
            #将 true_dist 中填充词的位置的概率设置为 0
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        #保存真实分布 true_dist
        self.true_dist = true_dist
        #使用交叉熵损失函数计算标签平滑后的损失
        return self.criterion(x, true_dist.clone().detach())

# Example of label smoothing.


def example_label_smoothing():
    crit = LabelSmoothing(5, 0, 0.4)
    predict = torch.FloatTensor(
        [
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
        ]
    )
    crit(x=predict.log(), target=torch.LongTensor([2, 1, 0, 3, 3]))
    LS_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "target distribution": crit.true_dist[x, y].flatten(),
                    "columns": y,
                    "rows": x,
                }
            )
            for y in range(5)
            for x in range(5)
        ]
    )

    return (
        alt.Chart(LS_data)
        .mark_rect(color="Blue", opacity=1)
        .properties(height=200, width=200)
        .encode(
            alt.X("columns:O", title=None),
            alt.Y("rows:O", title=None),
            alt.Color(
                "target distribution:Q", scale=alt.Scale(scheme="viridis")
            ),
        )
        .interactive()
    )


show_example(example_label_smoothing)

标签平滑实际上在模型对某一选择非常自信时，开始对其进行惩罚。 

#在 PyTorch 中可视化标签平滑的惩罚效果
#loss() 函数计算给定输入 x 和 LabelSmoothing 对象 crit 的标签平滑损失
#图表显示，标签平滑损失在训练步骤的初始阶段较高。这是因为模型最初对其预测非常有信心。随着训练的进行，模型逐渐学习区分正确答案和已添加到目标标签中的噪声。这导致标签平滑损失减小。
def loss(x, crit):
    d = x + 3 * 1
    predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])
    return crit(predict.log(), torch.LongTensor([1])).data


def penalization_visualization():
    crit = LabelSmoothing(5, 0, 0.1)
    loss_data = pd.DataFrame(
        {
            "Loss": [loss(x, crit) for x in range(1, 100)],
            "Steps": list(range(99)),
        }
    ).astype("float")

    return (
        alt.Chart(loss_data)
        .mark_line()
        .properties(width=350)
        .encode(
            x="Steps",
            y="Loss",
        )
        .interactive()
    )


show_example(penalization_visualization)

 A First Example

 我们可以从尝试一个简单的复制任务开始。给定一个小词库中的一组随机输入符号，目标是生成相同的符号。

Synthetic Data

#data_gen() 函数用于生成用于源-目标复制任务的随机数据，V: 词汇表的大小。batch_size: 每个批次的数据量。nbatches: 要生成的批次的数量。
def data_gen(V, batch_size, nbatches):
    "Generate random data for a src-tgt copy task."
    for i in range(nbatches):
        data = torch.randint(1, V, size=(batch_size, 10))#生成一个随机张量，该张量的形状为 (batch_size, 10)，其中的元素介于 1 和 V 之间
        data[:, 0] = 1#将 data 张量的第一列设置为 1。这是因为源序列的第一个元素总是起始符号
        #创建src，tgt张量，该张量是data张量的克隆，并且其requires_grad属性设置为False。这意味着src，tgt张量中的参数在训练过程中不会被更新
        src = data.requires_grad_(False).clone().detach()
        tgt = data.requires_grad_(False).clone().detach()
        #使用 yield 语句生成一个 Batch 对象，该对象包含 src 和 tgt 张量，以及一个目标长度为 0 的 len 张量
        yield Batch(src, tgt, 0)

Loss Computation 

损失计算和训练函数：

class SimpleLossCompute:
    "A simple loss compute and train function."

    def __init__(self, generator, criterion):
        self.generator = generator#生成器，用于生成模型的预测
        self.criterion = criterion#损失函数

    def __call__(self, x, y, norm):#x输入张量 y目标张量 norm归一化因子
        x = self.generator(x)#使用 generator 生成模型的预测
        #使用 criterion 计算预测和目标之间的损失 将 y 张量转换为 (seq_len,) 的形状。将 x 张量转换为 (batch_size * seq_len, 1) 的形状
        sloss = (
            self.criterion(
                x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
            )
            / norm
        )
        return sloss.data * norm, sloss

Greedy Decoding 

这段代码为了简化问题，使用贪心解码来预测翻译。

使用贪婪解码算法解码序列：

#model一个神经网络模型，src: 源语言序列。src_mask: 源语言序列掩码。max_len: 最大解码长度。start_symbol: 起始符号
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    memory = model.encode(src, src_mask)#编码器
    ys = torch.zeros(1, 1).fill_(start_symbol).type_as(src.data)#创建一个 ys 张量，包含起始符号
    for i in range(max_len - 1):
        out = model.decode(
            memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
        )#解码
        prob = model.generator(out[:, -1])#计算模型对下一个单词的预测概率。
        _, next_word = torch.max(prob, dim=1)#选择下一个要预测的单词
        next_word = next_word.data[0]
        ys = torch.cat(
            [ys, torch.zeros(1, 1).type_as(src.data).fill_(next_word)], dim=1
        )#将下一个单词添加到 ys 张量中
    return ys

# Train the simple copy task.


def example_simple_model():
    V = 11#词汇表大小
    criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)#创建一个 LabelSmoothing 对象，平滑因子为 0.0
    model = make_model(V, V, N=2)#创建一个 Transformer 模型，输入和输出词汇表的大小均为 11，编码层和解码层各包含 2 个自注意力模块。

    optimizer = torch.optim.Adam(
        model.parameters(), lr=0.5, betas=(0.9, 0.98), eps=1e-9
    )#创建一个 Adam 优化器，学习率为 0.5。
    lr_scheduler = LambdaLR(
        optimizer=optimizer,
        lr_lambda=lambda step: rate(
            step, model_size=model.src_embed[0].d_model, factor=1.0, warmup=400
        ),
    )#创建一个 LambdaLR 学习率调整器，使用 Noam 调度算法

    batch_size = 80#批量大小为 80
    for epoch in range(20):
        model.train()
        run_epoch(
            data_gen(V, batch_size, 20),
            model,
            SimpleLossCompute(model.generator, criterion),
            optimizer,
            lr_scheduler,
            mode="train",
        )
        model.eval()
        run_epoch(
            data_gen(V, batch_size, 5),
            model,
            SimpleLossCompute(model.generator, criterion),
            DummyOptimizer(),
            DummyScheduler(),
            mode="eval",
        )[0]

    model.eval()
    src = torch.LongTensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
    max_len = src.shape[1]
    src_mask = torch.ones(1, 1, max_len)
    print(greedy_decode(model, src, src_mask, max_len=max_len, start_symbol=0))


# execute_example(example_simple_model)

 Part 3: A Real World Example

 现在，我们来看一个实际使用Multi30k德英翻译任务的例子。这个任务比论文中考虑的WMT任务要小得多，但它说明了整个系统。我们还展示了如何使用多GPU处理来使其速度更快。

Data Loading

我们将使用torchtext和spacy进行分词并加载数据集。

加载德语和英语分词器模型：

# Load spacy tokenizer models, download them if they haven't been
# downloaded already

#load_tokenizers() 函数用于加载德语和英语的 spaCy 分词器模型。如果模型尚未下载，该函数将下载它们。
def load_tokenizers():
    #尝试加载德语 spaCy 分词器模型。如果模型尚未下载，该函数将使用 os.system() 函数下载它。
    try:
        spacy_de = spacy.load("de_core_news_sm")
    except IOError:
        os.system("python -m spacy download de_core_news_sm")
        spacy_de = spacy.load("de_core_news_sm")
    #尝试加载英语 spaCy 分词器模型。如果模型尚未下载，该函数将使用 os.system() 函数下载它。
    try:
        spacy_en = spacy.load("en_core_web_sm")
    except IOError:
        os.system("python -m spacy download en_core_web_sm")
        spacy_en = spacy.load("en_core_web_sm")

    return spacy_de, spacy_en

#text: 要分词的文本。tokenizer: 一个分词器对象。该函数返回一个分词后的单词列表。
def tokenize(text, tokenizer):
    return [tok.text for tok in tokenizer.tokenizer(text)]

#data_iter: 一个数据迭代器。tokenizer: 一个分词器对象。index: 数据迭代器中要分词的文本的下标。
def yield_tokens(data_iter, tokenizer, index):
    for from_to_tuple in data_iter:
        yield tokenizer(from_to_tuple[index])

def build_vocabulary(spacy_de, spacy_en):
    def tokenize_de(text):
        return tokenize(text, spacy_de)

    def tokenize_en(text):
        return tokenize(text, spacy_en)

    print("Building German Vocabulary ...")
    #使用 `Datasets` 库加载 Multi30k 数据集
    train, val, test = datasets.Multi30k(language_pair=("de", "en"))
    #使用 `build_vocab_from_iterator()` 函数构建词表。
    vocab_src = build_vocab_from_iterator(
        yield_tokens(train + val + test, tokenize_de, index=0),#生成词元列表
        min_freq=2,
        specials=["<s>", "</s>", "<blank>", "<unk>"],
    )

    print("Building English Vocabulary ...")
    train, val, test = datasets.Multi30k(language_pair=("de", "en"))
    vocab_tgt = build_vocab_from_iterator(
        yield_tokens(train + val + test, tokenize_en, index=1),#生成词元列表
        min_freq=2,
        specials=["<s>", "</s>", "<blank>", "<unk>"],
    )
    #设置默认索引。
    vocab_src.set_default_index(vocab_src["<unk>"])
    vocab_tgt.set_default_index(vocab_tgt["<unk>"])
    #返回词表
    return vocab_src, vocab_tgt


def load_vocab(spacy_de, spacy_en):
    #检查是否存在词表文件 `vocab.pt`
    if not exists("vocab.pt"):
        #如果不存在，则使用 `build_vocabulary()` 函数构建词表
        vocab_src, vocab_tgt = build_vocabulary(spacy_de, spacy_en)
        torch.save((vocab_src, vocab_tgt), "vocab.pt")
    else:
        #如果存在，则加载词表
        vocab_src, vocab_tgt = torch.load("vocab.pt")
    #打印词表大小
    print("Finished.\nVocabulary sizes:")
    print(len(vocab_src))
    print(len(vocab_tgt))
    ##返回词表
    return vocab_src, vocab_tgt


if is_interactive_notebook():#表示如果当前环境是交互式笔记本，则执行该代码块。
    # global variables used later in the script
    #加载德语和英语 spaCy 分词器模型
    spacy_de, spacy_en = show_example(load_tokenizers)
    #加载德语和英语词表。
    vocab_src, vocab_tgt = show_example(load_vocab, args=[spacy_de, spacy_en])

批处理对速度至关重要。我们希望将数据均匀地分成若干批次，并尽量减少填充。要做到这一点，我们需要在默认的torchtext批处理中进行一些修改。这段代码修改了它们的默认批处理，以确保我们在足够多的句子中搜索以找到紧凑的batches。 

 Iterators

对源语言文本和目标语言文本进行预处理，然后将其转换为 ID：

def collate_batch(
    batch,#列表，其中包含每个样本的输入和输出
    src_pipeline,#函数，用于对源语言文本进行预处理
    tgt_pipeline,#函数，用于对目标语言文本进行预处理
    src_vocab,#词典，其中包含源语言单词及其对应的 ID
    tgt_vocab,#词典，其中包含目标语言单词及其对应的 ID。
    device,#可以是 cpu 或 cuda
    max_padding=128,#最大填充长度
    pad_id=2,#填充 ID
):
    bs_id = torch.tensor([0], device=device)  # <s> token id开始符号的 ID
    eos_id = torch.tensor([1], device=device)  # </s> token id结束符号的 ID
    src_list, tgt_list = [], []#列表，用于存储源语言文本的 ID/目标语言文本的 ID
    for (_src, _tgt) in batch:
        #对源语言文本进行预处理，并将其转换为 ID
        processed_src = torch.cat(
            [
                bs_id,
                torch.tensor(
                    src_vocab(src_pipeline(_src)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        #对目标语言文本进行预处理，并将其转换为 ID
        processed_tgt = torch.cat(
            [
                bs_id,
                torch.tensor(
                    tgt_vocab(tgt_pipeline(_tgt)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        #将源语言文本的 ID 添加到 src_list 中
        src_list.append(
            # warning - overwrites values for negative values of padding - len
            pad(
                processed_src,
                (
                    0,
                    max_padding - len(processed_src),
                ),
                value=pad_id,
            )
        )
        #将目标语言文本的 ID 添加到 tgt_list 中
        tgt_list.append(
            pad(
                processed_tgt,
                (0, max_padding - len(processed_tgt)),
                value=pad_id,
            )
        )
    #将 src_list 和 tgt_list 转换为张量，并返回
    src = torch.stack(src_list)
    tgt = torch.stack(tgt_list)
    return (src, tgt)

def create_dataloaders(
    device,#设备，可以是 cpu 或 cuda
    vocab_src,#源语言的词汇表
    vocab_tgt,#目标语言的词汇表
    spacy_de,#用于德语分词的 SpaCy 模型
    spacy_en,#用于英文分词的 SpaCy 模型
    batch_size=12000,#批处理大小
    max_padding=128,#最大填充长度
    is_distributed=True,#是否使用分布式训练
):
    # def create_dataloaders(batch_size=12000):
    #函数，用于对德语文本进行分词
    def tokenize_de(text):
        return tokenize(text, spacy_de)
    #函数，用于对英文文本进行分词
    def tokenize_en(text):
        return tokenize(text, spacy_en)
    #对源语言文本和目标语言文本进行预处理，然后将其转换为 ID
    def collate_fn(batch):
        return collate_batch(
            batch,
            tokenize_de,
            tokenize_en,
            vocab_src,
            vocab_tgt,
            device,
            max_padding=max_padding,
            pad_id=vocab_src.get_stoi()["<blank>"],
        )
    #训练集 验证集 测试集
    train_iter, valid_iter, test_iter = datasets.Multi30k(
        language_pair=("de", "en")
    )
    #函数，用于将迭代器转换为字典格式
    train_iter_map = to_map_style_dataset(
        train_iter
    )  # DistributedSampler needs a dataset len()
    #采样器，用于训练集
    train_sampler = (
        DistributedSampler(train_iter_map) if is_distributed else None
    )
    #函数，用于将迭代器转换为字典格式
    valid_iter_map = to_map_style_dataset(valid_iter)
    ##采样器，用于验证集
    valid_sampler = (
        DistributedSampler(valid_iter_map) if is_distributed else None
    )
    #一个用于训练集的数据加载器
    train_dataloader = DataLoader(
        train_iter_map,
        batch_size=batch_size,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        collate_fn=collate_fn,
    )
    #一个用于验证集的数据加载器
    valid_dataloader = DataLoader(
        valid_iter_map,
        batch_size=batch_size,
        shuffle=(valid_sampler is None),
        sampler=valid_sampler,
        collate_fn=collate_fn,
    )
    return train_dataloader, valid_dataloader

 Training the System

def train_worker(
    gpu,#要使用的 GPU 设备号
    ngpus_per_node,#每个节点上的 GPU 设备数
    vocab_src,#源语言词汇表
    vocab_tgt,#目标语言词汇表
    spacy_de,#德语 spaCy 模型
    spacy_en,#英语 spaCy 模型
    config,#训练配置，包括 batch size、最大填充长度等
    is_distributed=False,#是否使用分布式训练
):
    print(f"Train worker process using GPU: {gpu} for training", flush=True)
    torch.cuda.set_device(gpu)#将当前工作进程的设备设置为指定的 GPU

    pad_idx = vocab_tgt["<blank>"]#获取目标语言词汇表中 <blank> 对应的索引，用于指示填充
    d_model = 512
    model = make_model(len(vocab_src), len(vocab_tgt), N=6)#创建模型
    model.cuda(gpu)#将模型移到指定的 GPU 上进行计算
    module = model
    is_main_process = True#表示当前进程是否为主进程
    if is_distributed:
        #初始化分布式训练环境，设置了通信后端为 "nccl"，并使用指定的方法进行初始化。
        dist.init_process_group(
            "nccl", init_method="env://", rank=gpu, world_size=ngpus_per_node
        )
        model = DDP(model, device_ids=[gpu])#使用 DDP 包装模型，以支持分布式训练
        module = model.module
        is_main_process = gpu == 0#更新 module和is_main_process变量 以标识当前进程是否为主进程
    #定义了一个平滑的交叉熵损失函数 criterion，用于计算模型预测与真实目标之间的损失。
    criterion = LabelSmoothing(
        size=len(vocab_tgt), padding_idx=pad_idx, smoothing=0.1
    )
    criterion.cuda(gpu)#将损失函数移动到指定的 GPU 上
    #创建训练和验证数据加载器
    train_dataloader, valid_dataloader = create_dataloaders(
        gpu,
        vocab_src,
        vocab_tgt,
        spacy_de,
        spacy_en,
        batch_size=config["batch_size"] // ngpus_per_node,
        max_padding=config["max_padding"],
        is_distributed=is_distributed,
    )
    #创建 Adam 优化器
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config["base_lr"], betas=(0.9, 0.98), eps=1e-9
    )
    #学习率调度器 lr_scheduler，使用了自定义的学习率变化函数 rate
    lr_scheduler = LambdaLR(
        optimizer=optimizer,
        lr_lambda=lambda step: rate(
            step, d_model, factor=1, warmup=config["warmup"]
        ),
    )
    #初始化了一个用于跟踪训练状态的对象 train_state
    train_state = TrainState()

    for epoch in range(config["num_epochs"]):#config["num_epochs"] 表示总共的训练周期数。
        #如果是分布式训练，则设置训练数据加载器和验证数据加载器的采样器的 epoch 为当前 epoch。这样可以确保在分布式环境中每个进程都使用相同顺序的数据进行训练
        if is_distributed:
            train_dataloader.sampler.set_epoch(epoch)
            valid_dataloader.sampler.set_epoch(epoch)
        #将模型设置为训练模式 这会启用模型中的训练相关的层，比如 Dropout 和 Batch Normalization。
        model.train()
        #打印当前 GPU 上的训练轮数信息，并使用 flush=True 立即刷新输出流，确保信息能够及时显示。
        print(f"[GPU{gpu}] Epoch {epoch} Training ====", flush=True)
        #计算每个 batch 的损失和梯度，并更新模型参数
        _, train_state = run_epoch(
            (Batch(b[0], b[1], pad_idx) for b in train_dataloader),
            model,
            SimpleLossCompute(module.generator, criterion),
            optimizer,
            lr_scheduler,
            mode="train+log",
            accum_iter=config["accum_iter"],
            train_state=train_state,
        )

        GPUtil.showUtilization()#显示当前 GPU 的利用率信息，可能是为了在训练过程中监控 GPU 使用情况
        if is_main_process:
            file_path = "%s%.2d.pt" % (config["file_prefix"], epoch)
            torch.save(module.state_dict(), file_path)#保存模型的参数到文件中
        torch.cuda.empty_cache()#清空 GPU 缓存以释放已使用的内存，确保在下个 epoch 开始时有足够的空间。

        print(f"[GPU{gpu}] Epoch {epoch} Validation ====", flush=True)
        #将模型设置为评估模式
        model.eval()
        #计算每个 batch 的损失
        sloss = run_epoch(
            (Batch(b[0], b[1], pad_idx) for b in valid_dataloader),
            model,
            SimpleLossCompute(module.generator, criterion),
            DummyOptimizer(),
            DummyScheduler(),
            mode="eval",
        )
        print(sloss)
        torch.cuda.empty_cache()#清空 GPU 缓存

    if is_main_process:
        file_path = "%sfinal.pt" % config["file_prefix"]
        torch.save(module.state_dict(), file_path)#保存最终的模型参数

#函数，用于在分布式环境下训练模型。
def train_distributed_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config):
    from the_annotated_transformer import train_worker

    ngpus = torch.cuda.device_count()#获取可用的 GPU 数量
    #设置了环境变量，指定了分布式训练的主节点地址和端口号。
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12356"
    print(f"Number of GPUs detected: {ngpus}")#打印检测到的可用 GPU 数量
    print("Spawning training processes ...")#打印提示信息，表示即将开始生成训练进程
    #使用 multiprocessing 模块的 spawn 方法生成多个训练进程，进行分布式训练
    mp.spawn(
        train_worker,
        nprocs=ngpus,
        args=(ngpus, vocab_src, vocab_tgt, spacy_de, spacy_en, config, True),
    )

#函数，根据配置选择是进行分布式训练还是单机训练。
def train_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config):
    if config["distributed"]:
        train_distributed_model(
            vocab_src, vocab_tgt, spacy_de, spacy_en, config
        )
    else:
        train_worker(
            0, 1, vocab_src, vocab_tgt, spacy_de, spacy_en, config, False
        )

#函数，用于加载已经训练好的模型。
def load_trained_model():
    #定义了一个包含模型训练参数的字典 config
    config = {
        "batch_size": 32,
        "distributed": False,
        "num_epochs": 8,
        "accum_iter": 10,
        "base_lr": 1.0,
        "max_padding": 72,
        "warmup": 3000,
        "file_prefix": "multi30k_model_",
    }
    model_path = "multi30k_model_final.pt"#指定了已经训练好的模型的保存路径
    if not exists(model_path):#如果已经训练好的模型文件不存在则调用train_model函数重新训练。
        train_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config)

    model = make_model(len(vocab_src), len(vocab_tgt), N=6)#创建了一个 Transformer 模型，其中 len(vocab_src) 和 len(vocab_tgt) 分别表示源语言和目标语言的词汇表大小，N=6 表示 Transformer 模型的层数。
    model.load_state_dict(torch.load("multi30k_model_final.pt"))#加载训练好的模型参数
    return model


if is_interactive_notebook():#在交互式笔记本中调用 load_trained_model 函数，加载训练好的模型
    model = load_trained_model()

一旦训练好模型，我们就可以解码模型以生成一组翻译。在这里，我们简单地翻译验证集中的第一句话。这个数据集很小，因此贪婪搜索的翻译相当准确。

Additional Components: BPE, Search, Averaging 

 BPE/分词：我们可以使用库将数据首先预处理成子词单元。请参考 Rico Sennrich 的 subword-nmt 实现。这些模型将把训练数据转换成以下形式

共享嵌入：当使用共享词汇表的BPE时，我们可以在源/目标/生成器之间共享相同的权重向量。要将此添加到模型中，只需执行以下操作： 

if False:
    model.src_embed[0].lut.weight = model.tgt_embeddings[0].lut.weight
    model.generator.lut.weight = model.tgt_embed[0].lut.weight

 束搜索: 这里介绍起来有点复杂，可参见OpenNMT-py中关于PyTorch的实现。

模型平均: 论文中提到将最后的k个检查点平均起来以产生集成效果。如果我们有一堆模型，可以在事后进行此操作。

def average(model, models):
    "Average models into model"
    for ps in zip(*[m.params() for m in [model] + models]):
        ps[0].copy_(torch.sum(*ps[1:]) / len(ps[1:]))

 Results

在WMT 2014年英语到德语翻译任务中，大型Transformer模型（表2中的Transformer（big））比之前报道的最佳模型（包括集成模型）高出2.0个BLEU分以上，建立了新的最先进的BLEU分数28.4。该模型的配置列在表3的最后一行。在8个P100 GPU上进行训练需要3.5天。即使我们的基础模型也超过了之前发布的所有模型和集成模型，其训练成本仅为任何竞争模型的四分之一。

在WMT 2014年英语到法语翻译任务中，我们的大型模型取得了BLEU分数41.0，超过了之前发布的所有单个模型，其训练成本不到之前最先进模型的四分之一。训练用于英语到法语的Transformer（big）模型使用了dropout rate Pdrop = 0.1，而不是0.3。

使用上一部分中的额外扩展，在EN-DE WMT上，OpenNMT-py复制达到26.9。我在我们的实现中加载了这些参数。

def check_outputs(
    valid_dataloader,
    model,
    vocab_src,
    vocab_tgt,
    n_examples=15,
    pad_idx=2,
    eos_string="</s>",
):
    results = [()] * n_examples# 创建一个长度为 n_examples 的空列表 results
    for idx in range(n_examples):
        print("\nExample %d ========\n" % idx)# 输出当前例子的序号
        b = next(iter(valid_dataloader))# 从验证数据加载器中获取下一个批次的数据
        rb = Batch(b[0], b[1], pad_idx)# 构造一个 Batch 对象 rb，用于处理批次数据
        greedy_decode(model, rb.src, rb.src_mask, 64, 0)[0]# 使用贪婪解码得到模型输出

        src_tokens = [
            vocab_src.get_itos()[x] for x in rb.src[0] if x != pad_idx
        ]# 从源语言词汇表中获取源语言词汇
        tgt_tokens = [
            vocab_tgt.get_itos()[x] for x in rb.tgt[0] if x != pad_idx
        ]# 从目标语言词汇表中获取目标语言词汇

        print(
            "Source Text (Input)        : "
            + " ".join(src_tokens).replace("\n", "")
        ) # 打印原始文本（输入）
        print(
            "Target Text (Ground Truth) : "
            + " ".join(tgt_tokens).replace("\n", "")
        )# 打印目标文本（正确结果）
        model_out = greedy_decode(model, rb.src, rb.src_mask, 72, 0)[0]# 再次使用贪婪解码得到模型输出
        model_txt = (
            " ".join(
                [vocab_tgt.get_itos()[x] for x in model_out if x != pad_idx]
            ).split(eos_string, 1)[0]
            + eos_string
        )# 处理模型输出文本
        print("Model Output               : " + model_txt.replace("\n", "")) # 打印模型输出
        results[idx] = (rb, src_tokens, tgt_tokens, model_out, model_txt)# 将当前例子的相关数据存入结果列表
    return results# 返回存储了每个例子相关数据的结果列表


def run_model_example(n_examples=5):
    global vocab_src, vocab_tgt, spacy_de, spacy_en# 声明引用全局变量

    print("Preparing Data ...")# 打印准备数据的信息
    _, valid_dataloader = create_dataloaders(# 创建验证数据加载器
        torch.device("cpu"),
        vocab_src,
        vocab_tgt,
        spacy_de,
        spacy_en,
        batch_size=1,
        is_distributed=False,
    )

    print("Loading Trained Model ...")# 打印加载训练好的模型的信息

    model = make_model(len(vocab_src), len(vocab_tgt), N=6)
    model.load_state_dict(# 加载之前训练好的模型参数
        torch.load("multi30k_model_final.pt", map_location=torch.device("cpu"))
    )

    print("Checking Model Outputs:") # 打印检查模型输出的信息
    example_data = check_outputs(# 调用 check_outputs 函数，对加载的模型进行测试，得到模型输出的例子数据
        valid_dataloader, model, vocab_src, vocab_tgt, n_examples=n_examples
    )
    return model, example_data# 返回加载的模型以及模型输出的例子数据