Transformer

前言

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

原论文：Attention is all you need

---

一、Transformer总体结构

将这个模型看成是一个黑箱操作。如在机器翻译中，输入的是一种语言，输出的是另一种语言。

图中输入端是法语，输出端是英语。
拆开这个模型可以看到它是由两大部分组成：encoders和decoders。

encoders部分是由一堆编码器（encoder）构成，论文中是6个，也可以是其他数字，decoders部分也是有相同数量的解码器（decoder）构成。

更具体的，其中，N=6时，就是上面的模型结构了

OK，我们知道了Transformer模型的总体结构了，接下来看看具体的encoder和decoder是由什么构成的。

二、Encoder

encoder就分成四部分来讲

1.输入部分

1.1 词嵌入

最开始输入一个句子比如是“我爱你…”，我们得将每个词其转为向量，首先应进行Embedding，如下图有12个字“我爱你…”，将每个字转为512维的字向量，可以使用Word2Vec、Glove等。
这样这个句子就变成了[12,512]的矩阵了，如果加上batch_size，那么就是[batch_size,12,512]的张量。

1.2位置嵌入

由于Transformer没有RNN网络的迭代操作，所以我们必须提供每个字的位置信息给Transformer，这样它才能识别出语言中的顺序关系。
论文中使用的是sin和cos函数的线性变换来提供给模型的位置信息：

上式中pos指一个字在句子中的位置，取值范围是[0，max_seq_length]，i指的是字向量的维度序号，上面我们字向量的维度embedding_dimension=512，i的取值范围是[0，embedding_dimension/2]，d_model就是embedding_dimension=512。

至于为何这样设计，可以参考这篇文章Transformer 中的 Positional Encoding。

知道字向量和位置向量怎么转换后，将句子中每个字的字向量和其位置向量直接相加，相加之后的结果就是encoder的输入。

2.多头注意力机制

在此之前，需要知道注意力机制是什么？

2.1注意力机制介绍

从注意力模型的命名方式来看，很明显借鉴了人类的注意力机制，因此，我们首先简单介绍人类视觉的选择性注意力机制。

上面这幅图展示了人类在看一幅图像时是怎样分配注意力的，颜色越深越被视觉系统关注，在上面这幅图的场景，很明显，人们会把更多的注意力放在婴儿的脸，文本的标题以及文章的首句等位置。那么换一个场景呢？

The animal didn’t cross the street because it was too tired

上面这句话中的it到底是指animal还是street呢？对于我们来说，判断出来很容易，但是对于机器来说却很难，self-attention就能让机器把it和animal联系起来。
当模型处理序列的每个单词时，self-attention会关注整个序列的所有单词，帮助模型对当前这个单词更好的进行编码。

2.2注意力机制计算

接下来看看如何使用向量来计算自注意力：
公式：

第一步： 得到Q、K、V，我们通过三个参数矩阵WQ、WK、WV（训练的权重矩阵）,来得到Q、K、V向量。

上图可以看出，我们先将一个句子中的每个单词Thinking和Machines先进行嵌入转化为向量（输入部分已讲过），然后每个词向量X1、X2都分别与WQ、WK、WV矩阵相乘，得到Q、K、V向量q1、k1、v1和q2、k2、v2（如果非要纠结Q、K、V的含义，那么你可以看看Q、K、V的含义，这里就不再描述）。

第二步： 计算scores，将Q与每个K向量进行点积(Q*KT)，如在给第一个词Thinking的进行打分时，q1 * k1T，q1 * k2T（T表示转置）。

这些scores决定了在编码单词“Thinking”时，有多重视句子中其他单词。

第三步： 先将这些scores进行scaled，除以8（根号d_k，论文中使用的d_k维度为64），再对进行softmax归一化。
如果你想知道为什么要除以根号d_k，你可以看下这篇文章attention为什么要scaled？

这个softmax分数决定了每个单词对当前编码的单词的重要程度，分数越高，编码单词的时候，越要关注分数高的单词。
第四步： 对第一步计算出的V进行加权求和（softmax与V相乘最后再相加）。

上图中的z1是“Thinking”的attention值，即z1 = 0.88v1 + 0.12v2，如果要计算出“Machines”得attention值，与“Thinking”类似，第一步先算出Q、K、V，第二步q2 * k1T、q2 * k2T（当前编码的词的Q向量与所有的K向量进行点乘）,得到scores，第三步除以8，之后softmax，最后进行加权求和，z2 = （q2 * k1）* v1 + （q2 * k2T）* v2

