Transformer工作原理及源码解析

Tansformer架构

网络结构

下面结合具体案例进行网络结构参数，其中案例以机器翻译为例：将“我有一只猫”翻译成“I have a cat <end>”。

输入Embedding

在Transformer中对于输入的序列“我 有 一只 猫”进行Embedding操作，得到输入序列的向量$X$，其中$X$由嵌入Word Embedding和位置嵌入Positional Embedding两个操作得到的向量进行点加而成。

• ** Word Embedding**: 词嵌入方式可以采用 Word2Vec、Glove 等算法预训练得到，也可以在 Transformer 中训练得到。其中自训练的方式实际上就是采用OneHot编码+一系列全连接层得到的。
• Positional Embedding: 位置编码采用的是以下公式，其中dmodel是向量的维度，从公式中可以看出，位置向量中偶数维为sin函数，奇数维为cos函数。

$$PE(pos,2i)=sin(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$$
$$PE(pos,2i+1)=cos(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$$

注意：

• 为什么需要Position Embedding?
  Transformer结构中采用了若干个注意力机制单元，而注意力机制单元与RNN网络单元而言，未考虑序列的位置信息，因此通过Positional Embedding操作，加入位置信息。

• 为什么采用上面公式进行位置编码？
  Positional Embedding不仅仅需要将位置转变成向量，同时相对位置信息通过编码的向量也能同样得到表征，如下图所示，$pos$相对位置$k$的位置向量可以通过$pos$和$k$的位置向量的线性组合得到，即 $5$的位置向量可以由$1$和$4$或者$2$和$3$的位置向量线性组合得到。

  公式中 $PE(pos+k,2i)=PE(pos,2i)\times PE(k,2i+1)+PE(pos,2i+1)\times PE(k,2i)$ 是$2i+1$的原因是，因为$cos()$函数对应的是奇数维度。
  

编码器Encoder

在Transformer中编码器由$N$个编码器单元组成，其中默认$N=6$。对于输入Embedding之后的词向量$X$，经过6个编码单元输出新的向量。
对于每一个编码单元包含两个核心层** Multi-Head Attention** 和 Feed Forward，两个核心层输出时都会进行残差连接和归一化，即Add & Norm。

Attention机制:

注意力机制：定义$Q$、$K$、$V$三个序列向量，计算注意力向量$V^{\prime }$。记序列中第$i$个词对应的$Q_i$、$K_i$、$V_i$，计算对应的注意力向量$V_i^{\prime }$，其中每一个词的计算是可以并行的。

---

以计算注意力向量$V_i^{\prime }$为例，具体步骤如下：

• Step1：对于$Q_i$与$K$计算关联度$A_i$，通常使用点积操作，例如：$A_i=(a_1,a_2,...,a_n)=(Q_i\times K_1,Q_i\times K_2,...Q_i\times K_n)$

• Step2: 对$A_i$进行缩放操作，如下，其中$d$为$Q,K,V$的维度：
  $$A_i=\frac{A_i}{\sqrt{d}}$$
  注意：缩放的目的是为了让$a_1,a_2,...,a_n$分别差异缩小一点，进而使得step3得到的概率差异不会太大。

• Step3：对$A_i$使用SoftMax进行归一化，如下：$$A_i=Softmax(A_i)$$

• Step4: 使用计算得到的权重，去组合得到新的$V_i$，其中计算方法如下:$$V_i=A_i\times V=a_1\times v_1+a_2\times v_2+...+a_n\times v_n$$

注意，新得到的$V_i$融合了$V$中所有向量的重要信息。

---

注意力机制结构图

注意力机制阶段分析结构图  注意力机制公式可视化图

• Q、K、V理解：上述提到了具体计算过程，但是是直接假设Q、K、V是存在的，那具体Q、K、V是如何来的呢，这个就衍生出来不同的注意力机制模型：
  • 自注意机制：当Q、K、V来源一致都是输入向量$X$时，就是自注意力机制

- 交叉注意力机制：Q 和 V 不同源，但是 K 和 V 同源
- cyd注意力机制：Q 和 V 同源，Q 和 K 不同源
- XXX注意力机制：Q 必须为 1，K 和 V 不同源（这个仅仅是假设）

Multi-Head Attention机制

通过$Q、K、V$设计的查询能够使得新的向量相比于原始的输入$X$有更好的表征，因此，设计多个$Q、K、V$对同一个输入$X$做自注意力机制，然后综合所有的多头信息，融合得到新的输出$X^{\prime }$。
具体结构如图：

具体操作步骤如下：

• 在transformer结构中，多头$h=8$，因此对于输入$X$经过8个自注意机制之后，得到8个输出结果，记为$Z_1,Z_2,Z3,...,Z_8$
  

• 接着对$Z_1$~$Z_8$进行拼接，然后经过Liner层(实际就是$W\times X+b$)，得到最终输出$X^{\prime }$
  

Add&Norm结构

Add&Norm结果本质上就是残差和归一化操作，残差就是为了信息恒等映射，即使增加网络深度，也能获取到前面的信息。归一化的目录是让输入均值为0，方差为1，这样可以加上收敛。具体公式为：$$LayerNorm(X+F(X))$$
在Transformer结构中，$F(X)$就两种，一种是$FeedForward(X)$，另一种就是$MultiHeadAttention(X)$

FeedForward结构

Feed Forward 层比较简单，是一个两层的全连接层，第一层的激活函数为 Relu，第二层不使用激活函数，对应的公式如下：

编码器编码流程

上面说个Transformer结果中有6个同样结构的编码器，其中单个编码器编码流程如图：

• step1:对于输入$X=(x_1,x_2)$，进行Positioinal Embedding编码操作，使得新的$x_1,x_2$包含位置信息。
• step2:对step1得到的$X$使用多头注意力机制，学习句法特征和语义特征，得到新的向量表征$Z=(z_1,z_2)$
• step3:使用Add和Norm操作，即$LayerNorm(X+Z)$
• step4:使用FeedForward操作对step3输入的$Z=(z_1,z_2)$进行激活
• step5:使用Add和Norm操作对step4的输出进行残差和归一化操作。

###解码器Encoder
与编码器同理，在Transformer中解码器由$N$个解码器单元组成，其中默认$N=6$。对于编码器输出的高维语义向量$C$，会分别输入给6个解码单元，经过一系列解码处理之后输出高维语义向量$Z$。
对于每一个编码单元包含三个核心层** Mask Multi-Head Attention**、Multi-Head Attention 和 Feed Forward，三个核心层输出时都会进行残差连接和归一化，即Add & Norm。

Mask Self Attention机制

与Self Attention机制及其类似，仅仅在自注意机制计算权重$a$时，进行Mask操作，即当前词的$Q$只能与当前词之前的一系列对应的$K$计算关联度。因此，对应的关联度矩阵如下：

即，I第一次注意力权重计算时，只有I自己；当hava计算注意力权重时，可以与I和hava一起计算，依次类推。这个原理本质上就是解码过程。在训练时，整体最终生成的标签我们是知道的，但是实际上推理阶段，最终生成的内容是依次逐步生成的，当生成第一个词的时候，还没有后面的词，因此，只能与之前的词计算关联概率，因此引入Mask操作，这样就可以保证训练和推理是一样的。

Mask Multi Head Attention机制

这个在此不多讲，就是多头注意力机制中没有自注意模块采用掩码自注意力机制替换而已。

Multi Head Attention机制

该结构原理在编码模块已经详细阐述，在解码阶段，该结构的Q、K、V与编码器的不一样，这里K和V同源，都来自于编码器输出的编码信息矩阵$C$，而Q则来自于Mask Multi Head Attention的输出，这儿很好理解，每一次解码生成下一次词时，帮之前生成的词的向量当做Q，编码信息矩阵$C$中查询重要的信息，去生成下一个词。

其他

在解码器中的 Feed Forward和Add&Norm等结构和编码器中原理一致，不在阐述。

解码器解码流程

• step1:将之前解码输出的内容当做Mask Multi Head Attention单元的输入记录为$O$，对于输入$O=(o_1,o_2)$，进行Positioinal Embedding编码操作，使得新的$o_1,o_2$包含位置信息。
• step2:对step1得到的$O$使用多头掩码注意力机制，学习句法特征和语义特征，得到新的向量表征$M=(m_1,m_2)$。
• step3:使用Add和Norm操作，即$LayerNorm(Y+M)$
• step4:将编码器输出的信息矩阵$C$，与$W_K$和$W_V$相乘，得到$K$和$V$，把step3的输出乘上$W_Q$得到$Q$，然后输入到Multi-Head Attention单元，得到新的向量表征$Z$。
• step5:使用Add和Norm操作对step4的输出进行残差和归一化操作。
• step6:使用FeedForward操作对step5输入的$Z=(z_1,z_2)$进行激活
• step7:使用Add和Norm操作对step6的输出进行残差和归一化操作。

输出

对最后一个解码器的输出的高维度向量进行Liner操作，即$W\times X+b$，这样可以转维度，同时融合语义，最后通过Softmax进行概率计算，选取概率最大的词作为输出。

注意：

• Transformer是支持并行的，这儿的并行是只当一个编码单元或者解码单元计算一系列词时，词与词之间是并行的。而RNN系列的模型，并行计算能力很差。RNN并行计算的问题就出在这里，因为 T 时刻的计算依赖 T-1 时刻的隐层计算结果，而 T-1 时刻的计算依赖 T-2 时刻的隐层计算结果，如此下去就形成了所谓的序列依赖关系。

• Transformer的特征抽取能力比RNN系列的模型要好。

代码实现

基于Pytorch框架一步步实现Transformer结构：

#! -*- coding: utf-8 -*-
# import os
# os.environ['http_proxy'] = 'http://127.0.0.1:7890'
# os.environ['https_proxy'] = 'http://127.0.0.1:7890'

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

class EncoderDecoder(nn.Module):

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

    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)


def clones(module : nn.Module, N : int):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class LayerNorm(nn.Module):

    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

class Encoder(nn.Module):

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        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.size = size
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return self.norm(x + self.dropout(sublayer(x)))
        # return x + self.dropout(sublayer(self.norm(x)))

class EncoderLayer(nn.Module):

    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)

