5.4　使用PyTorch实现Transformer

Transformer的原理在前面已经分析得很详细了，下面将使用PyTorch 1.0+完整实现Transformer的整个架构，并用简单实例进行验证。代码参考哈佛大学OpenNMT团队针对Transformer实现的代码，该代码是用PyTorch 0.3.0实现的。

5.4.1　Transformer背景介绍

目前的主流神经序列转换模型大都基于Encoder-Decoder模型。所谓序列转换模型就是把一个输入序列转换成另外一个输出序列，它们的长度很可能是不同的。比如基于神经网络的机器翻译，输入是中文句子，输出是英语句子，这就是一个序列转换模型。类似的，文本摘要、对话等问题都可以看作序列转换问题。虽然这里主要关注机器翻译，但是任何输入是一个序列，输出是另外一个序列的问题都可以使用Encoder-Decoder模型。

Encoder把输入序列（x 1，…，x n）映射（或编码）成一个连续的序列z=（z 1，…，z n）。而Decoder根据Z来解码得到输出序列y 1，…，y m。Decoder是自回归的（Auto-Regressive），会把前一个时刻的输出作为当前时刻的输入。Encoder-Decoder模型架构的代码分析如下。

5.4.2　构建Encoder-Decoder模型

1. 导入需要的库

导入Torch及Python中用于可视化的一些库，如matplotlib、seaborn等。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time

import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")

%matplotlib inline

2. 创建Encoder-Decoder类

定义EncoderDecoder类，该类继承nn.Module。

class EncoderDecoder(nn.Module):
"""
    这是一个标准的 Encoder-Decoder模型
"""
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        # encoder和decoder都是构造的时候传入的，这样会非常灵活
        self.encoder = encoder
        self.decoder = decoder
        # 输入和输出的embedding
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        #Decoder部分最后的Linear+softmax
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask):
        #接收并处理屏蔽src和目标序列
        #首先调用encode方法对输入进行编码，然后调用decode方法进行解码
        return self.decode(self.encode(src, src_mask), src_mask,tgt, tgt_mask)

    def encode(self, src, src_mask):
        #传入参数包括src的embedding和src_mask
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        #传入的参数包括目标的embedding、Encoder的输出memory，及两种掩码
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

从上述代码可以看出，Encoder和Decoder都使用了掩码（如src_mask、tgt_mask）。

3. 创建Generator类

Decoder的输出将传入一个全连接层，而后经过log_softmax函数的作用，成为概率值。

class Generator(nn.Module):
    """定义标准的一个全连接（linear）+ softmax
    根据Decoder的隐状态输出一个词
    d_model是Decoder输出的大小，vocab是词典大小"""
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)
    #全连接再加上一个softmax
    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)

5.4.3　构建Encoder

Encoder由N个相同结构的EncoderLayer堆积（stack）而成，而每个Encoder层又有两个子层。第一个是一种多头部的自注意力机制，第二个比较简单，是按位置全连接的前馈网络。其间还有LayerNorm及残差连接等。

1. 定义复制模块的函数

首先定义clones函数，用于克隆相同的Encoder层。

def clones(module, N):
    "克隆N个完全相同的子层，使用了copy.deepcopy"
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

这里使用了nn.ModuleList。ModuleList就像一个普通的Python的List，我们可以使用下标来访问它，好处是传入的ModuleList的所有Module都会注册到PyTorch里，这样Optimizer就能找到其中的参数，从而用梯度下降进行更新。但是nn.ModuleList并不是Module（的子类），因此它没有forward等方法，通常会被放到某个Module里。接下来创建Encoder。

2. 创建Encoder

创建Encoder，代码如下：

class Encoder(nn.Module):
    "Encoder是N个EncoderLayer的堆积而成"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        #layer是一个SubLayer，我们clone N个
        self.layers = clones(layer, N)
        #再加一个LayerNorm层
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "把输入(x,mask)被逐层处理"
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x) #N个EncoderLayer处理完成之后还需要一个LayerNorm

由代码可知，Encoder就是N个子层的栈，最后加上一个LayerNorm。下面我们来构建LayerNorm。

3. 构建LayerNorm

构建LayerNorm模型，代码如下：

class LayerNorm(nn.Module):
    "构建一个LayerNorm模型"
    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

具体处理过程如下：

x ->x+self-attention(x) ->layernorm(x+self-attention(x)) => y
y-> dense(y) ->y+dense(y) ->layernorm(y+dense(y)) => z(输入下一层)

这里把Layernorm层放到前面了，即处理过程如下：