上面的操作就是计算self-attention的具体过程，然而实际中，这些计算都是以矩阵形式完成的，接下来看看如何使用矩阵来完成的。（过程与上面用向量计算是一样的）

先计算出每个词向量的的q、k、v。这边的ai表示句子中第i个词向量（a1、a2就类似于上面的词向量x1、x2）。

然后计算出所有的scores，对scores除以根号d_k（图中省略了这步），再进行softmax得到scaled scores。

最后scaled scores与V矩阵相乘得到attention矩阵。

2.3 多头注意力机制

至此注意力机制讲完，接下来是多头注意力机制。
论文中进一步完善自注意力层，加入了一种多头注意力机制（Multi-Headed Attention）。从字面上来看，多个自注意力机制一起进行，就是说不仅仅使用一组Q、K、V矩阵，而是初始化多组，transformer使用了8组，所以最后得到了8个attention矩阵

上面的最是其中的两个self-attention，使用不同的权重矩阵进行8次attention计算，会得到8个不同的Z矩阵。

上面的Multi-Headed Attention得到8个Z矩阵，但是前馈层只需要一个矩阵，所以我们得把这8个矩阵压缩成一个，该怎么做呢？

我们将8个Z矩阵拼接起来，然后乘上一个额外的权重矩阵WO，得到一个唯一的Z矩阵。
下面这幅图就是全过程的描述

3.残差连接和LayerNormalization

我们在上一步得到了Attention(Q, K, V)，之后要进行残差连接。

红色框框是encoder中的残差+LayerNormalization

3.1为什么使用残差连接？

随着深度网络层数增加，带来一系列的问题，梯度消失、梯度爆炸、过拟合等问题。
针对这些问题，已有一些解决方案，如dropout层用来防止过拟合，Relu层主要用来防止梯度消失，BN（Batch-Nomalization）避免了消失，减少过拟合。

下图是残差连接结构图

f(x)是第二个weight layer的输出，下图是上图的简易结构

对XAout 即上面的X（identity）根据反向传播的链式法则进行求导

梯度消失一般情况下是因为连乘，上图中的最后一步括号中的连乘即使再多，即使变为0，但始终有个1在，确保了梯度不会为0，缓解了梯度消失。

3.2LayerNormalization

LayerNormalization对数据进行归一化
为什么使用LayerNormalization而不使用BatchNormalization可以看看这篇文章BatchNormalization、LayerNormalization、InstanceNorm、GroupNorm、SwitchableNorm总结

4.前馈层Feed Forward

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

X是输入，Feed Forward最终的输出维度与X是一致的。

三、Decoder

decoder结构与encoder类似，残差连接和LayerNormalization还有FeedForward跟encoder的是一致的，Masked Multi-Head Attention与Multi-Head Attention其实也是一样的，只不过多了Masked。decoder中间部分的Multi-Head Attention与encoder的略微有些不同。decoder还多了一个Linear+Softmax

3.1输入部分

图中decoder也有输入，在做Translation时，这个输入就会用上

3.2Masked Self-Attention

Masked Self-Attention是Decoder self-attention，相比encoder多个一个mask机制
传统的Seq2Seq中Decoder使用的是RNN，因此在训练的过程中输入t时刻的词，模型是看不到未来时刻的词。而Transformer Decoder抛弃了RNN，改成了Self-Attention，由此产生了一个问题，整个Decoder的输入都是暴露的，这显然不行，因此对Decoder的输入进行了一些处理，即Mask。
Mask操作是我们计算self-attention进行softmax之前，对Scaled Scores进行Mask。
举个例子， " I am fine"，我们输入“I”的时候，模型目前只知道当前输入“I”和之前的所有词的信息，我们要做的就是不要看到后面词的信息，该怎么做呢？
Mask很简单，首先生成一个下三角全为0，上三角全为负无穷的mask矩阵，然后再与Scaled Scores矩阵相加得到Masked Scores矩阵，如下图：

将得到的Masked Scores进行softmax，矩阵中为负无穷的位置经过softmax会变为0（不了解softmax函数的可以去看看softmax）

为什么要进行mask呢？
前面已经讲过了，对于当前输入的词，我不希望让decoder知道后面的词，那么当前输入的词计算attention时，不应将注意力放在后面的词身上，也就是上图中Scaled Scores矩阵我们应该将上三角全部置为负无穷，为什么要置为负无穷呢？因为这样下一步对Scaled Scores进行softmax时，负无穷会变为0。

3.3Encoder-Decoder Attention