class Decoder(nn.Module):

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

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

    def __init__(self, size, self_attn, memory_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.memory_attn = memory_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(self.size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        # 采用Masked Self-Attention，即只能看到当前字符之前的字符，通过自注意力机制来学习新生成的字符对应的上下文，在将这个上下文与memory进行attention
        x = self.sublayer[0](x, lambda x:self.self_attn(x, x, x, tgt_mask))
        # 基于新生成的字符的对应的上下文，与memory进行attention，在生成新的字符
        x = self.sublayer[1](x, lambda x:self.memory_attn(x, memory, memory, src_mask))
        return self.sublayer[2](x, self.feed_forward)


def subsequent_mask(size):
    attn_shape = (1, size, size)
    #生成一个上三角矩阵，k=1表示对角线上移一个位置，
    #假设size=5，那么生成的矩阵如下：
    # [[0, 1, 1, 1, 1],
    #  [0, 0, 1, 1, 1],
    #  [0, 0, 0, 1, 1],
    #  [0, 0, 0, 0, 1],
    #  [0, 0, 0, 0, 0]】
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 通过将矩阵转换为torch.Tensor，然后通过==0来生成一个mask矩阵
    # 此时mask矩阵为下三角矩阵，即对角线以下（包含对角线）的元素为1，对角线以上的元素为0
    # 假设size=5，那么生成的矩阵如下：
    # [[1, 0, 0, 0, 0],
    #  [1, 1, 0, 0, 0],
    #  [1, 1, 1, 0, 0],
    #  [1, 1, 1, 1, 0],
    #  [1, 1, 1, 1, 1]】
    return torch.from_numpy(subsequent_mask) == 0

def attention(query, key, value, mask=None, dropout=None):
    d_model = query.size(-1)
    # 计算query和key的点积，作为attention的分数，即关联度，通过除以根号d_model来缩放
    scores = torch.matmul(query, key.transpose(-2, -1)) / np.sqrt(d_model)
    if mask is not None:
        # 将mask为0的位置的分数设置为一个很大的负数，这样在softmax之后，这个位置的值就会接近于0
        scores = scores.masked_fill(mask == 0, -1e9)
    # 通过softmax来计算权重
    p_attn = F.softmax(scores, dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    # 通过权重和value的乘积来计算attention的结果
    return torch.matmul(p_attn, value), p_attn


class MultiHeadSelfAttention(nn.Module):

    def __init__(self, d_model, h, dropout=0.1):
        d_h = d_model // h
        assert d_h * h == d_model
        super(MultiHeadSelfAttention, self).__init__()

        self.h = h
        self.d_h = d_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):
        if mask is not None:
            # mask 扩张到多头的维度
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        # query, key, value维度为(batch, seq_len, num_head * d_h)
        # 为了计算多头注意力，需要将输入的query, key, value进行维度变换，将输入的query, key, value转换为(batch, seq_len, num_head, d_h)
        # 在计算维度变换之前，需要将输入的query, key, value通过线性变换，将输入的query, key, value映射到d_model维度
        query, key, value = [l(x).view(nbatches, -1, self.h, self.d_h).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))]

        # 计算多头注意力
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 将多头注意力的结果进行拼接，然后通过线性变换，将多头注意力的结果映射到d_model维度
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_h)
        return self.linears[-1](x)


