变压器型号代码和应用

原论文：Attention Is All You Need

在这篇博客中，transformer模型主要用于翻译，中译英，但是这里主要注重的是模型的代码。

具体应用的例子放在最后，大家有兴趣的可以看看。

首先，让我们来了解一下transformer的两个部分，也就是encoder和decoder

  在这里，我们的transformer模型是由6个编码器（encoder）和6个解码器（decoder）组成的，他们的数量可以自己定义，论文中用的是6个，所以我们在这里也是用的6个，

  需要注意的是，这6个encoder的结构虽然是一摸一样的，但是它的参数是不一样的，它们每一个的参数都是独立训练的；同时，encoder和decoder的结构也是不同的。

  下面，我们来看一下，原论文中的transformer模型：

  左边是encoder，右边是decoder，可以看出来，他们的内部结构是不同的。

准备部分

这部分是为了让后面的代码成功运行，不报错。

import math
import numpy as np
import torch
from torch.utils.data import Dataset,DataLoader
from torch import nn

# --------------------------手动代码构造句子----------------------------------

# 训练集
train_data=[
    ['我 有 一 个 好 朋 友 P', 'S I have a good friend .', 'I have a good friend . E'],
    ['我 有 零 个 女 朋 友 P', 'S I have zero girl friend .', 'I have zero girl friend . E'],
    ['我 有 一 个 男 朋 友 P', 'S I have a boy friend .', 'I have a boy friend . E']
]

# 建立源词库，用于将中文转化为数字序列
source= {'P': 0, '我': 1, '有': 2, '一': 3,
        '个': 4, '好': 5, '朋': 6, '友': 7, '零': 8, '女': 9, '男': 10}
# 字典推导式，i为索引(整数)，w为对应的中文字符或特殊符号
source_index={i:w for i,w in enumerate(source)}
source_size=len(source)

# 建立目标词库，用于将英文转化为数字序列，并在解码时将数字序列转为英语
target={'P': 0, 'I': 1, 'have': 2, 'a': 3, 'good': 4,
        'friend': 5, 'zero': 6, 'girl': 7, 'boy': 8, 'S': 9, 'E': 10, '.': 11}
target_index={i:w for i,w in enumerate(target)}
target_size=len(target)

source_len=8
target_len=7

device='cuda'

# ---------------------------------将数字序列转为tensor----------------------------

def to_tensor(train_data):
    enc_inputs,dec_inputs,dec_outputs=[],[],[]
    for i in range(len(train_data)):
        enc_input = [[source[n] for n in train_data[i][0].split()]]
        dec_input = [[target[n] for n in train_data[i][1].split()]]
        decr_output = [[target[n] for n in train_data[i][2].split()]]
        enc_inputs.extend(enc_input)
        dec_inputs.extend(dec_input)
        dec_outputs.extend(decr_output)

    # print("编码器输入的数字序列：{}".format(encoder_inputs))
    # print("解码器输入的数字序列：{}".format(decoder_inputs))
    # print("解码器输出的数字序列：{}".format(decoder_outputs))

    # print(torch.tensor(encoder_inputs))
    # print(torch.LongTensor(encoder_inputs))

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

enc_inputs,dec_inputs,dec_outputs=to_tensor(train_data)

# ----------------------------------自定义数据集----------------------------------

class Data_set(Dataset):
    def __init__(self,enc_inputs,dec_inputs,dec_outputs):
        super(Data_set,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, item):
        return self.enc_inputs[item],self.dec_inputs[item],self.dec_outputs[item]

data_loader=DataLoader(Data_set(enc_inputs,dec_inputs,dec_outputs),2,True)
# 参数设置
d_model = 512  # Embedding Size（token embedding和position编码的维度）
# FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_ff = 2048
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（有几个头）

encoder

我们可以看出，encoder主要分为3部分，即input（输入），multi-head attention（多头注意力机制）和feed forward（前馈神经网络），我们分开来解析。

positional encoding 位置编码

在这里，我们首先要知道为什么要使用位置编码。

原理

假如我们的输入这么几个字：1班2号，位置编码就能够识别每个字的位置，就能知道，哦，你是1班2号。

如果没有位置编码呢？我们的输入就是相当于几个独立的，没有顺序的字，既然如此，1班2号和2班1号是没有区别的。要知道，transformer经常被运用于机器翻译、文本生成等，元素的顺序是非常关键的。单词在句子中的位置以及排列顺序不仅是句子语法结构的组成部分，更是表达语义的重要概念。

那么transformer是怎么实现它得呢？

这主要利用了下面的公式：

 pos是字，2i和2i+1分别代表偶数和奇数位置，在偶数使用sin，奇数使用cos

举个栗子：班这个字为pos，它有512个维度，在0，2，4 .....510这些偶数位置使用的就是sin，而1，3.....511这些位置使用的就是cos

