Transformer架构详解

Google 2017年论文Attention is all you need提出了Transformer模型，完全基于Attention mechanism，抛弃了传统的CNN和RNN。

1. Transformer架构

Transformer

解释下这个结构图。首先，Transformer模型也是使用经典的encoder-decoder架构，由encoder和decoder两部分组成。

上图左侧用Nx框出来的，就是我们encoder的一层。encoder一共有6层这样的结构。

上图右侧用Nx框出来的，就是我们decoder的一层。decoder一共有6层这样的结构。

输入序列经过word embedding和positional embedding相加后，输入到encoder中。

输出序列经过word embedding和positional embedding相加后，输入到decoder中。

最后，decoder输出的结果，经过一个线性层，然后计算softmax。

2. Encoder

encoder由6层相同的层组成，每一层分别由两部分组成：

  • 第一部分是multi-head self-attention mechanism
  • 第二部分是position-wise feed-forward network，是一个全连接层。

两部分，都有一个残差连接(residual connection)，然后接着一个Layer Normalization。

3. Decoder

与encoder类似，decoder也是由6个相同层组成，每一个层包括以下3个部分：

  • 第一部分是multi-head self-attention mechanism
  • 第二部分是multi-head context-attention mechanism
  • 第三部分是position-wise feed-forward network

同样，上面三部分中每一部分，都有一个残差连接(residual connection)，后接着一个Layer Normalization。

4. Attention机制

Attention是指对于某个时刻的输出y，它在输入x上各个部分的注意力。这个注意力可以理解为权重。

attention机制有很多计算方式，下面是一张比较全面的表格：

image.png

seq2seq模型中，使用的是加性注意力(addtion attention)较多。

为什么这种attention叫做addtion attention呢？很简单，对于输入序列隐状态h_i和输出序列的隐状态s_t，它的处理方式很简单，直接合并为[s_t;h_i]

但是transformer模型使用的不是这种attention机制，使用的是另一种，叫做乘性注意力(multiplicative attention)。

那么这种乘性注意力机制是怎么样的呢？从上表中的公式也可以看出来：两个隐状态进行点积！

4.1 Self-attention是什么？

上面我们说的attention机制的时候，都会提到两个隐状态，分别是h_i和s_t，前者是输入序列第i个位置产生的隐状态，后者是输出序列在第t个位置产生的隐状态。

所谓self-attention实际上就是输出序列就是输入序列，因此计算自己的attention得分，就叫做self-attention!

4.2 Context-attention是什么？

context-attention是encoder和decoder之间的attention！，所以，也可以成为encoder-decoder attention！

不管是self-attention还是context-attention，它们计算attention分数的时候，可以选择很多方式，比如上面表中提到的：

  • additive attention
  • local-base
  • general
  • dot-product
  • scaled dot-product

那么Transformer模型，采用的是哪种呢？答案是：scaled dot-product attention。

4.3 Scaled dot-product attention是什么？

论文Attention is all you need里面对于attention机制的描述是这样的：

An attention function can be described as a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility of the query with the corresponding key.

这句话描述得很清楚了。翻译过来就是：通过确定Q和K之间的相似程度来选择V！

用公式来描述更加清晰：
Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V\tag{4.3.1}

scaled dot-product attention和dot-product attention唯一区别是，scaled dot-product attention有一个缩放因子\frac{1}{\sqrt{d_k}}。

上面公式中d_k表示的是K的维度，在论文中，默认是64。

那么为什么需要加上这个缩放因子呢？论文中给出了解释：对于d_k很大时，点积得到的结果维度很大，使得结果处理softmax函数梯度很小的区域。

我们知道，梯度很小时，这对反向传播不利。为了克服这个负面影响，除以一个缩放因子，在一定程度上减缓这种情况。

为什么是\frac{1}{\sqrt{d_k}}呢？论文没有进一步说明。个人觉得你可以使用其他缩放因子，看看模型效果有没有提升。

论文中也提供了一张很清晰的结果图，供大家参考：

image.png

