《The Annotated Transformer》 Transformer论文解读与代码注释 中英对照 By Harvard NLP

原文链接：
http://nlp.seas.harvard.edu/2018/04/03/attention.html
Motivation：
之前在网上看到的许多翻译版本由于截图较多以及水印问题不方便进行编辑
为了方便使用自己进行了整理
Tips：
英文部分照搬原文，保留公式格式与引用链接，部分链接需要科学上网
在原文机翻基础上进行修改，可能有错误，后续再进行补充
图片均来自原文外链，无水印
代码部分可复制

---

from IPython.display import Image
Image(filename='images/aiayn.png')

The Transformer from “Attention is All You Need” has been on a lot of people’s minds over the last year. Besides producing major improvements in translation quality, it provides a new architecture for many other NLP tasks. The paper itself is very clearly written, but the conventional wisdom has been that it is quite difficult to implement correctly.
《Attention is All You Need》一文中的Transformer在过去一年里得到了许多关注。除了显著提升翻译质量外，它还为许多其他 NLP 任务提供了新的架构。论文本身写得很清楚，但传统观点认为它很难正确实现。

In this post I present an “annotated” version of the paper in the form of a line-by-line implementation. I have reordered and deleted some sections from the original paper and added comments throughout. This document itself is a working notebook, and should be a completely usable implementation. In total there are 400 lines of library code which can process 27,000 tokens per second on 4 GPUs.
在这篇文章中，我以逐行实现的形式展示了论文的“注释”版本。我已经重新排序并删除了原始论文中的一些部分，并在整个过程中添加了评论。该文档本身就是一个Notebook，应该是一个完全可用的实现。总共有 400 行库代码，可以在 4 个 GPU 上每秒处理 27,000 个tokens。

To follow along you will first need to install PyTorch. The complete notebook is also available on github or on Google Colab with free GPUs.
要继续学习，您首先需要安装 PyTorch。完整的Notebook也可以在 github或带有免费 GPU的 Google Colab上找到。

Note this is merely a starting point for researchers and interested developers. The code here is based heavily on our OpenNMT packages. (If helpful feel free to cite.) For other full-sevice implementations of the model check-out Tensor2Tensor (tensorflow) and Sockeye (mxnet).
请注意，这只是研究人员和感兴趣的开发人员的起点。这里的代码很大程度上基于我们的OpenNMT包。（如果有帮助，请随意引用，BibTex格式在文章末尾）。对于模型的其他框架的实现，请查看 Tensor2Tensor (tensorflow) 和 Sockeye (mxnet)。

Alexander Rush (@harvardnlp or [email protected]), with help from Vincent Nguyen and Guillaume Klein

Prelims 准备

# !pip install http://download.pytorch.org/whl/cu80/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl numpy matplotlib spacy torchtext seaborn 

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
from torch.autograd import Variable
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")
%matplotlib inline

My comments are blockquoted. The main text is all from the paper itself.
作者的评论会用引用形式表示。正文全部来自论文本身。

Background 背景

The goal of reducing sequential computation also forms the foundation of the Extended-Neural GPU, ByteNet and ConvS2S, all of which use convolutional neural networks as basic building block, computing hidden representations in parallel for all input and output positions. In these models, the number of operations required to relate signals from two arbitrary input or output positions grows in the distance between positions, linearly for ConvS2S and logarithmically for ByteNet. This makes it more difficult to learn dependencies between distant positions. In the Transformer this is reduced to a constant number of operations, albeit at the cost of reduced effective resolution due to averaging attention-weighted positions, an effect we counteract with Multi-Head Attention.
Extended Neural GPU、ByteNet和ConvS2S的出现是为了减少序列计算量，他们都使用卷积神经网络作为基本构建块，并行计算所有输入和输出位置的隐藏表示。在这些模型中，关联来自两个任意输入或输出位置的信号所需的操作数量随着位置之间的距离而增长，对于ConvS2S呈线性增长，对于ByteNet则呈对数增长。 这使得学习远距离位置之间的依赖关系变得更加困难。在Transformer中，这被减少到恒定数量的操作，尽管由于平均注意力加权位置而降低了有效分辨率，我们用多头注意力来抵消这种影响。

Self-attention, sometimes called intra-attention is an attention mechanism relating different positions of a single sequence in order to compute a representation of the sequence. Self-attention has been used successfully in a variety of tasks including reading comprehension, abstractive summarization, textual entailment and learning task-independent sentence representations. End- to-end memory networks are based on a recurrent attention mechanism instead of sequence aligned recurrence and have been shown to perform well on simple- language question answering and language modeling tasks.
自注意力，有时称为内部注意力，是一种将单个序列的不同位置关联起来以计算序列表示的注意力机制。自注意力已成功用于各种任务，包括阅读理解、抽象摘要、文本蕴涵和学习任务无关的句子表示。端到端记忆网络基于循环注意机制而不是序列对齐循环，并且已被证明在简单语言问答和语言建模任务中表现良好。
To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequence aligned RNNs or convolution.
然而，据我们所知，Transformer 是第一个 完全依赖自注意力来计算其输入和输出表示而不使用序列对齐RNN或卷积的 转换模型。

Model Architecture 模型架构

Most competitive neural sequence transduction models have an encoder-decoder structure (cite). Here, the encoder maps an input sequence of symbol representations $(x_1,\dots ,x_n)$ to a sequence of continuous representations $\mathbf{z}=(z_1,\dots ,z_n)$. Given $\mathbf{z}$, the decoder then generates an output sequence $(y_1,\dots ,y_m)$ of symbols one element at a time. At each step the model is auto-regressive (cite), consuming the previously generated symbols as additional input when generating the next.
大多数常用神经序列转换模型都有编码器-解码器结构（引用）。这里，编码器将符号表示的输入序列$(x_1，，x_n)$映射为连续表示的序列$\mathbf{z}=(z_1,\dots ,z_n)$。给定 $\mathbf{z}$，解码器然后一次一个元素地生成符号的输出序列 $(y_1,\dots ,y_m)$ 。 在每个步骤中，模型是自动回归的（引用），在生成下一个时，把先前生成的符号作为附加输入。

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)
    
    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__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)