然后将字向量的512维度和位置编码的512个维度相加，最终就得到了transformer的输入。

代码实现

class Position_encoding(nn.Module):
    # dropout减少过拟合，增强鲁棒性，max_len为数字序列的最大长度
    def __init__(self,d_module,dropout=0.1,max_len=5000):
        super(Position_encoding,self).__init__()
        self.dropout=nn.Dropout(p=dropout)

        # pe是输入序列
        pe=torch.zeros(max_len,d_module)
        # 位置索引张量
        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))
        # 偶数位置编码
        position[:,0::2]=torch.sin(position*div_term)
        # 奇数位置编码
        position[:,1::2]=torch.cos(position*div_term)
        # 将位置编码的维度与输入序列的维度对齐
        pe=pe.unsqueeze(0).transpose(0,1)
        # buffer通常用于存储那些不需要在训练过程中更新的数据
        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)

用于生成并添加位置编码到输入序列的嵌入中，接收的参数为d_module 

position=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1)

max_len 为了便于演示，我这里取12，实际上可能是512，即输入序列的维度

div_term=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))

对应的数学公式如下：

 

multi-head attention 多头注意力 

原理

注意力机制（Attention Mechanism）是一种模拟人类注意力分配的计算模型，用于在机器学习和深度学习中提高模型性能。它源于对人类视觉的研究，并模仿了人类在处理信息时选择性地关注重要部分、忽略其他信息的机制。注意力机制的核心思想是让模型能够有选择性地关注输入序列中的不同部分，为输入序列的各个部分分配不同的权重，以此来突出对任务更关键的信息。

让我们举个栗子：

 在这张图片中，颜色越深，代表人们的注意力关注的越多。

当我们看这个图片的时候，我们会先注意到婴儿的脸，然后是标题和首个句子。

下面是公式注意力机制的公式：

 首先是Q,K,V三个矩阵，分别是query查询矩阵，key键矩阵，value值矩阵。那么他们是怎么得到的呢？通过对输入序列的每个元素（例如，词嵌入向量）分别与权重矩阵WQ,WK,WV进行线性变换得到。

接下来是Q*K_T,他们是在干嘛呢，他们其实是在判断查询矩阵和键矩阵的相似性，做了点积。

点积的公式：

cos的图像我们也知道，这两个矩阵越相似，点积结果越大。

为什么会有除呢？我们要知道，这两个点积，他们的结果可能很大，然后进行softmax时，他的梯度就可能很小，造成梯度消失。之所以是除以根号d_k是为了保证方差为1

然后进行softmax归一化，得到概率值， 也就是相似度。然后和value相乘，最终相加，得到加权和。

上面就是注意力机制。理解了注意力机制之后，多头注意力就很好理解了。

在多头注意力中，我们会用多个权重矩阵，主要有以下作用。

一、增强模型表达能力

1. 捕捉不同子空间信息：每个权重矩阵都代表了一个不同的线性变换，这些变换能够将输入数据投射到不同的表示子空间上。这样，每个注意力头都能够学习到输入数据的不同方面，捕捉到更丰富的特征信息。
2. 多样化特征提取：不同的权重矩阵允许每个注意力头专注于输入数据的不同特征，如语义信息、句法结构等。这种多样化特征提取方式有助于模型更全面地理解输入数据。

二、提高计算效率和并行性

1. 并行计算：多头注意力机制通过并行处理多个注意力头，可以同时计算多个子空间的注意力权重，从而显著提高计算效率。
2. 减少计算复杂度：虽然使用了多个权重矩阵，但每个权重矩阵的维度通常较小，这使得单个头的计算量相对较小。同时，由于并行计算的优势，整体计算效率得以提升。

三、优化模型训练

1. 权重共享与独立：虽然每个注意力头都使用独立的权重矩阵，但在某些实现中，这些权重矩阵的某些部分可能会共享参数，以减少模型参数量并避免过拟合。同时，独立的权重矩阵也使得每个头能够学习到更加独特和专注的特征。
2. 灵活性：多个权重矩阵的引入增加了模型的灵活性，使得模型能够在训练过程中学习到更复杂的特征变换和关系。这种灵活性有助于模型更好地适应不同的任务和数据集。

四、特征融合与变换

1. 特征融合：在多头注意力机制的最后一步，通常会将多个头的输出进行拼接，并通过一个额外的权重矩阵进行线性变换。这个权重矩阵的作用是将不同头学习到的特征有效地融合在一起，形成更加丰富和全面的特征表示。
2. 维度控制：通过调整这个额外权重矩阵的维度，可以控制最终输出特征矩阵的大小，使其与输入特征矩阵保持一致或满足其他特定要求。

 代码

首先，我们先敲出多头注意力公式的代码

