Attention Is All You Need （Transformer）阅读笔记

这篇来自Transformer的开篇之作，以下将通过论文顺序进行介绍，同时穿插该有的讲解

Abstract

The dominant sequence transduction models are based on complex recurrent or convolutional neural networks that include an encoder and a decoder. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train.

论文在摘要就提出，现在的传统模型大多使用encoder-decoder，如RNN和CNN，要么就是加一个注意力机制 （做序列到序列的生成）。这篇文章提出简单架构Transforms，仅仅利用自注意力机制 没有依赖于RNN或CNN，并且能很好的应用于其他任务场景。

Introduction

Introduction大多就是对作者提出的Transformer进行一个补充介绍：

• Transformer 和 传统 RNN/LSTM 模型的最大区别，就是 LSTM 的训练是迭代的、串行的，必须要等当前字处理完，才可以处理下一个字。
• 而 Transformer 的训练时并行的，即所有字是同时训练的，这样就大大增加了计算效率以及更高的并行度。
• Transformer 使用了位置嵌入 (Positional Encoding) 来理解语言的顺序，使用自注意力机制（Self Attention Mechanism）和全连接层进行计算，这些后面会讲到

Background

背景知识也是强调了一下之前的不足，以及我们为什么如此设计，之前很多模型都采用卷积来替代传统的RNN，但是当我们有很多个像素点或者两个像素点相隔很远的时候，我们需要通过很多个卷积核，才能将两个特征融合。
而本文提出的self-attention，我们一次可以看到所有像素，通过一次矩阵计算即可解决。
另外作者提到，我们为了能够像不同卷积核抽取不同特征，我们提出了MutiHead Self-Attetion（在模型中介绍），从而来达到我们抽取不同特征的目的。

Model

Transformer 模型主要分为两大部分，分别是 Encoder 和 Decoder。Encoder 负责把输入（语言序列）隐射成隐藏层（下图中第 2 步用九宫格代表的部分），然后解码器再把隐藏层映射为自然语言序列。例如下图机器翻译的例子（Decoder 输出的时候，是通过 N 层 Decoder Layer 才输出一个 token，并不是通过一层 Decoder Layer 就输出一个 token）

下面将按照我自我的顺序来进行逐一讲解

Encoder

首先解释 Encoder 部分，即把自然语言序列映射为隐藏层的数学表达的过程。理解了 Encoder 的结构，再理解 Decoder 就很简单了

Positional Encoding

由于 Transformer 模型没有循环神经网络的迭代操作，所以我们必须提供每个字的位置信息给 Transformer，这样它才能识别出语言中的顺序关系。
现在定义一个位置嵌入的概念，也就是 Positional Encoding，位置嵌入的维度与词向量嵌入的维度是相同的，都是 embedding_dimension。
注意，我们一般以字为单位训练 Transformer 模型。首先初始化字编码的大小为 [vocab_size, embedding_dimension]，vocab_size 为字库中所有字的数量，embedding_dimension 为字向量的维度，对应到 PyTorch 中，其实就是 nn.Embedding(vocab_size, embedding_dimension)

论文中使用了 sin 和 cos 函数的线性变换来提供给模型位置信息，这一位置嵌入函数是写死的，而不像BERT中可以进行自我学习:
$$\begin{array}{rl}P{E}_{(pos,2i)}& =\sin \left(pos/10000^{2i/d_{\text{model }}}\right)\\ P{E}_{(pos,2i+1)}& =\cos \left(pos/10000^{2i/d_{\text{model }}}\right)\end{array}$$
具体的公式含义在这里就不解析了，有兴趣可以参考这篇文章：Posititon Encoding
核心思想：利用sin cos 的周期，将任何一个值，用长为512的向量（也就是词嵌入的长度，512这里是论文中规定）表示其位置信息。

Multi-Head Self Attention Mechanism

对于输入的句子，通过 WordEmbedding 得到该句子中每个字的字向量，同时通过 Positional Encoding 得到所有字的位置向量，将其相加（维度相同，可以直接相加），得到该字真正的向量表示，接着传入Self-Attention。
具体Self-Attention的知识这里不赘述，如有疑问查看：Self-Attention
这里稍微说明的一下是，论文中为什么要除以 根号Dk:

