【精读】Transformer模型深度解读

英语原文：The Annotated Transformer

翻译：雷锋字幕组（Icarus、）

---

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

 

 

在过去的一年里，《注意力就是你所需要的》中的Transformer被很多人所关注。除了在翻译质量上产生重大改进外，它还为许多其他NLP任务提供了一个新的架构。这篇论文本身写得非常清楚，但传统的观点是，它的正确实现相当困难。

在这篇文章中，我以逐行实现的形式呈现了论文的 "注释 "版本。我对原论文中的一些章节进行了重新排序和删除，并在全文中添加了注释。这个文档本身就是一个工作笔记，应该是一个完全可用的实现。总共有400行库代码，可以在4个GPU上每秒处理27000个token。

要想跟上，你首先需要安装PyTorch。完整的笔记本也可以在github或Google Colab上使用免费的GPU。

请注意，这仅仅是研究人员和感兴趣的开发人员的一个起点。这里的代码主要基于我们的OpenNMT包。(如果有帮助，请自由引用。)对于其他模型的全服务实现，请查看Tensor2Tensor (tensorflow)和Sockeye (mxnet)。

Alexander Rush (@harvardnlp 或 [email protected])，得到Vincent Nguyen和Guillaume Klein的帮助。

预先设置

# !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

目录

• Prelims

• 背景资料

• 模型结构

        编码器

        解码器

        注意事项

        注意力在我们的模型中的应用

        编码器和解码器堆栈

        位置导向的前馈网络

        嵌入和Softmax

        位置编码

        完整模型

• 训练

        标签平滑化

        批量和遮蔽

        训练回路

        硬件和时间表

        优化器

        正规化

• 第一个例子

        合成数据

        损失计算

        greedy解码

• 一个现实例子

        加载数据

        迭代器

        多GPU训练

        训练系统

        附加组件：BPE、搜索、平均值

• 结果

        注意力可视化

• 结论

背景资料

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

自注意，有时也被称为内注意，是一种将单个序列的不同位置联系起来以计算序列的表示的注意机制。自我注意已被成功地应用于各种任务中，包括阅读理解、抽象概括、文本内涵和学习任务无关的句子表征。端到端记忆网络是基于循环注意机制而不是序列对齐的循环，并且已经被证明在简单语言问题回答和语言建模任务上表现良好。

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

模型结构

大多数竞争性神经序列转导模型都有一个编码器-解码器结构(cite)。在这里，编码器将输入的符号表示序列(x1，...,xn)映射为连续表示序列z=(z1，...,zn)。在给定z的情况下，解码器每次生成一个符号的输出序列(y1,...,ym)。在每一步中，该模型都是自动递减的(cite)，在生成下一个符号时，消耗之前生成的符号作为额外的输入。

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)

 

 

Transformer沿用了这种整体架构，编码器和解码器都采用堆叠式自关注和点对点的全连接层，分别如图1的左半部分和右半部分所示。

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

 

编码器和解码器堆栈

编码器

编码器是由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)

 

我们在两个子层周围分别采用残差连接(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

 

也就是说，每个子层的输出是LayerNorm(x+Sublayer(x))，其中Sublayer(x)是子层自己实现的函数。我们在每个子层的输出中应用dropout(引用)，然后再加到子层的输入中并进行归一化。

为了方便这些残差连接，模型中的所有子层以及嵌入层都会产生维度为dmodel=512dmodel=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)))

 

每层有两个子层。第一层是一个多头自注意机制，第二层是一个简单的、基于位置的全连接前馈网络。

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)

 

解码器

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

 

除了每个编码器层中的两个子层外，解码器还插入了第三个子层，它在编码器堆栈的输出上执行多头关注。与编码器类似，我们在每个子层周围采用残余连接，然后进行层归一化。

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)

 

我们还修改了解码器堆栈中的自注意力子层，以防止位置对后续位置的关注。这种遮蔽，加上输出嵌入偏移这一个位置的事实，确保位置ii的预测只能依赖于小于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

 

笔者注：注意力掩码下方显示了每个tgt词（行）被允许看的位置（列）。在训练过程中，单词会被屏蔽，以便关注之后的单词。

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

 

注意事项

注意函数可以描述为将一个查询和一组键值对映射到一个输出，其中查询、键、值和输出都是向量。输出是以值的加权和来计算的，其中分配给每个值的权重是由查询与对应键的兼容性函数计算出来的。我们把我们的特殊注意力称为 "Scaled Dot-Product注意力"。输入由维度为dk的查询和键，以及维度为dv的值组成。我们计算查询与所有键的点积，每一个点积除以√dk，然后应用一个softmax函数来获得值的权重。

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

 

在实际应用中，我们对一组查询同时计算关注函数，打包成一个矩阵QQ。键和值也一起打包成矩阵KK和VV。我们计算输出的矩阵为：

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

 

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