The Transformer follows this overall architecture using stacked self-attention and point-wise, fully connected layers for both the encoder and decoder, shown in the left and right halves of Figure 1, respectively.
Transformer 遵循这种整体架构，对编码器和解码器使用堆叠的自注意力和逐点全连接层，分别如图 1 的左半部分和右半部分所示。

Encoder and Decoder Stacks 堆叠的编码器和解码器

Encoder 编码器

The encoder is composed of a stack of $N=6$ identical layers.
编码器由$N=6$个相同的层组成。

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        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."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

We employ a residual connection (cite) around each of the two sub-layers, followed by layer normalization (cite).
我们在两个子层中的每一个周围使用残差连接（cite），接着进行层归一化（cite）。

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        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

That is, the output of each sub-layer is $\mathrm{L}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}\mathrm{N}\mathrm{o}\mathrm{r}\mathrm{m}(x+\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x))$, where $\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x)$ is the function implemented by the sub-layer itself. We apply dropout (cite) to the output of each sub-layer, before it is added to the sub-layer input and normalized.
也就是说，每个子层的输出都是$\mathrm{L}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}\mathrm{N}\mathrm{o}\mathrm{r}\mathrm{m}(x+\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x))$，其中$\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x)$是子层本身实现的函数。我们将dropout (引用)应用于每个子层的输出，然后将其添加到子层的输入并进行归一化。

To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension $d_{\text{model}}=512$.
为了促进这些残差连接，模型中的所有子层以及嵌入层都产生了维度 $d_{\text{model}}=512$的输出。

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."
        return x + self.dropout(sublayer(self.norm(x)))

Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed- forward network.
每层有两个子层。第一个是多头自注意机制，第二个是简单的逐位置全连接前馈网络。

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)
        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 解码器

The decoder is also composed of a stack of $N=6$ identical layers.
解码器也是由$N=6$个相同的层组成。

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):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization.
除了每个编码器层中的两个子层之外，解码器还插入了第三个子层，该子层对堆叠的编码器的输出执行多头注意力。与编码器类似，我们在每个子层周围使用残差连接，然后进行层归一化。

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)

We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position $i$ can depend only on the known outputs at positions less than $i$.
我们还修改了解码器中的自注意力子层，以防止位置对后续位置的影响。这种mask操作，结合输出嵌入偏移一个位置的事实，确保了位置 $i$ 的预测只能依赖于位置小于 $i$ 的已知输出。

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0

Below the attention mask shows the position each tgt word (row) is allowed to look at (column). Words are blocked for attending to future words during training.
作者注释：下方的注意力mask图显示了每个目标（target）单词（行）允许查看的位置（列）。阻止了单词在训练期间关注未来的单词。

plt.figure(figsize=(5,5))
plt.imshow(subsequent_mask(20)[0])
None

Attention 注意力机制

An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
注意力函数可以描述为将查询（Query）和一组键值对（Key-Value）映射到输出，其中 Query, Key, Value和输出都是向量。输出计算为Value的加权和，其中分配给每个Value的权重由Query与相应Key的兼容函数计算。
We call our particular attention “Scaled Dot-Product Attention”. The input consists of queries and keys of dimension $d_k$, and values of dimension $d_v$. We compute the dot products of the query with all keys, divide each by $\sqrt{d_k}$, and apply a softmax function to obtain the weights on the values.
我们把这种特别的注意力叫做“缩放点积注意力”。输入包括维度为$d_k$的Query和Key，维度为$d_v$的Value。我们计算Query与所有Key的点积，将每个Key除以$\sqrt{d_k}$，并应用softmax函数来获得值的权重。

Image(filename='images/ModalNet-19.png')

In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix $Q$. The keys and values are also packed together into matrices $K$ and $V$. We compute the matrix of outputs as:
在实际应用中，我们同时计算一组Query的注意函数，这些Query被打包成一个矩阵$Q$。Key和Value也被打包成矩阵$K$和$V$。我们计算输出矩阵：
$$\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}(Q,K,V)=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(\frac{Q{K}^T}{\sqrt{d_k}})V$$

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-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 = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

The two most commonly used attention functions are additive attention (cite), and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\frac{1}{\sqrt{d_k}}$. Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
两个最常用的注意力机制是加性注意力(additive attention)、乘法(点积)注意力。点积注意力与我们的算法相同，除了比例因子为$\frac{1}{\sqrt{d_k}}$. 加性注意力使用具有单个隐藏层的前馈网络计算兼容函数。虽然两者在理论上的复杂性相似，但点积注意力在实践中更快且更节省空间，因为它可以使用高度优化的矩阵乘法代码来实现。

While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $d_k$ (cite). We suspect that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients (To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random variables with mean $0$ and variance $1$. Then their dot product, $q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$, has mean $0$ and variance $d_k$.). To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$.
当对于较小的$d_k$值，这两种机制的表现类似，对于$d_k$较大值，加性注意力优于无缩放的点积注意力(引用)。我们怀疑对于较大的$d_k$值，点积的数量级变得很大，将softmax函数推入梯度极小的区域(为了说明为什么点积会变大，我们假设$q$和$k$的分量是均值为0、方差为1的独立随机变量。那么它们的点积$q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$，均值为$0$，方差为$d_k$)。为了抵消这种影响，我们将点积按比例$\frac{1}{\sqrt{d_k}}$缩放。

Image(filename='images/ModalNet-20.png')

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
Where the projections are parameter matrices $W_i^Q\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$, $W_i^K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$, $W_i^V\in {\mathbb{R}}^{d_{\text{model}}\times d_v}$ and $W^O\in {\mathbb{R}}^{h{d}_v\times d_{\text{model}}}$. In this work we employ $h=8$ parallel attention layers, or heads. For each of these we use $d_k=d_v=d_{\text{model}}/h=64$. Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.
多头注意力允许模型共同关注来自不同位置的不同表示子空间的信息。对于单个注意力头，平均操作会抑制这一点。投影是参数矩阵$W_i^Q\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$, $W_i^K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$, $W_i^V\in {\mathbb{R}}^{d_{\text{model}}\times d_v}$ 和 $W^O\in {\mathbb{R}}^{h{d}_v\times d_{\text{model}}}$。在这项工作中，我们采用了8个平行力注意层或头。对于其中的每一个，我们使用$d_k=d_v=d_{\text{model}}/h=64$。由于每个头的维度减少了，总计算成本与全维度的单头注意力相似。

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
        # 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"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        
        # 1) Do all the linear projections in batch from d_model => h x d_k 
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]
        
        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