首先说明一下我们的K、Q、V是什么：

  • 在encoder的self-attention中，Q、K、V都来自同一个地方（相等），他们是上一层encoder的输出。对于第一层encoder，它们就是word embedding和positional encoding相加得到的输入。

  • 在decoder的self-attention中，Q、K、V都来自同一个地方（相等），他们是上一层decoder的输出。对于第一层decoder，它们就是word embedding和positional encoding相加得到的输入。但是对于decoder，我们不希望它能获得下一个time step，因此我们需要进行sequence masking。

  • 在encoder-decoder attention中，Q来自于decoder的上一层的输出，K和V来自于encoder的输出，K和V是一样的。

  • Q、K、V三者的维度一样，即d_q=d_k=d_v。

4.4 Scaled dot-product attention代码实现

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    """
    Scaled dot-product attention mechanism.
    """

    def __init__(self, attention_dropout=0.0):
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(attention_dropout)
        self.softmax = nn.Softmax(dim=2)

    def forward(self, q, k, v, scale=None, attn_mask=None):
        """
        前向传播

        args:
            q: Queries张量，形状[B, L_q, D_q]
            k: keys张量， 形状[B, L_k, D_k]
            v: Values张量，形状[B, L_v, D_v]
            scale: 缩放因子，一个浮点标量
            attn_mask: Masking张量，形状[B, L_q, L_k]
        returns:
            上下文张量和attention张量
        """
        attention = torch.bmm(q, k.transpose(1, 2))
        if scale:
            attention = attention * scale
        if attn_mask:
            # 给需要mask的地方设置一个负无穷
            attention = attention.masked_fill_(attn_mask, -np.inf)
        # 计算softmax
        attention = self.softmax(attention)
        # 添加dropout
        attention = self.dropout(attention)
        # 和V做点积
        context = torch.bmm(attention, v)

        return context, attention

5. Multi-head attention是什么呢？

理解了Scaled dot-product attention，Multi-head attention也很简单了。论文提到，他们发现将Q、K、V通过一个线性映射之后，分成h份，对每一份进行scaled dot-product attention效果更好。然后，把各个部分的结果合并起来，再次经过线性映射，得到最终的输出。这就是所谓的multi-head attention。上面的超参数h就是heads数量。论文默认是8。

multi-head attention的结构图如下：


image.png

值得注意的是，上面所说的分成h份是在d_k、d_q、d_v维度上面进行切分的。因此，进入到scaled dot-product attention的d_k实际上等于未进入之前的D_k/h。

Multi-head attention允许模型加入不同位置的表示子空间的信息。

Multi-head attention的公式如下：
MultiHead(Q,K,V)=Concat(head_1,...,head_h)W^O\tag{5.1}

其中，
head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)\tag{5.2}

论文中，d_{model}=512, h=8。所以scaled dot-product attention里面的
d_q=d_k=d_v=d_{model}/h=512/8=64

5.1 Multi-head attention代码实现