虽然对于dk的小值，两种机制的表现相似，但对于dk的大值，加法注意力的表现优于点积注意力，而没有缩放（引用）。我们怀疑，对于dk的大值，点积的幅度会变大，将softmax函数推到它的梯度极小的区域（为了说明点积为什么会变大，假设qq和kk的分量是均值00、方差11的独立随机变量。那么它们的点积q⋅k=∑dki=1qiki，其均值为00，方差为dk）。

为了抵消这种影响，我们将点乘积的比例为1/√dk。

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

 

多头注意使模型能够共同注意来自不同位置的不同表征子空间的信息。在单注意头的情况下，平均化会抑制这一点。

其中，投影为参数矩阵

和

在这项工作中，我们采用了h=8h=8个平行的注意层，或者说头。对于其中的每一个层，我们使用 ：

由于每个头的维度减少，总的计算成本与全维度的单头注意力相似。

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)

 

注意力在我们的模型中的应用

Transformer以三种不同的方式使用多头注意：

1. 在 "编码器-解码器关注 "层中，查询来自前一个解码器层，而内存键和值来自编码器的输出。这使得解码器中的每个位置都能在输入序列的所有位置上进行关注。这模仿了序列到序列模型中典型的编码器-解码器的注意机制，如（引用）。

2. 编码器包含自注意层。在自注意层中，所有的键、值和查询都来自同一个地方，在这种情况下，就是编码器中上一层的输出。编码器中的每个位置都可以参加编码器前一层的所有位置。

3. 同样，解码器中的自注意层允许解码器中的每个位置都可以参加解码器中的所有位置，直到并包括该位置。我们需要防止解码器中的左向信息流，以保持自动递减特性。我们通过屏蔽（设置为-∞）softmax的输入中所有对应非法连接的值，在scaled dot- product attention里面实现。

位置导向的前馈网络

除了注意力子层，我们的编码器和解码器中的每个层都包含一个完全连接的前馈网络，该网络分别和相同地应用于每个位置。这包括两个线性变换，中间有一个ReLU激活。

虽然不同位置的线性变换是相同的，但它们在层与层之间使用不同的参数。另一种描述方式是内核大小为1的两个卷积。输入和输出的维度为dmodel=512，内层的维度为dff=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))))

嵌入和Softmax

与其他序列转导模型类似，我们使用学习的嵌入将输入令牌和输出令牌转换为维数dmodel的向量。我们还使用通常的学习线性变换和softmax函数将解码器输出转换为预测的下一个标记概率。在我们的模型中，我们在两个嵌入层和预softmax线性变换之间共享相同的权重矩阵，类似于（引用）。在嵌入层中，我们将这些权重乘以√dmodel。

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)

 

位置编码

由于我们的模型不包含递归和卷积，为了使模型能够利用序列的顺序，我们必须注入一些关于序列中标记的相对或绝对位置的信息。为此，我们在编码器和解码器堆栈底部的输入嵌入中加入 "位置编码"。位置编码与嵌入的维度dmodel相同，因此两者可以相加。位置编码有很多选择，有学习的和固定的（引用）。

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

其中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.                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)

 

下面的位置编码会根据位置加入一个正弦波。每个维度的波的频率和偏移量是不同的。

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

 

我们还试验了使用学习的位置嵌入（cite）来代替，并发现这两个版本产生的结果几乎是相同的。我们选择了正弦版本，因为它可能允许模型外推到比训练过程中遇到的序列长度更长的序列。

完整模型

在这里，我们定义了一个函数，它可以接受超参数并产生一个完整的模型。

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

 

训练

本节介绍了我们模型的训练制度。

笔者注：我们停下来做一个快速的插曲，介绍一些训练标准编码器解码器模型所需的工具。首先，我们定义了一个批处理对象，它保存了用于训练的src和目标句子，以及构建掩码。

批量和屏蔽

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

 

笔者注：接下来我们创建一个通用的训练和评分函数来跟踪损失。我们传递了一个通用的损失计算函数，它也处理参数更新。

 

训练循环

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

 

训练数据和批处理

我们在标准的WMT 2014英德数据集上进行了训练，该数据集由大约450万句子对组成。句子使用字节对编码，其共享的源-目标词汇约为37000个tokens。对于英语-法语，我们使用了明显更大的WMT 2014英法数据集，该数据集由36M句子组成，并将tokens拆分为32000个词片词汇。

句子对被按近似序列长度分批在一起。每个训练批次都包含一组句子对，包含大约25000个源标记和25000个目标标记。  

笔者注：我们将使用火炬文本进行批处理。这将在下面详细讨论。在这里，我们在torchtext函数中创建批处理，以确保我们的批处理大小垫到最大batchsize不超过一个阈值（25000，如果我们有8个gpus）。

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)

 

硬件和时间表

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

优化器

我们使用了Adam优化器（引用），β1=0.9，β2=0.98，ϵ=10^-9。我们根据公式，在训练过程中改变学习率。