# ------------------------------------多头注意力公式-------------------------------------------
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)
        attn=nn.Softmax(dim=-1)(scores)
        # scores : [batch_size, n_heads, len_q, len_k] * V: [batch_size, n_heads, len_v(=len_k), d_v]
        # context: [batch_size, n_heads, len_q, d_v]
        context = torch.matmul(attn, V)
        return context,attn

这里的attn指的使注意力权重。 

然后是多头注意力整体的

# ------------------------------------MultiHeadAttention多头注意力-------------------------------
class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention,self).__init__()
        # 输入：seq_len*d_module
        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)
        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        self.fc=nn.Linear(d_v*n_heads,d_model,bias=False)

    def forword(self,input_Q,input_K,input_V,attn_mask):
        # input_Q: [batch_size, len_q, d_model]
        # 残差     批次
        residual,batch_size=input_Q,input_K.size(0)
        Q=self.W_Q(input_Q).view(batch_size,-1,n_heads,d_k).transpose(1,2)
        K=self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        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,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)
        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        output=self.fc(context)
        return nn.LayerNorm(d_model).to(device)(output+residual),attn

PoswiseFeedForwardNet前馈传播网络

很基础的一部分，所以我们这里直接上代码

# ---------------------------------PoswiseFeedForwardNet前馈传播网络--------------------------------------
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_module]
        residual=inputs
        outputs=self.fc(inputs)
        return nn.LayerNorm(d_model).to(device)(outputs+residual)

单层encoder

class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer,self).__init__()
        self.self_attention=MultiHeadAttention()
        self.feed_forward=PoswiseFeedForwardNet()
    def forward(self,enc_inputs,enc_outputs,enc_mask):
        # enc_inputs: [batch_size, source_len, d_model]
        # enc_mask: [batch_size, source_len, source_len] mask矩阵
        # 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_inputs to same Q, K, V（未线性变换前）,在 MultiHeadAttention中他们会被线性转换为不同的 Q,K,V
        enc_outputs,attn=self.self_attention(enc_inputs,enc_inputs,enc_inputs,enc_mask)
        enc_outputs=self.feed_forward(enc_outputs)
        return enc_outputs,attn

在单层的编码器中，我们需要将各部分按图所示组合起来：

就是将多头注意力和前馈传播网络组合起来。

一个注意力权重矩阵attn记录了每个头在每个位置上的注意力分布 ，我们将这部分代码与多头注意力部分联系起来就好懂了。

最终得到了在这一层编码器中的编码器输出和这部分的注意力权重。既然有了一层的，肯定会有多层的，但是我们不着急弄多层的，大家肯定对注意力掩码这部分有疑惑，我们趁此机会，先把注意力掩码的代码搞懂。

注意力机制的填充掩码padding mask

def get_attn_pad_mask(seq_q,seq_k):
    # seq_q: [batch_size, seq_len]
    # seq_k: [batch_size, seq_len]
    batch_size,len_q=seq_q.size()
    batch_size,len_k=seq_k.size()
    seq_k.data.eq(0)
    # 检查seq_k中的每个元素是否等于0（通常0表示填充）。这会返回一个布尔张量，其中True表示对应位置是填充。
    pad_attn_mask=seq_k.data.eq(0).unsquee(1)
    return pad_attn_mask.expand(batch_size,len_q,len_k)

在这里，我们的seq-q和seq-k只是用来表示两个可能不同的序列集合，它们的实际含义取决于函数被调用的上下文。

然后检查seq-k是否有0，返回布尔张量，然后增加一个维度，最后将它扩充成

(batch_size,len_q,len_k)

的形状，最终返回的是布尔张量

注意力掩码

def get_attn_subsequence_mask(seq):
    # seq: [batch_size, target_len]
    attn_shape=[seq.size(0),seq.size(1),seq.size(1)]
    # attn_shape: [batch_size, target_len,target_len]
    subsequence_mask=np.triu(np.ones(attn_shape),k=1)
    subsequence_mask=torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask

为什么我们要把attn-shape是[batch_size, target_len,target_len]不是[batch_size, target_len]呢

如果我们使用 [batch_size, target_len] 的形状，那么我们只能为每个序列位置存储一个单一的权重值，而不是一个与其他所有位置相关的权重向量。这将无法捕捉到序列中不同位置之间的相互作用，这是自注意力机制的核心功能之一。

知道了这个之后，我们就可以理解为什么是这个形状了：

第一个维度 batch_size 表示批次中的序列数量。

第二个维度 target_len 表示我们正在处理序列中的哪个位置（即查询位置）。

第三个维度 target_len（与第二个维度相同）表示我们正在计算与序列中哪个位置的注意力权重（即键位置）。

然后使用NumPy生成一个形状为attn_shape的全1矩阵，并通过np.triu函数将其转换为上三角矩阵（包括对角线以上的一行）。k=1表示对角线以上的第一行及以上的元素为1，对角线及以下的元素为0。