class MultiHeadAttention(nn.Module):

    def __init__(self, model_dim=512, num_heads=8, dropout=0.0):
        super(MultiHeadAttention, self).__init__()

        self.dim_per_head = model_dim / num_heads
        self.num_heads = num_heads
        self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads)
        self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads)
        self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads)

        self.dot_product_attention = ScaledDotProductAttention(dropout)
        self.linear_final = nn.Linear(model_dim, model_dim)
        self.dropout = nn.Dropout(dropout)
        # multi-head attention之后需要做layer norm
        self.layer_num = nn.LayerNorm(model_dim)

    def forward(self, query, key, value, attn_mask=None):
        # 残差连接
        residual = query

        batch_size = key.size(0)

        # linear projection
        query = self.linear_q(query) # [B, L, D]
        key = self.linear_k(key) # [B, L, D]
        value = self.linear_v(value) # [B, L, D]

        # split by head
        query = query.view(batch_size * num_heads, -1, dim_per_head) # [B * 8, , D / 8]
        key = key.view(batch_size * num_heads, -1, dim_per_head) # 
        value = value.view(batch_size * num_heads, -1, dim_per_head)

        if attn_mask:
            attn_mask = attn_mask.repeat(num_heads, 1, 1)
        # scaled dot product attention
        scale = (key.size(-1) // num_heads) ** -0.5
        context, attention = self.dot_product_attention(
            query, key, value, scale, attn_mask
        ) 

        # concat heads
        context = context.view(batch_size, -1, dim_per_head * num_heads)
        
        # final linear projection
        output = self.linear_final(context)

        # dropout
        output = self.dropout(output)

        # add residual and norm layer
        output = self.layer_num(residual + output)

        return output, attention

上面代码中出现了 Residual connection和Layer normalization。下面进行解释：

5.1.1 Residual connection是什么？

残差连接其实比较简单！看图就会比较清晰：

image.png

假设网络中某个层对输入x作用后的输出为F(x)，那么增加residual connection之后，变成：

F(x) + x \tag{5.2.1}

这个+x操作被称为shotcut。

残差结构因为增加了一项x，该层网络对x求偏导时，为常数项1！所以可以在反向传播过程中，梯度连乘，不会造成梯度消失！

5.1.2 Layer normalization是什么？

归一化层，主要有这几种方法，BatchNorm（2015年）、LayerNorm（2016年）、InstanceNorm（2016年）、GroupNorm（2018年）；
将输入的图像shape记为[N,C,H,W]，这几个方法主要区别是：

  • BatchNorm：batch方向做归一化，计算NHW的均值，对小batchsize效果不好；(BN主要缺点是对batchsize的大小比较敏感，由于每次计算均值和方差是在一个batch上，所以如果batchsize太小，则计算的均值、方差不足以代表整个数据分布)

  • LayerNorm：channel方向做归一化，计算CHW的均值；(对RNN作用明显)

  • InstanceNorm：一个batch，一个channel内做归一化。计算HW的均值，用在风格化迁移；(因为在图像风格化中，生成结果主要依赖于某个图像实例，所以对整个batch归一化不适合图像风格化中，因而对HW做归一化。可以加速模型收敛，并且保持每个图像实例之间的独立。)

  • GroupNorm：将channel方向分group，然后每个group内做归一化，算(C//G)HW的均值；这样与batchsize无关，不受其约束。

Normalization layers

6. Mask是什么？

mask顾名思义就是掩码，大概意思是对某些值进行掩盖，使其不产生效果.

需要说明的是，Transformer模型中有两种mask。分别是padding mask和sequence mask。其中，padding mask在所有的scaled dot-product attention里都需要用到，而sequence mask只在decoder的self-attention中用到。

所以，我们之前的ScaledDotProductAttention的forward方法里的参数attn_mask在不同的地方有不同的含义。

6.1 Padding mask

什么是padding mask呢？回想一下，我们的每个批次输入序列长度是不一样的！也就是说，我们要对输入序列进行对齐！具体来说，就是给较短序列后面填充0。因为这些填充位置，其实没有意义，所以我们的attention机制不应该把注意力放在这些位置上，所以我们需要进行一些处理。

具体做法是：把这些位置的值加上一个非常大的负数(可以是负无穷)，这样的话，经过softmax，这些位置的概率就会接近0。

而我们的padding mask实际上是一个张量，每个值都是一个Boolean，值为False的地方就是我们要进行处理的地方。

下面是代码实现：

def padding_mask(seq_q, seq_k):
    # seq_k和seq_q的形状都是[B,L]
    len_q = seq_q.size(1)
    # `PAD` is 0
    pad_mask = seq_k.eq(0)
    pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1) # shape [B,L_q,L_k]

[B,L]->[B,1,L]->[B,L,L]

F F T T
F F T T
F F T T
F F T T

6.2 Sequence mask

sequence mask是为了使得decoder不能看到未来的信息。也就是对于一个序列，在time step为t的时刻，我们的解码输出只能依赖于t时刻之前的输出，而不能依赖t之后的输出。因此我们需要想一个办法，把t之后的信息给隐藏起来。

那具体如何做呢？也很简单：产生一个上三角矩阵，上三角矩阵的值全为1，下三角的值全为0，对角线值也为0。把这个矩阵作用在每一个序列上，就可以达到我们的目的。

具体代码如下：

def sequence_mask(seq):
    batch_size, seq_len = seq.size()
    mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8),
                    diagonal=1)
    mask = mask.unsqueeze(0).expand(batch_size, -1, -1)  # [B, L, L]
    return mask