这相当于对第一个warmupstepswarmupsteps训练步数线性增加学习率，此后按步数的倒平方根按比例减少。我们使用warmupsteps=4000。

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

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))

 

该模型在不同模型大小和优化超参数下的曲线示例。

# 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

 

正规化

标签平滑化

在训练过程中，我们采用了价值ϵls=0.1ϵls=0.1的标签平滑（引用）。这伤害了迷惑性，因为模型学会了更多的不确定，但提高了准确性和BLEU得分。

笔者注：我们使用KL div loss实现标签平滑。而不是使用一个一热的目标分布，我们创建了一个分布，有信心的正确的单词和其余的平滑质量分布在整个词汇。

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))

 

在这里，我们可以看到一个例子，说明质量是如何根据置信度分配给单词的。

 

# 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

 

标签平滑实际上是开始惩罚模型，如果它对给定的选择非常自信的话。

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

 

一个例子

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

综合数据

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)

 

损失计算

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

# 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

 

这个代码为了简单起见，使用贪婪的解码来预测翻译。

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]

 

一个真实的例子

笔者注：现在我们考虑使用IWSLT德英翻译任务的一个真实世界的例子。这个任务比论文中考虑的WMT任务小得多，但它说明了整个系统。我们还展示了如何使用多GPU处理来使其真正快速。

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

 

数据加载

笔者注：我们将使用 torchtext 和 spacy 加载数据集进行标记化。

# 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)

 

批量对速度有很大的影响。我们希望有非常均匀的批次，绝对最小的填充。要做到这一点，我们必须对默认的 torchtext 批量进行一些改进。这段代码修补了他们的默认批处理，以确保我们搜索到的句子足够多，从而找到紧凑的批处理。

迭代器

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)

 

多CPU训练

笔者注：最后为了真正针对快速训练，我们将使用多gpu。这段代码实现了多gpu词的生成。它不是专门针对transformer的，所以我不会说得太详细。我们的想法是在训练时将单词生成分割成块，在许多不同的gpus上并行处理。我们使用 pytorch 并行基元来实现这一目的。

• replicate - 将模块分割到不同的gpus上。

• scatter - 分割批次到不同的gpus上。

• parallel_apply - 将模块应用到不同的gpus上的批次。

• gather - 把分散的数据拉回到一个gpu上。

• 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

 

现在我们创建我们的模型、准则、优化器、数据迭代器和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

 

现在我们训练模型。我会稍微玩一下热身步骤，但其他一切都使用默认参数。在一个有4个特斯拉V100的AWS p3.8xlarge上，这以每秒约27,000个代币的速度运行，批次大小为12,000个。

训练系统

#!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")

 

笔者注：一旦训练好了，我们就可以对模型进行解码，产生一组翻译。这里我们只需翻译验证集中的第一句话。这个数据集相当小，所以通过贪婪搜索的翻译是相当准确的。

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 .

附加组件：BPE、搜索、平均

笔者注：所以这主要是涵盖了变压器模型本身。有四个方面我们没有明确地覆盖。我们还在OpenNMT-py中实现了所有这些附加功能。

1. BPE/字片。我们可以使用一个库先把数据预处理成子字单元。参见Rico Sennrich的subword- nmt实现。这些模型会将训练数据转化成这样的样子。

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

2. 共享嵌入。当使用共享词汇的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. 光束搜索:这有点太复杂了，这里就不多说了。参见OpenNMT- py中的pytorch实现。

4. 模型平均化:本文对最后的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:]))

 

结果

在WMT 2014英译德任务上，大变压器模型（Transformer（big）在表2中）比之前报道的最佳模型（包括合集）的表现要好2.0 BLEU以上，建立了新的最先进的BLEU得分28.4。该模型的配置列在表3的底行。训练在8个P100 GPU上花了3.5天。即使我们的基础模型也超越了所有之前发布的模型和合集，而训练成本只是任何竞争模型的一小部分。

在WMT 2014英译法任务上，我们的大模型实现了41.0的BLEU得分，超越了之前发布的所有单体模型，而训练成本不到之前最先进模型的1/4。为英译法训练的Transformer（大）模型使用的辍学率Pdrop=0.1，而不是0.3。

 

我们在这里写的代码是基础模型的一个版本。这里有这个系统的完全训练版本（示例模型）。

通过上一节中的附加扩展，OpenNMT-py复制在EN-DE 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)

 

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

 

注意力的可视化

即使是用greedy解码，翻译出来的效果也很不错。我们可以进一步将其可视化，看看在注意力的每一层发生了什么。

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

Encoder Self Layer 2

Encoder Src Layer 4

Encoder Self Layer 6

Decoder Src Layer 6

结论

笔者注：希望这段代码对以后的研究有用。如果你有任何问题，请联系我们。如果你觉得这段代码有帮助，也可以看看我们其他的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} }