最后将NumPy数组转换为PyTorch张量，并使用.byte()方法将其数据类型转换为torch.uint8（即布尔类型在PyTorch中的表示）。

encoder整体

指的是下图部分：

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder,self).__init__()
        self.source_embedding=nn.Embedding(source_size,d_model)
        # transformer中的位置编码是固定的，不需要学习
        self.pos_embedding=Position_encoding(d_model)
        # 初始化一个编码器层（EncoderLayer）的堆叠
        self.layers=nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self,encoder_inputs):
        # encoder_inputs: [batch_size, source_len]

        # 该层会为每个索引查找对应的嵌入向量。由于每个序列有source_len个词汇或子词，并且每个词汇或子词都被转换为一个d_model维的向量，
        # 因此输出的形状就是[batch_size, source_len, d_model]。
        enc_outputs=self.source_embedding(encoder_inputs)  # [batch_size, source_len, d_model]

        #Position_encoding要求的输入形状为[seq_len, batch_size, d_model]，所以我们要先转换为要求形状，然后再转换回来。
        enc_outputs=self.pos_embedding(enc_outputs.transpose(0,1)).transpose(0,1)  # [batch_size, source_len, d_model]

        enc_self_attn_pad_mask=get_attn_pad_mask(encoder_inputs,enc_outputs) # [batch_size, source_len, source_len]

        enc_self_attns=[] # 在计算中不需要用到，它主要用来保存你接下来返回的attention的值（这个主要是为了你画热力图等，用来看各个词之间的关系

        for layer in self.layers:
            # 上一个block的输出enc_outputs作为当前block的输入
            # enc_outputs: [batch_size, source_len, d_model], enc_self_attn: [batch_size, n_heads, source_len, source_len]
            # 传入的enc_outputs其实是input，传入mask矩阵是因为要做self attention
            enc_outputs,enc_self_attn=layer(enc_outputs,enc_self_attn_pad_mask)

            enc_self_attns.append(enc_self_attn) #可视化用
        return enc_outputs,enc_self_attns

ok，现在encoder部分的代码就差不多弄好了

下面就是decoder部分啦

decoder

让我们回忆一下decoder的结构

单层decoder

代码

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

    # 解码器自注意力掩码dec_self_attn_mask
    def forward(self,dec_inputs,enc_outputs,dec_self_attn_mask,dec_enco_attn_mask):
        # dec_inputs: [batch_size, target_len, d_model]
        # enc_outputs: [batch_size, source_len, d_model]
        # dec_self_attn_mask: [batch_size, target_len, target_len]
        # dec_enco_attn_mask: [batch_size, target_len, source_len]

        # dec_outputs: [batch_size, target_len, d_model], decoder_self_attn: [batch_size, n_heads, target_len, target_len]
        # 这里的Q,K,V全是Decoder自己的输入
        dec_outputs,dec_self_attn=self.dec_self_attn(dec_inputs,dec_inputs,dec_inputs,dec_enco_attn_mask)

        # dec_outputs: [batch_size, target_len, d_model], dec_enc_attn: [batch_size, h_heads, target_len, source_len]
        # Attention层的Q(来自decoder) 和 K,V(来自encoder)
        dec_outputs,dec_enc_attn=self.dec_enco_attn(dec_outputs,enc_outputs,enc_outputs,dec_enco_attn_mask)

        # [batch_size, target_len, d_model]
        dec_outputs=self.pos_ffn(dec_outputs)
        # dec_self_attn, dec_enc_attn这两个是为了可视化的
        return dec_outputs,dec_self_attn,dec_enc_attn

 解析

这里，dec_self_attn指的是解码器自注意力机制，而dec_enco_attn指的是编码器-解码器注意力机制

自注意力机制允许解码器中的每个位置都能关注到输入序列的所有位置（包括它自己）。

编码器-解码器注意力机制允许解码器中的每个位置都能关注到编码器输出的所有位置，这是理解输入序列并生成相应输出序列的关键。

dec_outputs,dec_self_attn=self.dec_self_attn(dec_inputs,dec_inputs,dec_inputs,dec_enco_attn_mask)

dec_outputs,dec_enc_attn=self.dec_enco_attn(dec_outputs,enc_outputs,enc_outputs,dec_enco_attn_mask)

我在学习代码的过程中被这两行迷惑了，就感觉有点迷，所以这里放一下它们两个的具体解析

第一行是通过解码器自注意力机制（self.dec_self_attn）处理解码器的输入，得到处理后的输出dec_outputs和自注意力权重dec_self_attn。这里的查询（Q）、键（K）、值（V）都来自解码器的输入。对应图中的Msked Muti-head attention

第二行是通过编码器-解码器注意力机制（self.dec_enc_attn）处理上一步的输出，得到进一步处理后的输出dec_outputs和编码器-解码器注意力权重dec_enc_attn。这里的查询（Q）来自解码器，而键（K）和值（V）来自编码器的输出。对应着Muti-head attention。