在Transforms中我们用到的Dk都比较大，是512维度，这样我们在做self-attention的点积时候，会导致结果越来越大，从而导致将softmax结果推到梯度较小的位置，因此会训练不动。所以出除一个这个是一个不错的选择。

至于多头注意力，作者提到，我们为了能够像不同卷积核抽取不同特征，我们提出了MutiHead Self-Attetion（在模型中介绍），从而来达到我们抽取不同特征的目的。最后把每个注意力函数的输出并在一起，再投影到512维度，从而得到最终输出

最后，作者在文中提到：

Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.

多头注意力并没有给我们带来很高的计算复杂度，因为 每个头的维度 * 头数量 = 总维度
因为有残差连接，所以这里投影的维度等于：最后输出维度/多少个头

Residual connection

我们在上一步得到了经过 self-attention 加权之后输出，也就是 $\text{ Self-Attention }(Q,K,V)$ 是通过 Euler integral，然后把他们加起来做残差连接

$$X_{\text{embedding }}+\text{ Self-Attention }(Q,K,V)$$

Layer Normalization

什么是Layer Normalization：
把每个样本变为均值为0 ，方差为1，而batchnorm则是对每个特征
相当于把 数据转置 放到 batchnorm 再转置回去 == layernorm

Layer Normalization好处：不会受到句子长短不一的影响，而batchnorm在句子长短不一的时候，均值和方差抖动较大。

Position-wise Feed-Forward Networks

In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.

$$X_{\text{hidden }}=\mathrm{Linear}\left(\mathrm{ReLU}\left(\mathrm{Linear}\left(X_{\text{attention }}\right)\right)\right)$$

对每一个词，作用同一个MLP 权重一样，前面attention起的作用就是整个序列信息抓取出来，做一次汇聚，attention的每一个输出都含有了我整个序列了历史信息（全局信息）
Position-wise：之所以可以每个点独立做MLP，因为attention每个输出已经含有了整个历史信息

其经过两次投影，首先从512维 到 2048维 再回到512维（因为有残差连接）所以必须回到512维度。

Transformer Encoder Overall structure

1. 字向量与位置编码
   $$X=\text{ Embedding-Lookup }(X)+\text{ Positional-Encoding }$$
2. 自注意力机制
   $$\begin{array}{c}Q={\mathrm{Linear}}_q(X)=X{W}_Q\\ K={\mathrm{Linear}}_k(X)=X{W}_K\\ V={\mathrm{Linear}}_v(X)=X{W}_V\\ X_{attention}=\mathrm{Self}\text{ Attention }(Q,K,V)\end{array}$$
3. self-attention 残差连接与 Layer Normalization
   $$\begin{array}{c}{X}_{\text{attention }}=X+X_{\text{attention }}\\ X_{\text{attention }}=\mathrm{LayerNorm}\left(X_{\text{attention }}\right)\end{array}$$
4. 进行 Encoder block 结构图中的第 4 部分，也就是 FeedForward，其实就是两层线性映射并用激活函数激活
   $$X_{\text{hidden }}=\mathrm{Linear}\left(\mathrm{ReLU}\left(\mathrm{Linear}\left(X_{\text{attention }}\right)\right)\right)$$
5. FeedForward 残差连接与 Layer Normalization
   $$\begin{array}{c}{X}_{\text{hidden }}=X_{\text{attention }}+X_{\text{hidden }}\\ X_{\text{hidden }}=\text{ LayerNorm }\left(X_{\text{hidden }}\right)\end{array}$$

Decoder

和 Encoder 一样，上面三个部分的每一个部分，都有一个残差连接，后接一个 Layer Normalization。Decoder 的中间部件并不复杂，大部分在前面 Encoder 里我们已经介绍过了，但是 Decoder 由于其特殊的功能，因此在训练时会涉及到一些细节

Masked Self-Attention