class PositionwiseFeedForward(nn.Module):

    def __init__(self, d_model, d_ff, dropout=0.1):

        super(PositionwiseFeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(F.relu(self.linear1(x))))


class Embedding(nn.Module):

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

    def forward(self, x):
        # 通过可学习的Embedding将输入的序列转换为d_model维度的向量
        # 此时输出维度为：（batch，seq_len, d_model）
        return self.lut(x) * math.sqrt(self.d_model)


class PositionalEmbedding(nn.Module):

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

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        # pe[:, 0::2] = torch.sin(position / torch.pow(10000, 2 * torch.arange(0, d_model, 2) / d_model))
        # pe[:, 1::2] = torch.cos(position / torch.pow(10000, 2 * torch.arange(1, d_model, 2) / d_model))

        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 这次self.pe是提前计算的，设置的长度为max_len=50000，所以这儿需要根据输入的序列长度进行截取
        # 其次，之前说过可训练的PositionalEmbedding和固定的PositionalEmbedding效果近似，所以这儿直接使用固定的PositionalEmbedding，不需要训练，故requires_grad=False
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        # 这儿的dropout是为了防止过拟合
        return self.dropout(x)

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h = 8, dropout=0.1):
    c = copy.deepcopy

    attn = MultiHeadSelfAttention(d_model, h)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEmbedding(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(Embedding(d_model, src_vocab), c(position)),
        nn.Sequential(Embedding(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab)
    )

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


class Batch:

    def __init__(self, src, trg=None, pad=0):
        self.src = src
        # 计算mask
        self.src_mask = (src != pad).unsqueeze(-2)

        if trg is not None:
            # 生成过程中，假设生成的序列是这样的： <sos> a b c d <eos>
            # 那么输入的序列是 <sos> a b c d
            # 输出的序列是 a b c d <eos>
            # 所以，对于模型动态输入序列，需要去掉最后一个字符，同时，使用mask来标记输入序列的有效长度，即每次只能看到当前字符之前的字符
            # 对于输入序列的标签，trg_y是去掉第一个字符<sos>，即 a b c d <eos>
            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):
        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):
    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 == 0:
            print("Epoch Step: %d Loss: %f Tokens per Sec: %f" % (i, loss / batch.ntokens, tokens / (time.time() - start)))

    return total_loss / total_tokens