[B,L,L]

0 1 1 1
0 0 1 1
0 0 0 1
0 0 0 0

哈佛大学的文章The Annotated Transformer有一张效果图:

image.png

值得注意的是，本来mask只需要二维矩阵即可，但是考虑到我们的输入序列都是批量的，所以我们需要把原本二维矩阵扩张成3维张量。上面代码中，已经做了处理。

回到本节开始的问题，attn_mask参数有几种情况？分别是什么意思？

  • 对于decoder的self-attention，里面使用的scaled dot-product attention，同时需要padding mask和sequence mask作为attn_mask，具体实现就是两个mask相加作为attn_mask。
  • 其它情况，attn_mask都等于padding mask。

7. Positional encoding是什么？

就目前而言，Transformer架构似乎少了点东西。没错，那就是它对序列的顺序没有约束！我们知道序列的顺序是一个很重要的信息，如果缺失了这个信息，可能我们的结果就是：所有词语都对了，但是无法组成有意义的语句。

为了解决这个问题，论文中提出了positional encoding。一句话概括就是：对序列中的词语出现的位置进行编码！如果对位置进行编码，那么我们的模型就可以捕捉顺序信息。

那么具体怎么做呢？论文的实现是使用正余弦函数。公式如下：
PF(pos,2i)=sin(pos/10000^{2i/d_{model}})\tag{7.1}

PF(pos,2i+1)=cos(pos/10000^{2i/d_{model}})\tag{7.2}

其中，pos是指词语在序列中的位置。可以看出，在偶数位置，使用正弦编码，在奇数位置，使用余弦编码。

上面公式中的d_{model}是模型的维度，论文默认是512。

这个编码公式的意思就是：给定词语的位置pos，我们可以把它编码成d_{model}维的向量！也就是说，位置编码的每一个维度对应正弦曲线，波长构成了从2\pi到10000*2\pi的等比序列。

Postional encoding是对词汇的位置编码。

7.1 Positional encoding代码实现

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, max_seq_len):
        """
        初始化

        args:
            d_model: 一个标量。模型的维度，论文默认是512
            max_seq_len: 一个标量。文本序列的最大长度
        """
        super(PositionalEncoding, self).__init__()

        # 根据论文给出的公式，构造出PE矩阵
        position_encoding = np.array([
            [pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)]
            for pos in range(max_seq_len)
        ])
        # 偶数列使用sin，奇数列使用cos
        position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2])
        position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2])

        # 在PE矩阵的一次行，加上一个全是0的向量，代表这`PAD`的positional_encoding
        # 在word embedding中也会经常加上`UNK`，代表位置单词的word embedding，两者十分类似
        # 那么为什么需要这个额外的PAD的编码呢？很简单，因为文本序列的长度不易，我们需要对齐，
        # 短的序列我们使用0在结尾不全，我们也需要这些补全位置的编码，也就是`PAD`对应的位置编码
        pad_row = torch.zeros([1, d_model])
        position_encoding = torch.cat((pad_row, position_encoding))

        # 嵌入操作，+1是因为增加了`PAD`这个补全位置的编码
        # word embedding中如果词典增加`UNK`，我们也需要+1。
        self.position_encoding = nn.Embedding(max_seq_len+1, d_model)
        self.position_encoding.weight = nn.Parameter(position_encoding, requires_grad=False)

    def forward(self, input_len):
        """
        神经网络前向传播

        args:
            input_len: 一个张量，形状为[BATCH_SIZE, 1]。每一个张量的值代表这一批文本序列中对应的长度。

        returns:
            返回这一批序列的位置编码，进行了对齐。
        """

        # 找出这一批序列的最大长度
        max_len = torch.max(input_len)
        # 对每一个序列的位置进行对齐，在原序列位置的后面补上0
        # 这里range从1开始也是因为要避开PAD(0)的位置
        input_pos = torch.LongTensor(
            [list(range(1, len+1)) + [0] * (max_len-len) for len in input_len]
        )
        return self.position_encoding(input_pos)

8. Word embedding是什么？

Word embedding是对序列中的词汇的编码，把每一个词汇编码成d_{model}维的向量！它实际上就是一个二维浮点矩阵，里面的权重是可训练参数，我们只需要把这个矩阵构建出来就完成了word embedding的工作。

embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0)

上面vocab_size是词典大小，embedding_size是词嵌入的维度大小，论文里面就是等于d_{model}=512。所以word embedding矩阵就是一个vocab_size*embedding_size的二维张量。

9. Position-wise Feed-Forward netword是什么？

这是一个全连接网络，包含连个线性变换和一个非线性函数（ReLU）。公式如下：
FFN(x)=max(0,xW_1+b_1)W2+b2\tag{9.1}

这个线性变换在不同的位置都是一样的，并且在不同的层之间使用不同的参数。

论文提到，这个公式还可以用两个核大小为1的一维卷积来解释，卷积的输入输出都是d_{model}=512，中间层维度是d_{ff}=2048。

代码如下：

class PositionalWiseFeedForward(nn.Module):

    def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0):
        super(PositionalWiseFeedForward, self).__init__()
        self.w1 = nn.Conv1d(model_dim, ffn_dim, 1)
        self.w2 = nn.Conv2d(model_dim, ffn_dim, 1)
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(model_dim)

    def forward(self, x):
        output = x.transpose(1, 2)
        output = self.w2(F.relu(self.w1(output)))
        output = self.dropout(output.transpose(1, 2))

        # add residual and norm layer
        output = self.layer_norm(x + output)
        return output

10. 完整代码

至此，所有的细节都解释完了。

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    """
    Scaled dot-product attention mechanism.
    """

    def __init__(self, attention_dropout=0.0):
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(attention_dropout)
        self.softmax = nn.Softmax(dim=2)

    def forward(self, q, k, v, scale=None, attn_mask=None):
        """
        前向传播

        args:
            q: Queries张量，形状[B, L_q, D_q]
            k: keys张量， 形状[B, L_k, D_k]
            v: Values张量，形状[B, L_v, D_v]
            scale: 缩放因子，一个浮点标量
            attn_mask: Masking张量，形状[B, L_q, L_k]
        returns:
            上下文张量和attention张量
        """
        attention = torch.bmm(q, k.transpose(1, 2))
        if scale:
            attention = attention * scale
        if attn_mask:
            # 给需要mask的地方设置一个负无穷
            attention = attention.masked_fill_(attn_mask, -np.inf)
        # 计算softmax
        attention = self.softmax(attention)
        # 添加dropout
        attention = self.dropout(attention)
        # 和V做点积
        context = torch.bmm(attention, v)

        return context, attention

class MultiHeadAttention(nn.Module):

    def __init__(self, model_dim=512, num_heads=8, dropout=0.0):
        super(MultiHeadAttention, self).__init__()

        self.dim_per_head = model_dim / num_heads
        self.num_heads = num_heads
        self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads)
        self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads)
        self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads)

        self.dot_product_attention = ScaledDotProductAttention(dropout)
        self.linear_final = nn.Linear(model_dim, model_dim)
        self.dropout = nn.Dropout(dropout)
        # multi-head attention之后需要做layer norm
        self.layer_num = nn.LayerNorm(model_dim)

    def forward(self, query, key, value, attn_mask=None):
        # 残差连接
        residual = query

        batch_size = key.size(0)

        # linear projection
        query = self.linear_q(query) # [B, L, D]
        key = self.linear_k(key) # [B, L, D]
        value = self.linear_v(value) # [B, L, D]

        # split by head
        query = query.view(batch_size * num_heads, -1, dim_per_head) # [B * 8, , D / 8]
        key = key.view(batch_size * num_heads, -1, dim_per_head) # 
        value = value.view(batch_size * num_heads, -1, dim_per_head)

        if attn_mask:
            attn_mask = attn_mask.repeat(num_heads, 1, 1)
        # scaled dot product attention
        scale = (key.size(-1) // num_heads) ** -0.5
        context, attention = self.dot_product_attention(
            query, key, value, scale, attn_mask
        ) 

        # concat heads
        context = context.view(batch_size, -1, dim_per_head * num_heads)
        
        # final linear projection
        output = self.linear_final(context)

        # dropout
        output = self.dropout(output)

        # add residual and norm layer
        output = self.layer_num(residual + output)

        return output, attention

def padding_mask(seq_q, seq_k):
    # seq_k和seq_q的形状都是[B,L]
    len_q = seq_q.size(1)
    # `PAD` is 0
    pad_mask = seq_k.eq(0)
    pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1) # shape [B,L_q,L_k]