具体来说，传统 Seq2Seq 中 Decoder 使用的是 RNN 模型，因此在训练过程中输入 时刻的词，模型无论如何也看不到未来时刻的词，因为循环神经网络是时间驱动的，只有当 t 时刻运算结束了，才能看到 t + 1 时刻的词。而 Transformer Decoder 抛弃了 RNN，改为 Self-Attention，由此就产生了一个问题，在训练过程中，整个 ground truth 都暴露在 Decoder 中，这显然是不对的，我们需要对 Decoder 的输入进行一些处理，该处理被称为 Mask

Similarly, self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot-product attention by masking out (setting to−∞) all values in the input of the softmax which correspond to illegal connections

也就是说，我们要做到一种自回归的方法来让他self-attention，因此我们要将没有看到的词设置为负无穷，这样在计算attention时候他的分数趋近于0。

具体如何实现：
道理很简单，我们做预测的时候是按照顺序一个字一个字的预测，怎么能这个字都没预测完，就已经知道后面字的信息了呢？Mask 非常简单，首先生成一个下三角全 0，上三角全为负无穷的矩阵，然后将其与 Scaled Scores 相加即可，

之后再做 softmax，就能将 - inf 变为 0，得到的这个矩阵即为每个字之间的权重

Multi-Head Self-Attention 无非就是并行的对上述步骤多做几次，前面 Encoder 也介绍了，这里就不多赘述了

Encoder-Decoder Attention

其实这一部分的计算流程和前面 Masked Self-Attention 很相似，结构也一摸一样，唯一不同的是这里的 $\text{K},V$为 Encoder 的输出， $\text{Q}$为 Decoder 中 Masked Self-Attention 的输出。

Training&Results

同学们自己下去看吧，浅显易懂

Conclusion

我们用一张图展示其完整结构

可以看到，模型首先会顺序执行完六个Encoder，再讲最后一个Encoder的输出作为 K，V 传入每一个Decoder中，而Decoder也是顺序执行的。

Q&A

1. Transformer 为什么需要进行 Multi-head Attention？
   原论文中说到进行 Multi-head Attention 的原因是将模型分为多个头，形成多个子空间，可以让模型去关注不同方面的信息，最后再将各个方面的信息综合起来。其实直观上也可以想到，如果自己设计这样的一个模型，必然也不会只做一次 attention，多次 attention 综合的结果至少能够起到增强模型的作用，也可以类比 CNN 中同时使用多个卷积核的作用，直观上讲，多头的注意力有助于网络捕捉到更丰富的特征 / 信息

2. Transformer 相比于 RNN/LSTM，有什么优势？为什么？
   (1) RNN 系列的模型，无法并行计算，因为 T 时刻的计算依赖 T-1 时刻的隐层计算结果，而 T-1 时刻的计算依赖 T-2 时刻的隐层计算结果
   (2) Transformer 的特征抽取能力比 RNN 系列的模型要好

3. 为什么说 Transformer 可以代替 seq2seq？
   这里用代替这个词略显不妥当，seq2seq 虽已老，但始终还是有其用武之地，seq2seq 最大的问题在于将 Encoder 端的所有信息压缩到一个固定长度的向量中，并将其作为 Decoder 端首个隐藏状态的输入，来预测 Decoder 端第一个单词 (token) 的隐藏状态。在输入序列比较长的时候，这样做显然会损失 Encoder 端的很多信息，而且这样一股脑的把该固定向量送入 Decoder 端，Decoder 端不能够关注到其想要关注的信息。Transformer 不但对 seq2seq 模型这两点缺点有了实质性的改进 (多头交互式 attention 模块)，而且还引入了 self-attention 模块，让源序列和目标序列首先 “自关联” 起来，这样的话，源序列和目标序列自身的 embedding 表示所蕴含的信息更加丰富，而且后续的 FFN 层也增强了模型的表达能力，并且 Transformer 并行计算的能力远远超过了 seq2seq 系列模型

Transforms’ Code

注释详细

import math
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data

# S: Symbol that shows starting of decoding input
# E: Symbol that shows starting of decoding output
# P: Symbol that will fill in blank sequence if current batch data size is short than time steps
sentences = [
    # enc_input           dec_input         dec_output
    ['ich mochte ein bier P', 'S i want a beer .', 'i want a beer . E'],
    ['ich mochte ein cola P', 'S i want a coke .', 'i want a coke . E']
]