def run_batch(batch, model, loss_compute):
    out = model.forward(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
    loss = loss_compute(out, batch.trg_y, batch.ntokens)
    return loss, batch.ntokens


global max_src_in_batch, max_tgt_in_batch

def batch_size_fn(new, count, sofar):
    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.tgt) + 2)
    src_elements = count * max_src_in_batch
    tgt_elements = count * max_tgt_in_batch
    return max(src_elements, tgt_elements)

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._ratio = 0

    def step(self):
        self._step += 1
        lr = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = lr
        self._ratio = lr
        self.optimizer.step()

    def rate(self, step=None):
        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))


# 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"])
# plt.show()

class SmoothLabel(nn.Module):

    def __init__(self, size, padding_idx, smoothing=0.0):
        super(SmoothLabel, self).__init__()
        # KLDivLoss是KL散度损失函数
        # 输入的是log_softmax的结果
        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中生成一个新的张量，维度与X一致，
        # 其中每个类别的概率值都是self.smoothing / (self.size - 2)
        # 真实标签位置的概率值是self.confidence
        # 然后将padding_idx位置的概率值设置为0
        # 最后将target中为padding_idx的位置的概率值设置为0
        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))


# crit = SmoothLabel(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(F.log_softmax(predict)),
#          Variable(torch.LongTensor([2, 1, 0])))
# print(v)
# plt.imshow(crit.true_dist)
# plt.show()