Multi-Head Attention计算流程其实和decoder的多头self-attention很相似，encoder的Q、K、V都是encoder的，但是这层的K、V都是Encoder那边传过来的，而Q才是Decoder的，所以这一层又叫交互注意力层

3.4Linear+Softmax

到这流程基本上快结束了，Linear就是一个全连接层，将Decoder的输出进行映射，如果translation时我们的词典有1w个词，那么就会映射成1w维，最后经过softmax会输出1w个词的概率，概率值最大的就是我们最终的结果。

四、Mask

虽然前面讲Masked Self-Attention时提到了mask，这里再进行补充
mask表示掩码，对某些值进行掩盖，使其在参数更新时不产生效果，Transformer里涉及两种mask，一种是padding mask另一种是sequence mask

4.1padding mask

为什么要进行padding mask呢？
每个batch输入的句子长度有可能不一样，但是我们得确定输入句子的长度，该怎么做呢？假设我们指定输入句子的长度为input_len，那么输入句子长度大于input_len的，直接截断超过input_len部分，如果句子长度小于input_len的，进行padding，如用“pad”填充。

input_len = 5 # 指定输入句子长度为5（词的个数）
s1 = 'i want a beer' # 长度为4
s2 = 'i want to eat a apple' # 长度为6

上面s1长度小于input_len，需要padding，s2大于input_len，需要截断。

s1 = 'i want a beer pad' # 长度为5
s2 = 'i want to eat a' # 长度为5

在计算attention时，由于这些“pad”是没有意义的，注意力不应该放在这部分，计算也与之前讲的mask一样，构建mask矩阵的时候，这些被padding的位置，将会置为负无穷，其他位置为0，之后将其与Scaled Scores矩阵相加得到Masked Scores矩阵，然后将Masked Scores矩阵进行softmax，矩阵中为负无穷的位置经过softmax会变为0，是不是更之前讲的mask类似？

4.2sequence mask

sequence mask就是上面讲Masked Self-Attention时提到了mask，是为了使得decoder不能看见后面单词的信息。

五、代码

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

device = 'cpu'
# device = 'cuda'

# transformer epochs
epochs = 10
# epochs = 1000

# 这里我没有用什么大型的数据集，而是手动输入了两对德语→英语的句子
# 还有每个字的索引也是我手动硬编码上去的，主要是为了降低代码阅读难度
# 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_idx2word = {i: w for i, w in enumerate(src_vocab)}
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（token embedding和position编码的维度）
d_ff = 2048  # FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_k = d_v = 64  # dimension of K(=Q), V（Q和K的维度需要相同，这里为了方便让K=V）
n_layers = 6  # number of Encoder of Decoder Layer（Block的个数）
n_heads = 8  # number of heads in Multi-Head Attention（有几套头）


# ==============================================================================================
# 数据构建


def make_data(sentences):
    """把单词序列转换为数字序列"""
    enc_inputs, dec_inputs, dec_outputs = [], [], []
    for i in range(len(sentences)):
        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]]

        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):
    """自定义DataLoader"""

    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)


# ====================================================================================================
# Transformer模型

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, 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, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-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).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)


"""
如果是encoder端：
    为了保证输入模型的input的size是一致的，而句子的长度不一定都相同，所以得让句子的长度完全一样记为src_len，如果长度大于src_len
    进行截断，如果小于，则进行pad，但是这些pad部分的是没有意义的，因此attention时，注意力不应该放在这部分，故需要将这部分mask掉，
    也就是需要将query与padding部分对应的key的相似度度量统一转化为一个很小的数，这样经过softmax的时候，这部分的权重就会很接近0
    然后这里的seq_q、seq_k都为输入序列batch*src_len
"""
def get_attn_pad_mask(seq_q, seq_k):
    # pad mask的作用：在对value向量加权平均的时候，可以让pad对应的alpha_ij=0，这样注意力就不会考虑到pad向量
    """这里的q,k表示的是两个序列（跟注意力机制的q,k没有关系），例如encoder_inputs (x1,x2,..xm)和encoder_inputs (x1,x2..xm)
    encoder和decoder都可能调用这个函数，所以seq_len视情况而定
    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()  # 这个seq_q只是用来expand维度的
    batch_size, len_k = seq_k.size()
    # eq(zero) is PAD token
    # 例如:seq_k = [[1,2,3,4,0], [1,2,3,5,0]]
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, len_k], True is masked
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size, len_q, len_k] 构成一个立方体(batch_size个这样的矩阵)


def get_attn_subsequence_mask(seq):
    """建议打印出来看看是什么的输出（一目了然）
    seq: [batch_size, tgt_len]
    """
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    # attn_shape: [batch_size, tgt_len, tgt_len]
    subsequence_mask = np.triu(np.ones(attn_shape), k=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]
        说明：在encoder-decoder的Attention层中len_q(q1,..qt)和len_k(k1,...km)可能不同
        """
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)  # scores : [batch_size, n_heads, len_q, len_k]
        # mask矩阵填充scores（用-1e9填充scores中与attn_mask中值为1位置相对应的元素）
        scores.masked_fill_(attn_mask, -1e9)  # Fills elements of self tensor with value where mask is True.

        attn = nn.Softmax(dim=-1)(scores)  # 对最后一个维度(v)做softmax
        # scores : [batch_size, n_heads, len_q, len_k] * V: [batch_size, n_heads, len_v(=len_k), d_v]
        context = torch.matmul(attn, V)  # context: [batch_size, n_heads, len_q, d_v]
        # context：[[z1,z2,...],[...]]向量, attn注意力稀疏矩阵（用于可视化的）
        return context, attn