x ->layernorm(x) -> self-attention(layernorm(x)) -> x + self-attention(layernorm(x)) => y
y ->layernorm(y) -> dense(layernorm(y)) ->y+dense(layernorm(y)) =>z(输入下一层)

PyTorch中各层权重的数据类型是nn.Parameter，而不是Tensor。故需对初始化后的参数（Tensor型）进行类型转换。每个Encoder层又有两个子层，每个子层通过残差把每层的输入转换为新的输出。不管是自注意力层还是全连接层，都先是LayerNorm，然后是Self-Attention/Dense，接着是Dropout层，最后是残差连接层。接下来把它们封装成SublayerConnection。

4. 构建SublayerConnection

构建SublayerConnection模型，代码如下：

class SublayerConnection(nn.Module):
    """
    LayerNorm + sublayer(Self-Attenion/Dense) + dropout + 残差连接
    为了简单，把LayerNorm放到了前面，这和原始论文稍有不同，原始论文LayerNorm在最后
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        #将残差连接应用于具有相同大小的任何子层
        return x + self.dropout(sublayer(self.norm(x)))

5. 构建EncoderLayer

有了以上这些代码，构建EncoderLayer就很简单了，代码如下。

class EncoderLayer(nn.Module):
    "Encoder由self_attn 和 feed_forward构成"
    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):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

为了复用，这里把self_attn层和feed_forward层作为参数传入，只构造两个子层。forward调用sublayer[0]（这是sublayer对象），最终会调到它的forward方法，而这个方法需要两个参数，一个是输入Tensor，一个是对象或函数。在Python中，类似的实例可以像函数一样，可以被调用。而self_attn函数需要4个参数，即query的输入、key的输入、value的输入和mask，因此，使用lambda的技巧是把它变成一个参数x的函数（mask可以看成已知的数）。

5.4.4　构建Decoder

前文提到过，Decoder也是N个Decoder层的堆叠，参数layer是Decoder层数，它也是一个调用对象，最终会调用DecoderLayer.forward方法，这个方法需要4个参数，输入x、Encoder层的输出memory、输入Encoder的mask（src_mask）和输入Decoder的mask（tgt_mask）。所有这里的Decoder的forward方法也需要这4个参数。

1. 创建Decoder类

创建Decoder类，代码如下：

class Decoder(nn.Module):
    "构建N个完全相同的Decoder层"
    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)

2. 创建DecoderLayer类

创建DecoderLayer类，具体实现如下：

class DecoderLayer(nn.Module):
    "Decoder包括self_attn、src_attn和feed_forward层"
    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):
        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)

由代码可知，DecoderLayer类比EncoderLayer类多一个src-attn层，这是解码器关注编码器的输出（memory）。src-attn和self-attn的实现是一样的，只是query、key和value的输入不同。普通注意力（src-attn）的query来自下层的输入（即self-attn的输出），key和value来自Encoder最后一层的输出memory；而Self-Attention的query、key和value都是来自下层的输入。

3. 定义subsequent_mask函数

Decoder和Encoder有一个关键的不同：Decoder在解码第t个时刻的时候只能使用1…t时刻的输入，而不能使用t+1时刻及其之后的输入。因此我们需要一个函数来生成一个Mask矩阵，代码如下：

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

我们看一下这个函数生成的一个简单样例，假设语句长度为6，其运行结果如图5-30所示。

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

运行结果：

ensor([[ True, False, False, False, False, False],
        [ True,  True, False, False, False, False],
        [ True,  True,  True, False, False, False],
        [ True,  True,  True,  True, False, False],
        [ True,  True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True,  True]])

由以上结果可知，输出的是一个方阵，对角线和左下方都是True。其中，第一行只有第一列是True，它的意思是时刻1只能关注输入1；第三行说明时刻3可以关注{1，2，3}而不能关注{4，5，6}的输入，因为在真正解码的时候，这属于预测的信息。知道了这个函数的用途之后，上面的代码就很容易理解了。

5.4.5　构建MultiHeadedAttention

下面讲解如何构建MultiHeadedAttention。

1. 定义Attention

Attention（包括Self-Attention和普通的Attention）可以看作一个函数，它的输入是query、key、value和mask，输出是一个张量（Tensor）。其中输出是value的加权平均，而权重由query和key计算得出。具体的计算如5.18式所示：

def attention(query, key, value, mask=None, dropout=None):
    "计算 '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

在上面的代码实现中，和公式里的稍微不同，这里的Q和K都是4维张量，包括

批量和头维度。torch.matmul会对query和key的最后两维进行矩阵乘法，这样效率更高，如果我们用标准的矩阵（2维张量）乘法来实现，则需要遍历batch维度和head维度。

用一个具体例子跟踪一些不同张量的形状变化，然后对照公式会更容易理解。比如Q是（30, 8, 33, 64)，其中30是批量个数，8是头个数，33是序列长度，64是每个时刻的特征数。K和Q的shape必须相同，而V可以不同，但是这里的shape也是相同的。接下来是scores.masked_fill(mask == 0, -1e9)，用于把mask是0的得分变成一个很小的数，这样后面经过Softmax函数计算之后的概率就会接近零。self_attention中的mask主要是padding格式，与Decoder中的mask不同。

接下来对score进行Softmax函数计算，把得分变成概率p_attn，如果有dropout，则对p_attn进行Dropout（原论文没有dropout）。最后把p_attn和value相乘。p_attn是（30, 8, 33, 33），value是（30, 8, 33, 64），我们只看后两维，（33x33）x（33x64），最终得到33x64。

2. 定义MultiHeadedAttention

前面可视化部分介绍了如何将输入变成Q、K和V，即对于每一个head，都使用三个矩阵wQ、wK、wV把输入转换成Q、K和V。然后分别用每一个head进行自注意力计算，把N个head的输出拼接起来，与矩阵wO相乘。MultiHeadedAttention的计算公式如下：

MultiHead（Q，K，V）=concat（head 1，head 2，…，head h)WO　（5.19）

这里映射是参数矩阵，其中。

假设head个数为8，即h=8，。详细的计算过程如下：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "传入head个数及model的维度."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # 这里假设d_v=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:
            # 相同的mask适应所有的head.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) 首先使用线性变换，然后把d_model分配给h个Head，每个head为d_k=d_model/h
        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) 使用attention函数计算缩放点积注意力
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) 实现Multi-head attention，用view函数把8个head的64维向量拼接成一个512的向量。
        #然后再使用一个线性变换(512,521)，shape不变
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

其中zip(self.linears, (query, key, value))是把（self.linears[0], self.linears[1], self.linears[2]）和（query, key, value）放到一起遍历。这里我们只看一个self.linears[0] (query)。根据构造函数的定义，self.linears[0]是一个（512, 512）的矩阵，而query是（batch, time, 512），相乘之后得到的新query还是512（d_model）维的向量，然后用view函数把它变成（batch, time, 8, 64）。接着转换成（batch, 8, time, 64），这是attention函数要求的shape。分别对应8个head，每个head的query都是64维。

key和value的运算完全相同，因此我们也分别得到8个head的64维的key和64维的value。接下来调用attention函数，得到x和self.attn。其中x的shape是（batch, 8, time, 64），而attn是（batch, 8, time, time），x.transpose（1, 2）把x变成（batch, time, 8, 64），然后用view函数把它变成（batch, time, 512），也就是把最后8个64维的向量拼接成512的向量。最后使用self.linears[-1]对x进行线性变换，self.linears[-1]是（512, 512）的，因此最终的输出还是（batch, time, 512）。我们最初构造了4个（512, 512）的矩阵，前3个矩阵用于对query、key和value进行变换，而最后一个矩阵对8个head拼接后的向量再做一次变换。

MultiHeadedAttention的应用主要有以下几种：

• Encoder的自注意力层。query、key和value都是相同的值，来自下层的输入，Mask都是1（当然padding的不算）。
• Decoder的自注意力层。query、key和value都是相同的值，来自下层的输入，但是Mask使得它不能访问未来的输入。
• Encoder-Decoder的普通注意力层。query来自下层的输入，key和value相同，是Encoder最后一层的输出，而Mask都是1。

5.4.6　构建前馈网络层

除了需要注意子层之外，还需要注意编码器和解码器中的每个层都包含一个完全连接的前馈网络，该网络应用于每层的对应位置。这包括两个线性转换，中间有一个ReLU激活函数，具体公式为：

FFN（x）=max（0，xW 1+b 1）W 2+b 2　（5.21）

全连接层的输入和输出的d_model都是512，中间隐单元的个数d_ff为2048，具体代码如下：

class PositionwiseFeedForward(nn.Module):
    "实现FFN函数"
    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))))

5.4.7　预处理输入数据

输入的词序列都是ID序列，所以我们对其进行预处理。源语言和目标语言都需要嵌入，此外我们还需要一个线性变换把隐变量变成输出概率，这可以通过前面的Generator类来实现。Transformer模型的注意力机制并没有包含位置信息，也就是说，即使一句话中词语在不同的位置，但其在Transformer中是没有区别的，这显然不符合实际。因此，在Transformer中引入位置信息对于CNN、RNN等模型有非常重要的作用。笔者添加位置编码的方法是：构造一个跟输入嵌入维度一样的矩阵，然后跟输入嵌入相加得到多头注意力的输入。对输入的处理过程如图5-31所示。

1. 将输入数据转换为Embedding

先把输入数据转换为Embedding，具体代码如下：

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)

2. 添加位置编码

位置编码的公式如下：

PE（pos，2i）=sin（pos/10002i/d model）　（5.22）

PE（pos，2i+1）=cos（pos/10002i/d model））　（5.23）

具体实现代码如下：

class PositionalEncoding(nn.Module):
    "实现PE函数"
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 计算位置编码
        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 + self.pe[:, :x.size(1)].clone().detach()
        return self.dropout(x)

注意这里调用了Module.register_buffer函数。该函数的作用是创建一个buffer缓冲区，比如这里把pe保存下来。register_buffer通常用于保存一些模型参数之外的值，比如在BatchNorm中，我们需要保存running_mean（平均位移），它不是模型的参数（不是通过迭代学习的参数），但是模型会修改它，而且在预测的时候也要用到它。这里也是类似的，pe是一个提前计算好的常量，在forward函数会经常用到。在构造函数里并没有把pe保存到self参数里，但是在forward函数调用时却可以直接使用它（self.pe）。如果保存（序列化）模型到磁盘，则PyTorch框架将把缓冲区里的数据保存到磁盘，这样反序列化的时候才能恢复它们。

3. 可视化位置编码

假设输入是长度为10的ID序列，如果输入Embedding之后是（10, 512），那么位置编码的输出也是（10, 512）。对应pos就是位置（0～9），512维的偶数维使用sin函数，而奇数维使用cos函数。这种位置编码的好处是：PEpos+k可以表示成PEpos的线性函数，这样前馈网络就能很容易地学习到相对位置的关系。图5-32就是这样一个示例，向量大小d_model=20，这里画出来的是第4、5、6和7维（下标从零开始）维的图像，最大的位置是100。可以看到它们都是正弦（余弦）函数，而且周期越来越长。

##语句长度为100，这里假设d_model=20
plt.figure(figsize=(15, 5))
pe = PositionalEncoding(20, 0)
y = pe.forward(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]])

4. 简单示例

生成位置编码的简单示例如下：

d_model, dropout, max_len=512,0,5000
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)
print(pe.shape)
pe = pe.unsqueeze(0)
print(pe.shape)

5.4.8　构建完整网络

把前面创建的各网络层整合成一个完整网络，实现代码如下。

def make_model(src_vocab,tgt_vocab,N=6,d_model=512, d_ff=2048, h=8, dropout=0.1):
    "构建模型"
    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))

    # 随机初始化参数，非常重要，这里用Glorot/fan_avg
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

首先把copy.deepcopy命名为c，这样可以使下面的代码简洁一点。然后构造Multi-HeadedAttention、PositionwiseFeedForward和PositionalEncoding对象。接着构造Encoder-Decoder对象，该对象需要5个参数：Encoder、Decoder、src-embed、tgt-embed和Generator。

我们先看后面三个简单的参数：Generator可以直接构造，它的作用是把模型的隐单元变成输出词的概率；而src-embed是一个嵌入层和一个位置编码层c；tgt-embed与之类似。

最后我们来看Decoder（Encoder和Decoder类似）。Decoder由N个DecoderLayer组成，而DecoderLayer需要传入self-attn、src-attn、连接层和Dropout层。因为所有的Multi-HeadedAttention都是一样的，因此我们直接深度复制（deepcopy）就行；同理，所有的PositionwiseFeedForward也是一样的，我们可以深度复制而不用再构造一个。

实例化这个类，可以看到模型包含哪些组件，如下所示。

#测试一个简单模型，输入、目标语句长度分别为10，Encoder、Decoder各2层。
tmp_model = make_model(10, 10, 2)
tmp_model

5.4.9　训练模型

1）训练前，先介绍一个便于批次训练的Batch类，具体定义如下：

class Batch:
    "在训练期间，构建带有掩码的批量数据"
    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 & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.
            data).clone().detach()
        return tgt_mask

Batch构造函数的输入是src、trg和pad，其中trg的默认值为None。刚预测的时候是没有tgt的。为便于理解，这里用一个例子来说明。假设这是训练阶段的一个Batch，src的维度为（40，20），其中40是批量大小，而20是最长的句子长度，其他不够长的都填充成20。而trg的维度为（40，25），表示翻译后的最长句子是25个词，不足的也已填充对齐。

那么，src_mask要如何实现呢？注意，表达式（src != pad）是指把src中大于0的时刻置为1，以表示它已在关注的范围内。然后unsqueeze(-2)把src_mask变成（40/batch, 1, 20/time）。它的用法可参考前面的attention函数。

对于训练（以教师强迫模式训练，教师强迫模型（Teacher Forcing）是一种用来训练有关序列模型的方法，该方法以上一时刻的输出作为下一时刻的输入）来说，Decoder有一个输入和一个输出。比如对于句子“<sos> it is a good day <eos>”，输入会变成”<sos> it is a good day”，而输出为”it is a good day<eos>”。对应到代码里，self.trg就是输入，而self.trg_y就是输出。接着对输入self.trg进行掩码操作，使得自注意力不能访问未来的输入。这是通过make_std_mask函数实现的，这个函数会调用我们之前详细介绍过的subsequent_mask函数，最终得到的trg_mask的shape是（40/batch，24，24），表示24个时刻的掩码矩阵。该矩阵在前面已经介绍过，这里不再赘述。

注意，src_mask的shape是（batch，1，time），而trg_mask是（batch，time，time）。这是因为src_mask的每一个时刻都能关注所有时刻（填充的时间除外），一次只需要一个向量就行了，而trg_mask需要一个矩阵。

2）构建训练迭代函数，具体代码如下：

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

它遍历epoch（epoch指整个训练集被训练的次数）次数据，然后调用forward函数，接着用loss_compute函数计算梯度，更新参数并且返回loss。

3）对数据进行批量处理，代码如下：

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)

4）定义优化器，实现代码如下：

class NoamOpt:
    "包括优化学习率的优化器"
    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):
        "更新参数及学习率"
        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))

5）下面来看学习率在不同场景下的变化情况，以下代码是一个简单示例，其运行结果如图5-33所示。

# 超参数学习率的3个场景.
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"])

6）正则化。对标签做正则化平滑处理，可以提高模型的准确性和BLEU分数[1]，代码如下。

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.criterion = nn.KLDivLoss(reduction='sum')
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None

    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, true_dist.clone().detach())

这里先定义实现标签平滑处理的类，该类使用KL散度损失（nn.KLDivLoss）实现标签平滑。创建一个分布，该分布具有对正确单词的置信度，而其余平滑质量分布在整个词汇表中。

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(predict.log().clone().detach(), torch.LongTensor([2, 1, 0])
    .clone().detach())
plt.imshow(crit.true_dist)

运行结果如图5-34所示。

通过图5-34可以看到如何基于置信度将质量分配给单词。

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],])
    return crit(predict.log().clone().detach(),torch.LongTensor([1]).clone().
        detach()).item()
plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])

从图5-35可以看出，如果标签平滑化对于给定的选择非常有信心，那么标签平滑处理实际上已开始对模型造成不利影响。

5.4.10　实现一个简单实例

下面以一个简单实例帮助读者加深理解。

1）生成合成数据。

def data_gen(V, batch, nbatches):
    "Generate random data for a src-tgt copy task."
    for i in range(nbatches):
        #把torch.Embedding的输入类型改为LongTensor。
        data = torch.from_numpy(np.random.randint(1, V, size=(batch, 10))).long()
        data[:, 0] = 1

src = data.clone().detach()
tgt = data.clone().detach()
        yield Batch(src, tgt, 0)

2）定义损失函数。

class SimpleLossCompute:
    "一个简单的计算损失的函数."
    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.item() * norm

3）训练简单任务。

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: 1.249925 Tokens per Sec: 1429.082397
Epoch Step: 1 Loss: 0.460243 Tokens per Sec: 1860.120972
tensor(0.3935)
Epoch Step: 1 Loss: 0.966166 Tokens per Sec: 1433.039185
Epoch Step: 1 Loss: 0.198598 Tokens per Sec: 1917.530884
tensor(0.1874)

4）为了简单起见，此代码使用贪婪解码来预测翻译结果。

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):
        #add torch.tensor 202005
        out = model.decode(memory, src_mask,ys, subsequent_mask(torch.tensor(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 = torch.LongTensor([[1,2,3,4,5,6,7,8,9,10]])
src_mask = torch.ones(1, 1, 10)
print(greedy_decode(model, src, src_mask, max_len=10, start_symbol=1))

运行结果如下：

tensor([[ 1,  2,  3,  4,  4,  6,  7,  8,  9, 10]])

---

[1]BLEU（Bilingual Evaluation Understudy，双倍评估替换）分数，常用于说明候选文本与参考文本的相似程度，值越接近1，说明两个文本越相似。