最后通过前馈传播部分，得到最终的dec-output。

decoder整体

代码

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder,self).__init__()
        self.target_embedding=nn.Embedding(target_size,d_model)
        self.pos_embedding=Position_encoding(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_inputs: [batch_size, src_len]
        # enc_outputs: [batch_size, src_len, d_model]  # 用在Encoder-Decoder Attention层
        dec_outputs=self.target_embedding(dec_inputs)  # [batch_size,target_len,d_model]
        dec_outputs=self.pos_embedding(dec_outputs.transpose(0,1)).transpose(0,1)    # [batch_size,target_len,d_model]

        # Decoder输入序列的pad mask矩阵
        dec_self_attn_pad_mask=get_attn_pad_mask(dec_inputs,dec_inputs).to(device) # [batch_size,target_len,target_len]

        # Masked self_Attention 当前时刻是看不到未来的信息的
        dec_self_attn_subsequence_mask=get_attn_subsequence_mask(dec_inputs).to(device) # [batch_size,target_len,target_len]

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

        # 这个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) # [batch_size,target_len,source_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

解析 

首先，通过目标嵌入层和位置编码层处理解码器输入。

然后，生成两种mask：

dec_self_attn_pad_mask：用于屏蔽解码器自注意力层中的pad位置。

dec_self_attn_subsequence_mask：用于确保解码器在生成每个词时只能关注到当前词之前的词（即屏蔽未来信息）。

将这两种mask相加并通过torch.gt转换为布尔mask，其中True表示需要关注的位置，False表示需要屏蔽的位置。

dec_enc_attn_mask：用于编码器-解码器注意力层，屏蔽编码器输出中的pad位置。

遍历解码器层，每层都会处理输入并返回解码器自注意力、编码器-解码器注意力的输出以及解码器的输出。

dec_outputs,dec_self_attn,dec_enc_attn=layer(dec_outputs,enc_outputs,dec_self_attn_mask,dec_enc_attn_mask)

我相信不只我一个人对为什么用这些参数有疑问吧~让我们来看看传递的这些参数是怎么用的。

将这些参数传递给解码器层后，该层会执行以下操作：

• 应用自注意力机制（可能带有掩码）来处理 dec_outputs，生成新的表示。
• 应用编码器-解码器注意力机制（也可能带有掩码）来结合 enc_outputs 和自注意力的输出，进一步生成新的表示。
• 通过一个或多个前馈神经网络层来处理这些表示，最终输出更新后的 dec_outputs，以及可选的自注意力权重 dec_self_attn 和编码器-解码器注意力权重 dec_enc_attn。

这些输出随后会被传递给下一层解码器（如果有的话），或者作为解码器的最终输出（对于最后一层而言）。通过这种方式，解码器能够逐步构建目标序列的表示，并在每个步骤中利用编码器的输出和先前的目标词信息。

到这里，transformer模型的各个部分就组装好了。

下面就是整体的transformer的整体代码了

transformer整体

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,target_size,bias=False).to(device)
    def forward(self,enc_inputs,dec_inputs):
        # Transformers的输入：两个序列
        # enc_inputs: [batch_size, source_len]
        # dec_inputs: [batch_size, target_len]
        # tensor to store decoder outputs
        # outputs = torch.zeros(batch_size, target_len, tgt_vocab_size).to(self.device)

        # enc_outputs: [batch_size, source_len, d_model], enc_self_attns: [n_layers, batch_size, n_heads, source_len,source_len]
        # 经过Encoder网络后，得到的输出还是[batch_size, source_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

 完整模型

import math
import numpy as np
import torch
from torch.utils.data import Dataset,DataLoader
from torch import nn

# --------------------------手动代码构造句子----------------------------------

# 训练集
train_data=[
    ['我 有 一 个 好 朋 友 P', 'S I have a good friend .', 'I have a good friend . E'],
    ['我 有 零 个 女 朋 友 P', 'S I have zero girl friend .', 'I have zero girl friend . E'],
    ['我 有 一 个 男 朋 友 P', 'S I have a boy friend .', 'I have a boy friend . E']
]

# 建立源词库，用于将中文转化为数字序列
source= {'P': 0, '我': 1, '有': 2, '一': 3,
        '个': 4, '好': 5, '朋': 6, '友': 7, '零': 8, '女': 9, '男': 10}
# 字典推导式，i为索引(整数)，w为对应的中文字符或特殊符号
source_index={i:w for i,w in enumerate(source)}
source_size=len(source)

# 建立目标词库，用于将英文转化为数字序列，并在解码时将数字序列转为英语
target={'P': 0, 'I': 1, 'have': 2, 'a': 3, 'good': 4,
        'friend': 5, 'zero': 6, 'girl': 7, 'boy': 8, 'S': 9, 'E': 10, '.': 11}