class MultiHeadAttention(nn.Module):
    """这个Attention类可以实现:
    Encoder的Self-Attention
    Decoder的Masked Self-Attention
    Encoder-Decoder的Attention
    """

    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)  # q,k必须维度相同，不然无法做点积
        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: batch_size, S:seq_len, D: dim
        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, Head, W) -trans-> (B, Head, S, W)
        #           线性变换               拆成多头

        # Q: [batch_size, n_heads, len_q, d_k]
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # K: [batch_size, n_heads, len_k, d_k] # K和V的长度一定相同，维度可以不同
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # V: [batch_size, n_heads, len_v(=len_k), d_v]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)

        # 因为是多头，所以mask矩阵要扩充成4维的
        # attn_mask: [batch_size, seq_len, seq_len] -> [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], attn: [batch_size, n_heads, len_q, len_k]
        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)
        # 下面将不同头的输出向量拼接在一起
        # context: [batch_size, n_heads, len_q, d_v] -> [batch_size, len_q, n_heads * d_v]
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # 再做一个projection
        output = self.fc(context)  # [batch_size, len_q, d_model]
        return nn.LayerNorm(d_model).to(device)(output + residual), attn


# Pytorch中的Linear只会对最后一维操作，所以正好是我们希望的每个位置用同一个全连接网络
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)
        return nn.LayerNorm(d_model).to(device)(output + residual)  # [batch_size, seq_len, d_model]


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):
        """E
        enc_inputs: [batch_size, src_len, d_model]
        enc_self_attn_mask: [batch_size, src_len, src_len]  mask矩阵(pad mask or sequence mask)
        """
        # enc_outputs: [batch_size, src_len, d_model], attn: [batch_size, n_heads, src_len, src_len]
        # 第一个enc_inputs * W_Q = Q
        # 第二个enc_inputs * W_K = K
        # 第三个enc_inputs * W_V = 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（未线性变换前）
        enc_outputs = self.pos_ffn(enc_outputs)
        # enc_outputs: [batch_size, src_len, d_model]
        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]
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs,
                                                        dec_self_attn_mask)  # 这里的Q,K,V全是Decoder自己的输入
        # dec_outputs: [batch_size, tgt_len, d_model], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs,
                                                      dec_enc_attn_mask)  # Attention层的Q(来自decoder) 和 K,V(来自encoder)
        dec_outputs = self.pos_ffn(dec_outputs)  # [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attn, dec_enc_attn  # 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)  # token Embedding
        self.pos_emb = PositionalEncoding(d_model)  # Transformer中位置编码时固定的，不需要学习
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        """
        enc_inputs: [batch_size, src_len]
        """
        enc_outputs = self.src_emb(enc_inputs)  # [batch_size, src_len, d_model]
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # [batch_size, src_len, d_model]
        # Encoder输入序列的pad mask矩阵
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)  # [batch_size, src_len, src_len]
        enc_self_attns = []  # 在计算中不需要用到，它主要用来保存你接下来返回的attention的值（这个主要是为了你画热力图等，用来看各个词之间的关系
        for layer in self.layers:  # for循环访问nn.ModuleList对象
            # 上一个block的输出enc_outputs作为当前block的输入
            # enc_outputs: [batch_size, src_len, d_model], enc_self_attn: [batch_size, n_heads, src_len, src_len]
            enc_outputs, enc_self_attn = layer(enc_outputs,
                                               enc_self_attn_mask)  # 传入的enc_outputs其实是input，传入mask矩阵是因为你要做self attention
            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)  # Decoder输入的embed词表
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])  # Decoder的blocks

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        """
        dec_inputs: [batch_size, tgt_len]
        enc_inputs: [batch_size, src_len]
        enc_outputs: [batch_size, src_len, d_model]   # 用在Encoder-Decoder Attention层
        """
        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).to(
            device)  # [batch_size, tgt_len, d_model]
        # Decoder输入序列的pad mask矩阵（这个例子中decoder是没有加pad的，实际应用中都是有pad填充的）
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).to(device)  # [batch_size, tgt_len, tgt_len]
        # Masked Self_Attention：当前时刻是看不到未来的信息的
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).to(
            device)  # [batch_size, tgt_len, tgt_len]

        # Decoder中把两种mask矩阵相加（既屏蔽了pad的信息，也屏蔽了未来时刻的信息）
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask),
                                      0).to(device)  # [batch_size, tgt_len, tgt_len]; torch.gt比较两个矩阵的元素，大于则返回1，否则返回0

        # 这个mask主要用于encoder-decoder attention层
        # get_attn_pad_mask主要是enc_inputs的pad mask矩阵(因为enc是处理K,V的，求Attention时是用v1,v2,..vm去加权的，要把pad对应的v_i的相关系数设为0，这样注意力就不会关注pad向量)
        #                       dec_inputs只是提供expand的size的
        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]
            # Decoder的Block是上一个Block的输出dec_outputs（变化）和Encoder网络的输出enc_outputs（固定）
            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)
        # dec_outputs: [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attns, dec_enc_attns


