transformer算法解析

本文参考：

详解Transformer （Attention Is All You Need） - 知乎

Transformer 代码完全解读！_AI科技大本营的博客-CSDN博客

Transformer学习笔记一：Positional Encoding（位置编码） - 知乎

1、Transformer概述

Transformer抛弃了CNN和RNN，整个网络结构完全由self-Attention和Feed Forward Neural Network机制组成。同时它有很好的并行性。

Attention机制的公式如下：

对于Self-Attention则Q=K=V

Feed Forward Neural Network（FFN）的公式如下：

全连接有两层，第一层的激活函数是ReLU，第二层是一个线性激活函数。

Transformer的整体架构如下：

训练阶段过程概述（针对seq2seq场景）：

（1）输入input，首先经过input embedding和positional embedding得到了带位置信息的embedding信息

（2）encoder阶段，在每个encoder子层中包含多头注意力机制（Mutli-Head Attention）和全连接（Feed Forward），而每个Multi-Head Attention和Feed Forward又都是残差结构，最后输出包含inputs特征的memory信息

（3）输入output，output的首位是起始符，同样经过input embedding和positional embedding得到了带位置信息的embedding信息

（4）decoder阶段，在每个decoder子层中包含两个Multi-Head Attention和全连接（Feed Forward），其中第一个Multi-Head Attention得到了output自身的特征，第二个Multi-Head Attention融合了encoder输出的memory后优化了output的特征信息

（5）decoder输出的特征经过Linear和softmax之后得到了在所有输出字符的预测分布信息，与输出字符的实际分布进行对比计算loss

推理预测阶段过程概述（针对seq2seq场景）：

（1）输入inputs通过embedding模块和encoder模块得到输入特征memory

（2）将输入decoder模块结合memory预测出第一个字符，然后将和第一个字符再输入decoder模块预测出第二个字符，直到预测出结束字符

2、Transformer详述

2.1、位置编码

2.1.1 编码公式

使用sin和cos函数进行固定位置的编码，编码公式如下：

2.1.2 位置编码的演变过程

（1）用整型值标记位置

用1、2、3…标记位置，会带来以下问题：

• 模型可能会遇到比训练序列更长的序列，不利于模型的泛化
• 模型的位置值不断变大是无界的

（2）用[0,1]范围标记位置

用0表示第一个字符，用1表示最后一个字符。但是会带来以下问题：

当序列长度不同时，字符间的相对距离是不一样的。

比如序列长度为3时某两个字符的相对距离是0.5，而序列长度为4时某两个字符的相对距离是0.33。

（3）用二进制向量标记位置

用一个和input embedding维度一样的向量来表示位置，形式如下：

但是也存在问题：这样编码出来的位置向量是离散的而非连续的。

比如有4个位置需要编码，我们得到了[0,0], [0,1], [1,0], [1,1]

第2个位置[0,1]和第3个位置[1,0]到第1个位置[0,0]的相对距离是一样的。

（4）用周期函数sin来表示位置

把位置向量当中的每一个元素都用一个sin函数来表示，则第t个字符的位置可以表示为：

为了避免出现位置的冲突，可以把所有频率都设置成同一个非常小的值。公式修改为：

以上公式的特点为：

• 每个字符的向量唯一（每个sin函数的频率足够小）
• 位置向量有界且连续，模型在处理位置向量时更容易泛化，即更好处理长度和训练数据分布不一致的序列。

存在的问题：不同的位置不能通过线性变化转换得到。

（5）用sin和cos交替来表示位置

理想位置表示公式：

通过旋转的线性变化可以得到如下：

融合周期函数sin和cos得到最终版本：

2.2、多头注意力机制

2.2.1、注意力机制

注意力机制启发：人们在观察事物时不能同时仔细观察眼前的一切，只能聚焦到某一个局部。对于神经网络的应用中，就是通过一种方式得到权重再更新特征向量，使得重要的特征更突出，不重要的特征被打压。公式如下：