target_index={i:w for i,w in enumerate(target)}
target_size=len(target)

source_len=8
target_len=7

device='cuda'

# ---------------------------------将数字序列转为tensor----------------------------

def to_tensor(train_data):
    enc_inputs,dec_inputs,dec_outputs=[],[],[]
    for i in range(len(train_data)):
        enc_input = [[source[n] for n in train_data[i][0].split()]]
        dec_input = [[target[n] for n in train_data[i][1].split()]]
        decr_output = [[target[n] for n in train_data[i][2].split()]]
        enc_inputs.extend(enc_input)
        dec_inputs.extend(dec_input)
        dec_outputs.extend(decr_output)

    # print("编码器输入的数字序列：{}".format(encoder_inputs))
    # print("解码器输入的数字序列：{}".format(decoder_inputs))
    # print("解码器输出的数字序列：{}".format(decoder_outputs))

    # print(torch.tensor(encoder_inputs))
    # print(torch.LongTensor(encoder_inputs))

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

enc_inputs,dec_inputs,dec_outputs=to_tensor(train_data)

# ----------------------------------自定义数据集----------------------------------

class Data_set(Dataset):
    def __init__(self,enc_inputs,dec_inputs,dec_outputs):
        super(Data_set,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, item):
        return self.enc_inputs[item],self.dec_inputs[item],self.dec_outputs[item]

data_loader=DataLoader(Data_set(enc_inputs,dec_inputs,dec_outputs),2,True)

# -------------------------------------------------transformer-------------------------------------------------------

# 参数设置
d_model = 512  # Embedding Size（token embedding和position编码的维度）
# FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_ff = 2048
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（有几个头）

# --------------------------------------------位置编码------------------------------------------
class Position_encoding(nn.Module):
    # dropout减少过拟合，增强鲁棒性，max_len为数字序列的最大长度
    def __init__(self,d_module,dropout=0.1,max_len=5000):
        super(Position_encoding,self).__init__()
        self.dropout=nn.Dropout(p=dropout)

        # pe是输入序列
        pe=torch.zeros(max_len,d_module)
        # 位置索引张量
        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))
        # 偶数位置编码
        position[:,0::2]=torch.sin(position*div_term)
        # 奇数位置编码
        position[:,1::2]=torch.cos(position*div_term)
        # 将位置编码的维度与输入序列的维度对齐
        pe=pe.unsqueeze(0).transpose(0,1)
        # buffer通常用于存储那些不需要在训练过程中更新的数据
        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)
# ------------------------------------多头注意力公式-------------------------------------------
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)
        attn=nn.Softmax(dim=-1)(scores)
        # scores : [batch_size, n_heads, len_q, len_k] * V: [batch_size, n_heads, len_v(=len_k), d_v]
        # context: [batch_size, n_heads, len_q, d_v]
        context = torch.matmul(attn, V)
        return context,attn
 # ------------------------------------MultiHeadAttention多头注意力-------------------------------
class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention,self).__init__()
        # 输入：seq_len*d_module
        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)
        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        self.fc=nn.Linear(d_v*n_heads,d_model,bias=False)

    def forword(self,input_Q,input_K,input_V,attn_mask):
        # input_Q: [batch_size, len_q, d_model]
        # 残差     批次
        residual,batch_size=input_Q,input_K.size(0)
        Q=self.W_Q(input_Q).view(batch_size,-1,n_heads,d_k).transpose(1,2)
        K=self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        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,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)
        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        output=self.fc(context)
        return nn.LayerNorm(d_model).to(device)(output+residual),attn

# ---------------------------------PoswiseFeedForwardNet前馈传播网络--------------------------------------
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_module]
        residual=inputs
        outputs=self.fc(inputs)
        return nn.LayerNorm(d_model).to(device)(outputs+residual)

# --------------------------------------EncoderLayer--------------------------------
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer,self).__init__()
        self.self_attention=MultiHeadAttention()
        self.feed_forward=PoswiseFeedForwardNet()
    def forward(self,enc_inputs,enc_outputs,enc_mask):
        # enc_inputs: [batch_size, source_len, d_model]
        # enc_mask: [batch_size, source_len, source_len] mask矩阵
        # 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_inputs to same Q, K, V（未线性变换前）,在 MultiHeadAttention中他们会被线性转换为不同的 Q,K,V
        enc_outputs,attn=self.self_attention(enc_inputs,enc_inputs,enc_inputs,enc_mask)
        enc_outputs=self.feed_forward(enc_outputs)
        return enc_outputs,attn


# ----------------------------------------注意力机制的填充掩码-------------------------------------