Applications of Attention in our Model 注意力在模型中的应用

The Transformer uses multi-head attention in three different ways:
Transformer 以三种不同的方式使用多头注意力：

1. In “encoder-decoder attention” layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This allows every position in the decoder to attend over all positions in the input sequence. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models such as (cite).
   1）在“编码器-解码器注意力”层中，Query来自前一个解码器层，记忆的Key和Value来自编码器的输出。这允许解码器中的每个位置参与输入序列中的所有位置。这模仿了Seq2Seq模型中典型的编码器-解码器注意力机制，例如 (cite)。

2. The encoder contains self-attention layers. In a self-attention layer all of the keys, values and queries come from the same place, in this case, the output of the previous layer in the encoder. Each position in the encoder can attend to all positions in the previous layer of the encoder.
   2）编码器包含自注意力层。在自注意力层中，所有的Key、Value和Query都来自同一个地方。本模型中是编码器中前一层的输出。编码器中的每个位置都可以关注编码器上一层中的所有位置。

3. Similarly, self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot- product attention by masking out (setting to $-\infty$) all values in the input of the softmax which correspond to illegal connections.
   3)类似地，解码器中的自注意层允许解码器中的每个位置都能关注到：该位置之前及包括该位置在内的解码器中的所有位置。我们需要防止解码器中的信息向左流动，以保持自回归特性。我们通过对softmax输入中与非法连接对应的所有值进行屏蔽mask(设置为$-\infty$) 来实现缩放点乘注意力。

Position-wise Feed-Forward Networks 逐位置的前馈网络

In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.
除了注意力子层之外，我们的编码器和解码器中的每一层都包含一个全连接前馈网络，该网络分别且相同地应用于每个位置。这包括两个线性变换，中间使用ReLU激活函数。
$$\mathrm{F}\mathrm{F}\mathrm{N}(x)=\max (0,x{W}_1+b_1)W_2+b_2$$

While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way of describing this is as two convolutions with kernel size 1. The dimensionality of input and output is $d_{\text{model}}=512$, and the inner-layer has dimensionality $d_{ff}=2048$.
虽然线性变换在不同位置上是相同的，但它们在层与层之间使用不同的参数。另一种描述方式是内核大小为 1 的两个卷积。输入和输出的维度为$d_{\text{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):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

Embeddings and Softmax 词嵌入和Softmax

Similarly to other sequence transduction models, we use learned embeddings to convert the input tokens and output tokens to vectors of dimension $d_{\text{model}}$. We also use the usual learned linear transformation and softmax function to convert the decoder output to predicted next-token probabilities. In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to (cite). In the embedding layers, we multiply those weights by $\sqrt{d_{\text{model}}}$.

与其他序列转导模型类似，我们使用学习过的embedding将输入、输出token转换为$d_{\text{model}}$维的向量。我们还使用通常学到的线性变换和softmax函数将解码器输出转换为下一个token的预测概率。在我们的模型中，我们在两个embedding层和pre-softmax线性变换之间共享相同的权重矩阵，类似于 (cite)。在embedding层中，我们将这些权重乘以$\sqrt{d_{\text{model}}}$。

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

Positional Encoding 位置信息编码

Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence. To this end, we add “positional encodings” to the input embeddings at the bottoms of the encoder and decoder stacks. The positional encodings have the same dimension $d_{\text{model}}$ as the embeddings, so that the two can be summed. There are many choices of positional encodings, learned and fixed (cite).
由于我们的模型不包含递归和卷积，为了让模型利用序列的顺序，我们必须注入一些关于序列中token的相对或绝对位置的信息。为此，我们在编码器和解码器底层的输入embedding中添加“位置信息编码”。位置信息编码与词嵌入层具有相同的维度$d_{\text{model}}$，因此可以将两者相加。位置编码有很多选择：动态学习或者和固定的 （引用）。

In this work, we use sine and cosine functions of different frequencies:
本文我们使用不同频率的正弦和余弦函数:$$P{E}_{(pos,2i)}=sin(pos/10000^{2i/d_{\text{model}}})$$$$P{E}_{(pos,2i+1)}=cos(pos/10000^{2i/d_{\text{model}}})$$ where $pos$ is the position and $i$ is the dimension. That is, each dimension of the positional encoding corresponds to a sinusoid. The wavelengths form a geometric progression from $2\pi$ to $10000\cdot 2\pi$. We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset $k$, $P{E}_{pos+k}$ can be represented as a linear function of $P{E}_{pos}$.
其中$pos$是位置，$i$是维数。也就是说，位置编码的每个维度都对应一个正弦信号。波长以几何级数形式从$2\pi$到$10000\cdot 2\pi$。我们选择这个函数是因为我们假设它可以让模型轻松地学会使用相对位置，因为对于任何固定偏移$k$， $P{E}_{pos+k}$可以表示为$P{E}_{pos}$的线性函数。

In addition, we apply dropout to the sums of the embeddings and the positional encodings in both the encoder and decoder stacks. For the base model, we use a rate of $P_{drop}=0.1$.
此外，我们将 dropout 应用于编码器和解码器堆叠中词嵌入和位置编码的总和。对于基本模型，我们使用的比率为$P_{drop}=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.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        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[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], 
                         requires_grad=False)
        return self.dropout(x)

Below the positional encoding will add in a sine wave based on position. The frequency and offset of the wave is different for each dimension.
下面的位置编码将根据位置添加一个正弦波。波的频率和偏移对于每个维度都是不同的。