# Padding Should be Zero
src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4, 'cola': 5}
src_vocab_size = len(src_vocab)

tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'coke': 5, 'S': 6, 'E': 7, '.': 8}
idx2word = {i: w for i, w in enumerate(tgt_vocab)}
tgt_vocab_size = len(tgt_vocab)

src_len = 5  # enc_input max sequence length
tgt_len = 6  # dec_input(=dec_output) max sequence length

# Transformer Parameters
d_model = 512  # Embedding Size 论文中 的参数
d_ff = 2048  # FeedForward dimension
d_q = d_k = d_v = 64  # dimension of K(=Q), V
n_layers = 6  # number of Enco der of Decoder La yer
# 多头注意力
n_heads = 8  # number of heads in Multi-Head Attent ion


def make_data(sentences):
    enc_inputs, dec_inputs, dec_outputs = [], [], []
    for i in range(len(sentences)):
        # sentences[i][0] encoder_input 的位置
        enc_input = [[src_vocab[n] for n in sentences[i][0].split()]]  # [[1, 2, 3, 4, 0], [1, 2, 3, 5, 0]]
        dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]]  # [[6, 1, 2, 3, 4, 8], [6, 1, 2, 3, 5, 8]]
        dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]]  # [[1, 2, 3, 4, 8, 7], [1, 2, 3, 5, 8, 7]]

        # extend 函数表示在末尾追加列表 把每句话的 加进去 [[1, 2, 3, 4, 0], [1, 2, 3, 5, 0]]
        enc_inputs.extend(enc_input)
        dec_inputs.extend(dec_input)
        dec_outputs.extend(dec_output)

    return torch.LongTensor(enc_inputs), torch.LongTensor(dec_inputs), torch.LongTensor(dec_outputs)


enc_inputs, dec_inputs, dec_outputs = make_data(sentences)


class MyDataSet(Data.Dataset):
    def __init__(self, enc_inputs, dec_inputs, dec_outputs):
        super(MyDataSet, self).__init__()
        self.enc_inputs = enc_inputs
        self.dec_inputs = dec_inputs
        self.dec_outputs = dec_outputs

    def __len__(self):
        return self.enc_inputs.shape[0]

    def __getitem__(self, idx):
        return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx]


loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # 论文中 提到了 在position 中 也使用了 dropout
        self.dropout = nn.Dropout(p=dropout)

        # 根据论文中的公式 3.5
        pe = torch.zeros(max_len, d_model)
        # 增加一个维度 变成两位
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # 0::2 每隔两步
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        '''
        x: [seq_len, batch_size, d_model]
        '''
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


# 把pad这个字符 mask掉 因为我们不需要计算pad和其他单词的注意力
# 注意 这里的 pad mask  并没有 mask 掉 最后一行（pad 与 所有词） 因为不太重要
# 但一定要 保证 mask 掉 最后一列
def get_attn_pad_mask(seq_q, seq_k):
    '''
    seq_q: [batch_size, seq_len]
    seq_k: [batch_size, seq_len]
    seq_len could be src_len or it could be tgt_len
    seq_len in seq_q and seq_len in seq_k maybe not equal

    '''
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    # eq(zero) is PAD token
    # eq 相同为 0 不同为 1
    # seq_k = [[1,2,3,4,0],[1,2,3,5,0]]
    # pad_attn_mask = [[F,F,F,F,T],[F,F,F,F,T]]
    # 增加一个维度
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, len_k], False is masked
    # 并进行扩充维度
    # 返回的是 true fasle矩阵
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size, len_q, len_k]

# Decoder 的 词 mask
def get_attn_subsequence_mask(seq):
    '''
    seq: [batch_size, tgt_len]
    '''
    # 创建一个三维全1矩阵 矩阵大小 如上
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    # triu 用于返回上三角 k = 1 表示对角线上移动
    subsequence_mask = np.triu(np.ones(attn_shape), k=1)  # Upper triangular matrix 上三角矩阵
    # 转换成张量 .byte() 表示 非零返回 1 零返回 0  最后 上三角全是 1
    subsequence_mask = torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask  # [batch_size, tgt_len, tgt_len]