def get_attn_pad_mask(seq_q,seq_k):
    # seq_q: [batch_size, seq_len]
    # seq_k: [batch_size, seq_len]
    batch_size,len_q=seq_q.size()
    batch_size,len_k=seq_k.size()
    # 检查seq_k中的每个元素是否等于0（通常0表示填充）。这会返回一个布尔张量，其中True表示对应位置是填充。
    pad_attn_mask=seq_k.data.eq(0).unsquee(1)
    return pad_attn_mask.expand(batch_size,len_q,len_k)

# ----------------------------------------注意力机制的掩码-------------------------------------

def get_attn_subsequence_mask(seq):
    # seq: [batch_size, target_len]
    attn_shape=[seq.size(0),seq.size(1),seq.size(1)]
    # attn_shape: [batch_size, target_len,target_len]
    subsequence_mask=np.triu(np.ones(attn_shape),k=1)
    subsequence_mask=torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask

# -------------------------------------------编码层--------------------------------------
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder,self).__init__()
        self.source_embedding=nn.Embedding(source_size,d_model)
        # transformer中的位置编码是固定的，不需要学习
        self.pos_embedding=Position_encoding(d_model)
        # 初始化一个编码器层（EncoderLayer）的堆叠
        self.layers=nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self,encoder_inputs):
        # encoder_inputs: [batch_size, source_len]

        # 该层会为每个索引查找对应的嵌入向量。由于每个序列有source_len个词汇或子词，并且每个词汇或子词都被转换为一个d_model维的向量，
        # 因此输出的形状就是[batch_size, source_len, d_model]。
        enc_outputs=self.source_embedding(encoder_inputs)  # [batch_size, source_len, d_model]

        #Position_encoding要求的输入形状为[seq_len, batch_size, d_model]，所以我们要先转换为要求形状，然后再转换回来。
        enc_outputs=self.pos_embedding(enc_outputs.transpose(0,1)).transpose(0,1)  # [batch_size, source_len, d_model]

        enc_self_attn_pad_mask=get_attn_pad_mask(encoder_inputs,enc_outputs) # [batch_size, source_len, source_len]

        enc_self_attns=[] # 在计算中不需要用到，它主要用来保存你接下来返回的attention的值（这个主要是为了你画热力图等，用来看各个词之间的关系

        for layer in self.layers:
            # 上一个block的输出enc_outputs作为当前block的输入
            # enc_outputs: [batch_size, source_len, d_model], enc_self_attn: [batch_size, n_heads, source_len, source_len]
            # 传入的enc_outputs其实是input，传入mask矩阵是因为要做self attention
            enc_outputs,enc_self_attn=layer(enc_outputs,enc_self_attn_pad_mask)

            enc_self_attns.append(enc_self_attn) #可视化用
        return enc_outputs,enc_self_attns


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

    # 解码器自注意力掩码dec_self_attn_mask
    def forward(self,dec_inputs,enc_outputs,dec_self_attn_mask,dec_enco_attn_mask):
        # dec_inputs: [batch_size, target_len, d_model]
        # enc_outputs: [batch_size, source_len, d_model]
        # dec_self_attn_mask: [batch_size, target_len, target_len]
        # dec_enco_attn_mask: [batch_size, target_len, source_len]

        # dec_outputs: [batch_size, target_len, d_model], decoder_self_attn: [batch_size, n_heads, target_len, target_len]
        # 这里的Q,K,V全是Decoder自己的输入
        dec_outputs,dec_self_attn=self.dec_self_attn(dec_inputs,dec_inputs,dec_inputs,dec_enco_attn_mask)

        # dec_outputs: [batch_size, target_len, d_model], dec_enc_attn: [batch_size, h_heads, target_len, source_len]
        # Attention层的Q(来自decoder) 和 K,V(来自encoder)
        dec_outputs,dec_enc_attn=self.dec_enco_attn(dec_outputs,enc_outputs,enc_outputs,dec_enco_attn_mask)

        # [batch_size, target_len, d_model]
        dec_outputs=self.pos_ffn(dec_outputs)
        # dec_self_attn, dec_enc_attn这两个是为了可视化的
        return dec_outputs,dec_self_attn,dec_enc_attn