plt.figure(figsize=(15, 5))
pe = PositionalEncoding(20, 0)
y = pe.forward(Variable(torch.zeros(1, 100, 20)))
plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
plt.legend(["dim %d"%p for p in [4,5,6,7]])
None

We also experimented with using learned positional embeddings (cite) instead, and found that the two versions produced nearly identical results. We chose the sinusoidal version because it may allow the model to extrapolate to sequence lengths longer than the ones encountered during training.
我们还尝试在实验中使用动态学习的位置编码方法(cite)替代，发现这两个版本产生了几乎相同的结果。我们选择了正弦版本，因为它可以让模型推断出比训练期间遇到的序列长度更长的序列长度。

Full Model 完整模型

Here we define a function that takes in hyperparameters and produces a 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(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), 
                             c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab))
    
    # This was important from their code. 
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model

# Small example model.
tmp_model = make_model(10, 10, 2)
None

Training 训练

This section describes the training regime for our models.
本节描述了我们模型的训练机制。

We stop for a quick interlude to introduce some of the tools needed to train a standard encoder decoder model. First we define a batch object that holds the src and target sentences for training, as well as constructing the masks.
我们停下来做个简短的插曲，介绍一些训练标准编码器解码器模型所需的工具。首先，我们定义一个批处理对象，其中包含用于训练的source和目标句子，以及构建掩码。

Batches and Masking Batch和Mask

class Batch:
    "Object for holding a batch of data with mask during training."
    def __init__(self, src, trg=None, pad=0):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if trg is not None:
            self.trg = trg[:, :-1]
            self.trg_y = trg[:, 1:]
            self.trg_mask = \
                self.make_std_mask(self.trg, pad)
            self.ntokens = (self.trg_y != pad).data.sum()
    
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & Variable(
            subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
        return tgt_mask

Next we create a generic training and scoring function to keep track of loss. We pass in a generic loss compute function that also handles parameter updates.
接下来，我们创建一个通用的训练和评分函数来跟踪损失。我们传入一个通用的损失计算函数，它也处理参数更新。

Training Loop 训练Epoch

def run_epoch(data_iter, model, loss_compute):
    "Standard Training and Logging Function"
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    for i, batch in enumerate(data_iter):
        out = model.forward(batch.src, batch.trg, 
                            batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)
        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens
        if i % 50 == 1:
            elapsed = time.time() - start
            print("Epoch Step: %d Loss: %f Tokens per Sec: %f" %
                    (i, loss / batch.ntokens, tokens / elapsed))
            start = time.time()
            tokens = 0
    return total_loss / total_tokens

Training Data and Batching 训练数据和批处理

We trained on the standard WMT 2014 English-German dataset consisting of about 4.5 million sentence pairs. Sentences were encoded using byte-pair encoding, which has a shared source-target vocabulary of about 37000 tokens. For English- French, we used the significantly larger WMT 2014 English-French dataset consisting of 36M sentences and split tokens into a 32000 word-piece vocabulary.
我们在由大约450万个句子对组成的标准 WMT 2014 English-German数据集上进行了训练。句子使用字节对编码进行编码，该编码具有大约37000个token的source-target词汇。对于英语-法语，我们使用了更大的 WMT 2014 English- French数据集，该数据集由3600万个句子组成，并将标记拆分为32000个单词的词汇表。
Sentence pairs were batched together by approximate sequence length. Each training batch contained a set of sentence pairs containing approximately 25000 source tokens and 25000 target tokens.
句子对按近似的序列长度划分batch。每个batch包含一组句子对，其中包含大约25000个source token和25000个target token。

We will use torch text for batching. This is discussed in more detail below. Here we create batches in a torchtext function that ensures our batch size padded to the maximum batchsize does not surpass a threshold (25000 if we have 8 gpus).
我们将使用torch text进行批处理。这将在下面更详细地讨论。在这里，我们在 torchtext 函数中创建批次，以确保我们padding到最大batchsize的batchsize不超过阈值（如果我们有 8 个 GPU，则为 25000）。

global max_src_in_batch, max_tgt_in_batch
def batch_size_fn(new, count, sofar):
    "Keep augmenting batch and calculate total number of tokens + padding."
    global max_src_in_batch, max_tgt_in_batch
    if count == 1:
        max_src_in_batch = 0
        max_tgt_in_batch = 0
    max_src_in_batch = max(max_src_in_batch,  len(new.src))
    max_tgt_in_batch = max(max_tgt_in_batch,  len(new.trg) + 2)
    src_elements = count * max_src_in_batch
    tgt_elements = count * max_tgt_in_batch
    return max(src_elements, tgt_elements)

Hardware and Schedule 硬件和时间

We trained our models on one machine with 8 NVIDIA P100 GPUs. For our base models using the hyperparameters described throughout the paper, each training step took about 0.4 seconds. We trained the base models for a total of 100,000 steps or 12 hours. For our big models, step time was 1.0 seconds. The big models were trained for 300,000 steps (3.5 days).
我们在一台配备 8 个 NVIDIA P100 GPU 的机器上训练我们的模型。对于我们使用整篇论文中描述的超参数的基础模型，每个训练步骤大约需要0.4秒。我们对基础模型进行了总共100,000步或12小时的训练。对于我们的大模型，每个步骤时间为 1.0 秒。大型模型训练了300,000步（3.5 天）。

Optimizer 优化器

We used the Adam optimizer (cite) with ${\beta }_1=0.9$, ${\beta }_2=0.98$ and $ϵ=10^{-9}$. We varied the learning rate over the course of training, according to the formula: $$lrate=d_{\text{model}}^{-0.5}\cdot \min ({step\_num}^{-0.5},step\_num\cdot {warmup\_steps}^{-1.5})$$ This corresponds to increasing the learning rate linearly for the first $warmup\_steps$ training steps, and decreasing it thereafter proportionally to the inverse square root of the step number. We used $warmup\_steps=4000$.

我们使用了 Adam 优化器（引用），其中 ${\beta }_1=0.9$、${\beta }_2=0.98$、$ϵ=10^{-9}$。我们根据以下公式在训练过程中改变学习率： $$lrate=d_{\text{model}}^{-0.5}\cdot \min ({step\_num}^{-0.5},step\_num\cdot {warmup\_steps}^{-1.5})$$
这对应于第一个$warmup\_steps$训练步骤的学习速率线性增加，然后与步数的平方根的倒数成比例地减少。$warmup\_steps=4000$。

Note: This part is very important. Need to train with this setup of the model.
注意：这部分非常重要。需要使用此模型设置进行训练。

class NoamOpt:
    "Optim wrapper that implements rate."
    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0
        
    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()
        
    def rate(self, step = None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * \
            (self.model_size ** (-0.5) *
            min(step ** (-0.5), step * self.warmup ** (-1.5)))
        
def get_std_opt(model):
    return NoamOpt(model.src_embed[0].d_model, 2, 4000,
            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

Example of the curves of this model for different model sizes and for optimization hyperparameters.
该模型针对不同模型大小和优化超参数的曲线示例。

# Three settings of the lrate hyperparameters.
opts = [NoamOpt(512, 1, 4000, None), 
        NoamOpt(512, 1, 8000, None),
        NoamOpt(256, 1, 4000, None)]
plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])
plt.legend(["512:4000", "512:8000", "256:4000"])
None

Regularization 正则化

Label Smoothing 标签平滑

During training, we employed label smoothing of value ${ϵ}_{ls}=0.1$ (cite). This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.
在训练期间，我们使用了标签平滑值${ϵ}_{ls}=0.1$（引用）。这会损害困惑度，因为模型会变得更加不确定，但会提高准确性和BLEU分数。

We implement label smoothing using the KL div loss. Instead of using a one-hot target distribution, we create a distribution that has confidence of the correct word and the rest of the smoothing mass distributed throughout the vocabulary.
我们使用KL散度实现标签平滑。我们没有使用独热目标分布，而是创建了一个分布，该分布具有confidence正确的单词，其余的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(size_average=False)
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None
        
    def forward(self, x, target):
        assert x.size(1) == self.size
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (self.size - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, Variable(true_dist, requires_grad=False))

Here we can see an example of how the mass is distributed to the words based on confidence.
在这里，我们可以看到质量如何基于置信度分配给单词的示例。

# Example of 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]])
v = crit(Variable(predict.log()), 
         Variable(torch.LongTensor([2, 1, 0])))

# Show the target distributions expected by the system.
plt.imshow(crit.true_dist)
None

Label smoothing actually starts to penalize the model if it gets very confident about a given choice.
如果模型对给定的选择非常有信心，标签平滑实际上会开始惩罚模型。

crit = LabelSmoothing(5, 0, 0.1)
def loss(x):
    d = x + 3 * 1
    predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d],
                                 ])
    #print(predict)
    return crit(Variable(predict.log()),
                 Variable(torch.LongTensor([1]))).data[0]
plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])
None

A First Example 第一个例子

We can begin by trying out a simple copy-task. Given a random set of input symbols from a small vocabulary, the goal is to generate back those same symbols.
我们可以从尝试一个简单的复制任务开始。给定来自小词汇表的一组随机输入符号，目标是生成相同的符号。

Synthetic Data 综合数据

def data_gen(V, batch, nbatches):
    "Generate random data for a src-tgt copy task."
    for i in range(nbatches):
        data = torch.from_numpy(np.random.randint(1, V, size=(batch, 10)))
        data[:, 0] = 1
        src = Variable(data, requires_grad=False)
        tgt = Variable(data, requires_grad=False)
        yield Batch(src, tgt, 0)

Loss Computation 损失计算

class SimpleLossCompute:
    "A simple loss compute and train function."
    def __init__(self, generator, criterion, opt=None):
        self.generator = generator
        self.criterion = criterion
        self.opt = opt
        
    def __call__(self, x, y, norm):
        x = self.generator(x)
        loss = self.criterion(x.contiguous().view(-1, x.size(-1)), 
                              y.contiguous().view(-1)) / norm
        loss.backward()
        if self.opt is not None:
            self.opt.step()
            self.opt.optimizer.zero_grad()
        return loss.data[0] * norm

Greedy Decoding 贪心搜索解码