# 文中的 自注意力机制
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        '''
        Q: [batch_size, n_heads, len_q, d_k]
        K: [batch_size, n_heads, len_k, d_k]
        V: [batch_size, n_heads, len_v(=len_k), d_v]
        attn_mask: [batch_size, n_heads, seq_len, seq_len]
        '''
        # 计算注意力分数
        # K.transpose(-1, -2) 根据线性代数矩阵计算规则 必须把最后两维度换一下位置 才能计算
        # scores : [batch_size, n_heads, len_q, len_k]
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)

        # pad mask  把 pad填充 mask掉 根据前面计算的 TF 矩阵
        # 这样 softmax 出来就是 0 也就是pad 和我们任何单词都没有相关性
        scores.masked_fill_(attn_mask, -1e9)  # Fills elements of self tensor with value where mask is True.

        # 真正的权重
        attn = nn.Softmax(dim=-1)(scores)
        # 和 V  相乘
        # [batch_size, n_heads, len_q, d_v]
        context = torch.matmul(attn, V)
        return context, attn


class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        # 输入的维度是　d_model 　那么我们输出维度　根据论文中要求　是　dk　维度　由于　有多个　head　
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, input_Q, input_K, input_V, attn_mask):
        '''
        input_Q: [batch_size, len_q, d_model]
        input_K: [batch_size, len_k, d_model]
        input_V: [batch_size, len_v(=len_k), d_model]
        attn_mask: [batch_size, seq_len, seq_len]
        '''
        # 残差
        residual, batch_size = input_Q, input_Q.size(0)
        # 下面三行代码就执行的这个过程
        # 原始维度  映射 分开（分成多头） 维度调换
        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, H, W) -trans-> (B, H, S, W)
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # Q: [batch_size, n_heads, len_q, d_k]
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # K: [batch_size, n_heads, len_k, d_k]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)  # V: [batch_size, n_heads, len_v(=len_k), d_v]

        # 把 T F 矩阵 也变成 一样的维度 用于 pad mask
        # attn_mask : [batch_size, n_heads, seq_len, seq_len]
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)

        # context: [batch_size, n_heads, len_q, d_v] 经过 Q K V 计算的自注意力
        # attn: [batch_size, n_heads, len_q, len_k]
        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)  # 没有在 self里面实例化 所以这样传参
        # context: [batch_size, len_q, n_heads * d_v] 四维度 变成 三维
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # 但是这样还不能进行 输出 因为根据论文中的图 我们需要 d_model 这样的输出 如此才能残差连接
        output = self.fc(context)  # [batch_size, len_q, d_model]
        # 在输出的时候 加上残差 并且使用 layernorm nn.LayerNorm(d_model) 填入最后一个维度
        return nn.LayerNorm(d_model).cuda()(output + residual), attn


# pointerWise 前馈神经网路  每个从attention出来的 单独做 MLP
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False)
        )

    def forward(self, inputs):
        '''
        inputs: [batch_size, seq_len, d_model]
        '''
        # 同样要加入残差
        residual = inputs
        output = self.fc(inputs)
        # nn.LayerNorm(d_model) 填入最后一个维度
        return nn.LayerNorm(d_model).cuda()(output + residual)  # [batch_size, seq_len, d_model]


# Encoder
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, enc_inputs, enc_self_attn_mask):
        '''
        enc_inputs: [batch_size, src_len, d_model]
        # 两两词之间的 pad 矩阵 所以肯定是 src_len * src*len
        enc_self_attn_mask: [batch_size, src_len, src_len]
        '''
        # enc_outputs: [batch_size, src_len, d_model]
        # attn: [batch_size, n_heads, src_len, src_len]
        # enc_inputs 分别乘以 Wq Wk Wv 生成不同的 Q K　V　矩阵
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs,
                                               enc_self_attn_mask)  # enc_inputs to same Q,K,V
        # 第二个 sublayer 做一个 pointerwise 特征提取
        # enc_outputs: [batch_size, src_len, d_model]
        # enc_outputs 和 enc_inputs 维度相等
        enc_outputs = self.pos_ffn(enc_outputs)
        return enc_outputs, attn