在seq2seq中，首先计算每个时间步的系数，是通过当前时间步的query和其他时间步对应的key做内积得到，最后用该系数乘以每个时间步的特征向量value得到注意力计算结果。

2.2.2、多头注意力机制

多头注意力机制，即多个注意力模块组合在一起。

让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义拥有来自更多元表达。

举例说明：bank是银行的意思，如果只有一个注意力模块，那么它大概率会学习去关注类似money、loan贷款这样的词。如果我们使用多个多头机制，那么不同的头就会去关注不同的语义，比如bank还有一种含义是河岸，那么可能有一个头就会去关注类似river这样的词汇，这时多头注意力的价值就体现出来了。

举例说明如下：

（1）输入X：batch_size=32, seq_len=15(字符串的长度), feature_size=512

（2）分别通过Wq、Wk、Wv三个权重矩阵得到query、key和value

（3）将特征拆分为多头，512=8*64，将feature_size=512拆分为8头，每一头对应64位的特征，变为(32, 8, 15, 64)

（4）并行计算score分数，batch_size * head_num=32*8=256同时进行score分数计算，计算的是seq_len中两两之间的score分数。

（5）计算attention值。(32, 8, 15, 15) 与(32, 8, 15, 64)进行矩阵相乘，得到(32, 8, 15, 64)。

（6）通过view操作再将多头特征重新组装为一个特征，得到(32, 15, 512)

（7）通过512 -> 512的权重矩阵得到输出X为(32, 15, 512)。

2.2.3 Decoder中的注意力机制

（1）假设此时输入Y为(32, 14, 512)。区别于Encoder中seq_len=15，此时Decoder中的输入Y的seq_len=14。经过第1次多头注意力机制后得到(32, 14, 512)，整个过程与Encoder中的多头流程完全一致。

（2）Decoder阶段的输入变query(32, 14, 512)拆分为多头后为(32, 8, 14, 64)，encoder阶段的memory拆分为多头后为key(32, 8, 15, 64)和value(32, 8, 15, 64)。通过query和key的转置相乘后得到score(32, 8, 14, 15)，然后score再乘以value得到(32, 8, 14, 64)，经过后续操作后最终转化为(32, 14, 512)。

2.3、掩码mask机制

在transformer中掩码主要有两个作用，一个是屏蔽掉无效的padding区域，一个是屏蔽掉来自”未来”的信息。Encoder中的掩码主要是起到第一个作用，Decoder中的掩码则同时发挥着两种作用。

屏蔽掉无效的padding区域：我们训练需要组batch进行，就以机器翻译任务为例，一个batch中不同样本的输入长度很可能是不一样的，此时我们要设置一个最大句子长度，然后对空白区域进行padding填充，而填充的区域无论在Encoder还是Decoder的计算中都是没有意义的，因此需要用mask进行标识，屏蔽掉对应区域的响应。

屏蔽掉来自未来的信息：我们已经学习了attention的计算流程，它是会综合所有时间步的计算的，那么在解码的时候，就有可能获取到未来的信息，这是不行的。因此，这种情况也需要我们使用mask进行屏蔽。如下图所示：

在训练阶段预测出最后一个字“人“即可终止，所以mask是一个倒三角的结构，通过这种mask机制可以保证并行进行训练。

2.4 整体架构流程详情手稿

 3、示例代码实现

3.1 transformer.py实现transformer网络结构

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



class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        """
        类的初始化
        :param d_model: 词嵌入的维度
        :param vocab: 词表的大小
        """
        super(Embeddings, self).__init__()
        # 调用nn中的预定义层Embedding，获得一个词嵌入对象self.lut
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        """
        Embedding层的前向传播逻辑
        :param x: 输入给模型的单词文本通过词表映射后的one-hot向量
        :return: 将x传给self.lut并与根号下self.d_model相乘作为结果返回
        """
        embedds = self.lut(x)
        return embedds * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        """
        位置编码器类的初始化函数
        :param d_model: 词嵌入维度
        :param dropout: dropout比率
        :param max_len: 每个句子的最大长度
        """
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # compute the positional encodings
        # 下面代码的计算方式与公式中给出的不同但等价，这样计算是为了避免中间的数值计算结果超出float的范围
        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)  # 变量级别控制gradient
        return self.dropout(x)