# ---------------------------------------------------解码层-------------------------------------------------
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder,self).__init__()
        self.target_embedding=nn.Embedding(target_size,d_model)
        self.pos_embedding=Position_encoding(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_inputs: [batch_size, src_len]
        # enc_outputs: [batch_size, src_len, d_model]  # 用在Encoder-Decoder Attention层
        dec_outputs=self.target_embedding(dec_inputs)  # [batch_size,target_len,d_model]
        dec_outputs=self.pos_embedding(dec_outputs.transpose(0,1)).transpose(0,1)    # [batch_size,target_len,d_model]

        # Decoder输入序列的pad mask矩阵
        dec_self_attn_pad_mask=get_attn_pad_mask(dec_inputs,dec_inputs).to(device) # [batch_size,target_len,target_len]

        # Masked self_Attention 当前时刻是看不到未来的信息的
        dec_self_attn_subsequence_mask=get_attn_subsequence_mask(dec_inputs).to(device) # [batch_size,target_len,target_len]

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

        # 这个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) # [batch_size,target_len,source_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,target_size,bias=False).to(device)
    def forward(self,enc_inputs,dec_inputs):
        # Transformers的输入：两个序列
        # enc_inputs: [batch_size, source_len]
        # dec_inputs: [batch_size, target_len]
        # tensor to store decoder outputs
        # outputs = torch.zeros(batch_size, target_len, tgt_vocab_size).to(self.device)

        # enc_outputs: [batch_size, source_len, d_model], enc_self_attns: [n_layers, batch_size, n_heads, source_len,source_len]
        # 经过Encoder网络后，得到的输出还是[batch_size, source_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

应用

 例子及代码的出处：Transformer 代码详解(Pytorch版)_transformer pytorch-****博客

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 = 'cuda'

# transformer epochs
epochs = 100

# S: 代表解码器输入序列的开始
# E: 代表解码器输出序列的结束
# P: 这个符号用于填充。在批处理中，如果某个批次中的句子长度不一致，但模型需要固定长度的输入或输出，那么较短的句子会使用这个符号来填充到所需的长度。
# 这有助于模型在训练时保持输入/输出序列的维度一致。

# 训练集
sentences = [
    # 中文和英语的单词个数不要求相同
    # 中文句子的编码输入              解码器的输入序列           期望的解码器输出序列
    ['我 有 一 个 好 朋 友 P', 'S I have a good friend .', 'I have a good friend . E'],
    ['我 有 零 个 女 朋 友 P', 'S I have zero girl friend .', 'I have zero girl friend . E'],
    ['我 有 一 个 男 朋 友 P', 'S I have a boy friend .', 'I have a boy friend . E']
]

# Padding Should be Zero
src_vocab = {'P': 0, '我': 1, '有': 2, '一': 3,
             '个': 4, '好': 5, '朋': 6, '友': 7, '零': 8, '女': 9, '男': 10}
src_idx2word = {i: w for i, w in enumerate(src_vocab)}
src_vocab_size = len(src_vocab)

tgt_vocab = {'P': 0, 'I': 1, 'have': 2, 'a': 3, 'good': 4,
             'friend': 5, 'zero': 6, 'girl': 7, 'boy': 8, 'S': 9, 'E': 10, '.': 11}
idx2word = {i: w for i, w in enumerate(tgt_vocab)}
tgt_vocab_size = len(tgt_vocab)

src_len = 8  # （源句子的长度）enc_input max sequence length
tgt_len = 7  # dec_input(=dec_output) max sequence length

# Transformer Parameters
d_model = 512  # Embedding Size（token embedding和position编码的维度）
# FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_ff = 2048
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()]]
        dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]]
        dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]]

        # [[1, 2, 3, 4, 5, 6, 7, 0], [1, 2, 8, 4, 9, 6, 7, 0], [1, 2, 3, 4, 10, 6, 7, 0]]
        enc_inputs.extend(enc_input)
        # [[9, 1, 2, 3, 4, 5, 11], [9, 1, 2, 6, 7, 5, 11], [9, 1, 2, 3, 8, 5, 11]]
        dec_inputs.extend(dec_input)
        # [[1, 2, 3, 4, 5, 11, 10], [1, 2, 6, 7, 5, 11, 10], [1, 2, 3, 8, 5, 11, 10]]
        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)


# ----------------------------------------注意力机制的填充（padding）掩码（mask）-------------------------------------

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]]
    # [batch_size, 1, len_k], True is masked
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)
    # [batch_size, len_q, len_k] 构成一个立方体(batch_size个这样的矩阵)
    return pad_attn_mask.expand(batch_size, len_q, len_k)


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位置相对应的元素）
        # Fills elements of self tensor with value where mask is True.
        scores.masked_fill_(attn_mask, -1e9)

        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: [batch_size, n_heads, len_q, d_v]
        context = torch.matmul(attn, V)
        # context：[[z1,z2,...],[...]]向量, attn注意力稀疏矩阵（用于可视化的）
        return context, attn


class MultiHeadAttention(nn.Module):
    """这个Attention类可以实现:
    Encoder的Self-Attention
    Decoder的Masked Self-Attention
    Encoder-Decoder的Attention
    输入：seq_len x d_model
    输出：seq_len x d_model
    """

    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)
        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        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)

        # 这个全连接层可以保证多头attention的输出仍然是seq_len x d_model
        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)
        # [batch_size, seq_len, d_model]
        return nn.LayerNorm(d_model).to(device)(output + residual)


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)
        # [batch_size, tgt_len, d_model]
        dec_outputs = self.pos_ffn(dec_outputs)
        # dec_self_attn, dec_enc_attn这两个是为了可视化的
        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)  # 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