class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention()
        self.dec_enc_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        '''
        dec_inputs: [batch_size, tgt_len, d_model]
        enc_outputs: [batch_size, src_len, d_model]
        dec_self_attn_mask: [batch_size, tgt_len, tgt_len]
        dec_enc_attn_mask: [batch_size, tgt_len, src_len]
        '''
        # dec_outputs: [batch_size, tgt_len, d_model]
        # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]
        # 每一次来自上一层的 Decoder的第一个Sublayer 并且传入 dec_self_attn_mask 的 pad mask
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs,
                                                        dec_self_attn_mask)

        # dec_outputs: [batch_size, tgt_len, d_model]
        # dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
        # 这里是把 encoder的输出 拿过来 当作 K V  然后 Q 是来自 Decoder 的第一个 Sublayer
        # dec_enc_attn_mask 传入 dec_enc 的padmask  TF 矩阵
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs,
                                                      dec_enc_attn_mask)

        # point-wise 前馈神经网路 提取特征信息
        dec_outputs = self.pos_ffn(dec_outputs)  # [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attn, dec_enc_attn


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)
        self.pos_emb = PositionalEncoding(d_model)
        # 像列表一样 增加模型 嵌套八个 encoder
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        '''
        enc_inputs: [batch_size, src_len]
        '''
        # word_embedding
        enc_outputs = self.src_emb(enc_inputs)  # [batch_size, src_len, d_model]
        # 加入位置encoding
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # [batch_size, src_len, d_model]
        # 对 pad进行 mask
        # [batch_size, src_len, src_len] 后面两个维度 都是单词的个数
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)
        # 用来存储返回的attention的值
        enc_self_attns = []
        for layer in self.layers:
            # enc_outputs 就是通过了 wordembedding 和 posEncoding
            # enc_self_attn_mask 一个 pad mask 的 T F 矩阵
            # enc_outputs: [batch_size, src_len, d_model]
            # enc_self_attn: [batch_size, n_heads, src_len, src_len]
            # 从 第二层 encoder开始 就不需要 进行 embedded 和 posEncoding
            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)
            # 注意力加入进来 没啥用
            enc_self_attns.append(enc_self_attn)
        return enc_outputs, enc_self_attns


class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        '''
        dec_inputs: [batch_size, tgt_len]
        enc_intpus: [batch_size, src_len]
        enc_outputs 用在第二个sublayer
        enc_outputs: [batsh_size, src_len, d_model]
        '''
        # [batch_size, tgt_len, d_model]
        dec_outputs = self.tgt_emb(dec_inputs)
        # [batch_size, tgt_len, d_model]
        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).cuda()
        # [batch_size, tgt_len, tgt_len] decoder 一样要计算一个 pad 的 TF  矩阵
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).cuda()

        # [batch_size, tgt_len, tgt_len] 用于当前时刻 不能看到未来时刻的信息
        # 使用的mask 方法 不同于上一个 padmask
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).cuda()

        # 把 pad mask 的 TF 矩阵 和 submask 非0即1 矩阵 相加
        # 比较两个矩阵 大于为 1  小于为 0
        # dec_self_attn_mask 即屏蔽了 pad 也屏蔽了 未来时刻的信息
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask),
                                      0).cuda()  # [batch_size, tgt_len, tgt_len]

        '''
        在进行 encode-decode-attention 之前，是不是还做了一次 add & norm，要命的就是这个 add。
        由于进行了残差连接，使得原本为 0 的地方，又产生了值

        此时接下来又要进行一次 attention（即 encoder 的输出和 decoder 进行的 attention），
        如果这个时候，不重新 mask 一次，是不是未来时刻的信息仍然被用上了，因此要进行 mask
        '''
        # 这个 pad mask  用于 Decoder 中第二个 sublayer
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)  # [batc_size, tgt_len, src_len]

        dec_self_attns, dec_enc_attns = [], []
        for layer in self.layers:
            # dec_outputs: [batch_size, tgt_len, d_model]
            # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]
            # dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask,
                                                             dec_enc_attn_mask)

            # 把两个注意力矩阵加进入 只为了打印画图
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        return dec_outputs, dec_self_attns, dec_enc_attns