# Train the simple copy task.
V = 11
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
model = make_model(V, V, N=2)
model_opt = NoamOpt(model.src_embed[0].d_model, 1, 400,
        torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

for epoch in range(10):
    model.train()
    run_epoch(data_gen(V, 30, 20), model, 
              SimpleLossCompute(model.generator, criterion, model_opt))
    model.eval()
    print(run_epoch(data_gen(V, 30, 5), model, 
                    SimpleLossCompute(model.generator, criterion, None)))

Epoch Step: 1 Loss: 3.023465 Tokens per Sec: 403.074173
Epoch Step: 1 Loss: 1.920030 Tokens per Sec: 641.689380
1.9274832487106324
Epoch Step: 1 Loss: 1.940011 Tokens per Sec: 432.003378
Epoch Step: 1 Loss: 1.699767 Tokens per Sec: 641.979665
1.657595729827881
Epoch Step: 1 Loss: 1.860276 Tokens per Sec: 433.320240
Epoch Step: 1 Loss: 1.546011 Tokens per Sec: 640.537198
1.4888023376464843
Epoch Step: 1 Loss: 1.682198 Tokens per Sec: 432.092305
Epoch Step: 1 Loss: 1.313169 Tokens per Sec: 639.441857
1.3485562801361084
Epoch Step: 1 Loss: 1.278768 Tokens per Sec: 433.568756
Epoch Step: 1 Loss: 1.062384 Tokens per Sec: 642.542067
0.9853351473808288
Epoch Step: 1 Loss: 1.269471 Tokens per Sec: 433.388727
Epoch Step: 1 Loss: 0.590709 Tokens per Sec: 642.862135
0.5686767101287842
Epoch Step: 1 Loss: 0.997076 Tokens per Sec: 433.009746
Epoch Step: 1 Loss: 0.343118 Tokens per Sec: 642.288427
0.34273059368133546
Epoch Step: 1 Loss: 0.459483 Tokens per Sec: 434.594030
Epoch Step: 1 Loss: 0.290385 Tokens per Sec: 642.519464
0.2612409472465515
Epoch Step: 1 Loss: 1.031042 Tokens per Sec: 434.557008
Epoch Step: 1 Loss: 0.437069 Tokens per Sec: 643.630322
0.4323212027549744
Epoch Step: 1 Loss: 0.617165 Tokens per Sec: 436.652626
Epoch Step: 1 Loss: 0.258793 Tokens per Sec: 644.372296
0.27331129014492034

This code predicts a translation using greedy decoding for simplicity.
为简单起见，此代码使用Greedy Search来预测翻译。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)
    for i in range(max_len-1):
        out = model.decode(memory, src_mask, 
                           Variable(ys), 
                           Variable(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.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
    return ys

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

    1     2     3     4     5     6     7     8     9    10
[torch.LongTensor of size 1x10]

A Real World Example 一个真实的例子

Now we consider a real-world example using the IWSLT German-English Translation task. This task is much smaller than the WMT task considered in the paper, but it illustrates the whole system. We also show how to use multi-gpu processing to make it really fast.
现在我们考虑一个使用 IWSLT 德语-英语翻译任务的真实示例。这个任务比论文中考虑的 WMT 任务要小得多，但它说明了整个系统。我们还展示了如何使用多 GPU 处理使其真正快速。

#!pip install torchtext spacy
#!python -m spacy download en
#!python -m spacy download de

Data Loading 数据加载

We will load the dataset using torchtext and spacy for tokenization.
我们将使用 torchtext 和 spacy 加载数据集以进行tokenization。

# For data loading.
from torchtext import data, datasets

if True:
    import spacy
    spacy_de = spacy.load('de')
    spacy_en = spacy.load('en')

    def tokenize_de(text):
        return [tok.text for tok in spacy_de.tokenizer(text)]

    def tokenize_en(text):
        return [tok.text for tok in spacy_en.tokenizer(text)]

    BOS_WORD = ''
    EOS_WORD = ''
    BLANK_WORD = ""
    SRC = data.Field(tokenize=tokenize_de, pad_token=BLANK_WORD)
    TGT = data.Field(tokenize=tokenize_en, init_token = BOS_WORD, 
                     eos_token = EOS_WORD, pad_token=BLANK_WORD)

    MAX_LEN = 100
    train, val, test = datasets.IWSLT.splits(
        exts=('.de', '.en'), fields=(SRC, TGT), 
        filter_pred=lambda x: len(vars(x)['src']) <= MAX_LEN and 
            len(vars(x)['trg']) <= MAX_LEN)
    MIN_FREQ = 2
    SRC.build_vocab(train.src, min_freq=MIN_FREQ)
    TGT.build_vocab(train.trg, min_freq=MIN_FREQ)

Batching matters a ton for speed. We want to have very evenly divided batches, with absolutely minimal padding. To do this we have to hack a bit around the default torchtext batching. This code patches their default batching to make sure we search over enough sentences to find tight batches.
批处理对速度很重要。我们希望batch划分的非常均匀，且padding绝对最小。为此，我们必须对默认的 torchtext 批处理进行一些修改。这段代码修补了它们的默认批处理，以确保我们搜索足够多的句子来找到紧密的批处理。

Iterators 迭代器

class MyIterator(data.Iterator):
    def create_batches(self):
        if self.train:
            def pool(d, random_shuffler):
                for p in data.batch(d, self.batch_size * 100):
                    p_batch = data.batch(
                        sorted(p, key=self.sort_key),
                        self.batch_size, self.batch_size_fn)
                    for b in random_shuffler(list(p_batch)):
                        yield b
            self.batches = pool(self.data(), self.random_shuffler)
            
        else:
            self.batches = []
            for b in data.batch(self.data(), self.batch_size,
                                          self.batch_size_fn):
                self.batches.append(sorted(b, key=self.sort_key))

def rebatch(pad_idx, batch):
    "Fix order in torchtext to match ours"
    src, trg = batch.src.transpose(0, 1), batch.trg.transpose(0, 1)
    return Batch(src, trg, pad_idx)

Multi-GPU Training 多GPU训练

Finally to really target fast training, we will use multi-gpu. This code implements multi-gpu word generation. It is not specific to transformer so I won’t go into too much detail. The idea is to split up word generation at training time into chunks to be processed in parallel across many different gpus. We do this using pytorch parallel primitives:
最后，为了真正以快速训练为目标，我们将使用多GPU。这段代码实现了多GPU的词的生成。它不是特定用于Transformer的，所以我不会详细介绍。这个想法是将训练时的单词生成分成块，以便在许多不同的GPU上并行处理。我们使用pytorch并行原语来做到这一点：

• replicate - split modules onto different gpus.
• replicate - 将模块拆分到不同的 GPU 上
• scatter - split batches onto different gpus
• scatter - 将批次拆分到不同的 GPU 上
• parallel_apply - apply module to batches on different gpus
• parallel_apply - 将模块应用于不同 GPU 上的批次
• gather - pull scattered data back onto one gpu.
• gather - 将分散的数据拉回一个 GPU
• nn.DataParallel - a special module wrapper that calls these all before evaluating.
• nn.DataParallel - 一个特殊的模块包装器，在评估之前调用这些

# Skip if not interested in multigpu.
class MultiGPULossCompute:
    "A multi-gpu loss compute and train function."
    def __init__(self, generator, criterion, devices, opt=None, chunk_size=5):
        # Send out to different gpus.
        self.generator = generator
        self.criterion = nn.parallel.replicate(criterion, 
                                               devices=devices)
        self.opt = opt
        self.devices = devices
        self.chunk_size = chunk_size
        
    def __call__(self, out, targets, normalize):
        total = 0.0
        generator = nn.parallel.replicate(self.generator, 
                                                devices=self.devices)
        out_scatter = nn.parallel.scatter(out, 
                                          target_gpus=self.devices)
        out_grad = [[] for _ in out_scatter]
        targets = nn.parallel.scatter(targets, 
                                      target_gpus=self.devices)

        # Divide generating into chunks.
        chunk_size = self.chunk_size
        for i in range(0, out_scatter[0].size(1), chunk_size):
            # Predict distributions
            out_column = [[Variable(o[:, i:i+chunk_size].data, 
                                    requires_grad=self.opt is not None)] 
                           for o in out_scatter]
            gen = nn.parallel.parallel_apply(generator, out_column)

            # Compute loss. 
            y = [(g.contiguous().view(-1, g.size(-1)), 
                  t[:, i:i+chunk_size].contiguous().view(-1)) 
                 for g, t in zip(gen, targets)]
            loss = nn.parallel.parallel_apply(self.criterion, y)

            # Sum and normalize loss
            l = nn.parallel.gather(loss, 
                                   target_device=self.devices[0])
            l = l.sum()[0] / normalize
            total += l.data[0]

            # Backprop loss to output of transformer
            if self.opt is not None:
                l.backward()
                for j, l in enumerate(loss):
                    out_grad[j].append(out_column[j][0].grad.data.clone())

        # Backprop all loss through transformer.            
        if self.opt is not None:
            out_grad = [Variable(torch.cat(og, dim=1)) for og in out_grad]
            o1 = out
            o2 = nn.parallel.gather(out_grad, 
                                    target_device=self.devices[0])
            o1.backward(gradient=o2)
            self.opt.step()
            self.opt.optimizer.zero_grad()
        return total * normalize

Now we create our model, criterion, optimizer, data iterators, and paralelization
现在我们创建模型、标准、优化器、数据迭代器和并行化

# GPUs to use
devices = [0, 1, 2, 3]
if True:
    pad_idx = TGT.vocab.stoi[""]
    model = make_model(len(SRC.vocab), len(TGT.vocab), N=6)
    model.cuda()
    criterion = LabelSmoothing(size=len(TGT.vocab), padding_idx=pad_idx, smoothing=0.1)
    criterion.cuda()
    BATCH_SIZE = 12000
    train_iter = MyIterator(train, batch_size=BATCH_SIZE, device=0,
                            repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)),
                            batch_size_fn=batch_size_fn, train=True)
    valid_iter = MyIterator(val, batch_size=BATCH_SIZE, device=0,
                            repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)),
                            batch_size_fn=batch_size_fn, train=False)
    model_par = nn.DataParallel(model, device_ids=devices)