# 定义一个clones函数，来更方便的将某个结构复制若干份
def clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class Encoder(nn.Module):
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        # 调用时会将编码器传进来，简单克隆N份叠加在一起，组成完整的Encoder
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class SublayerConnection(nn.Module):
    """
    实现子层连接结构的类
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        sublayer_out = sublayer(x)
        sublayer_out = self.dropout(sublayer_out)
        x_norm = x + self.norm(sublayer_out)
        return x_norm

class EncoderLayer(nn.Module):
    "EncoderLayer is made up of two sublayer: self-attn and 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 # embedding's dimentionof model 512

    def forward(self, x, mask):
        # attention sub layer
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        # feed forward sub layer
        z = self.sublayer[1](x, self.feed_forward)
        return z

def attention(query, key, value, mask=None, dropout=None):
    "Compute Scaled Dot Product Attention"

    # 首先取query的最后一维的大小，对应词嵌入维度
    d_k = query.size(-1)
    # 按照注意力公式，将query和key的转置相乘，这里面key是将最后两个维度进行转置，再除以缩放系数得到注意力得分张量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 接着判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法，将掩码张量和scores张量每个位置一一比较，如果掩码张量则对应的scores张量用-1e9这个值来替换
        scores = scores.masked_fill(mask == 0, -1e9)

    # 对scores的最后一维进行softmax操作，使用F.softmax方法，这样获得最终的注意力张量
    p_attn = F.softmax(scores, dim=-1)

    # 之后判断是否使用dropout进行随机置0
    if dropout is not None:
        p_attn = dropout(p_attn)

    # 最后根据公式将p_attn与value张量相乘获得最终的query注意力表示，同时返回注意力张量
    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        # 在类的初始化时，会传入三个参数，h代表头数，d_model代表词嵌入的维数，dropout代表进行置0比率
        super(MultiHeadedAttention, self).__init__()
        # 在函数中，首先判断h是否能被d_model整除，这是因为我们之后要给每个头分配等量的词特征，也就是embedding_dim/head个
        assert d_model % h == 0
        # 得到每个头获得的分割词向量维度d_k
        self.d_k = d_model // h
        # 传入头数h
        self.h = h

        # 创建Linear层，通过nn的Linear实例化，它的内部变换矩阵是embedding_dim * embedding_dim.
        # 为什么是四个呢？这是因为在多头注意力中，Q/K/V各需要一个，最后拼接的矩阵还需要一个，因此一共是四个
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        # self.attn为None，它代表最后得到的注意力张量，现在还没有结果所以为None
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        # 前向逻辑函数，它输入参数有四个，前三个就是注意力机制需要的Q，K，V。最后一个是注意力机制中可能需要的mask掩码张量，默认是None
        if mask is not None:
            # Same mask applied to all h heads.使用unsqueeze扩展维度，代表多头中的第n头
            mask = mask.unsqueeze(1)

        # 接着获得一个batch_size的变量，它是query尺寸的第1个数字，代表有多少条样本
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch for d_model => h * d_k
        # 首先利用zip将输入QKV与三个线性层组到一起，
        # 然后利用for循环将输入QKV分别传到线性层中
        # 接下来为每个头分割输入，这里使用view方法对线性变换的结构进行维度重塑，多加了一个维度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> Apply attention on all the projected vaectors in batch
        # 得到每个头的输入后，接下来就是将他们传入到attention中，这里直接调用我们之前实现的attention函数，同时也将mask和dropout传入其中
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 3) "concat" using a view and apply a final linear
        # 通过多头注意力计算后，我们就得到了每个头计算结果组成的4维张量，我们需要将其转换为输入的形状以便后续的计算
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)

        # 最后使用线性层列表中的最后一个线性变换得到最终的多头注意力结构的输出
        return self.linears[-1](x)

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        """
        FFN前馈全连接层，我们希望输入通过前馈全连接层后输入和输出的维度不变
        """
        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))))

class LayerNorm(nn.Module):
    # 规范层网络，类似于nn.BatchNorm
    def __init__(self, feature_size, eps=1e-6):
        """
        :param feature_size: 词嵌入的维度
        :param eps: 它是一个足够小的数，在规范化公式的分母中出现，防止分母为0，默认是1e-6
        """
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(feature_size))
        self.b_2 = nn.Parameter(torch.zeros(feature_size))
        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

def subsequent_mask(size):
    # 生成向后遮掩的掩码张量，参数size是掩码张量最后两个维度的大小，它最后两维形成一个方阵

    attn_shape = (1, size, size)
    # 然后使用np.ones方法向这个形状中添加1元素，形成上三角阵
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0


class Decoder(nn.Module):
    # 根据编码器的结果以及上一次预测的结果，输出序列的下一个结果
    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):
        """
        :param x: 目标数据的嵌入表示
        :param memory: 编码器层的输出
        :param src_mask: 源数据的掩码张量
        :param tgt_mask: 目标数据的掩码张量
        :return:
        """
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

# 使用DecoderLayer的类实现解码器层
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        """
        :param size: 词嵌入的维度大小，也代表解码器的尺寸
        :param self_attn: 多头自注意力对象，要求Q=K=V
        :param src_attn: 多头注意力对象，要求Q!=K=V
        :param feed_forward: 前馈全连接层对象
        :param dropout:
        """
        super(DecoderLayer, self).__init__()
        self.size  = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        # 根据结构图使用clones函数克隆三个子层连接对象
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        """
        :param x: 上一层的输入x
        :param memory: 编码器语音存储变量memory
        :param src_mask:
        :param tgt_mask:
        :return:
        """
        m = memory
        # 将x输入第一个子层结构，第一个子层结构的输入分别是x和self_attn函数，因为是自注意力机制，所以Q,K,V都是x，
        # 最后一个参数时目标数据掩码张量，这时要对目标数据进行遮掩，因为此时模型可能还没有生成任何目标数据
        # 比如在解码器准备生成一个字符或词汇时，我们其实已经传入了第一个字符以便计算损失，但是我们不希望在生成第一个字符时模型能利用这个信息，因此我们会将其遮掩
        # 同样生成第二个字符或词汇时，模型只能使用第一个字符或词汇信息，第二个字符以及之后的信息都不允许被模型使用
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        # 接着进入第二个子层，这个子层中常规的注意力机制，q是输入x; k,v是编码器层输出memory，同样也传入source_mask，
        # 但是进行源数据遮掩并非是抑制信息泄露，而是遮掩对结果没有意义的padding
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))

        # 最后一个子层就是前馈连接子层，经过它的处理后就可以返回结果，这就是我们的解码器结构
        return self.sublayer[2](x, self.feed_forward)

# 将线性层和softmax计算层一起实现，因为二者的共同目标是生成最后的结构，因此把类的名字叫做Generator生成器类
class Generator(nn.Module):
    def __init__(self, d_model, vocab):
        """
        :param d_model: 词嵌入维度
        :param 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)