# crit = SmoothLabel(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(F.log_softmax(predict)),
#                  Variable(torch.LongTensor([1]))).item()
#
# # print(loss(5))
# plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])
# plt.show()


def data_gen(V, batch, nbatches):
    # 生成随机的数据，每个数据的长度是10
    for i in range(nbatches):
        data = torch.from_numpy(np.random.randint(1, V, size=(batch, 10)))

        data[:, 0] = 1
        # 生成的数据中，每个数据的第一个字符是1，表示start_symbol，即<sos>
        # 这儿src和tgt一样，模拟测试数据
        src = Variable(data, requires_grad=False)
        tgt = Variable(data, requires_grad=False)
        yield Batch(src, tgt, 0)


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


def greedy_decode(model, src, src_mask, max_len, start_symbol, end_symbol=None):
    # 对原始输入序列进行编码，得到memory向量
    memory = model.encode(src, src_mask)

    # 初始化输出序列，即<sos>
    ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)

    # 逐个生成新的字符，这儿设置了最大长度，即生成的字符长度不能超过max_len，同时，生成字符<eos>应该停止
    # TODO 原作者这儿的实现是有问题的，应该是生成<eos>也停止
    for i in range(max_len - 1):
        # 对于当前生成的字符序列，进行解码，得到新生成的字符
        # 在实际预测过程中，解码理论将上个时间步的上下文记录下来，就可以每次只是计算当前时间步的上下文，而不是每次都计算所有时间步的上下文
        # TODO 不知道HuggingFace的transformer库是不是这样实现的
        out = model.decode(memory, src_mask, Variable(ys), Variable(subsequent_mask(ys.size(1))))
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.item()
        if end_symbol is not None and next_word == end_symbol:
            break
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
    return ys


import torchtext.data as data



# V = 11
# criterion = SmoothLabel(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)))

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

import torchtext
from torchtext import data, datasets
from torchtext.vocab import build_vocab_from_iterator


from torch.utils.data import DataLoader, Dataset