class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder().to(device)
        self.decoder = Decoder().to(device)
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).to(device)

    def forward(self, enc_inputs, dec_inputs):
        """Transformers的输入：两个序列
        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]
        # 经过Encoder网络后，得到的输出还是[batch_size, src_len, d_model]
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)
        # dec_outputs: [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)
        # dec_outputs: [batch_size, tgt_len, d_model] -> dec_logits: [batch_size, tgt_len, tgt_vocab_size]
        dec_logits = self.projection(dec_outputs)
        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns


model = Transformer().to(device)
# 这里的损失函数里面设置了一个参数 ignore_index=0，因为 "pad" 这个单词的索引为 0，这样设置以后，就不会计算 "pad" 的损失（因为本来 "pad" 也没有意义，不需要计算）
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)  # 用adam的话效果不好

# ====================================================================================================
for epoch in range(epochs):
    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.to(device), dec_inputs.to(device), dec_outputs.to(device)
        # outputs: [batch_size * tgt_len, tgt_vocab_size]
        outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)
        loss = criterion(outputs, dec_outputs.view(-1))  # dec_outputs.view(-1):[batch_size * tgt_len * tgt_vocab_size]
        print('Epoch:', '%d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))

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


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序列会一点点变长（每次添加一个新预测出来的单词）
        dec_input = torch.cat([dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
                              -1)
        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]  # 拿出当前预测的单词(数字)。我们用x'_t对应的输出z_t去预测下一个单词的概率，不用z_1,z_2..z_{t-1}
        next_symbol = next_word
        if next_symbol == tgt_vocab["E"]:
            terminal = True
        # print(next_word)

    # greedy_dec_predict = torch.cat(
    #     [dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
    #     -1)
    greedy_dec_predict = dec_input[:, 1:]
    return greedy_dec_predict


# ==========================================================================================
# 预测阶段
enc_inputs, _, _ = next(iter(loader))
for i in range(len(enc_inputs)):
    greedy_dec_predict = greedy_decoder(model, enc_inputs[i].view(1, -1).to(device), start_symbol=tgt_vocab["S"])
    print(enc_inputs[i], '->', greedy_dec_predict.squeeze())
    print([src_idx2word[t.item()] for t in enc_inputs[i]], '->',
          [idx2word[n.item()] for n in greedy_dec_predict.squeeze()])

---

参考文章

图解Transformer（完整版）
BERT大火却不懂Transformer？读这一篇就够了 原版 可视化机器学习 可视化神经网络 可视化深度学习
Transformer从零详细解读(可能是你见过最通俗易懂的讲解)
台大李宏毅21年机器学习课程 self-attention和transformer
Transformer代码从零解读(Pytorch）
transformer中的attention为什么scaled?
Batch Normalization详解
transformer面试题的简单回答
Transformer 详解
Transformer 的 PyTorch 实现
Transformer 模型详解
一文读懂BERT(原理篇)
nlp中的Attention注意力机制+Transformer详解