# 使用EncoderDecoder类实现编码器-解码器结构
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        """
        :param encoder:  编码器对象
        :param decoder:  解码器对象
        :param src_embed: 源数据嵌入函数
        :param tgt_embed: 目标数据嵌入函数
        :param 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 encode(self, src, src_mask):
        src_embedds = self.src_embed(src)
        return self.encoder(src_embedds, src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        target_embedds = self.tgt_embed(tgt)
        return self.decoder(target_embedds, memory, src_mask, tgt_mask)

    def forward(self, src, tgt, src_mask, tgt_mask):
        memory = self.encode(src, src_mask)
        res = self.decode(memory, src_mask, tgt, tgt_mask)
        return res

# full model
def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
    """
    :param src_vocab:
    :param tgt_vocab:
    :param N: 编码器和解码器堆叠基础模块的个数
    :param d_model: 模型中embedding的size，512
    :param d_ff: FeedForward Layer层中embedding的size，2048
    :param h: MultiHeadAttention中多头的个数，必须被d_model整除
    :param dropout:
    :return:
    """
    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))

    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    return model

if __name__ == '__main__':
    print("\n------------------------")
    print("test subsequect_mask")
    temp_mask = subsequent_mask(4)
    print(temp_mask)

    print("\n------------------------")
    print("test build model")
    tmp_model = make_model(10, 10, 2)
    print(tmp_model)

3.2  train_demo.py实现训练和推理

import time
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from transformer import make_model, subsequent_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]  # decoder的输入（即期望输出除了最后一个token以外的部分）
            self.trg_y = trg[:, 1:]  # decoder的期望输出（trg基础上再删去句子起始符）
            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.
        padd 和 future words 均在mask中用0表示
        """
        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 data_gen(V, slen, batch, nbatches, device):
    """
    generate random data for a src-tgt copy task
    :param V: 词典数量，取值范围[0, V-1], 约定0作为特殊符号使用代表padding
    :param slen: 生成的序列数据的长度
    :param batch:  batch_size
    :param nbatches: number of batches to generate
    :param device:
    :return:
    """
    for i in range(nbatches):
        data = torch.from_numpy(np.random.randint(2, V, size=(batch, slen))).long()
        # 约定输出为输入除去序列第一个元素，即向后平移一位进行输出，同时输出数据要在第一个时间步添加一个起始符
        tgt_data = data.clone()
        tgt_data[:, 0] = 1 # 将序列的第一个时间步置为1（即约定的起始符），即可完成GT数据的构造
        src = Variable(data, requires_grad=False)
        tgt = Variable(tgt_data, requires_grad=False)
        if device == "cuda":
            src = src.cuda()
            tgt = tgt.cuda()
        yield Batch(src, tgt, 0)