def sequence_mask(seq):
    batch_size, seq_len = seq.size()
    mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8),
                    diagonal=1)
    mask = mask.unsqueeze(0).expand(batch_size, -1, -1)  # [B, L, L]
    return mask

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, max_seq_len):
        """
        初始化

        args:
            d_model: 一个标量。模型的维度，论文默认是512
            max_seq_len: 一个标量。文本序列的最大长度
        """
        super(PositionalEncoding, self).__init__()

        # 根据论文给出的公式，构造出PE矩阵
        position_encoding = np.array([
            [pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)]
            for pos in range(max_seq_len)
        ])
        # 偶数列使用sin，奇数列使用cos
        position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2])
        position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2])

        # 在PE矩阵的一次行，加上一个全是0的向量，代表这`PAD`的positional_encoding
        # 在word embedding中也会经常加上`UNK`，代表位置单词的word embedding，两者十分类似
        # 那么为什么需要这个额外的PAD的编码呢？很简单，因为文本序列的长度不易，我们需要对齐，
        # 短的序列我们使用0在结尾不全，我们也需要这些补全位置的编码，也就是`PAD`对应的位置编码
        pad_row = torch.zeros([1, d_model])
        position_encoding = torch.cat((pad_row, position_encoding))

        # 嵌入操作，+1是因为增加了`PAD`这个补全位置的编码
        # word embedding中如果词典增加`UNK`，我们也需要+1。
        self.position_encoding = nn.Embedding(max_seq_len+1, d_model)
        self.position_encoding.weight = nn.Parameter(position_encoding, requires_grad=False)

    def forward(self, input_len):
        """
        神经网络前向传播

        args:
            input_len: 一个张量，形状为[BATCH_SIZE, 1]。每一个张量的值代表这一批文本序列中对应的长度。

        returns:
            返回这一批序列的位置编码，进行了对齐。
        """

        # 找出这一批序列的最大长度
        max_len = torch.max(input_len)
        # 对每一个序列的位置进行对齐，在原序列位置的后面补上0
        # 这里range从1开始也是因为要避开PAD(0)的位置
        input_pos = torch.LongTensor(
            [list(range(1, len+1)) + [0] * (max_len-len) for len in input_len]
        )
        return self.position_encoding(input_pos)

# embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0)
# 获得输入的词嵌入编码
# seq_embedding = seq_embedding(inputs) * np.sqrt(d_model)

class PositionalWiseFeedForward(nn.Module):

    def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0):
        super(PositionalWiseFeedForward, self).__init__()
        self.w1 = nn.Conv1d(model_dim, ffn_dim, 1)
        self.w2 = nn.Conv2d(model_dim, ffn_dim, 1)
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(model_dim)

    def forward(self, x):
        output = x.transpose(1, 2)
        output = self.w2(F.relu(self.w1(output)))
        output = self.dropout(output.transpose(1, 2))

        # add residual and norm layer
        output = self.layer_norm(x + output)
        return output

class EncoderLayer(nn.Module):
    """Encoder的一层。"""
    def __init__(self, model_dim=512, num_heads=8, ffn_dim=2048, dropout=0.0):
        super(EncoderLayer, self).__init__()

        self.attention = MultiHeadAttention(model_dim, num_heads, dropout)
        self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)

    def forward(self, inputs, attn_mask=None):
        # self attention
        context, attention = self.attention(inputs, inputs, inputs, attn_mask)

        # feed forward network
        output = self.feed_forward(context)

        return output, attention


class Encoder(nn.Module):
    """多层EncoderLayer组成的Encoder"""
    def __init__(self,
                vocab_size,
                num_layers=6,
                model_dim=512,
                num_heads=8,
                ffn_dim=2048,
                dropout=0.0):
        super(Encoder, self).__init__()

        self.encoder_layers = nn.ModuleList(
            [EncoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in range(num_layers)]
        )

        self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)
        self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)

    def forward(self, inputs, inputs_len):
        output = self.seq_embedding(inputs)
        output += self.pos_embedding(inputs_len)

        self_attention_mask = padding_mask(inputs, inputs)

        attentions = []
        for encoder in self.encoder_layers:
            output, attention = encoder(output, self_attention_mask)
            attentions.append(attention)

        return output, attentions