class MyDataSet(Dataset):
    def __init__(self, json_file, max_token_len=200, min_freq=2):
        super(Dataset, self).__init__()
        import json
        import spacy
        self.min_freq = min_freq
        self.en = spacy.load('en_core_web_sm')
        self.max_token_len = max_token_len
        json_data = json.load(open(json_file))
        self.en_data = []
        self.zh_data = []

        for en_zh in json_data:
            en = en_zh[0]
            zh = en_zh[1]
            en = self.en_toeknizer(en)
            zh = self.zh_tokenizer(zh)
            if len(en) > self.max_token_len or (len(zh) + 2) > self.max_token_len:
                continue
            for i in range(self.max_token_len - len(en)):
                en.append('<pad>')
            zh.insert(0, '<sos>')
            zh.append('<eos>')
            for i in range(self.max_token_len - len(zh)):
                zh.append('<pad>')
            self.en_data.append(en)
            self.zh_data.append(zh)
        self.en_vocab = build_vocab_from_iterator(self.en_data, min_freq=self.min_freq, specials=['<sos>', '<eos>', '<unk>', '<pad>'])
        self.zh_vocab = build_vocab_from_iterator(self.zh_data, min_freq=self.min_freq, specials=['<sos>', '<eos>', '<unk>', '<pad>'])


    def en_toeknizer(self, text):
        return [tok.text for tok in self.en.tokenizer(text)]

    def zh_tokenizer(self, text):
        return [ch for ch in text]

    def __len__(self):
        return len(self.en_data)

    def __getitem__(self, idx):
        src_text = self.en_data[idx]
        trg_text = self.zh_data[idx]
        src_token = [self.en_vocab[token] if token in self.en_vocab else self.en_vocab['<unk>'] for token in src_text]
        trg_token = [self.zh_vocab[token] if token in self.zh_vocab else self.zh_vocab['<unk>'] for token in trg_text]
        src_token = torch.tensor(src_token, dtype=torch.long)
        trg_token = torch.tensor(trg_token, dtype=torch.long)
        return src_token, trg_token



device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


train_dataset = MyDataSet("/home/fengyuejin/下载/WMT18-English-Chinese-Machine-Translation-master/data/dataset/dev.json")
dev_dataset = MyDataSet("/home/fengyuejin/下载/WMT18-English-Chinese-Machine-Translation-master/data/dataset/dev.json")



en_vocab_size = len(train_dataset.en_vocab)
zh_vocab_size = len(train_dataset.zh_vocab)
criterion = SmoothLabel(size=zh_vocab_size, padding_idx=train_dataset.zh_vocab['<pad>'], smoothing=0.1)
model = make_model(en_vocab_size, zh_vocab_size, 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))


data_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=0)
pad_id = train_dataset.en_vocab['<pad>']



epoch = 10
for i in range(0, epoch):
    total_tokens = 0
    total_loss = 0
    start_time = time.time()
    total_step = len(data_loader)
    model.train()
    for i, batch in enumerate(data_loader):
        en_token, zh_token = batch
        src = Variable(en_token, requires_grad=False)
        tgt = Variable(zh_token, requires_grad=False)
        input_batch = Batch(src, tgt, pad_id)
        loss, tokens = run_batch(input_batch, model, SimpleLossCompute(model.generator, criterion, model_opt))
        total_tokens += tokens
        total_loss += loss
        if i % 50 == 0:
            print("Epoch Step: %d/%d Loss: %f Tokens per Sec: %f" % (i, total_step, loss / tokens, tokens / (time.time() - start_time)))
    print("Epoch %d Loss: %f" % (i, total_loss / total_tokens))

    model.eval()
    for i in range(len(dev_dataset)):
        src_token, trg_token = dev_dataset[i]
        src = Variable(src_token, requires_grad=False)
        src = src.unsqueeze(0)
        src_mask = (src != pad_id).unsqueeze(-2)
        predict_idx = greedy_decode(model, src, src_mask, max_len=200, start_symbol=train_dataset.zh_vocab['<sos>'], end_symbol=train_dataset.zh_vocab['<eos>'])
        predict_text = train_dataset.zh_vocab.lookup_tokens(torch.squeeze(predict_idx).tolist())
        src_text = train_dataset.en_vocab.lookup_tokens(torch.squeeze(src_token).tolist())
        src_text = " ".join([ch for ch in src_text if ch != '<pad>'])
        text_str = "".join(predict_text)
        print("src: %s, predict: %s" % (src_text, text_str))