def run_epoch(data_iter, model, loss_compute, device=None):
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    model = model.to(device)

    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

class LabelSmoothing(nn.Module):
    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))

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):
        """
        :param norm: loss的归一化系数，用batch中所有有效token数即可
        :return:
        """
        x = self.generator(x)
        x_ = x.contiguous().view(-1, x.size(-1))
        y_ = y.contiguous().view(-1)
        loss = self.criterion(x_, y_)
        loss /= norm
        loss.backward()
        if self.opt is not None:
            self.opt.step()
            self.opt.zero_grad()
        return loss.item() * norm

# -----------------------------------
# A Easy Example
# -----------------------------------
device = "cuda"
nrof_epochs = 40
batch_size = 32
V = 11 # 词典的数量
sequence_len = 15 # 生成的序列数据的长度
nrof_batch_train_epoch = 30 # 训练时每个epoch多少个batch
nrof_batch_valid_epoch = 10 # 验证时每个epoch多少个batch
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
model = make_model(V, V, N=2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9)
if device == "cuda":
    model.cuda()

for epoch in range(nrof_epochs):
    print("traing...{%d}" % (epoch + 1 ))
    model.train()
    data_iter = data_gen(V, sequence_len, batch_size, nrof_batch_train_epoch, device)
    loss_compute = SimpleLossCompute(model.generator, criterion, optimizer)
    train_mean_loss = run_epoch(data_iter, model, loss_compute, device)
    print("valid...")
    model.eval()
    valid_data_iter = data_gen(V, sequence_len, batch_size, nrof_batch_valid_epoch, device)
    valid_loss_compute = SimpleLossCompute(model.generator, criterion, None)
    valid_mean_loss = run_epoch(valid_data_iter, model, valid_loss_compute, device)
    print(f"valid loss: {valid_mean_loss}")

# greedy decode
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    memory = model.encode(src, src_mask)
    # ys代表目前已生成的序列，最初为仅包含一个起始符的序列，不断将预测结果追加到序列最后
    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

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