class DecoderLayer(nn.Module):
    def __init__(self, model_dim, num_heads=8, ffn_dim=2048, dropout=0.0):
        super(DecoderLayer, self).__init__()

        self.attention = MultiHeadAttention(model_dim, num_heads, dropout)
        self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)

    def forward(self,
                dec_inputs,
                enc_outputs,
                self_attn_mask=None,
                context_attn_mask=None):
        # self attention, all inputs are decoder inputs
        dec_output, self_attention = self.attention(dec_inputs, dec_inputs, dec_inputs, self_attn_mask)

        # context attention
        # query is decoder's outputs, key and value are encoder's inputs
        dec_output, context_attention = self.attention(dec_output, enc_outputs, enc_outputs, context_attn_mask)

        # decoder's output, or context
        dec_output = self.feed_forward(dec_output)

        return dec_output, self_attention, context_attention

class Decoder(nn.Module):
    def __init__(self,
                vocab_size,
                max_seq_len,
                num_layers=6,
                model_dim=512,
                num_heads=8,
                ffn_dim=2048,
                dropout=0.0):
        super(Decoder).__init__()

        self.num_layers = num_layers

        self.decoder_layers = nn.ModuleList(
            [DecoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in range(num_layers)]
        )
        
        self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)
        self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)

    def forward(self, inputs, inputs_len, enc_output, context_attn_mask=None):
        output = self.seq_embedding(inputs)
        output += self.pos_embedding(inputs_len)

        self_attention_padding_mask = padding_mask(inputs, inputs)
        seq_mask = sequence_mask(inputs)
        self_attn_mask = torch.gt((self_attention_padding_mask + seq_mask), 0)

        self_attentions = []
        context_attentions = []
        for decoder in self.decoder_layers:
            output, self_attn, context_attn = decoder(
            output, enc_output, self_attn_mask, context_attn_mask)
            self_attentions.append(self_attn)
            context_attentions.append(context_attn)

        return output, self_attentions, context_attentions

    
class Transformer(nn.Module):
    def __init__(self,
                src_vocab_size,
                src_max_len,
                tgt_vocab_size,
                tgt_max_len,
                num_layers=6,
                model_dim=512,
                num_heads=8,
                ffn_dim=2048,
                dropout=0.0):
        super(Transformer).__init__()

        self.encoder = Encoder(src_vocab_size, src_max_len, num_layers, model_dim, num_heads, ffn_dim, dropout)
        self.decoder = Decoder(tgt_vocab_size, tgt_max_len, num_layers, model_dim, num_heads, ffn_dim, dropout)

        self.linear = nn.Linear(model_dim, tgt_vocab_size, bias=False)
        self.softmax = nn.Softmax()

    def forward(self, src_seq, src_len, tgt_seq, tgt_len):
        context_attn_mask = padding_mask(tgt_seq, src_seq)

        output, enc_self_attn = self.encoder(src_seq, src_len)

        output, dec_self_attn, ctx_attn = self.decoder(tgt_seq, tgt_len, output, context_attn_mask)

        output = self.linear(output)
        output = self.softmax(output)

        return output, enc_self_attn, dec_self_attn, ctx_attn
最后编辑于 ：
©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 206,214评论 6赞 481
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 88,307评论 2赞 382
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 152,543评论 0赞 341
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 55,221评论 1赞 279
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 64,224评论 5赞 371
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 49,007评论 1赞 284
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 38,313评论 3赞 399
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 36,956评论 0赞 259
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 43,441评论 1赞 300
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,925评论 2赞 323
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,018评论 1赞 333
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 33,685评论 4赞 322
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 39,234评论 3赞 307
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,240评论 0赞 19
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,464评论 1赞 261
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 45,467评论 2赞 352
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 42,762评论 2赞 345

推荐阅读更多精彩内容