None

Now we train the model. I will play with the warmup steps a bit, but everything else uses the default parameters. On an AWS p3.8xlarge with 4 Tesla V100s, this runs at ~27,000 tokens per second with a batch size of 12,000
现在我们训练模型。我会稍微玩一下热身步骤，但其他一切都使用默认参数。在具有 4 个 Tesla V100 的 AWS p3.8xlarge 上，它以每秒约 27,000 个token的速度运行，批量大小为 12,000

Training the System 训练

#!wget https://s3.amazonaws.com/opennmt-models/iwslt.pt

if False:
    model_opt = NoamOpt(model.src_embed[0].d_model, 1, 2000,
            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
    for epoch in range(10):
        model_par.train()
        run_epoch((rebatch(pad_idx, b) for b in train_iter), 
                  model_par, 
                  MultiGPULossCompute(model.generator, criterion, 
                                      devices=devices, opt=model_opt))
        model_par.eval()
        loss = run_epoch((rebatch(pad_idx, b) for b in valid_iter), 
                          model_par, 
                          MultiGPULossCompute(model.generator, criterion, 
                          devices=devices, opt=None))
        print(loss)
else:
    model = torch.load("iwslt.pt")

Once trained we can decode the model to produce a set of translations. Here we simply translate the first sentence in the validation set. This dataset is pretty small so the translations with greedy search are reasonably accurate.
一旦经过训练，我们就可以对模型进行解码以产生一组翻译。这里我们简单翻译验证集中的第一句话。这个数据集非常小，所以贪婪搜索的翻译相当准确。

for i, batch in enumerate(valid_iter):
    src = batch.src.transpose(0, 1)[:1]
    src_mask = (src != SRC.vocab.stoi[""]).unsqueeze(-2)
    out = greedy_decode(model, src, src_mask, 
                        max_len=60, start_symbol=TGT.vocab.stoi[""])
    print("Translation:", end="\t")
    for i in range(1, out.size(1)):
        sym = TGT.vocab.itos[out[0, i]]
        if sym == "": break
        print(sym, end =" ")
    print()
    print("Target:", end="\t")
    for i in range(1, batch.trg.size(0)):
        sym = TGT.vocab.itos[batch.trg.data[i, 0]]
        if sym == "": break
        print(sym, end =" ")
    print()
    break

Translation:	  . In my language , that means , thank you very much . 
Gold:	  . It means in my language , thank you very much . 

Additional Components: BPE, Search, Averaging 附加组件：BPE、搜索、平均

So this mostly covers the transformer model itself. There are four aspects that we didn’t cover explicitly. We also have all these additional features implemented in OpenNMT-py.
所以这主要涵盖了Transformer模型本身。我们没有明确涵盖四个方面。我们还在OpenNMT-py中实现了所有这些附加功能。

1. BPE/ Word-piece: We can use a library to first preprocess the data into subword units. See Rico Sennrich’s subword- nmt implementation. These models will transform the training data to look like this:
   BPE（byte pair encoder）字节对编码/ Word-piece：我们可以使用库首先将数据预处理成子词单元。请参阅 Rico Sennrich 的子词 nmt实现。这些模型会将训练数据转换为如下所示：

▁Die ▁Protokoll datei ▁kann ▁ heimlich ▁per ▁E - Mail ▁oder ▁FTP ▁an ▁einen ▁bestimmte n ▁Empfänger ▁gesendet ▁werden .

2. Shared Embeddings: When using BPE with shared vocabulary we can share the same weight vectors between the source / target / generator. See the (cite) for details. To add this to the model simply do this:
   共享嵌入：当使用具有共享词汇表的 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

3. Beam Search: This is a bit too complicated to cover here. See the OpenNMT- py for a pytorch implementation.
   Beam Search：这有点太复杂了，无法在这里介绍。有关 pytorch 实现，请参阅OpenNMT-py 。

4. Model Averaging: The paper averages the last k checkpoints to create an ensembling effect. We can do this after the fact if we have a bunch of models:
   模型平均：论文对最后k个检查点进行平均，以创建一个集成效果。如果我们有一堆模型，我们可以在事后这样做：

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

Results 结果

On the WMT 2014 English-to-German translation task, the big transformer model (Transformer (big) in Table 2) outperforms the best previously reported models (including ensembles) by more than 2.0 BLEU, establishing a new state-of-the-art BLEU score of 28.4. The configuration of this model is listed in the bottom line of Table 3. Training took 3.5 days on 8 P100 GPUs. Even our base model surpasses all previously published models and ensembles, at a fraction of the training cost of any of the competitive models.
在 WMT 2014 英德翻译任务中，大 Transformer 模型（表 2 中的 Transformer (big)）比之前报道的最佳模型（包括集成）高出 2.0 BLEU 以上，建立了一个新的最先进的 BLEU 得分为 28.4。该模型的配置列于表 3 的最后一行。在 8 个 P100 GPU 上训练耗时 3.5 天。甚至我们的基础模型也超过了所有先前发布的模型和集成，其训练成本只是任何有竞争力的模型的一小部分。

On the WMT 2014 English-to-French translation task, our big model achieves a BLEU score of 41.0, outperforming all of the previously published single models, at less than 1/4 the training cost of the previous state-of-the-art model. The Transformer (big) model trained for English-to-French used dropout rate Pdrop = 0.1, instead of 0.3.
在 WMT 2014 英法翻译任务中，我们的大模型达到 41.0 的 BLEU 分数，优于之前发布的所有单一模型，其训练成本不到之前最先进技术模型的1/4。用于英语到法语训练的Transformer大模型使用的dropout率 Pdrop = 0.1，而不是 0.3。

Image(filename="images/results.png")

The code we have written here is a version of the base model. There are fully trained version of this system available here (Example Models).
我们在这里编写的代码是基本模型的一个版本。此处提供了该系统的完全训练版本 （示例模型）。
With the addtional extensions in the last section, the OpenNMT-py replication gets to 26.9 on EN-DE WMT. Here I have loaded in those parameters to our reimplemenation.
通过上一节中的附加扩展，OpenNMT-py 复制在 英-德 WMT 上达到 26.9。在这里，我已将这些参数加载到我们的重新实现中。

!wget https://s3.amazonaws.com/opennmt-models/en-de-model.pt

model, SRC, TGT = torch.load("en-de-model.pt")

model.eval()
sent = "▁The ▁log ▁file ▁can ▁be ▁sent ▁secret ly ▁with ▁email ▁or ▁FTP ▁to ▁a ▁specified ▁receiver".split()
src = torch.LongTensor([[SRC.stoi[w] for w in sent]])
src = Variable(src)
src_mask = (src != SRC.stoi[""]).unsqueeze(-2)
out = greedy_decode(model, src, src_mask, 
                    max_len=60, start_symbol=TGT.stoi[""])
print("Translation:", end="\t")
trans = " "
for i in range(1, out.size(1)):
    sym = TGT.itos[out[0, i]]
    if sym == "": break
    trans += sym + " "
print(trans)

Translation:	 ▁Die ▁Protokoll datei ▁kann ▁ heimlich ▁per ▁E - Mail ▁oder ▁FTP ▁an ▁einen ▁bestimmte n ▁Empfänger ▁gesendet ▁werden . 

Attention Visualization 注意力可视化

Even with a greedy decoder the translation looks pretty good. We can further visualize it to see what is happening at each layer of the attention
即使使用贪婪搜索的解码器，翻译看起来也很不错。我们可以进一步将其可视化，以查看注意力的每一层发生了什么

tgt_sent = trans.split()
def draw(data, x, y, ax):
    seaborn.heatmap(data, 
                    xticklabels=x, square=True, yticklabels=y, vmin=0.0, vmax=1.0, 
                    cbar=False, ax=ax)
    
for layer in range(1, 6, 2):
    fig, axs = plt.subplots(1,4, figsize=(20, 10))
    print("Encoder Layer", layer+1)
    for h in range(4):
        draw(model.encoder.layers[layer].self_attn.attn[0, h].data, 
            sent, sent if h ==0 else [], ax=axs[h])
    plt.show()
    
for layer in range(1, 6, 2):
    fig, axs = plt.subplots(1,4, figsize=(20, 10))
    print("Decoder Self Layer", layer+1)
    for h in range(4):
        draw(model.decoder.layers[layer].self_attn.attn[0, h].data[:len(tgt_sent), :len(tgt_sent)], 
            tgt_sent, tgt_sent if h ==0 else [], ax=axs[h])
    plt.show()
    print("Decoder Src Layer", layer+1)
    fig, axs = plt.subplots(1,4, figsize=(20, 10))
    for h in range(4):
        draw(model.decoder.layers[layer].self_attn.attn[0, h].data[:len(tgt_sent), :len(sent)], 
            sent, tgt_sent if h ==0 else [], ax=axs[h])
    plt.show()

Encoder Layer 2

Encoder Layer 4

Encoder Layer 6

Decoder Self Layer 2

Decoder Src Layer 2

Decoder Self Layer 4

Decoder Src Layer 4

Decoder Self Layer 6

Decoder Src Layer 6

Conclusion 结论

Hopefully this code is useful for future research. Please reach out if you have any issues. If you find this code helpful, also check out our other OpenNMT tools.
希望这段代码对未来的研究有用。如果您有任何问题，请与我们联系。如果您发现此代码有帮助，还请查看我们的其他 OpenNMT 工具。

@inproceedings{opennmt,
  author    = {Guillaume Klein and
               Yoon Kim and
               Yuntian Deng and
               Jean Senellart and
               Alexander M. Rush},
  title     = {OpenNMT: Open-Source Toolkit for Neural Machine Translation},
  booktitle = {Proc. ACL},
  year      = {2017},
  url       = {https://doi.org/10.18653/v1/P17-4012},
  doi       = {10.18653/v1/P17-4012}
}