class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder().cuda()
        self.decoder = Decoder().cuda()
        # 转换维度 把decoder的输出 是512维度 转换为 target_size --- 求概率
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).cuda()

    def forward(self, enc_inputs, dec_inputs):
        '''
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        '''
        # tensor to store decoder outputs
        # outputs = torch.zeros(batch_size, tgt_len, tgt_vocab_size).to(self.device)

        # enc_outputs: [batch_size, src_len, d_model] 维度不变
        # enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len] 注意力矩阵 没啥作用 只为画图
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)

        # dec_outpus: [batch_size, tgt_len, d_model]
        # dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len]
        # dec_enc_attn: [n_layers, batch_size, tgt_len, src_len]
        dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)

        # 因为要变换到 我们 target size
        # dec_logits: [batch_size, tgt_len, tgt_vocab_size]
        dec_logits = self.projection(dec_outputs)

        # -1 第一个 -1 把 batch_size tgt_len 拼接起来  这里也就是把两个batch 拼接起来了
        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns


model = Transformer().cuda()
# 损失函数设置了一个参数 ignore_index=0，因为 "pad" 这个单词的编码为 0，计算损失时候 无需计算编码为 0 的损失
# 这样设置以后，就不会计算 "pad" 的损失（因为本来 "pad" 也没有意义，不需要计算）
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)
# optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(1000):
    for enc_inputs, dec_inputs, dec_outputs in loader:
        '''
      enc_inputs: [batch_size, src_len]
      dec_inputs: [batch_size, tgt_len]
      dec_outputs: [batch_size, tgt_len]
      '''
        enc_inputs, dec_inputs, dec_outputs = enc_inputs.cuda(), dec_inputs.cuda(), dec_outputs.cuda()
        # outputs: [batch_size * tgt_len, tgt_vocab_size]
        outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)

        # dec_outputs.view(-1) 也把两个batch 拼接起来 和 模型的返回值一样
        # 并没有通过 softmax 求概率最大值  而是直接求损失
        loss = criterion(outputs, dec_outputs.view(-1))
        print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


# 为了简单起见，当K=1时，贪婪解码器是Beam搜索。
# 这对于推理是必要的，因为我们不知道 目标序列的输入。因此，我们试图逐字生成目标输入，然后将其送入变换器。
def greedy_decoder(model, enc_input, start_symbol):
    """
    For simplicity, a Greedy Decoder is Beam search when K=1. This is necessary for inference as we don't know the
    target sequence input. Therefore we try to generate the target input word by word, then feed it into the transformer.
    Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding
    :param model: Transformer Model
    :param enc_input: The encoder input
    :param start_symbol: The start symbol. In this example it is 'S' which corresponds to index 4
    :return: The target input
    """
    enc_outputs, enc_self_attns = model.encoder(enc_input)
    dec_input = torch.zeros(1, 0).type_as(enc_input.data)
    terminal = False
    next_symbol = start_symbol
    while not terminal:
        # 自回归方式
        dec_input = torch.cat([dec_input.detach(), torch.tensor([[next_symbol]], dtype=enc_input.dtype).cuda()], -1)
        # infer 只用decoder
        dec_outputs, _, _ = model.decoder(dec_input, enc_input, enc_outputs)
        projected = model.projection(dec_outputs)
        # 从最后 线性层出来 求一个最大值 也就是最可能的那个值
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        # 每次拿出
        next_word = prob.data[-1]
        next_symbol = next_word
        # 终止条件
        if next_symbol == tgt_vocab["."]:
            terminal = True
        print(next_word)
    return dec_input


# Test
enc_inputs, _, _ = next(iter(loader))
enc_inputs = enc_inputs.cuda()
for i in range(len(enc_inputs)):
    # 粗糙的
    greedy_dec_input = greedy_decoder(model, enc_inputs[i].view(1, -1), start_symbol=tgt_vocab["S"])
    # 在传入模型进行预测
    predict, _, _, _ = model(enc_inputs[i].view(1, -1), greedy_dec_input)
    predict = predict.data.max(1, keepdim=True)[1]
    print(enc_inputs[i], '->', [idx2word[n.item()] for n in predict.squeeze()])