各位朋友大家好,欢迎来到月来客栈。今天要和大家介绍的一篇论文是谷歌2017年所发表的一篇论文,名字叫做”Attention is all you need“[1]。以下为文章目录,大家可以快速定位到自己关注部分的内容。

1. 多头注意力机制原理

1.1 动机

虽然,网上已经有了大量的关于这篇论文的解析,不过好菜不怕晚笔者在这里也会谈谈自己对于它的理解以及运用。按照我们一贯解读论文的顺序,首先让我们先一起来看看作者当时为什么要提出Transformer这个模型?需要解决什么样的问题?现在的模型有什么样的缺陷?

1.1.1 面临问题

在论文的摘要部分作者提到,现在主流的序列模型都是基于复杂的循环神经网络或者是卷积神经网络构造而来的Encoder-Decoder模型,并且就算是目前性能最好的序列模型也都是基于注意力机制下的Encoder-Decoder架构。为什么作者会不停的提及这些传统的Encoder-Decoder模型呢?接着,作者在介绍部分谈到,由于传统的Encoder-Decoder架构在建模过程中,下一个时刻的计算过程会依赖于上一个时刻的输出,而这种固有的属性就限制了传统的Encoder-Decoder模型就不能以并行的方式进行计算,如图1-1所示。

This inherently sequential nature precludes parallelization within training examples, which becomes critical at longer sequence lengths, as memory constraints limit batching across examples.

4万字50余图3个实战示例一网打尽Transformer
::: hljs-center
图 1-1. 循环神经网络编码图
:::

随后作者谈到,尽管最新的研究工作已经能够使得传统的循环神经网络在计算效率上有了很大的提升,但是本质的问题依旧没有得到解决。

Recent work has achieved significant improvements in computational efficiency through factorization tricks and conditional computation, while also improving model performance in case of the latter. The fundamental constraint of sequential computation, however, remains.

1.1.2 解决思路

因此,在这篇论文中,作者首次提出了一种全新的Transformer架构来解决这一问题,如图1-2所示。当然,Transformer架构的优点在于它完全摈弃了传统的循环结构,取而代之的是只通过注意力机制来计算模型输入与输出的隐含表示,而这种注意力的名字就是大名鼎鼎的自注意力机制(self-attention),也就是图1-2中的Multi-Head Attention模块。

To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequence- aligned RNNs or convolution.

请添加链接描述
::: hljs-center
图 1-2. Transformer网络结构图
:::

总体来说,所谓自注意力机制就是通过某种运算来直接计算得到句子在编码过程中每个位置上的注意力权重;然后再以权重和的形式来计算得到整个句子的隐含向量表示。最终,Transformer架构就是基于这种的自注意力机制而构建的Encoder-Decoder模型。

1.2 技术手段

在介绍完整篇论文的提出背景后,下面就让我们一起首先来看一看自注意力机制的庐山真面目,然后再来探究整体的网络架构。

1.2.1 什么是self-Attention

首先需要明白一点的是,所谓的自注意力机制其实就是论文中所指代的“Scaled Dot-Product Attention“。在论文中作者说道,注意力机制可以描述为将query和一系列的key-value对映射到某个输出的过程,而这个输出的向量就是根据query和key计算得到的权重作用于value上的权重和。

An attention function can be described as mapping 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 function of the query with the corresponding key.

不过想要更加深入的理解query、key和value的含义,得需要结合Transformer的解码过程,这部分内容将会在后续进行介绍。 具体的,自注意力机制的结构如图1-3所示。

请添加链接描述

::: hljs-center
图 1-3. 自注意力机制结构图
:::

从图1-3可以看出,自注意力机制的核心过程就是通过Q和K计算得到注意力权重;然后再作用于V得到整个权重和输出。具体的,对于输入Q、K和V来说,其输出向量的计算公式为:
$$
\text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})V\;\;\;\;\;\;(1.1)
$$
其中Q、K和V分别为3个矩阵,且其(第2个)维度分别为$d_q,d_k,d_v$ (从后面的计算过程其实可以发现$d_q=d_v)。而公式$(1.1)​中除以$\sqrt{d_k}$的过程就是图1-3中所指的Scale过程。

之所以要进行缩放这一步是因为通过实验作者发现,对于较大的$d_k$来说在完成$QK^T$后将会得到很大的值,而这将导致在经过sofrmax操作后产生非常小的梯度,不利于网络的训练。

We suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients.

如果仅仅只是看着图1-3中的结构以及公式$(1.1)$中的计算过程显然是不那么容易理解自注意力机制的含义,例如初学者最困惑的一个问题就是图1-3中的Q、K和V分别是怎么来的?下面,我们来看一个实际的计算示例。现在,假设输入序列“我 是 谁”,且已经通过某种方式得到了1个形状为$3\times4$的矩阵来进行表示,那么通过图1-3所示的过程便能够就算得到Q、K以及V[2]。

请添加链接描述
::: hljs-center

图 1-4. Q、K和V计算过程图
:::

从图1-4的计算过程可以看出,Q、K和V其实就是输入X​分别乘以3个不同的矩阵计算而来(但这仅仅局限于Encoder和Decoder在各自输入部分利用自注意力机制进行编码的过程,Encoder和Decoder交互部分的Q、K和V另有指代)。此处对于计算得到的Q、K、V,你可以理解为这是对于同一个输入进行3次不同的线性变换来表示其不同的3种状态。在计算得到Q、K、V之后,就可以进一步计算得到权重向量,计算过程如图1-5所示。
请添加链接描述
::: hljs-center
图 1-5. 注意力权重计算图(已经经过scale和softmax操作)
:::

如图1-5所示,在经过上述过程计算得到了这个注意力权重矩阵之后我们不禁就会问到,这些权重值到底表示的是什么呢?对于权重矩阵的第1行来说,0.7表示的就是“我”与“我”的注意力值;0.2表示的就是“我”与”是”的注意力值;0.1表示的就是“我”与“谁”的注意力值。换句话说,在对序列中的“我“进行编码时,应该将0.7的注意力放在“我”上,0.2的注意力放在“是”上,将0.1的注意力放在谁上。

同理,对于权重矩阵的第3行来说,其表示的含义就是,在对序列中”谁“进行编码时,应该将0.2的注意力放在“我”上,将0.1的注意力放在“是”上,将0.7的注意力放在“谁”上。从这一过程可以看出,通过这个权重矩阵模型就能轻松的知道在编码对应位置上的向量时,应该以何种方式将注意力集中到不同的位置上。

不过从上面的计算结果还可以看到一点就是,模型在对当前位置的信息进行编码时,会过度的将注意力集中于自身的位置(虽然这符合常识)而可能忽略了其它位置[2]。因此,作者采取的一种解决方案就是采用多头注意力机制(MultiHeadAttention),这部分内容我们将在稍后看到。

It expands the model’s ability to focus on different positions. Yes, in the example above, $z_1$ contains a little bit of every other encoding, but it could be dominated by the the actual word itself.

在通过图1-5示的过程计算得到权重矩阵后,便可以将其作用于V ,进而得到最终的编码输出,计算过程如图1-6所示。

4万字50余图3个实战示例一网打尽Transformer
::: hljs-center
图 1-6. 权重和编码输出图
:::

根据如图1-6所示的过程,我们便能够得到最后编码后的输出向量。当然,对于上述过程我们还可以换个角度来进行观察,如图1-7所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-7. 编码输出计算图
:::

从图1-7可以看出,对于最终输出“是”的编码向量来说,它其实就是原始“我 是 谁”3个向量的加权和,而这也就体现了在对“是”进行编码时注意力权重分配的全过程。

当然,对于整个图1-5到图1-6的过程,我们还可以通过如图1-8所示的过程来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-8. 自注意力机制计算过程图
:::

可以看出通过这种自注意力机制的方式确实解决了作者在论文伊始所提出的“传统序列模型在编码过程中都需顺序进行的弊端”的问题,有了自注意力机制后,仅仅只需要对原始输入进行几次矩阵变换便能够得到最终包含有不同位置注意力信息的编码向量。

对于自注意力机制的核心部分到这里就介绍完了,不过里面依旧有很多细节之处没有进行介绍。例如Encoder和Decoder在进行交互时的Q、K、V是如何得到的?在图1-3中所标记的Mask操作是什么意思,什么情况下会用到等等?这些内容将会在后续逐一进行介绍。

下面,让我们继续进入到MultiHeadAttention机制的探索中。

1.2.2 为什么要MultiHeadAttention

经过上面内容的介绍,我们算是在一定程度上对于自注意力机制有了清晰的认识,不过在上面我们也提到了自注意力机制的缺陷就是:模型在对当前位置的信息进行编码时,会过度的将注意力集中于自身的位置,因此作者提出了通过多头注意力机制来解决这一问题。同时,使用多头注意力机制还能够给予注意力层的输出包含有不同子空间中的编码表示信息,从而增强模型的表达能力。

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.

在说完为什么需要多头注意力机制以及使用多头注意力机制的好处之后,下面我们就来看一看到底什么是多头注意力机制。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-9. 多头注意力机制结构图
:::

如图1-9所示,可以看到所谓的多头注意力机制其实就是将原始的输入序列进行多组的自注意力处理过程;然后再将每一组自注意力的结果拼接起来进行一次线性变换得到最终的输出结果。具体的,其计算公式为:
$$
\text{MultiHead}(Q,K,V)=\text{Concat}(\text{head}_1,...,\text{head}_h)W^O\
\;\;\;\;\;\;\;\text{where}\;\;\text{head}_i=\text{Attention}(QW_i^Q,KW_i^K,VW_i^V)\;\;\;\;\;\;\;\;\;(1.2)
$$
其中
$$
W^Qi\in\mathbb{R}^{d{model}\times d_k},W^Ki\in\mathbb{R}^{d{model}\times d_k},W^Vi\in\mathbb{R}^{d{model}\times d_v},W^O\in\mathbb{R}^{hdv\times d{model}}
$$
同时,在论文中,作者使用了$h=8$个并行的自注意力模块(8个头)来构建一个注意力层,并且对于每个自注意力模块都限定了$d_k=dv=d{model}/h=64$。从这里其实可以发现,论文中所使用的多头注意力机制其实就是将一个大的高维单头拆分成了$h$个多头。因此,整个多头注意力机制的计算过程我们可以通过如图1-10所示的过程来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-10. 多头注意力机制计算过程图
:::

注意:图中的dm​就是指$d{model}$

如图1-10所示,根据输入序列X和$W^Q_1,W^K_1,W^V_1$ 我们就计算得到了$Q_1,K_1,V_1$,进一步根据公式$(1.1)$就得到了单个自注意力模块的输出$Z_1$;同理,根据X和$W^Q_2,W^K_2,W^V_2$就得到了另外一个自注意力模块输出$Z_2$。最后,根据公式$(1.2)$将$Z_1,Z_2$水平堆叠形成$Z$,然后再用$Z$乘以$W^O$便得到了整个多头注意力层的输出。同时,根据图1-9中的计算过程,还可以得到$d_q=d_k=d_v$。

到此,对于整个Transformer的核心部分,即多头注意力机制的原理就介绍完了。

1.2.3 同维度中的单头与多头的区别

在多头注意力中,对于初学者来说一个比较经典的问题就是,在相同维度下使用单头和多头的区别是什么?这句话什么意思呢?以图1-10中示例为例,此时的自注意力中使用了两个头,每个头的维度为$d_q$,即采用了多头的方式。另外一种做法就是,只是用一个头,但是其维度为$2d_q$,即采用单头的方式。那么在这两种情况下有什么区别呢?

首先,从论文中内容可知,作者在头注意力机制与多头个数之间做了如下的限制
$$
d_q=d_k=dv=\frac{d{model}}{h}\;\;\;\;\;\;\;\;\;\;(1.3)
$$
从式$(1.3)$可以看出,单个头注意力机制的维度$dk$乘上多头的个数$h$就等于模型的维度$d{model}$。

注意:后续的d_m,$dm$以及$d{model}$都是指代模型的维度。

同时,从图1-10中可以看出,这里使用的多头数量$h=2$,即$d_{model}=2\times d_q$。此时,对于第1个头来说有:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-11. 头1注意力计算过程
:::

对于第2个头来说有:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-12. 头2注意力计算过程
:::

最后,可以将$Z_1,Z_2$在横向堆叠起来进行一个线性变换得到最终的$Z$。因此,对于图1-10所示的计算过程,我们还可以通过图1-13来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-13. 多头注意力合并计算过程图
:::

从图1-13可知,在一开始初始化$W^Q,W^K,W^V$这3个权重矩阵时,可以直接同时初始化$h$个头的权重,然后再进行后续的计算。而且事实上,在真正的代码实现过程中也是采用的这样的方式,这部分内容将在3.3.2节中进行介绍。因此,对图1-13中的多头计算过程,还可以根据图1-14来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-14. 多头注意力计算过程图
:::

说了这么多,终于把铺垫做完了。此时,假如有如图1-15所示的头注意力计算过程:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-15. 头注意力计算过程图
:::

如图1-15所示,该计算过程采用了头注意力机制来进行计算,且头的计算过程还可通过图1-16来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-16. 头注意力机制计算过程题
:::

那现在的问题是图1-16中的$Z$能够计算得到吗?答案是不能。为什么?因为我没有告诉你这里的$h$等于多少。如果我告诉你多头$h=2$,那么毫无疑问图1-16的计算过程就等同于图1-14的计算过程,即

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-17. 当h=2时注意力计算过程图
:::

且此时$d_k=d_m/2$。但是如果我告诉你多头$h=3$,那么图1-16的计算过程会变成

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-18. 当h=3时注意力计算过程图
:::

那么此时$d_k$则为$d_m/3$。

现在回到一开始的问题上,根据上面的论述我们可以发现,在$d_m$固定的情况下,不管是使用单头还是多头的方式,在实际的处理过程中直到进行注意力权重矩阵计算前,两者之前没有任何区别。当进行进行注意力权重矩阵计算时,$h$越大那么$Q,K,V$就会被切分得越小,进而得到的注意力权重分配方式越多,如图1-19所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 1-19. 注意力机制分配图
:::

从图1-19可以看出,如果$h=1$,那么最终可能得到的就是一个各个位置只集中于自身位置的注意力权重矩阵;如果$h=2$,那么就还可能得到另外一个注意力权重稍微分配合理的权重矩阵;$h=3$同理如此。因而多头这一做法也恰好是论文作者提出用于克服模型在对当前位置的信息进行编码时,会过度的将注意力集中于自身的位置的问题。这里再插入一张真实场景下同一层的不同注意力权重矩阵可视化结果图:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图1-20. 注意力机制分配图
:::

同时,当$h$不一样时,$d_k$的取值也不一样,进而使得对权重矩阵的scale的程度不一样。例如,如果$d_m=768$,那么当$h=12$时,则$d_k=64$;当$h=1$时,则$d_k=768$。

所以,当模型的维度$d_m$确定时,一定程度上$h$越大整个模型的表达能力越强,越能提高模型对于注意力权重的合理分配。

2. 位置编码与编码解码过程

2.1 Embedding机制

在正式介绍Transformer的网络结构之前,我们先来一起看看Transformer如何对字符进行Embedding处理。

2.1.1 Token Embedding

熟悉文本处理的读者可能都知道,在对文本相关的数据进行建模时首先要做的便是对其进行向量化。例如在机器学习中,常见的文本表示方法有one-hot编码、词袋模型以及TF-IDF等。不过在深度学习中,更常见的做法便是将各个词(或者字)通过一个Embedding层映射到低维稠密的向量空间。因此,在Transformer模型中,首先第一步要做的同样是将文本以这样的方式进行向量化表示,并且将其称之为Token Embedding,也就是深度学习中常说的词嵌入(Word Embedding)如图2-1所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-1. Token Embedding
:::

如果是换做之前的网络模型,例如CNN或者RNN,那么对于文本向量化的步骤就到此结束了,因为这些网络结构本身已经具备了捕捉时序特征的能力,不管是CNN中的n-gram形式还是RNN中的时序形式。但是这对仅仅只有自注意力机制的网络结构来说却不行。为什么呢?根据自注意力机制原理的介绍我们知道,自注意力机制在实际运算过程中不过就是几个矩阵来回相乘进行线性变换而已。因此,这就导致即使是打乱各个词的顺序,那么最终计算得到的结果本质上却没有发生任何变换,换句话说仅仅只使用自注意力机制会丢失文本原有的序列信息。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-2. 自注意力机制弊端图(一)
:::

如图2-2所示,在经过词嵌入表示后,序列“我 在 看 书”经过了一次线性变换。现在,我们将序列变成“书 在 看 我”,然后同样以中间这个权重矩阵来进行线性变换,过程如图2-3所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-3. 自注意力机制弊端图(二)
:::

根据图2-3中的计算结果来看,序列在交换位置前和交换位置后计算得到的结果在本质上并没有任何区别,仅仅只是交换了对应的位置。因此,基于这样的原因,Transformer在原始输入文本进行Token Embedding后,又额外的加入了一个Positional Embedding来刻画数据在时序上的特征。

Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence.

2.1.2 Positional Embedding

说了这么多,那到底什么又是Positional Embedding呢?数无形时少直觉,下面我们先来通过一幅图直观看看经过Positional Embedding处理后到底产生了什么样的变化。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-4. Positional Embedding
:::

如图2-4所示,横坐标表示输入序列中的每一个Token,每一条曲线或者直线表示对应Token在每个维度上对应的位置信息。在左图中,每个维度所对应的位置信息都是一个不变的常数;而在右图中,每个维度所对应的位置信息都是基于某种公式变换所得到。换句话说就是,左图中任意两个Token上的向量都可以进行位置交换而模型却不能捕捉到这一差异,但是加入右图这样的位置信息模型却能够感知到。例如位置20这一处的向量,在左图中无论你将它换到哪个位置,都和原来一模一样;但在右图中,你却再也找不到与位置20处位置信息相同的位置。

下面,笔者通过两个实际的示例来进行说明。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-5. 常数Positional Embedding(一)
:::

如图2-5所示,原始输入在经过Token Embedding后,又加入了一个常数位置信息的的Positional Embedding。在经过一次线性变换后便得到了图2-5左右边所示的结果。接下来,我们再交换序列的位置,并同时进行Positional Embedding观察其结果。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-6. 常数Positional Embedding(二)
:::

如图2-6所示,在交换序列位置后,采用同样的Positional Embedding进行处理,并且进行线性变换。可以发现,其计算结果同图2-5中的计算结果本质上也没有发生变换。因此,这就再次证明,如果Positional Embedding中位置信息是以常数形式进行变换,那么这样的Positional Embedding是无效的。

在Transformer中,作者采用了如公式$(2.1)$所示的规则来生成各个维度的位置信息,其可视化结果如图2-4右所示。
$$
PE{pos,2i}=sin(pos/10000^{2i/d{model}})\;\;\;\;\;\;\
PE{pos,2i+1}=cos(pos/10000^{2i/d{model}})\;\;\;\;\;\;(2.1)
$$
其中$PE$就是这个Positional Embedding矩阵,$pos\in[0,max_len)$表示具体的某一个位置,$i\in[0,d_{model}/2)$表示具体的某一维度。

最终,在融入这种非常数的Positional Embedding位置信息后,便可以得到如图2-7所示的对比结果。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-7. 非常数Positional Embedding
:::

从图2-7可以看出,在交换位置前与交换位置后,与同一个权重矩阵进行线性变换后的结果截然不同。因此,这就证明通过Positional Embedding可以弥补自注意力机制不能捕捉序列时序信息的缺陷。

说完Transformer中的Embedding后,接下来我们再来继续探究Transformer的网络结构。

2.2 Transformer网络结构

如图2-8所示便是一个单层Transformer网络结构图。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-8. 单层Transformer网络结构图
:::

如图2-8所示,整个Transformer网络包含左右两个部分,即Encoder和Decoder。下面,我们就分别来对其中的各个部分进行介绍。

2.2.1 Encoder层

首先,对于Encoder来说,其网络结构如图2-8左侧所示(尽管论文中是以6个这样相同的模块堆叠而成,但这里我们先以堆叠一层来进行介绍,多层的Transformer结构将在稍后进行介绍)。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-9. Encoder网络结构图
:::

如图2-9所示,对于Encoder部分来说其内部主要由两部分网络所构成:多头注意力机制和两层前馈神经网络。

The encoder is composed of a stack of N = 6 identical layers. Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position- wise fully connected feed-forward network.

同时,对于这两部分网络来说,都加入了残差连接,并且在残差连接后还进行了层归一化操作。这样,对于每个部分来说其输出均为$\text{LayerNorm(x+Sublayer(x))}$,并且在都加入了Dropout操作。

We apply dropout to the output of each sub-layer, before it is added to the sub-layer input and normalized.

进一步,为了便于在这些地方使用残差连接,这两部分网络输出向量的维度均为$d_{model}=512$。

对于第2部分的两层全连接网络来说,其具体计算过程为
$$
\text{FFN}(x)=\text{max}(0,xW_1+b_1)W_2+b2\;\;\;\;\;\;(2.2)
$$
其中输入$x$的维度为$d
{model}=512$,第1层全连接层的输出维度为$d{ff}=2048$,第2层全连接层的输出为$d{model}=512$,且同时对于第1层网络的输出还运用了Relu激活函数。

到此,对于单层Encoder的网络结构就算是介绍完了,接下来让我们继续探究Decoder部分的网络结构。

2.2.2 Decoder层

同Encoder部分一样,论文中也采用了6个完全相同的网络层堆叠而成,不过这里我们依旧只是先看1层时的情况。对于Decoder部分来说,其整体上与Encoder类似,只是多了一个用于与Encoder输出进行交互的多头注意力机制,如图2-10所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-10. Decoder网络结构图
:::

不同于Encoder部分,在Decoder中一共包含有3个部分的网络结构。最上面的和最下面的部分(暂时忽略Mask)与Encoder相同,只是多了中间这个与Encoder输出(Memory)进行交互的部分,作者称之为“Encoder-Decoder attention”。对于这部分的输入,Q来自于下面多头注意力机制的输出,K和V均是Encoder部分的输出(Memory)经过线性变换后得到。而作者之所以这样设计也是在模仿传统Encoder-Decoder网络模型的解码过程。

In "encoder-decoder attention" layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models

为了能够更好的理解这里Q、K、V的含义,我们先来看看传统的基于Encoder-Decoder的Seq2Seq翻译模型是如何进行解码的,如图2-11所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-11. 传统的Seq2Seq网络模型图
:::

如图2-11所示是一个经典的基于Encoder-Decoder的机器翻译模型。左下边部分为编码器,右下边部分为解码器,左上边部分便是注意力机制部分。在图2-11中,$\overline{h_i}$表示的是在编码过程中,各个时刻的隐含状态,称之为每个时刻的Memory;$h_t$表示解码当前时刻时的隐含状态。此时注意力机制的思想在于,希望模型在解码的时刻能够参考编码阶段每个时刻的记忆。

因此,在解码第一个时刻"<s>"时,$h_t$会首先同每个记忆状态进行相似度比较得到注意力权重。这个注意力权重所蕴含的意思就是,在解码第一个时刻时应该将$50\%$的注意力放在编码第一个时刻的记忆上(其它的同理),最终通过加权求和得到4个Memory的权重和,即context vector。同理,在解码第二时刻"我"时,也会遵循上面的这一解码过程。可以看出,此时注意力机制扮演的就是能够使得Encoder与Decoder进行交互的角色。

回到Transformer的Encoder-Decoder attention中,K和V均是编码部分的输出Memory经过线性变换后的结果(此时的Memory中包含了原始输入序列每个位置的编码信息),而Q是解码部分多头注意力机制输出的隐含向量经过线性变换后的结果。在Decoder对每一个时刻进行解码时,首先需要做的便是通过Q与 K进行交互(query查询),并计算得到注意力权重矩阵;然后再通过注意力权重与V进行计算得到一个权重向量,该权重向量所表示的含义就是在解码时如何将注意力分配到Memory的各个位置上。这一过程我们可以通过如图2-12和图2-13所示的过程来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-12. 解码过程Q、K、V计算过程图
:::

如图2-12所示,待解码向量和Memory分别各自乘上一个矩阵后得到Q、K、V。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-13. 解码第1个时刻输出向量计算过程
:::

如图2-13所示,在解码第1个时刻时,首先Q通过与K进行交互得到权重向量,此时可以看做是Q(待解码向量)在K(本质上也就是Memory)中查询Memory中各个位置与Q有关的信息;然后将权重向量与V进行运算得到解码向量,此时这个解码向量可以看作是考虑了Memory中各个位置编码信息的输出向量,也就是说它包含了在解码当前时刻时应该将注意力放在Memory中哪些位置上的信息。

进一步,在得到这个解码向量并经过图2-10中最上面的两层全连接层后,便将其输入到分类层中进行分类得到当前时刻的解码输出值。

2.2.3 Decoder预测解码过程

当第1个时刻的解码过程完成之后,解码器便会将解码第1个时刻时的输入,以及解码第1个时刻后的输出均作为解码器的输入来解码预测第2个时刻的输出。整个过程可以通过如图2-14所示的过程来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-14. Decoder多时刻解码过程图(图片来自[3])
:::

如图2-14所示,Decoder在对当前时刻进行解码输出时,都会将当前时刻之前所有的预测结果作为输入来对下一个时刻的输出进行预测。假设现在需要将"我 是 谁"翻译成英语"who am i",且解码预测后前两个时刻的结果为"who am",接下来需要对下一时刻的输出"i"进行预测,那么整个过程就可以通过图2-15和图2-16来进行表示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-15. 解码过程中Q、K、V计算过程图
:::

如图2-15所示,左上角的矩阵是解码器对输入"<s> who am"这3个词经过解码器中自注意力机制编码后的结果;左下角是编码器对输入"我 是 谁"这3个词编码后的结果(同图2-12中的一样);两者分别在经过线性变换后便得到了Q、K和V这3个矩阵。此时值得注意的是,左上角矩阵中的每一个向量在经过自注意力机制编码后,每个向量同样也包含了其它位置上的编码信息。

进一步,Q与K作用和便得到了一个权重矩阵;再将其与V进行线性组合便得到了Encoder-Decoder attention部分的输出,如图2-16所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-16. 解码第3个时刻输出向量计算过程
:::

如图2-16所示,左下角便是Q与K作用后的权重矩阵,它的每一行就表示在对Memory(这里指图2-16中的V)中的每一位置进行解码时,应该如何对注意力进行分配。例如第3行$[0.6,0.2,0.2]$的含义就是在解码当前时刻时应该将$60\%$的注意力放在Memory中的"我"上,其它同理。这样,在经过解码器中的两个全连接层后,便得到了解码器最终的输出结果。接着,解码器会循环对下一个时刻的输出进行解码预测,直到预测结果为"<e>"或者达到指定长度后停止。

同时,这里需要注意的是,在通过模型进行实际的预测时,只会取解码器输出的其中一个向量进行分类,然后作为当前时刻的解码输出。例如图2-16中解码器最终会输出一个形状为[3,tgt_vocab_len]的矩阵,那么只会取其最后一个向量喂入到分类器中进行分类得到当前时刻的解码输出。具体细节见后续代码实现。

2.2.4 Decoder训练解码过程

在介绍完预测时Decoder的解码过程后,下面就继续来看在网络在训练过程中是如何进行解码的。从2.2.3小节的内容可以看出,在真实预测时解码器需要将上一个时刻的输出作为下一个时刻解码的输入,然后一个时刻一个时刻的进行解码操作。显然,如果训练时也采用同样的方法那将是十分费时的。因此,在训练过程中,解码器也同编码器一样,一次接收解码时所有时刻的输入进行计算。这样做的好处,一是通过多样本并行计算能够加快网络的训练速度;二是在训练过程中直接喂入解码器正确的结果而不是上一时刻的预测值(因为训练时上一时刻的预测值可能是错误的)能够更好的训练网络。

例如在用平行预料"我 是 谁"<==>"who am i"对网络进行训练时,编码器的输入便是"我 是 谁",而解码器的输入则是"&lt;s&gt; who am i",对应的正确标签则是"who am i &lt;e&gt;"

假设现在解码器的输入"&lt;s&gt; who am i"在分别乘上一个矩阵进行线性变换后得到了Q、K、V,且Q与K作用后得到了注意力权重矩阵(此时还未进行softmax操作),如图2-17所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-17. 解码器输入权重矩阵计算过程图
:::

从图2-17可以看出,此时已经计算得到了注意力权重矩阵。由第1行的权重向量可知,在解码第1个时刻时应该将$20\%$(严格来说应该是经过softmax后的值)的注意力放到"&lt;s&gt;"上,$30\%$的注意力放到"who"上等等。不过此时有一个问题就是,在2.2.3节中笔者介绍到,模型在实际的预测过程中只是将当前时刻之前(包括当前时刻)的所有时刻作为输入来预测下一个时刻,也就是说模型在预测时是看不到当前时刻之后的信息。因此,Transformer中的Decoder通过加入注意力掩码机制来解决了这一问题。

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

如图2-18所示,左边依旧是通过Q和K计算得到了注意力权重矩阵(此时还未进行softmax操作),而中间的就是所谓的注意力掩码矩阵,两者在相加之后再乘上矩阵V便得到了整个自注意力机制的输出,也就是图2-10中的Masked Multi-Head Attention。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 2-18. 注意力掩码计算过程图
:::

那为什么注意力权重矩阵加上这个注意力掩码矩阵就能够达到这样的效果呢?以图2-18中第1行权重为例,当解码器对第1个时刻进行解码时其对应的输入只有"&lt;s&gt;",因此这就意味着此时应该将所有的注意力放在第1个位置上(尽管在训练时解码器一次喂入了所有的输入),换句话说也就是第1个位置上的权重应该是1,而其它位置则是0。从图2-17可以看出,第1行注意力向量在加上第1行注意力掩码,再经过softmax操作后便得到了一个类似$[1,0,0,0,0]$的向量。那么,通过这个向量就能够保证在解码第1个时刻时只能将注意力放在第1个位置上的特性。同理,在解码后续的时刻也是类似的过程。

到此,对于整个单层Transformer的网络结构以及编码解码过程就介绍完了,更多细节内容见后续代码实现。

2.2.5 位置编码与Attention Mask

在刚接触Transformer的时候,有的人会认为在Decoder中,既然已经有了Attention mask那么为什么还需要Positional Embedding呢?如图2-18所示,持这种观点的朋友认为,Attention mask已经有了使得输入序列依次输入解码器的能力,因此就不再需要Positional Embedding了。这样想对吗?

根据2.2.4节内容的介绍可以知道,Attention mask的作用只有一个,那就是在训练过程中掩盖掉当前时刻之后所有位置上的信息,而这也是在模仿模型在预测时只能看到当前时刻及其之前位置上的信息。因此,持有上述观点的朋友可能是把“能看见”和“能看见且有序”混在一起了。

虽然看似有了Attention mask这个掩码矩阵能够使得Decoder在解码过程中可以有序地看到当前位置之前的所有信息,但是事实上没有Positional Embedding的Attention mask只能做到看到当前位置之前的所有信息,而做不到有序。前者的“有序”指的是喂入解码器中序列的顺序,而后者的“有序”指的是序列本身固有的语序。

如果不加Postional Embedding的话,那么以下序列对于模型来说就是一回事:

<s> → 北 → 京 → 欢 → 迎 → 你 → <e>
<s> → 北 → 京 → 迎 → 欢 → 你 → <e>
<s> → 北 → 京 → 你 → 迎 → 欢 → <e>

虽然此时Attention mask具有能够让上述序列一个时刻一个时刻的按序喂入到解码器中,但是它却无法识别出这句话本身固有的语序。

2.2.6 原始Q、K、V来源

在Transformer中各个部分的Q、K、V到底是怎么来的一直以来都是初学者最大的一个疑问,并且这部分内容在原论文中也没有进行交代,只是交代了如何根据Q、K、V来进行自注意力机制的计算。虽然在第2部分的前面几个小节已经提及过了这部分内容,但是这里再给大家进行一次总结。

根据图2-8(Transformer结构图)可知,在整个Transformer中涉及到自注意力机制的一共有3个部分:Encoder中的Multi-Head Attention;Decoder中的Masked Multi-Head Attention;Encoder和Decoder交互部分的Multi-Head Attention。

① 对于Encoder中的Multi-Head Attention来说,其原始q、k、v均是Encoder的Token输入经过Embedding后的结果。q、k、v分别经过一次线性变换(各自乘以一个权重矩阵)后得到了Q、K、V(也就是图1-4中的示例),然后再进行自注意力运算得到Encoder部分的输出结果Memory。

② 对于Decoder中的Masked Multi-Head Attention来说,其原始q、k、v均是Decoder的Token输入经过Embedding后的结果。q、k、v分别经过一次线性变换后得到了Q、K、V,然后再进行自注意力运算得到Masked Multi-Head Attention部分的输出结果,即待解码向量。

对于Encoder和Decoder交互部分的Multi-Head Attention,其原始q、k、v分别是上面的带解码向量、Memory和Memory。q、k、v分别经过一次线性变换后得到了Q、K、V(也就是图2-12中的示例),然后再进行自注意力运算得到Decoder部分的输出结果。之所以这样设计也是在模仿传统Encoder-Decoder网络模型的解码过程。

3. 网络结构与自注意力实现

在通过前面几部分内容详细介绍完Transformer网络结构的原理后,接下来就让我们来看一看如何借用Pytorch框架来实现MultiHeadAttention这一结构。同时,需要说明的一点是,下面所有的实现代码都是笔者直接从Pytorch 1.4版本中torch.nn.Transformer模块里摘取出来的简略版,目的就是为了让大家对于整个实现过程有一个清晰的认识。并且为了使得大家在阅读完以下内容后也能够对Pytorch中的相关模块有一定的了解,所以下面的代码在变量名方面也与Pytorch保持了一致。

3.1 多层Transformer

在第2部分中,笔者详细介绍了单层Transformer网络结构中的各个组成部分。尽管多层Transformer就是在此基础上堆叠而来,不过笔者认为还是有必要在这里稍微提及一下。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-1. 单层Transformer网络结构图
:::

如图3-1所示便是一个单层Transformer网络结构图,左边是编码器右边是解码器。而多层的Transformer网络就是在两边分别堆叠了多个编码器和解码器的网络模型,如图3-2所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-2. 多层Transformer网络结构图
:::

如图3-2所示便是一个多层的Transformer网络结构图(原论文中采用了6个编码器和6个解码器),其中的每一个Encoder都是图3-1中左边所示的网络结构(Decoder同理)。可以发现,它真的就是图3-1堆叠后的形式。不过需要注意的是其整个解码过程。

在多层Transformer中,多层编码器先对输入序列进行编码,然后得到最后一个Encoder的输出Memory;解码器先通过Masked Multi-Head Attention对输入序列进行编码,然后将输出结果同Memory通过Encoder-Decoder Attention后得到第1层解码器的输出;接着再将第1层Decoder的输出通过Masked Multi-Head Attention进行编码,最后再将编码后的结果同Memory通过Encoder-Decoder Attention后得到第2层解码器的输出,以此类推得到最后一个Decoder的输出。

值得注意的是,在多层Transformer的解码过程中,每一个Decoder在Encoder-Decoder Attention中所使用的Memory均是同一个。

3.2 Transformer中的掩码

由于在实现多头注意力时需要考虑到各种情况下的掩码,因此在这里需要先对这部分内容进行介绍。在Transformer中,主要有两个地方会用到掩码这一机制。第1个地方就是在2.2.4节中介绍到的Attention Mask,用于在训练过程中解码的时候掩盖掉当前时刻之后的信息;第2个地方便是对一个batch中不同长度的序列在Padding到相同长度后,对Padding部分的信息进行掩盖。下面分别就这两种情况进行介绍。

3.2.1 Attention Mask

如图3-3所示,在训练过程中对于每一个样本来说都需要这样一个对称矩阵来掩盖掉当前时刻之后所有位置的信息。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-3. 注意力掩码计算过程图
:::

从图3-3可以看出,这个注意力掩码矩阵的形状为[tgt_len,tgt_len],其具体Mask原理在2.2.4节中笔者已经介绍过l,这里就不再赘述。在后续实现过程中,我们将通过generate_square_subsequent_mask方法来生成这样一个矩阵。同时,在后续多头注意力机制实现中,将通过attn_mask这一变量名来指代这个矩阵。

3.2.2 Padding Mask

在Transformer中,使用到掩码的第2个地方便是Padding Mask。由于在网络的训练过程中同一个batch会包含有多个文本序列,而不同的序列长度并不一致。因此在数据集的生成过程中,就需要将同一个batch中的序列Padding到相同的长度。但是,这样就会导致在注意力的计算过程中会考虑到Padding位置上的信息。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-4. Padding时注意力计算过程图
:::

如图3-4所示,P表示Padding的位置,右边的矩阵表示计算得到的注意力权重矩阵。可以看到,此时的注意力权重对于Padding位置山的信息也会加以考虑。因此在Transformer中,作者通过在生成训练集的过程中记录下每个样本Padding的实际位置;然后再将注意力权重矩阵中对应位置的权重替换成负无穷,经softmax操作后对应Padding位置上的权重就变成了0,从而达到了忽略Padding位置信息的目的。这种做法也是Encoder-Decoder网络结构中通用的一种办法。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-5. Padding掩码计算过程图
:::

如图3-5所示,对于"我 是 谁 P P"这个序列来说,前3个字符是正常的,后2个字符是Padding后的结果。因此,其Mask向量便为[True, True, True, False, False]。通过这个Mask向量可知,需要将权重矩阵的最后两列替换成负无穷,在后续我们会通过torch.masked_fill这个方法来完成这一步,并且在实现时将使用key_padding_mask来指代这一向量。

到此,对于Transformer中所要用到Mask的地方就介绍完了,下面正式来看如何实现多头注意力机制。

3.3 实现多头注意力机制

3.3.1 多头注意力机制

根据前面的介绍可以知道,多头注意力机制中最为重要的就是自注意力机制,也就是需要前计算得到Q、K和V,如图3-6所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-6. Q、K和V计算过程
:::

然后再根据Q、K、V来计算得到最终的注意力编码,如图3-7所示:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-7. 注意力编码计算图
:::

同时,为了避免单个自注意力机制计算得到的注意力权重过度集中于当前编码位置自己所在的位置(同时更应该关注于其它位置),所以作者在论文中提到通过采用多头注意力机制来解决这一问题,如图3-8所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 3-8. 多头注意力计算图(2个头)
:::

3.3.2 定义类MyMultiHeadAttention

综上所述,我们可以给出类MyMultiHeadAttentiond的定义为

class MyMultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0., bias=True):
        super(MyMultiheadAttention, self).__init__()
        """
        :param embed_dim:   词嵌入的维度,也就是前面的d_model参数,论文中的默认值为512
        :param num_heads:   多头注意力机制中多头的数量,也就是前面的nhead参数, 论文默认值为 8
        :param bias:        最后对多头的注意力(组合)输出进行线性变换时,是否使用偏置
        """
        self.embed_dim = embed_dim  # 前面的d_model参数
        self.head_dim = embed_dim // num_heads  # head_dim 指的就是d_k,d_v
        self.kdim = self.head_dim
        self.vdim = self.head_dim
        self.num_heads = num_heads  # 多头个数
        self.dropout = dropout
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim 除以 num_heads必须为整数"
        # 上面的限制条件就是论文中的  d_k = d_v = d_model/n_head 条件
        self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  
        # embed_dim = kdim * num_heads
        # 这里第二个维度之所以是embed_dim,实际上这里是同时初始化了num_heads个W_q堆叠起来的, 也就是num_heads个头
        self.k_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  
        # W_k,  embed_dim = kdim * num_heads
        self.v_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  
        # W_v,  embed_dim = vdim * num_heads
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        # 最后将所有的Z组合起来的时候,也是一次性完成, embed_dim = vdim * num_heads

在上述代码中,embed_dim表示模型的维度(图3-8中的d_m);num_heads表示多头的个数;bias表示是否在多头线性组合时使用偏置。同时,为了使得实现代码更加高效,所以Pytorch在实现的时候是多个头注意力机制一起进行的计算,也就上面代码的第17-22行,分别用来初始化了多个头的权重值(这一过程从图3-8也可以看出)。当多头注意力机制计算完成后,将会得到一个形状为[src_len,embed_dim]的矩阵,也就是图3-8中多个$z_i$水平堆叠后的结果。因此,第24行代码将会初始化一个线性层来对这一结果进行一个线性变换。

3.3.3 定义前向传播过程

在定义完初始化函数后,便可以定义如下所示的多头注意力前向传播的过程

    def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
        """
        在论文中,编码时query, key, value 都是同一个输入, 
        解码时 输入的部分也都是同一个输入, 
        解码和编码交互时 key,value指的是 memory, query指的是tgt
        :param query: # [tgt_len, batch_size, embed_dim], tgt_len 表示目标序列的长度
        :param key:  #  [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param value: # [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param attn_mask: # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
                一般只在解码时使用,为了并行一次喂入所有解码部分的输入,所以要用mask来进行掩盖当前时刻之后的位置信息
        :param key_padding_mask: [batch_size, src_len], src_len 表示源序列的长度
        :return:
        attn_output: [tgt_len, batch_size, embed_dim]
        attn_output_weights: # [batch_size, tgt_len, src_len]
        """
        return multi_head_attention_forward(query, key, value, self.num_heads,
                                            self.dropout, self.out_proj.weight, self.out_proj.bias,
                                            training=self.training,
                                            key_padding_mask=key_padding_mask,
                                            q_proj_weight=self.q_proj_weight,
                                            k_proj_weight=self.k_proj_weight,
                                            v_proj_weight=self.v_proj_weight,
                                            attn_mask=attn_mask)

在上述代码中,querykeyvalue指的并不是图3-6中的Q、K和V,而是没有经过线性变换前的输入。例如在编码时三者指的均是原始输入序列src;在解码时的Mask Multi-Head Attention中三者指的均是目标输入序列tgt;在解码时的Encoder-Decoder Attention中三者分别指的是Mask Multi-Head Attention的输出、Memory和Memory。key_padding_mask指的是编码或解码部分,输入序列的Padding情况,形状为[batch_size,src_len]或者[batch_size,tgt_len];attn_mask指的就是注意力掩码矩阵,形状为[tgt_len,src_len]它只会在解码时使用。

注意,在上面的这些维度中,tgt_len本质上指的其实是query_lensrc_len本质上指的是key_len。只是在不同情况下两者可能会是一样,也可能会是不一样。

3.3.4 多头注意力计算过程

在定义完类MyMultiHeadAttentiond后,就需要定义出多头注意力的实际计算过程。由于这部分代码较长,所以就分层次进行介绍。

def multi_head_attention_forward(
    query,  # [tgt_len,batch_size, embed_dim]
    key,  # [src_len, batch_size, embed_dim]
    value,  # [src_len, batch_size, embed_dim]
    num_heads,
    dropout_p,
    out_proj_weight, # [embed_dim = vdim * num_heads, embed_dim]
    out_proj_bias,
    training=True,
    key_padding_mask=None,  # [batch_size,src_len/tgt_len]
    q_proj_weight=None,  # [embed_dim,kdim * num_heads]
    k_proj_weight=None,  # [embed_dim, kdim * num_heads]
    v_proj_weight=None,  # [embed_dim, vdim * num_heads]
    attn_mask=None,  # [tgt_len,src_len] 
  ):
    # 第一阶段: 计算得到Q、K、V
    q = F.linear(query, q_proj_weight)
    #  [tgt_len,batch_size,embed_dim] x [embed_dim,kdim * num_heads] 
    #  = [tgt_len,batch_size,kdim * num_heads]
    k = F.linear(key, k_proj_weight)
    # [src_len, batch_size,embed_dim] x [embed_dim,kdim * num_heads]
    # = [src_len,batch_size,kdim * num_heads]
    v = F.linear(value, v_proj_weight)
    # [src_len, batch_size,embed_dim] x [embed_dim,vdim * num_heads] 
    # = [src_len,batch_size,vdim * num_heads]

在上述代码中,第17-23行所做的就是根据输入进行线性变换得到图3-6中的Q、K和V。

    # 第二阶段: 缩放,以及attn_mask维度判断
        tgt_len, bsz, embed_dim = query.size()  # [tgt_len,batch_size, embed_dim]
    src_len = key.size(0)
    head_dim = embed_dim // num_heads  # num_heads * head_dim = embed_dim
    scaling = float(head_dim) ** -0.5
    q = q * scaling  # [query_len,batch_size,kdim * num_heads]

    if attn_mask is not None:  
      # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
        if attn_mask.dim() == 2:
            attn_mask = attn_mask.unsqueeze(0)  # [1, tgt_len,src_len] 扩充维度
            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
                raise RuntimeError('The size of the 2D attn_mask is not correct.')
        elif attn_mask.dim() == 3:
            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
                raise RuntimeError('The size of the 3D attn_mask is not correct.')
        # 现在 atten_mask 的维度就变成了3D

接着,在上述代码中第5-6行所完成的就是图3-7中的缩放过程;第8-16行用来判断或修改attn_mask的维度,当然这几行代码只会在解码器中的Masked Multi-Head Attention中用到。

        # 第三阶段: 计算得到注意力权重矩阵
        q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    # [batch_size * num_heads,tgt_len,kdim]
    # 因为前面是num_heads个头一起参与的计算,所以这里要进行一下变形,以便于后面计算。 且同时交换了0,1两个维度
    k = k.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
    #[batch_size * num_heads,src_len,kdim]
    v = v.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
    #[batch_size * num_heads,src_len,vdim]
    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
    # [batch_size * num_heads,tgt_len,kdim] x [batch_size * num_heads, kdim, src_len]
    # =  [batch_size * num_heads, tgt_len, src_len]  这就num_heads个QK相乘后的注意力矩阵

继续,在上述代码中第1-7行所做的就是交换Q、K、V中的维度,以便于多个样本同时进行计算;第9行代码便是用来计算注意力权重矩阵;其中上contiguous()方法是将变量放到一块连续的物理内存中;bmm的作用是用来计算两个三维矩阵的乘法操作[4]。

需要提示的是,大家在看代码的时候,最好是仔细观察一下各个变量维度的变化过程,笔者也在每次运算后进行了批注。

        # 第四阶段: 进行相关掩码操作
        if attn_mask is not None:
        attn_output_weights += attn_mask  # [batch_size * num_heads, tgt_len, src_len]
    if key_padding_mask is not None:
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        # 变成 [batch_size, num_heads, tgt_len, src_len]的形状
        attn_output_weights = attn_output_weights.masked_fill(
            key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf')) 
         # 扩展维度,从[batch_size,src_len]变成[batch_size,1,1,src_len]
        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len,src_len)  
        # [batch_size * num_heads, tgt_len, src_len]

进一步,在上述代码中第2-3行便是用来执行图3-3中的步骤;第4-8行便是用来执行图3-5中的步骤,同时还进行了维度扩充。

        attn_output_weights = F.softmax(attn_output_weights, dim=-1)
    # [batch_size * num_heads, tgt_len, src_len]
    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
    attn_output = torch.bmm(attn_output_weights, v)
    # Z = [batch_size * num_heads, tgt_len, src_len]  x  [batch_size * num_heads,src_len,vdim]
    # = # [batch_size * num_heads,tgt_len,vdim]
    # 这就num_heads个Attention(Q,K,V)结果

    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
    # 先transpose成 [tgt_len, batch_size* num_heads ,kdim]
    # 再view成 [tgt_len,batch_size,num_heads*kdim]
    attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)

    Z = F.linear(attn_output, out_proj_weight, out_proj_bias)
    # 这里就是多个z  线性组合成Z  [tgt_len,batch_size,embed_dim]
    return Z, attn_output_weights.sum(dim=1) / num_heads  # 将num_heads个注意力权重矩阵按对应维度取平均

最后,在上述代码中第1-4行便是用来对权重矩阵进行归一化操作,以及计算得到多头注意力机制的输出;第14行代码便是用来对多个注意力的输出结果进行线性组合;第16行代码用来返回线性组合后的结果,以及多个注意力权重矩阵的平均值。

3.3.5 示例代码

在实现完类MyMultiHeadAttention的全部代码后,便可以通过类似如下的方式进行使用。

if __name__ == '__main__':
    src_len = 5
    batch_size = 2
    dmodel = 32
    num_head = 1
    src = torch.rand((src_len, batch_size, dmodel))  # shape: [src_len, batch_size, embed_dim]
    src_key_padding_mask = torch.tensor([[True, True, True, False, False],
                                         [True, True, True, True, False]])  # shape: [src_len, src_len]

    my_mh = MyMultiheadAttention(embed_dim=dmodel, num_heads=num_head)
    r = my_mh(src, src, src,key_padding_mask = src_key_padding_mask)

在上述代码中,第6-11行其实也就是Encoder中多头注意力机制的实现过程。同时,在计算过程中还可以打印出各个变量的维度变化信息:

进入多头注意力计算:
     多头num_heads = 1, d_model=32, d_k = d_v = d_model/num_heads=32
     query的shape([tgt_len, batch_size, embed_dim]):torch.Size([5, 2, 32])
      W_q 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
       Q  的shape([tgt_len, batch_size,kdim * num_heads]):torch.Size([5, 2, 32])
    ----------------------------------------------------------------------
      key 的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
      W_k 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
       K  的shape([src_len,batch_size,kdim * num_heads]):torch.Size([5, 2, 32])
    ----------------------------------------------------------------------
     value的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
      W_v 的shape([embed_dim,vdim * num_heads]):torch.Size([32, 32])
       V  的shape([src_len,batch_size,vdim * num_heads]):torch.Size([5, 2, 32])
    ----------------------------------------------------------------------
     ***** 注意,这里的W_q, W_k, W_v是多个head同时进行计算的. 因此,Q,K,V分别也是包含了多个head的q,k,v堆叠起来的结果 *****
     多头注意力中,多头计算结束后的形状(堆叠)为([tgt_len,batch_size,num_heads*kdim])torch.Size([5, 2, 32])
     多头计算结束后,再进行线性变换时的权重W_o的形状为([num_heads*vdim, num_heads*vdim  ])torch.Size([32, 32])
     多头线性变化后的形状为([tgt_len,batch_size,embed_dim]) torch.Size([5, 2, 32])

4 Transformer的实现过程

在前面几部分内容中,笔者陆续介绍了多头注意力机制的原理、Transformer中编码器和解码器的工作流程以及多头注意力的实现过程等。接下来,笔者将会一步一步地来详细介绍如何通过Pytorch框架实现Transformer的整体网络结构,包括Token Embedding、Positional Embedding、编码器和解码器等。

下面,首先要介绍的就是对于Embedding部分的编码实现。

4.1 Embedding 实现

4.1.1 Token Embedding

这里首先要实现的便是最基础的Token Enbedding,也是字符转向量的一种常用做法,如下所示:

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens):
        """
        :param tokens: shape : [len, batch_size] 
        :return: shape: [len, batch_size, emb_size]
        """
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

如上代码所示便是TokenEmbedding的实现过程,由于这部分代码并不复杂所以就不再逐行进行介绍。注意,第12行代码对原始向量进行缩放是出自论文中3.4部分的描述。

4.1.2 Positional Embedding

在2.1.2节中笔者已经对Positional Embedding的原理做了详细的介绍,其每个位置的变化方式如式$(1)$所示。
$$
PE{pos,2i}=sin(pos/10000^{2i/d{model}})\;\;\;\;\;\;\
PE{pos,2i+1}=cos(pos/10000^{2i/d{model}})\;\;\;\;\;\;\;\;(4.1)
$$
进一步,我们还可以对式$(4.1)$中括号内的参数进行化简得到如式$(4.2)$中的形式。
$$
\frac{1}{10000^{2i/d{model}}}=\text{exp}{\text{log}(10000)^{\frac{-2i}{d{model}}}}=\text{exp}{\frac{-2i\cdot\text{log(10000)}}{d{model}}}=\text{exp}{2i\cdot(\frac{\text{-log}(10000)}{d{model}})}\;\;\;\;\;\;\;\;(4.2)
$$

由此,根据式$(4.1)(4.2)$便可以实现Positional Embedding部分的代码,如下所示:

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) # [max_len, d_model]
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # [max_len, 1]
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # [d_model/2]
        pe[:, 0::2] = torch.sin(position * div_term) # [max_len, d_model/2]
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)  # [max_len, 1, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        :param x: [x_len, batch_size, emb_size]
        :return: [x_len, batch_size, emb_size]
        """
        x = x + self.pe[:x.size(0), :]  # [x_len, batch_size, d_model]
        return self.dropout(x)

如上代码所示便是整个Positional Embedding的实现过程,其中第5行代码是用来初始化一个全0的位置矩阵(也就是图1中从左往右数第2个矩阵),同时还指定了一个序列的最大长度;第6-10行是用来计算每个维度(每一列)的相关位置信息;第19行代码首先是在位置矩阵中取与输入序列长度相等的前x_len行,然后在加上Token Embedding的结果;第20行是用来返回最后得到的结果并进行Dropout操作。同时,这里需要注意的一点便是,在输入x的维度中batch_size并不是第1个维度。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 4-1. Positional Embedding 计算过程图
:::

4.1.3 Embedding代码示例

在实现完这部分代码后,便可以通过如下方式进行使用:

if __name__ == '__main__':
    x = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], dtype=torch.long)
    x = x.reshape(5, 2)  # [src_len, batch_size]
    token_embedding = TokenEmbedding(vocab_size=11, emb_size=512)
    x = token_embedding(tokens=x)
    pos_embedding = PositionalEncoding(d_model=512)
    x = pos_embedding(x=x)
    print(x.shape) # torch.Size([5, 2, 512])

4.2 Transformer实现

在介绍完Embedding部分的编码工作后,下面就开始正式如何来实现Transformer网络结构。如图4-2所示,对于Transformer网络的实现一共会包含4个部分:TransformerEncoderLayerTransformerEncoderTransformerDecoderLayerTransformerDecoder,其分别表示定义一个单独编码层、构造由多个编码层组合得到的编码器、定义一个单独的解码层以及构造由多个解码层得到的解码器。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 4-2 Transformer实现结构图
:::

需要注意的是,图4-2中的一个EncoderLayer指的就是图3-2中的一个对应的Encoder,DecoderLayer同理。

4.2.1 编码层的实现

首先,我们需要实现最基本的编码层单元,也就是图4-2中的TransformerEncoderLayer,其内部结构为图4-3所示的前向传播过程(不包括Embedding部分)。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 4-3. 编码层前向传播过程
:::

对于这部分前向传播过程,可以通过如下代码来进行实现:

class MyTransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(MyTransformerEncoderLayer, self).__init__()
        """
        :param d_model:         d_k = d_v = d_model/nhead = 64, 模型中向量的维度,论文默认值为 512
        :param nhead:           多头注意力机制中多头的数量,论文默认为值 8
        :param dim_feedforward: 全连接中向量的维度,论文默认值为 2048
        :param dropout:         丢弃率,论文中的默认值为 0.1    
        """
        self.self_attn = MyMultiheadAttention(d_model, nhead, dropout=dropout)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.activation = F.relu

        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

在上述代码中,第10行用来定义一个多头注意力机制模块,并传入相应的参数;第11-20行代码便是用来定义其它层归一化和线性变换的模块。在完成类MyTransformerEncoderLayer的初始化后,便可以实现整个前向传播的forward方法:

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        """
        :param src: 编码部分的输入,形状为 [src_len,batch_size, embed_dim]
        :param src_mask:  编码部分输入的padding情况,形状为 [batch_size, src_len]
        :return: # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
        """
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask, )[0]  # 计算多头注意力
        # src2: [src_len,batch_size,num_heads*kdim] num_heads*kdim = embed_dim
        src = src + self.dropout1(src2)  # 残差连接
        src = self.norm1(src)  # [src_len,batch_size,num_heads*kdim]

        src2 = self.activation(self.linear1(src))  # [src_len,batch_size,dim_feedforward]
        src2 = self.linear2(self.dropout(src2))  # [src_len,batch_size,num_heads*kdim]
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src  # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]

在上述代码中,第7-8行便是用来实现图4-3中Multi-Head Attention部分的前向传播过程;第10-11行用来实现多头注意力后的Add&Norm部分;第13-16行用来实现图4-3中最上面的Feed Forward部分和Add&Norm部分。

这里再次提醒大家,在阅读代码的时候最好是将对应的维度信息带入以便于理解。

4.2.2 编码器实现

在实现完一个标准的编码层之后,便可以基于此来实现堆叠多个编码层,从而得到Transformer中的编码器。对于这部分内容,可以通过如下代码来实现:

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

class MyTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers, norm=None):
        super(MyTransformerEncoder, self).__init__()
        """
        encoder_layer: 就是包含有多头注意力机制的一个编码层
        num_layers: 克隆得到多个encoder layers 论文中默认为6
        norm: 归一化层
        """
        self.layers = _get_clones(encoder_layer, num_layers)  
        # 克隆得到多个encoder layers 论文中默认为6
        self.num_layers = num_layers
        self.norm = norm

在上述代码中,第1-2行是用来定义一个克隆多个编码层或解码层功能函数;第12行中的encoder_layer便是一个实例化的编码层,self.layers中保存的便是一个包含有多个编码层的ModuleList。在完成类MyTransformerEncoder的初始化后,便可以实现整个前向传播的forward方法:

    def forward(self, src, mask=None, src_key_padding_mask=None):
        """
        :param src: 编码部分的输入,形状为 [src_len,batch_size, embed_dim]
        :param mask:  编码部分输入的padding情况,形状为 [batch_size, src_len]
        :return:# [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
        """
        output = src
        for mod in self.layers:
            output = mod(output, src_mask=mask,
                         src_key_padding_mask=src_key_padding_mask)  
            # 多个encoder layers层堆叠后的前向传播过程
        if self.norm is not None:
            output = self.norm(output)
        return output  # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]

在上述代码中,第8-10行便是用来实现多个编码层堆叠起来的效果,并完成整个前向传播过程;第11-13行用来对多个编码层的输出结果进行层归一化并返回最终的结果。

4.2.3 编码器使用示例

在完成Transformer中编码器的实现过程后,便可以将其用于对输入序列进行编码。例如可以仅仅通过一个编码器对输入序列进行编码,然后将最后的输出喂入到分类器当中进行分类处理,这部分内容在后续也会进行介绍。下面先看一个使用示例。

if __name__ == '__main__':
    src_len = 5
    batch_size = 2
    dmodel = 32
    num_head = 3
    num_layers = 2
    src = torch.rand((src_len, batch_size, dmodel))  # shape: [src_len, batch_size, embed_dim]
    src_key_padding_mask = torch.tensor([[True, True, True, False, False],
        [True, True, True, True, False]])  # shape: [batch_size, src_len]

    my_transformer_encoder_layer = MyTransformerEncoderLayer(d_model=dmodel, nhead=num_head)
    my_transformer_encoder = MyTransformerEncoder(encoder_layer=my_transformer_encoder_layer,
                                                  num_layers=num_layers,
                                                  norm=nn.LayerNorm(dmodel))
    memory = my_transformer_encoder(src=src, mask=None, 
                                    src_key_padding_mask=src_key_padding_mask)
    print(memory.shape)  # torch.Size([5, 2, 32])

在上述代码中,第2-6行定义了编码器中各个部分的参数值;第11-12行则是首先定义一个编码层,然后再定义由多个编码层组成的编码器;第15-16行便是用来得到整个编码器的前向传播输出结果,并且需要注意的是在编码器中不需要掩盖当前时刻之后的位置信息,所以mask=None

4.2.4 解码层实现

在介绍完编码器的实现后,下面就开始介绍如何实现Transformer中的解码器部分。同编码器的实现流程一样,首先需要实现的依旧是一个标准的解码层,也就是图4-4所示的前向传播过程(不包括Embedding部分)。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 4-4. 解码层前向传播过程
:::

对于这部分前向传播过程,可以通过如下代码来进行实现:

class MyTransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(MyTransformerDecoderLayer, self).__init__()
        """
        :param d_model:         d_k = d_v = d_model/nhead = 64, 模型中向量的维度,论文默认值为 512
        :param nhead:           多头注意力机制中多头的数量,论文默认为值 8
        :param dim_feedforward: 全连接中向量的维度,论文默认值为 2048
        :param dropout:         丢弃率,论文中的默认值为 0.1    
        """
        self.self_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
        # 解码部分输入序列之间的多头注意力(也就是论文结构图中的Masked Multi-head attention)
        self.multihead_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
        # 编码部分输出(memory)和解码部分之间的多头注意力机制。
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = F.relu

在上述代码中,第10行代码用来定义图4-4中Masked Multi-head Attention部分的前向传播过程;第12行则是用来定义图4-4中编码器与解码器交互的多头注意力机制模块;第14-24行是用来定义剩余的全连接层以及层归一化相关操作。在完成类MyTransformerDecoderLayer的初始化后,便可以实现整个前向传播的forward方法:

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        """
        :param tgt:  解码部分的输入,形状为 [tgt_len,batch_size, embed_dim]
        :param memory: 编码部分的输出(memory), [src_len,batch_size,embed_dim]
        :param tgt_mask: 注意力Mask输入,用于掩盖当前position之后的信息, [tgt_len, tgt_len]
        :param memory_mask: 编码器-解码器交互时的注意力掩码,一般为None
        :param tgt_key_padding_mask: 解码部分输入的padding情况,形状为 [batch_size, tgt_len]
        :param memory_key_padding_mask: 编码部分输入的padding情况,形状为 [batch_size, src_len]
        :return:# [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
        """
        tgt2 = self.self_attn(tgt, tgt, tgt,  # [tgt_len,batch_size, embed_dim]
                              attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        # 解码部分输入序列之间的多头注意力(也就是论文结构图中的Masked Multi-head attention)
        tgt = tgt + self.dropout1(tgt2)  # 接着是残差连接
        tgt = self.norm1(tgt)  # [tgt_len,batch_size, embed_dim]

        tgt2 = self.multihead_attn(tgt, memory, memory,  # [tgt_len, batch_size, embed_dim]
                                   attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        # 解码部分的输入经过多头注意力后同编码部分的输出(memory)通过多头注意力机制进行交互
        tgt = tgt + self.dropout2(tgt2)  # 残差连接
        tgt = self.norm2(tgt)  # [tgt_len, batch_size, embed_dim]

        tgt2 = self.activation(self.linear1(tgt))  # [tgt_len, batch_size, dim_feedforward]
        tgt2 = self.linear2(self.dropout(tgt2))  # [tgt_len, batch_size, embed_dim]
        # 最后的两层全连接
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt  # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]

在上述代码中,第12-14行用来完成图4-4中Masked Multi-head Attention部分的前向传播过程,其中的tgt_mask就是在训练时用来掩盖当前时刻之后位置的注意力掩码;第16-17行用来完成图4-4中Masked Multi-head Attention之后Add&Norm部分的前向传播过程;第19-21行用来实现解码器与编码器之间的交互过程,其中memory_maskNonememory_key_padding_masksrc_key_padding_mask用来对编码器的输出进行(序列)填充部分的掩盖,这一点同编码器中的key_padding_mask原理一样;第23-31行便是用来实现余下的其它过程。

4.2.5 解码器实现

在实现完一个标准的解码层之后,便可以基于此来实现堆叠多个解码层,从而得到Transformer中的解码器。对于这部分内容,可以通过如下代码来实现:

class MyTransformerDecoder(nn.Module):
    def __init__(self, decoder_layer, num_layers, norm=None):
        super(MyTransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        """
        :param tgt: 解码部分的输入,形状为 [tgt_len,batch_size, embed_dim]
        :param memory: 编码部分最后一层的输出 [src_len,batch_size, embed_dim]
        :param tgt_mask: 注意力Mask输入,用于掩盖当前position之后的信息, [tgt_len, tgt_len]
        :param memory_mask: 编码器-解码器交互时的注意力掩码,一般为None
        :param tgt_key_padding_mask: 解码部分输入的padding情况,形状为 [batch_size, tgt_len]
        :param memory_key_padding_mask: 编码部分输入的padding情况,形状为 [batch_size, src_len]
        :return: # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
        """
        output = tgt  # [tgt_len,batch_size, embed_dim]
        for mod in self.layers:  # 这里的layers就是N层解码层堆叠起来的
            output = mod(output, memory,
                         tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)
        if self.norm is not None:
            output = self.norm(output)
        return output  # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]

在上述代码中,第4行用来克隆得到多个解码层;第20-25行用来实现多层解码层的前向传播过程;第28行便是用来返回最后的结果。

4.2.6 Transformer网络实现

在实现完Transformer中各个基础模块的话,下面就可以来搭建最后的Transformer模型了。总体来说这部分的代码也相对简单,只需要将上述编码器解码器组合到一起即可,具体代码如下所示:

class MyTransformer(nn.Module):
    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1):
        super(MyTransformer, self).__init__()

        """
        :param d_model:  d_k = d_v = d_model/nhead = 64, 模型中向量的维度,论文默认值为 512
        :param nhead:               多头注意力机制中多头的数量,论文默认为值 8
        :param num_encoder_layers:  encoder堆叠的数量,也就是论文中的N,论文默认值为6
        :param num_decoder_layers:  decoder堆叠的数量,也就是论文中的N,论文默认值为6
        :param dim_feedforward:     全连接中向量的维度,论文默认值为 2048
        :param dropout:             丢弃率,论文中的默认值为 0.1
        """
        #  ================ 编码部分 =====================
        encoder_layer = MyTransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
        encoder_norm = nn.LayerNorm(d_model)
        self.encoder = MyTransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
        # ================ 解码部分 =====================
        decoder_layer = MyTransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = MyTransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
        self._reset_parameters() # 初始化模型参数
        self.d_model = d_model
        self.nhead = nhead

在上述代码中,第15-17行是用来定义编码器部分;第19-21行是用来定义解码器部分;第22行用来以某种方式初始化Transformer中的权重参数,具体实现在稍后的内容中。在定义完类MyTransformer的初始化函数后,便可以来实现Transformer的整个前向传播过程,代码如下:

    def forward(self, src, tgt, src_mask=None, tgt_mask=None,
                memory_mask=None, src_key_padding_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        """
                :param src:   [src_len,batch_size,embed_dim]
        :param tgt:  [tgt_len, batch_size, embed_dim]
        :param src_mask:  None
        :param tgt_mask:  [tgt_len, tgt_len]
        :param memory_mask: None
        :param src_key_padding_mask: [batch_size, src_len]
        :param tgt_key_padding_mask: [batch_size, tgt_len]
        :param memory_key_padding_mask:  [batch_size, src_len]
        :return: [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
        """
        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
        # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
        output = self.decoder(tgt=tgt, memory=memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)
        return output # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]

在上述代码中,src表示编码器的输入;tgt表示解码器的输入;src_mask为空,因为编码时不需要对当前时刻之后的位置信息进行掩盖;tgt_mask用于掩盖解码输入中当前时刻以后的所有位置信息;memory_mask为空;src_key_padding_mask表示对编码输入序列填充部分的Token进行mask;tgt_key_padding_mask表示对解码输入序列填充部分的Token进行掩盖;memory_key_padding_mask表示对编码器的输出部分进行掩盖,掩盖原因等同于编码输入时的mask操作。

到此,对于整个Transformer的网络结构就算是搭建完毕了,不过这还没有实现论文中基于Transformer结构的翻译模型,而这部分内容笔者也将会在下一节中进行详细的介绍。当然,出了上述模块之外,Transformer中还有两个部分需要实现的就是参数初始化方法和注意力掩码矩阵生成方法,具体代码如下:

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)

    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask # [sz,sz]
4.2.7 Transfromer使用示例

在实现完Transformer的整个完了结构后,便可以通过如下步骤进行使用:

if __name__ == '__main__':
    src_len = 5
    batch_size = 2
    dmodel = 32
    tgt_len = 6
    num_head = 8
    src = torch.rand((src_len, batch_size, dmodel))  # shape: [src_len, batch_size, embed_dim]
    src_key_padding_mask = torch.tensor([[True, True, True, False, False],
                             [True, True, True, True, False]])  # shape: [batch_size, src_len]

    tgt = torch.rand((tgt_len, batch_size, dmodel))  # shape: [tgt_len, batch_size, embed_dim]
    tgt_key_padding_mask = torch.tensor([[True, True, True, False, False, False],
                      [True, True, True, True, False, False]])  # shape: [batch_size, tgt_len]

    my_transformer = MyTransformer(d_model=dmodel, nhead=num_head, num_encoder_layers=6,
                                   num_decoder_layers=6, dim_feedforward=500)
    tgt_mask = my_transformer.generate_square_subsequent_mask(tgt_len)
    out = my_transformer(src=src, tgt=tgt, tgt_mask=tgt_mask,
                         src_key_padding_mask=src_key_padding_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=src_key_padding_mask)
    print(out.shape) #torch.Size([6, 2, 32])

在上述代码中,第7-13行用来生成模拟的输入数据;第15-16行用来实例化类MyTransformer;第17行用来生成解码输入时的注意力掩码矩阵;第18-21行用来执行Transformer网络结构的前向传播过程。

5 基于Transformer的翻译模型

经过前面几节内容的介绍,相信各位读者对于Transformer的基本原理以及实现过程已经有了一个较为清晰的认识。不过想要对一个网络模型有更加深刻的认识,那么最好的办法便是从数据预处理到模型训练,自己完完全全的经历一遍。因此,为了使得大家能够更加透彻的理解Transformer的整个工作流程,在本节中笔者将继续带着大家一起来还原论文中的文本翻译模型。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 5-1. Transformer网络结构图
:::

如图5-1所示便是Transformer网络的整体结构图,对于这部分内容在上一节内容中总体上算是介绍完了,只是在数据预处理方面还未涉及。下面,笔者就以Multi30K[9]中的English-German平行语料为例进行介绍(注意这并不是论文中所用到数据集)。

本部分完整代码可参见[11]。

5.1 数据预处理

5.1.1 语料介绍

在这里,我们使用到的平行语料一共包含有6个文件train.detrain.enval.deval.entest_2016_flickr.detest_2016_flickr.en,其分别为德语训练语料、英语训练语料、德语验证语料、英语验证语料、德语测试语料和英语测试语料。同时,这三部分的样本量分别为29000、1014和1000条。

如下所示便是一条平行预料数据,其中第1行为德语,第2行为英语,后续我们需要完成的就是搭建一个翻译模型将德语翻译为英语。

Zwei junge weiße Männer sind im, Freien in der Nähe vieler Büsche.
Two young, White males are outside near many bushes.

5.1.2 数据集预览

在正式介绍如何构建数据集之前,我们先通过几张图来了解一下整个构建的流程,以便做到心中有数,不会迷路。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 5-2. 翻译模型数据集处理过程图(一)
:::

如图5-2所示,左边部分为原始输入,右边部分为目标输入。从图5-2可以看出,第1步需要完成的就是对原始语料进行tokenize操作。如果是对类似英文这样的语料进行处理,那就是直接按空格切分即可。但是需要注意的是要把其中的逗号和句号也给分割出来。第2步需要做的就是根据tokenize后的结果对原始输入和目标输入分别建立一个字典。第3步需要做的则是将tokenize后结果根据字典中的索引将其转换成token序列。第4步则是对同一个batch中的序列以最长的为标准其它样本进行padding,并且同时需要在目标输入序列的前后加上起止符(即'&lt;bos&gt;''&lt;eos&gt;') 。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 5-3. 翻译模型数据集处理过程图(二)
:::

如图5-3所示,在完成前面4个步骤后,对于目标序列来说第5步需要做的就是将第4步处理后的结果划分成tgt_inputtgt_output。从图5-3右侧可以看出,tgt_inputtgt_output是相互对应起来的。例如对于第1个样本来说,解码第1个时刻的输入应该是"2",而此时刻对应的正确标签就应该是tgt_output中的"8";解码第2个时刻的输入应该是tgt_input中的"8",而此时刻对应的正确标签就应该是tgt_output中的"45",以此类推下去。最后,第6步则是根据src_inputtgt_input各自的padding情况,得到一个padding mask向量(注意由于这里tgt_input中的两个样本长度一样,所以并不需要padding),其中"T"表示padding的位置。当然,这里的tgt_mask并没有画出。

同时,图5-3中各个部分的结果体现在Transformer网络中的情况如图5-4所示。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 5-4. 基于Transformer翻译模型的模型输入情况
:::

以上就是基于Transformer架构的翻译模型数据预处理的整个大致流程,下面我们开始正式来通过编码实现这一过程。

5.1.3 数据集构建

第1步:定义tokenize

如果是对类似英文这样的语料进行处理,大部分就是直接按空格切分即可。但是需要注意的是单词中的某些缩写也需要给分割出来,例如"you're"需要分割成"you"和"re'"。因此,这部分代码可以借助torchtext中的get_tokenizer方法来实现,具体代码如下:

from torchtext.data.utils import get_tokenizer
def my_tokenizer():
    tokenizer = {}
    tokenizer['de'] = get_tokenizer('spacy', language='de_core_news_sm')  # 德语
    tokenizer['en'] = get_tokenizer('spacy', language='en_core_web_sm')  # 英语
    return tokenizer

这里返回的是两个tokenizer,分别用于对德语和英语进行序列化。例如对于如下文本来说

s = "Moon Hotel, it's very interesting."

其tokenize后的结果为:

tokenizer = my_tokenizer()
print(tokenizer['en'](s))
['Moon', 'Hotel', ',', 'it', "'s", 'very', 'interesting', '.']

第2步:建立词表

在介绍完tokenize的实现方法后,我们就可以正式通过torchtext.vocab中的Vocab方法来构建词典了,代码如下:

def build_vocab(tokenizer, filepath, min_freq, specials=None):    
    if specials is None:
        specials = ['<unk>', '<pad>', '<bos>', '<eos>']
    counter = Counter()
    with open(filepath, encoding='utf8') as f:
        for string_ in f:
            counter.update(tokenizer(string_))
    return Vocab(counter, specials=specials, min_freq=min_freq)

在上述代码中,第3行代码用来指定特殊的字符;第5-7行代码用来遍历文件中的每一个样本(每行一个)并进行tokenize和计数,其中对于counter.update进行介绍可以参考[10];第8行则是返回最后得到词典。

在完成上述过程后,我们将得到两个Vocab类的实例化对象。

一个为原始序列的字典:

{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '.': 4, 'Männer': 5, 'ein': 6, 'in': 7, 'Ein': 8, 'Zwei': 9, 'und': 10, ',': 11, ......}

一个为目标序列的字典:

{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '.': 4, 'a': 5, 'are': 6, 'A': 7, 'Two': 8, 'in': 9, 'men': 10, ',': 11, 'Several': 12,......}

此时,我们就需要定义一个类,并在类的初始化过程中根据训练语料完成字典的构建,代码如下:

class LoadEnglishGermanDataset():
    def __init__(self, train_file_paths=None, tokenizer=None,
                 batch_size=2, min_freq=1):
        # 根据训练预料建立英语和德语各自的字典
        self.tokenizer = tokenizer()
        self.de_vocab = build_vocab(self.tokenizer['de'], 
                                    filepath=train_file_paths[0],   min_freq=min_freq)
        self.en_vocab = build_vocab(self.tokenizer['en'], 
                                    filepath=train_file_paths[1], min_freq=min_freq)
        self.specials = ['<unk>', '<pad>', '<bos>', '<eos>']
        self.PAD_IDX = self.de_vocab['<pad>']
        self.BOS_IDX = self.de_vocab['<bos>']
        self.EOS_IDX = self.de_vocab['<eos>']
        self.batch_size = batch_size

其中min_freq表示在构建词表时忽略掉出现次数小于该值的字。

第3步:转换为Token序列

在得到构建的字典后,便可以通过如下函数来将训练集、验证集和测试集转换成Token序列:

    def data_process(self, filepaths):
        """
        将每一句话中的每一个词根据字典转换成索引的形式
        :param filepaths:
        :return:
        """
        raw_de_iter = iter(open(filepaths[0], encoding="utf8))
        raw_en_iter = iter(open(filepaths[1], encoding="utf8))
        data = []
        for (raw_de, raw_en) in zip(raw_de_iter, raw_en_iter):
            de_tensor_ = torch.tensor([self.de_vocab[token] for token in
                                       self.tokenizer['de'](raw_de.rstrip("\n))], dtype=torch.long)
            en_tensor_ = torch.tensor([self.en_vocab[token] for token in
                                       self.tokenizer['en'](raw_en.rstrip("\n))], dtype=torch.long)
            data.append((de_tensor_, en_tensor_))
        return data

在上述代码中,第11-4行分别用来将原始序列和目标序列转换为对应词表中的Token形式。在处理完成后,就会得到类似如下的结果:

 [(tensor([9,37, 46, 5, 42, 36, 11, 16,7, 33, 24, 45, 13,4]),tensor([8,45, 11, 13,28, 6, 34,31, 30,16,  4])),  (tensor([22, 5, 40, 25, 30,  6, 12,  4]), tensor([12, 10,  9, 22, 23,  6, 33,  5, 20, 37, 41,  4])),  (tensor([8, 38, 23, 39,  7,  6, 26, 29, 19,  4]), tensor([ 7, 27, 21, 18, 24,  5, 44, 35,  4])),  (tensor([ 9,  5, 43, 27, 18, 10, 31, 14, 47,  4]), tensor([ 8, 10,  6, 14, 42, 40, 36, 19,  4]))  ]

其中左边的一列就是原始序列的Token形式,右边一列就是目标序列的Token形式,每一行构成一个样本。

第4步:padding处理

从上面的输出结果(以及图5-2中第③步后的结果)可以看到,无论是对于原始序列来说还是目标序列来说,在不同的样本中其对应长度都不尽相同。但是在将数据输入到相应模型时却需要保持同样的长度,因此在这里我们就需要对Token序列化后的样本进行padding处理。同时需要注意的是,一般在这种生成模型中,模型在训练过程中只需要保证同一个batch中所有的原始序列等长,所有的目标序列等长即可,也就是说不需要在整个数据集中所有样本都保证等长。

因此,在实际处理过程中无论是原始序列还是目标序列都会以每个batch中最长的样本为标准对其它样本进行padding,具体代码如下:

    def generate_batch(self, data_batch): 
        de_batch, en_batch = [], []
        for (de_item, en_item) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
            de_batch.append(de_item)  # 编码器输入序列不需要加起止符
            # 在每个idx序列的首位加上 起始token 和 结束 token
            en = torch.cat([torch.tensor([self.BOS_IDX]), en_item, torch.tensor([self.EOS_IDX])], dim=0)
            en_batch.append(en)
        # 以最长的序列为标准进行填充
        de_batch = pad_sequence(de_batch, padding_value=self.PAD_IDX)  # [de_len,batch_size]
        en_batch = pad_sequence(en_batch, padding_value=self.PAD_IDX)  # [en_len,batch_size]
        return de_batch, en_batch

在上述代码中,第6-7行用来在目标序列的首尾加上特定的起止符;第9-10行则是分别对一个batch中的原始序列和目标序列以各自当中最长的样本为标准进行padding(这里的pad_sequence导入自torch.nn.utils.rnn)。

第5步:构造mask向量

在处理完成图5-2中的第④步后,对于图5-3中的第⑤步来说就是简单的切片操作,因此就不作介绍。进一步需要根据src_inputtgt_input来构造相关的mask向量,具体代码如下:

    def generate_square_subsequent_mask(self, sz, device):
        mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def create_mask(self, src, tgt, device='cpu'):
        src_seq_len = src.shape[0]
        tgt_seq_len = tgt.shape[0]
        tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device)  # [tgt_len,tgt_len]
        # Decoder的注意力Mask输入,用于掩盖当前position之后的position,所以这里是一个对称矩阵
        src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
        # Encoder的注意力Mask输入,这部分其实对于Encoder来说是没有用的,所以这里全是0
        src_padding_mask = (src == self.PAD_IDX).transpose(0, 1)
        # 用于mask掉Encoder的Token序列中的padding部分,[batch_size, src_len]
        tgt_padding_mask = (tgt == self.PAD_IDX).transpose(0, 1)
        # 用于mask掉Decoder的Token序列中的padding部分,batch_size, tgt_len
        return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

在上述代码中,第1-4行是用来生成一个形状为[sz,sz]的注意力掩码矩阵,用于在解码过程中掩盖当前position之后的position;第6-17行用来返回Transformer中各种情况下的mask矩阵,其中src_mask在这里并没有作用。

第6步:构造DataLoade与使用示

经过前面5步的操作,整个数据集的构建就算是已经基本完成了,只需要再构造一个DataLoader迭代器即可,代码如下:

    def load_train_val_test_data(self, train_file_paths, val_file_paths, test_file_paths):
        train_data = self.data_process(train_file_paths)
        val_data = self.data_process(val_file_paths)
        test_data = self.data_process(test_file_paths)
        train_iter = DataLoader(train_data, batch_size=self.batch_size,
                                shuffle=True, collate_fn=self.generate_batch)
        valid_iter = DataLoader(val_data, batch_size=self.batch_size,
                                shuffle=True, collate_fn=self.generate_batch)
        test_iter = DataLoader(test_data, batch_size=self.batch_size,
                               shuffle=True, collate_fn=self.generate_batch)
        return train_iter, valid_iter, test_iter

在上述代码中,第2-4行便是分别用来将训练集、验证集和测试集转换为Token序列;第5-10行则是分别构造3个DataLoader,其中generate_batch将作为一个参数传入来对每个batch的样本进行处理。在完成类LoadEnglishGermanDataset所有的编码过程后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    train_filepath = ['data/train_.de',
                      'data/train_.en']

    data_loader = LoadEnglishGermanDataset(train_filepath, tokenizer=my_tokenizer, batch_size=2)
    train_iter, valid_iter, test_iter = data_loader.load_train_val_test_data(train_filepath,
                                                                             train_filepath,
                                                                             train_filepath)
    print(data_loader.PAD_IDX)
    for src, tgt in train_iter:
        tgt_input = tgt[:-1, :]
        tgt_out = tgt[1:, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = data_loader.create_mask(src, tgt_input)
        print("src shape:", src.shape)  # [de_tensor_len,batch_size]
        print("src_padding_mask shape (batch_size, src_len): ", src_padding_mask.shape)
        print("tgt input shape:", tgt_input.shape)
        print("tgt_padding_mask shape: (batch_size, tgt_len) ", tgt_padding_mask.shape)
        print("tgt output shape:", tgt_out.shape)
        print("tgt_mask shape (tgt_len,tgt_len): ", tgt_mask.shape)
        break

各位读者在阅读这部分代码时最好是能够结合图5-2到5-4进行理解,这样效果可能会更好。在介绍完数据集构建的整个过程后,下面就开始正式进入到翻译模型的构建中。

5.2 翻译模型

5.2.1 网络结构

总体来说,基于Transformer的翻译模型的网络结构其实就是图5-4所展示的所有部分,只是在前面介绍Transformer网络结构时笔者并没有把Embedding部分的实现给加进去。这是因为对于不同的文本生成模型,其Embedding部分会不一样(例如在诗歌生成这一情景中编码器和解码器共用一个TokenEmbedding即可,而在翻译模型中就需要两个),所以将两者进行了拆分。同时,待模型训练完成后,在inference过程中Encoder只需要执行一次,所以在此过程中也需要单独使用Transformer中的Encoder和Decoder。

首先,我们需要定义一个名为TranslationModel的类,其前向传播过程代码如下所示:

class TranslationModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, 
                 d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, 
                 dropout=0.1):
        super(TranslationModel, self).__init__()
        self.my_transformer = MyTransformer(d_model=d_model,
                                            nhead=nhead,
                                            num_encoder_layers=num_encoder_layers,
                                            num_decoder_layers=num_decoder_layers,
                                            dim_feedforward=dim_feedforward,
                                            dropout=dropout)
        self.pos_embedding = PositionalEncoding(d_model=d_model, dropout=dropout)
        self.src_token_embedding = TokenEmbedding(src_vocab_size, d_model)
        self.tgt_token_embedding = TokenEmbedding(tgt_vocab_size, d_model)
        self.classification = nn.Linear(d_model, tgt_vocab_size)

    def forward(self, src=None, tgt=None, src_mask=None, 
                tgt_mask=None, memory_mask=None, src_key_padding_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        """
        src: Encoder的输入 [src_len,batch_size]
        tgt: Decoder的输入 [tgt_len,batch_size]
        src_key_padding_mask: 用来Mask掉Encoder中不同序列的padding部分,[batch_size, src_len]
        tgt_key_padding_mask: 用来Mask掉Decoder中不同序列的padding部分 [batch_size, tgt_len]
        memory_key_padding_mask: 用来Mask掉Encoder输出的memory中不同序列的padding部分 [batch_size, src_len]
        :return:
        """
        src_embed = self.src_token_embedding(src)  # [src_len, batch_size, embed_dim]
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
        tgt_embed = self.tgt_token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
        outs = self.my_transformer(src=src_embed, tgt=tgt_embed, src_mask=src_mask, 
                                   tgt_mask=tgt_mask,memory_mask=memory_mask, 
                                   src_key_padding_mask=src_key_padding_mask,
                                   tgt_key_padding_mask=tgt_key_padding_mask,
                                   memory_key_padding_mask=memory_key_padding_mask)  
                                   # [tgt_len,batch_size,embed_dim]
        logits = self.classification(outs)  # [tgt_len,batch_size,tgt_vocab_size]
        return logits

在上述代码中,第7-12行便是用来定义一个Transformer结构;第13-16分别用来定义Positional Embedding、Token Embedding和最后的分类器;第29-39行便是用来执行整个前向传播过程,其中Transformer的整个前向传播过程在第4节中已经介绍过,在这里就不再赘述。

在定义完logits的前向传播过后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    src_len = 7
    batch_size = 2
    dmodel = 32
    tgt_len = 8
    num_head = 4
    src = torch.tensor([[4, 3, 2, 6, 0, 0, 0],
                        [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1) # 转换成 [src_len, batch_size]
    src_key_padding_mask = torch.tensor([[True, True, True, True, False, False, False],
                                         [True, True, True, True, True, False, False]])

    tgt = torch.tensor([[1, 3, 3, 5, 4, 3, 0, 0],
                        [1, 6, 8, 2, 9, 1, 0, 0]]).transpose(0, 1)
    tgt_key_padding_mask = torch.tensor([[True, True, True, True, True, True, False, False],
                                         [True, True, True, True, True, True, False, False]])

    trans_model = TranslationModel(src_vocab_size=10, tgt_vocab_size=15,
                                   d_model=dmodel,nhead=num_head,num_encoder_layers=6,
                                   num_decoder_layers=6, dim_feedforward=30, dropout=0.1)
    tgt_mask = trans_model.my_transformer.generate_square_subsequent_mask(tgt_len)
    logits = trans_model(src, tgt=tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=src_key_padding_mask)
    print(logits.shape) # torch.Size([8, 2, 15]) [tgt_len,batch_size,tgt_vocab_size]

接着,我们需要再定义一个EncoderDecoder在inference中使用,代码如下:

    def encoder(self, src):
        src_embed = self.src_token_embedding(src)  # [src_len, batch_size, embed_dim]
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
        memory = self.my_transformer.encoder(src_embed)
        return memory

    def decoder(self, tgt, memory, tgt_mask):
        tgt_embed = self.tgt_token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
        outs = self.my_transformer.decoder(tgt_embed, memory=memory,
                                 tgt_mask=tgt_mask)  # [tgt_len,batch_size,embed_dim]
        return outs

在上述代码中,第1-5行用于在inference时对输入序列进行编码并得到memory(只需要执行一次);第7-11行用于根据memory和当前解码时刻的输入对输出进行预测,需要循环执行多次,这部分内容详见模型预测部分。

5.2.2 模型训练

在定义完成整个翻译模型的网络结构后下面就可以开始训练模型了。由于这部分代码较长,所以下面笔者依旧以分块的形式进行介绍:

第1步:载入数据集

def train_model(config):
    data_loader = LoadEnglishGermanDataset(config.train_corpus_file_paths,
                                           batch_size=config.batch_size,
                                           tokenizer=my_tokenizer)
    train_iter, valid_iter, test_iter = \
        data_loader.load_train_val_test_data(config.train_corpus_file_paths,
                                             config.val_corpus_file_paths,
                                             config.test_corpus_file_paths)

首先我们可以根据前面的介绍,通过类LoadEnglishGermanDataset来载入数据集,其中config中定义了模型所涉及到的所有配置参数。

第2步:定义模型并初始化权重

    translation_model = TranslationModel(src_vocab_size=len(data_loader.de_vocab),
                                         tgt_vocab_size=len(data_loader.en_vocab),
                                         d_model=config.d_model,
                                         nhead=config.num_head,
                                         num_encoder_layers=config.num_encoder_layers,
                                         num_decoder_layers=config.num_decoder_layers,
                                         dim_feedforward=config.dim_feedforward,
                                         dropout=config.dropout)
    for p in translation_model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

在载入数据后,便可以定义一个翻译模型TranslationModel,并根据相关参数对其进行实例化;同时,可以对整个模型中的所有参数进行一个初始化操作。

第3步:定义损失学习率与优化器

    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=data_loader.PAD_IDX)
    learning_rate = CustomSchedule(config.d_model)
    optimizer = torch.optim.Adam(translation_model.parameters(),
                                 lr=config.warm_up_learning_rate,
                                 betas=(config.beta1, config.beta2), eps=config.epsilon)

在上述代码中,第1行是定义交叉熵损失函数,并同时指定需要忽略的索引ignore_index。因为根据图5-3的tgt_output可知,有些位置上的标签值其实是Padding后的结果,因此在计算损失的时候需要将这些位置给忽略掉。第2行代码则是论文中所提出来的动态学习率计算过程,其计算公式为:
$$
\text{lrate}=d_{model}^{-0.5}\cdot \text{min}(\text{step_num}^{-0.5},\text{step_num}\cdot \text{warmup_steps}^{-1.5})\;\;\;\;\;\;(5.1)
$$
具体实现代码为:

class CustomSchedule(nn.Module):
    def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()
        self.d_model = torch.tensor(d_model, dtype=torch.float32)
        self.warmup_steps = warmup_steps
        self.step = 1.

    def __call__(self):
        arg1 = self.step ** -0.5
        arg2 = self.step * (self.warmup_steps ** -1.5)
        self.step += 1.
        return (self.d_model ** -0.5) * min(arg1, arg2)

通过CustomSchedule,就能够在训练过程中动态的调整学习率。学习率随step增加而变换的结果如图5-5所示:

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 5-5. 动态学习率变化过程图
:::

从图5-5可以看出,在前warm_up个step中,学习率是线性增长的,在这之后便是非线性下降,直至收敛与0.0004。

第4步:开始训练

    for epoch in range(config.epochs):
        losses = 0
        start_time = time.time()
        for idx, (src, tgt) in enumerate(train_iter):
            src = src.to(config.device)  # [src_len, batch_size]
            tgt = tgt.to(config.device)
            tgt_input = tgt[:-1, :]  # 解码部分的输入, [tgt_len,batch_size]
            src_mask, tgt_mask, src_padding_mask, tgt_padding_mask \
                = data_loader.create_mask(src, tgt_input, config.device)
            logits = translation_model(
                src=src,  # Encoder的token序列输入,[src_len,batch_size]
                tgt=tgt_input,  # Decoder的token序列输入,[tgt_len,batch_size]
                src_mask=src_mask,  # Encoder的注意力Mask输入,这部分其实对于Encoder来说是没有用的
                tgt_mask=tgt_mask,
                # Decoder的注意力Mask输入,用于掩盖当前position之后的position [tgt_len,tgt_len]
                src_key_padding_mask=src_padding_mask,  # 用于mask掉Encoder的Token序列中的padding部分
                tgt_key_padding_mask=tgt_padding_mask,  # 用于mask掉Decoder的Token序列中的padding部分
                memory_key_padding_mask=src_padding_mask)  # 用于mask掉Encoder的Token序列中的padding部分
            # logits 输出shape为[tgt_len,batch_size,tgt_vocab_size]

            optimizer.zero_grad()
            tgt_out = tgt[1:, :]  # 解码部分的真实值  shape: [tgt_len,batch_size]
            loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
            # [tgt_len*batch_size, tgt_vocab_size] with [tgt_len*batch_size, ]
            loss.backward()
            lr = learning_rate()
            for p in optimizer.param_groups:
                p['lr'] = lr
            optimizer.step()
            losses += loss.item()
            acc, _, _ = accuracy(logits, tgt_out, data_loader.PAD_IDX)
            print(f"Epoch: {epoch}, Train loss :{loss.item():.3f}, Train acc: {acc})

在上述代码中,第5-9行是用来得到模型各个部分的输入;第10-18行是计算模型整个前向传播的过程;第21-25行则是执行损失计算与反向传播;第27-29则是将每个step更新后的学习率送入到模型中并更新参数;第31行是用来计算模型预测的准确率,具体过程将在后续文章中进行介绍。以下便是模型训练过程中的输出:

Epoch: 2, Train loss: 5.685, Train acc: 0.240947
Epoch: 2, Train loss: 5.668, Train acc: 0.241493
Epoch: 2, Train loss: 5.714, Train acc: 0.224682
Epoch: 2, Train loss: 5.660, Train acc: 0.235888
Epoch: 2, Train loss: 5.584, Train acc: 0.242052
Epoch: 2, Train loss: 5.611, Train acc: 0.243428

5.2.3 模型预测

在介绍完模型的训练过程后接下来就来看模型的预测部分。生成模型的预测部分不像普通的分类任务只需要将网络最后的输出做argmax操作即可,生成模型在预测过程中往往需要按时刻一步步进行来进行。因此,下面我们这里定义一个translate函数来执行这一过程,具体代码如下:

def translate(model, src, data_loader, config):
    src_vocab = data_loader.de_vocab
    tgt_vocab = data_loader.en_vocab
    src_tokenizer = data_loader.tokenizer
    model.eval()
        tokens = [src_vocab.stoi[tok] for tok in src_tokenizer(src)] # 构造一个样本
    num_tokens = len(tokens)
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1))  # 将src_len 作为第一个维度
    tgt_tokens = greedy_decode(model, src, max_len=num_tokens + 5,
                               start_symbol=data_loader.BOS_IDX, config=config,
                               data_loader=data_loader).flatten()  # 解码的预测结果
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", ").replace("<eos>", ")

在上述代码中,第6行是将待翻译的源序列进行序列化操作;第8-11行则是通过函数greedy_decode函数来对输入进行解码;第12行则是将最后解码后的结果由Token序列在转换成实际的目标语言。同时,greedy_decode函数的实现如下:

def greedy_decode(model, src, max_len, start_symbol, config, data_loader):
    src = src.to(config.device)
    memory = model.encoder(src)  # 对输入的Token序列进行解码翻译
    ys = torch.ones(1, 1).fill_(start_symbol). \
        type(torch.long).to(config.device)  # 解码的第一个输入,起始符号
    for i in range(max_len - 1):
        memory = memory.to(config.device)
        tgt_mask = (model.my_transformer.generate_square_subsequent_mask(ys.size(0))
                    .type(torch.bool)).to(config.device)  # 根据tgt_len产生一个注意力mask矩阵(对称的)
        out = model.decoder(ys, memory, tgt_mask)  # [tgt_len,tgt_vocab_size]
        out = out.transpose(0, 1)  # [tgt_vocab_size, tgt_len]
        prob = model.classification(out[:, -1])  # 只对对预测的下一个词进行分类
        _, next_word = torch.max(prob, dim=1)  # 选择概率最大者
        next_word = next_word.item()
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        # 将当前时刻解码的预测输出结果,同之前所有的结果堆叠作为输入再去预测下一个词。
        if next_word == data_loader.EOS_IDX:  # 如果当前时刻的预测输出为结束标志,则跳出循环结束预测。
            break
    return ys

在上述代码中,第3行是将源序列输入到Transformer的编码器中进行编码并得到Memory;第4-5行是初始化解码阶段输入的第1个时刻的,在这里也就是'&lt;sos&gt;';第6-18行则是整个循环解码过程,在下一个时刻为EOS_IDX或者达到最大长度后停止;第8-9行是根据当前解码器输入的长度生成注意力掩码矩阵tgt_mask;第10行是根据memory以及当前时刻的输入对当前时刻的输出进行解码;第12-14行则是分类得到当前时刻的解码输出结果;第15行则是将当前时刻的解码输出结果头当前时刻之前所有的输入进行拼接,以此再对下一个时刻的输出进行预测。

最后,我们只需要调用如下函数便可以完成对原始输入语言的翻译任务:

def translate_german_to_english(src, config):
    data_loader = LoadEnglishGermanDataset(config.train_corpus_file_paths,
                                           batch_size=config.batch_size,
                                           tokenizer=my_tokenizer)
    translation_model = TranslationModel(src_vocab_size=len(data_loader.de_vocab),
                                         tgt_vocab_size=len(data_loader.en_vocab),
                                         d_model=config.d_model,
                                         nhead=config.num_head,
                                         num_encoder_layers=config.num_encoder_layers,
                                         num_decoder_layers=config.num_decoder_layers,
                                         dim_feedforward=config.dim_feedforward,
                                         dropout=config.dropout)
    translation_model = translation_model.to(config.device)
    torch.load(config.model_save_dir + '/model.pkl')
    r = translate(translation_model, src, data_loader, config)
    return r

if __name__ == '__main__':
    srcs = ["Eine Gruppe von Menschen steht vor einem Iglu.",
            "Ein Mann in einem blauen Hemd steht auf einer Leiter und putzt ein Fenster."]
    tgts = ["A group of people are facing an igloo.",
            "A man in a blue shirt is standing on a ladder cleaning a window."]
    config = Config()
    for i, src in enumerate(srcs):
        r = translate_german_to_english(src, config)
        print(f"德语:{src})
        print(f"翻译:{r})
        print(f"英语:{tgts[i]})

在上述代码中,第5-14行是定义网络结构,以及恢复本地保存的网络权重;第15行则是开始执行翻译任务;第19-28行为翻译示例,其输出结果为:

德语:Eine Gruppe von Menschen steht vor einem Iglu.
翻译:A group of people standing in fraon of an igloo .
英语:A group of people are facing an igloo.
=========
德语:Ein Mann in einem blauen Hemd steht auf einer Leiter und putzt ein Fenster.
翻译:A man in a blue shirt is standing on a ladder cleaning a window.
英语:A man in a blue shirt is standing on a ladder cleaning a window.

其中第1句德语为训练集之外的数据。

6 基于Transformer的分类模型

经过前面几节内容的介绍,相信大家对于Transformer的原理应该有了一个比较清晰的认识。不过要想做到灵活运用Transformer结构,那就还得再看看其它情况下的运用场景。在这部分内容中,笔者将会以AG_News数据集为例,来搭建一个基于Transformer结构的文本分类模型。

4万字50余图3个实战示例一网打尽Transformer

::: hljs-center
图 6-1. 基于Transformer Encoder的文本分类网络结构图
:::

如图6-1所示便是一个基于Transformer结构的文本分类模型。不过准确的说应该只是一个基于Transformer中Encoder的文本分类模型。这是因为在文本分类任务中并没有解码这一过程,所以我们只需要将Encoder编码得到的向量输入到分类器中进行分类即可。同时需要注意的是,Encoder部分最后输出张量的形状为[batch_size,d_model,src_len](图6-1中Encoder输出的src_len为7),我们需要根据相应策略来进行下一步的处理,具体见后文。

这里值得一提的是,其实BERT模型的网络结构本质上就等同于图6-1所示的网络结构,看懂了图6-1所示的结果对于后续理解BERT问题就不大了。

本部分完整代码可参见[12]。

6.1 数据预处理

6.1.1 语料介绍

在正式介绍模型之前,我们还是先来看看后续所要用到的AG_News数据集。AG_News新闻主题分类数据集是通过从原始语料库中选择 4 个最大的类构建的。每个类包含 30000 个训练样本和 1900 个测试样本。训练样本总数为 120000条,测试总数为 7600条。AG_News原始数据大概长这样:

"3","Fears for T N pension after talks","Unions representing workers at Turner   Newall say they are 'disappointed' after talks with stricken parent firm Federal Mogul."
"4","The Race is On: Second Private Team Sets Launch Date for Human Spaceflight (SPACE.com)","SPACE.com - TORONTO, Canada -- A second\team of rocketeers competing for the  #36;10 million Ansari X Prize, a contest for\privately funded suborbital space flight, has officially announced the first\launch date for its manned rocket."
"4","Ky. Company Wins Grant to Study Peptides (AP)","AP - A company founded by a chemistry researcher at the University of Louisville won a grant to develop a method of producing better peptides, which are short chains of amino acids, the building blocks of proteins."

上述一共包含有3个样本,每1行为1个样本。同时,所有样本均使用逗号作为分隔符,一共包含有 3 列,分别对应类标(1 到 4)、标题和新闻描述。在本篇文中中,我们暂时只使用新闻描述作为输入(当然也可以用title作为输入来进行分类)。

对于该数据的载入,你可以使用Pytorch中的方法来下载并使用[13]:

from torchtext.datasets import AG_NEWS
train_iter = AG_NEWS(split='train')

也可以自己下载原始数据来进行处理。在这里为了延续使用与熟悉第5节中介绍的预处理代码,所以这里我们暂不使用Pytorch内置的代码。

6.1.2 数据集构建

由于分类模型数据集的构建过程并不复杂,其大部分工作在第5节中其实已经介绍过了,所以这里笔者就只是简单的介绍一下。

第1步:定义tokenize

对于英文语料的处理,可以继续沿用5.1.3节中对英文语料的类似处理方式,实现代码如下:

def my_tokenizer(s):
    tokenizer = get_tokenizer('basic_english')
    return tokenizer(s)

第2步:定义字符串清理

从6.1.1节中的示例语料中可以看到, 原始语料中有很多奇奇怪怪的字符,因此还需要对其稍微做一点处理。例如①只保留字母、数字、以及常用标点;②全部转换为小写字等。当然,你也可以自己再添加其它处理方式。具体代码如下:

def clean_str(string):
    string = re.sub("[^A-Za-z0-9\-\?\!\.\,]", " ", string).lower()
    return string

第3步:建立词表

在介绍完tokenize和字符串清理的实现方法后,我们就可以正式通过torchtext.vocab中的Vocab方法来构建词典了,代码如下:

def build_vocab(tokenizer, filepath, min_freq, specials=None):
    if specials is None:
        specials = ['<unk>', '<pad>']
    counter = Counter()
    with open(filepath, encoding='utf8') as f:
        for string_ in tqdm(f):
            string_ = string_.strip().split('","')[-1][:-1]  # 取标签和新闻描述
            counter.update(tokenizer(clean_str(string_)))
    return Vocab(counter, min_freq=min_freq, specials=specials)

在上述代码中,tokenizer表示所使用的分词器; min_freq表示最小字频,去掉小于min_freq的字。第3行代码用来指定特殊的字符;第5-8行代码用来遍历文件中的每一个样本(每行一个)并进行tokenize和计数,其中对于counter.update进行介绍可以参考[10];第9行则是返回最后得到词典。

在完成上述过程后,我们将得到一个Vocab类的实例化对象:

{'<unk>': 0, '<pad>': 1, 'the': 2, '.': 3, ',': 4, 'a': 5, 'to': 6, 'of': 7, 'in': 8, 'and': 9, 'on': 10, 's': 11, 'for': 12, '-': 13, '39': 14, 'that': 15,.....}

接下来,我们就需要定义一个类,并在类的初始化过程中根据训练语料完成字典的构建,代码如下:

class LoadSentenceClassificationDataset():
    def __init__(self, train_file_path=None,  # 训练集路径
                 tokenizer=None,
                 batch_size=20,
                 min_freq=1,  # 最小词频,去掉小于min_freq的词
                 max_sen_len='same'):  # 最大句子长度,默认设置其长度为整个数据集中最长样本的长度
        # max_sen_len = None时,表示按每个batch中最长的样本长度进行padding
        # 根据训练预料建立字典
        self.tokenizer = tokenizer
        self.min_freq = min_freq
        self.specials = ['<unk>', '<pad>']
        self.vocab = build_vocab(self.tokenizer,filepath=train_file_path,
                                 min_freq=self.min_freq,specials=self.specials)
        self.PAD_IDX = self.vocab['<pad>']
        self.UNK_IDX = self.vocab['<unk>']
        self.batch_size = batch_size
        self.max_sen_len = max_sen_len

第4步:转换为Token序列

在得到构建的字典后,便可以通过如下函数来将训练集和测试集转换成Token序列:

    def data_process(self, filepath):
        """
        将每一句话中的每一个词根据字典转换成索引的形式,同时返回所有样本中最长样本的长度
        :param filepath: 数据集路径
        :return:
        """

        raw_iter = iter(open(filepath, encoding="utf8"))
        data = []
        max_len = 0
        for raw in tqdm(raw_iter):
            line = raw.rstrip("\n").split('","')
            s, l = line[-1][:-1], line[0][1:]
            s = clean_str(s)
            tensor_ = torch.tensor([self.vocab[token] for token in
                                    self.tokenizer(s)], dtype=torch.long)
            l = torch.tensor(int(l) - 1, dtype=torch.long)# 标签
            max_len = max(max_len, tensor_.size(0))
            data.append((tensor_, l))
        return data, max_len

在上述代码中,第11-4行分别用来将原始输入序列转换为对应词表中的Token形式。在处理完成后,就会得到类似如下的结果:

[(tensor([   25,    65,    45,  1487,     5,  4062,  3291,    10,  2918, 20217,
            4,  1842,  4512,  1161,    15,   143,   140,  3658, 21658,  4762,
           40,   146,   409,    22,     8,    25,    65,     4,    16,     5,
          142,   287,    15,     4,   633,    39,   146,   409,    22,  5474,
            3,    40]), tensor(1)), .....]

第5步:padding处理

由于对于不同的样本来说其对应的序列长度通常来说都是不同的,但是在将数据输入到相应模型时却需要保持同样的长度。因此在这里我们就需要对Token序列化后的样本进行padding处理,具体代码如下:

def pad_sequence(sequences, batch_first=False, max_len=None, padding_value=0):
    max_size = sequences[0].size()
    trailing_dims = max_size[1:]
    length = max_len
    max_len = max([s.size(0) for s in sequences])
    if length is not None:
        max_len = max(length, max_len)
    if batch_first:
        out_dims = (len(sequences), max_len) + trailing_dims
    else:
        out_dims = (max_len, len(sequences)) + trailing_dims
    out_tensor = sequences[0].data.new(*out_dims).fill_(padding_value)
    for i, tensor in enumerate(sequences):
        length = tensor.size(0)
        # use index notation to prevent duplicate references to the tensor
        if batch_first:
            out_tensor[i, :length, ...] = tensor
        else:
            out_tensor[:length, i, ...] = tensor
    return out_tensor

在上述代码中,max_len 表示 最大句子长度,默认为None,即在每个batch中以最长样本的长度对其它样本进行padding;当然同样也可以指定max_len的值为整个数据集中最长样本的长度进行padding处理。padding处理后的结果类似如下:

tensor([[    2,   342,   578,  ...,     1,     1,     1],
        [   32,    13, 14585,  ...,     1,     1,     1],
        [ 1189,    11,   327,  ...,     1,     1,     1],...)

末尾的1即是padding的部分。

在定义完pad_sequence这个函数后,我们便可以通过它来对每个batch中的数据集进行padding处理:

    def generate_batch(self, data_batch):
        batch_sentence, batch_label = [], []
        for (sen, label) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
            batch_sentence.append(sen)
            batch_label.append(label)
        batch_sentence = pad_sequence(batch_sentence,  # [batch_size,max_len]
                                      padding_value=self.PAD_IDX,
                                      batch_first=False,
                                      max_len=self.max_sen_len)
        batch_label = torch.tensor(batch_label, dtype=torch.long)
        return batch_sentence, batch_label

第6步:构造DataLoade与使用示例

经过前面5步的操作,整个数据集的构建就算是已经基本完成了,只需要再构造一个DataLoader迭代器即可,代码如下:

    def load_train_val_test_data(self, train_file_paths, test_file_paths):
        train_data, max_sen_len = self.data_process(train_file_paths)  # 得到处理好的所有样本
        if self.max_sen_len == 'same':
            self.max_sen_len = max_sen_len
        test_data, _ = self.data_process(test_file_paths)
        train_iter = DataLoader(train_data, batch_size=self.batch_size,  # 构造DataLoader
                                shuffle=True, collate_fn=self.generate_batch)
        test_iter = DataLoader(test_data, batch_size=self.batch_size,
                               shuffle=True, collate_fn=self.generate_batch)
        return train_iter, test_iter

在上述代码中,第2-5行便是分别用来将训练集和测试集转换为Token序列;第6-9行则是分别构造2个DataLoader,其中generate_batch将作为一个参数传入来对每个batch的样本进行处理。在完成类LoadSentenceClassificationDataset所有的编码过程后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    path = "./data/ag_news_csv/test.csv"
    data_loader = LoadSentenceClassificationDataset(train_file_path=path,
                                                    tokenizer=my_tokenizer,
                                                    max_sen_len=None)
    train_iter, test_iter = data_loader.load_train_val_test_data(path, path)
    for sample, label in train_iter:
        print(sample.shape)  # [seq_len,batch_size]

最后,由于Encoder只会在padding部分有mask操作,所以每个样本的key_padding_mask向量我们在训练部分再生成即可。下面,我们正式进入到文本分类模型部分的介绍。

6.2 文本分类模型

6.2.1 网络结构

总体来说,基于Transformer文本分类模型的网络结构其实就是图6-1所展示的所有部分,当然你还可以使用多个Encoder进行堆叠。最后,只需要将Encdder的输出喂入到一个softmax分类器即可完成分类任务。不过这里有两个细节的地方需要大家注意:

①根据第4节的介绍可知,Encoder在编码结束后输出的形状为[src_len,batch_size,embed_dim](这里的src_len也可以理解为LSTM中time step的概念)。因此,在构造最后分类器的输入时就可以有多种不同的形式,例如只取最后一个位置上的向量、或者是取所有位置向量的平均(求和)等都可以。后面笔者也会将这3种方式都实现供大家参考。

②由于每个样本长度给不相同,因此在对样本进行padding的时候就有两种方式。一般来说在大多数模型中多需要保持所有的样本具有相同的长度,不过由于这里我们使用的是自注意力的编码机制,因此只需要保持同一个batch中的样本长度一致即可。不过后面笔者对这两种方式都进行了实现,只需要通过max_sen_len这个参数来控制即可。

首先,我们需要定义一个名为ClassificationModel的类,其前向传播过程代码如下所示:

class ClassificationModel(nn.Module):
    def __init__(self, vocab_size=None,
                 d_model=512, nhead=8,
                 num_encoder_layers=6,
                 dim_feedforward=2048,
                 dim_classification=64,
                 num_classification=4,
                 dropout=0.1):
        super(ClassificationModel, self).__init__()
        self.pos_embedding = PositionalEncoding(d_model=d_model, dropout=dropout)
        self.src_token_embedding = TokenEmbedding(vocab_size, d_model)
        encoder_layer = MyTransformerEncoderLayer(
                          d_model, nhead,dim_feedforward,dropout)
        encoder_norm = nn.LayerNorm(d_model)
        self.encoder = MyTransformerEncoder(encoder_layer,
                                            num_encoder_layers, encoder_norm)
        self.classifier = nn.Sequential(nn.Linear(d_model, dim_classification),
                                        nn.Dropout(dropout),
                                        nn.Linear(dim_classification, num_classification))

在上述代码中,第10-11行定义了Transformer中的Embedding操作;第12-13行定义了Transformer中的EncoderLayer;第14-16行用来定义Transformer中的Encoder;第17-19行用来定义一个分类器。最后,整个网络的前向传播过程如下:

    def forward(self,
                src,  # [src_len, batch_size]
                src_mask=None,
                src_key_padding_mask=None,  # [batsh_size, src_len]
                concat_type='sum'  # 解码之后取所有位置相加,还是最后一个位置作为输出
                ):
        src_embed = self.src_token_embedding(src)  # [src_len, batch_size, embed_dim]
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
        memory = self.encoder(src=src_embed,
                              mask=src_mask,
                              src_key_padding_mask=src_key_padding_mask)
        # [src_len,batch_size,embed_dim]
        if concat_type == 'sum':
            memory = torch.sum(memory, dim=0)
        elif concat_type == 'avg':
            memory = torch.sum(memory, dim=0) / memory.size(0)
        else:
            memory = memory[-1, ::]  # 取最后一个时刻
        # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
        out = self.classifier(memory)  # 输出logits
        return out  # [batch_size, num_class]

在上述代码中,第7-11行用来执行编码器的前向传播过程;第13-18行便是用来选择以何种方式来选择分类器的输入,经笔者实验后发现取各个位置的平均值效果最好;第20-21行便是将经过分类器后的输出进行返回。

在定义完logits的前向传播过后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    src_len = 7
    batch_size = 2
    dmodel = 32
    num_head = 4
    src = torch.tensor([[4, 3, 2, 6, 0, 0, 0],
                        [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1)  # 转换成 [src_len, batch_size]
    src_key_padding_mask = torch.tensor([[True, True, True, True, False, False, False],
                                         [True, True, True, True, True, False, False]])
    model = ClassificationModel(vocab_size=10, d_model=dmodel, nhead=num_head)
    logits = model(src, src_key_padding_mask=src_key_padding_mask)
    print(logits.shape) # torch.Size([2, 4])

6.2.2 模型训练

在定义完成整个分类模型的网络结构后下面就可以开始训练模型了。由于这部分代码较长,所以下面笔者依旧以分块的形式进行介绍:

第1步:载入数据集

def train_model(config):
    data_loader = LoadSentenceClassificationDataset(config.train_corpus_file_paths,
                                                    my_tokenizer,
                                                    batch_size=config.batch_size,
                                                    min_freq=config.min_freq,
                                                    max_sen_len=config.max_sen_len)
    train_iter, test_iter = data_loader.load_train_val_test_data(
        config.train_corpus_file_paths, config.test_corpus_file_paths)

首先我们可以根据前面的介绍,通过类LoadSentenceClassificationDataset来载入数据集,其中config中定义了模型所涉及到的所有配置参数。同时,可以通过max_sen_len参数来控制padding时保持所有样本一样还是仅在每个batch内部一样。

第2步:定义模型并初始化权重

    classification_model = ClassificationModel(vocab_size=len(data_loader.vocab),
                                               d_model=config.d_model,
                                               nhead=config.num_head,
                                               num_encoder_layers=config.num_encoder_layers,
                                               dim_feedforward=config.dim_feedforward,
                                               dim_classification=config.dim_classification,
                                               num_classification=config.num_class,
                                               dropout=config.dropout)

    for p in classification_model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

在载入数据后,便可以定义一个文本分类模型ClassificationModel,并根据相关参数对其进行实例化;同时,可以对整个模型中的所有参数进行一个初始化操作。

第3步:定义损失学习率与优化器

    loss_fn = torch.nn.CrossEntropyLoss()
    learning_rate = CustomSchedule(config.d_model)
    optimizer = torch.optim.Adam(classification_model.parameters(),
                                 lr=0.,
                                 betas=(config.beta1, config.beta2), 
                                 eps=config.epsilon)

在上述代码中,第1行是定义交叉熵损失函数;第2行代码则是论文中所提出来的动态学习率计算过程,其计算公式为(和公式$(5.1)$没有差别。):
$$
\text{lrate}=d_{model}^{-0.5}\cdot \text{min}(\text{step_num}^{-0.5},\text{step_num}\cdot \text{warmup_steps}^{-1.5})\;\;\;\;\;\;(6.1)
$$
具体实现代码同5.2.2节中的一样,这里就不再赘述。

第4步:开始训练

    for epoch in range(config.epochs):
        losses = 0
        start_time = time.time()
        for idx, (sample, label) in enumerate(train_iter):
            sample = sample.to(config.device)  # [src_len, batch_size]
            label = label.to(config.device)
            padding_mask = (sample == data_loader.PAD_IDX).transpose(0, 1)
            logits = classification_model(sample,
                                          src_key_padding_mask=padding_mask)  
                                                                        # [batch_size,num_class]
            optimizer.zero_grad()
            loss = loss_fn(logits, label)
            loss.backward()
            lr = learning_rate()
            for p in optimizer.param_groups:
                p['lr'] = lr
            optimizer.step()
            losses += loss.item()

            acc = (logits.argmax(1) == label).float().mean()
            if idx % 10 == 0:
                print(f"Epoch: {epoch}, Batch[{idx}/{len(train_iter)}], "
                      f"Train loss :{loss.item():.3f}, Train acc: {acc:.3f}")
        end_time = time.time()
        train_loss = losses / len(train_iter)
        print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Epoch time = {(end_time -                                             start_time):.3f}s")

在上述代码中,第7行代码用来生成每个样本对应的padding mask向量;第15-16行是将每个step更新后的学习率送入到模型中。以下便是模型训练过程中的输出:

    Epoch: 9, Batch: [410/469], Train loss 0.186, Train acc: 0.938
    Epoch: 9, Batch: [420/469], Train loss 0.150, Train acc: 0.938
    Epoch: 9, Batch: [430/469], Train loss 0.269, Train acc: 0.941
    Epoch: 9, Batch: [440/469], Train loss 0.197, Train acc: 0.925
    Epoch: 9, Batch: [450/469], Train loss 0.245, Train acc: 0.917
    Epoch: 9, Batch: [460/469], Train loss 0.272, Train acc: 0.902
    Accuracy on test 0.886

本部分完整代码可参见[12]。

7 基于Transformer的对联模型

经过前面6节内容的介绍,对于Transformer相信大家都应该理解得差不多了。不过要想做到灵活运用Transformer结构,那就还得多看看其它场景下的运用。在接下来的这部分内容中,笔者将会以一个含有70余万条的对联数据集为例,来搭建一个基于Transformer结构的对联生成模型。同时,这也是介绍Transformer结构系列文章的最后一节内容。

![]https://moonhotel.oss-cn-shanghai.aliyuncs.com/images/2105291232322424560.jpg)

::: hljs-center
图 7-1. Transformer网络结构图
:::

如图7-1所示便是一个基于Transformer结构的对联生成模型。可以看出,其实它与前面介绍的基于Transformer结构的翻译模型没有特别的变换,唯一不同的可能就是在对联生成模型中解码器和编码器共用了同一个词表(因为两者都是中文)。

7.1 数据预处理

7.1.1 语料介绍

按老规矩,在正式介绍模型搭建之前我们还是先来看看数据集都长什么样。本次所使用到的数据集是一个网上公开的对联数据集,在github中搜索"couplet-dataset”就能找到,其一共包含有770491条训练样本,4000条测试样本。和翻译数据集类似,对联数据集也包含上下两句:

# 上联 in.txt
风 弦 未 拨 心 先 乱 
花 梦 粘 于 春 袖 口 
晋 世 文 章 昌 二 陆 
# 下联 out.txt
夜 幕 已 沉 梦 更 闲 
莺 声 溅 落 柳 枝 头 
魏 家 词 赋 重 三 曹 

如上所示便是3条样本,分别存放在in.txtout.txt这两个文件中。可以看出,原始数据已经做了分字这步操作,所有后续我们只需要进行简单的split操作即可。

7.1.2 数据集构建

总体上来说对联生成模型的数据集构建过程和翻译模型的数据集构建过程基本上没有太大差别,主要步骤同样也是:①构建字典;②将文本中的每一个词(字)转换为Token序列;③对不同长度的样本序列按照某个标准进行padding处理;④构建DataLoader类。

第1步:定义tokenize

由于原始数据每个字已经被空格隔开了,所以这里tokenizer的定义只需要进行split操作即可,代码如下:

def my_tokenizer(s):
    return s.split()

可以看到,其实也非常简单。例如对于如下文本来说

腾 飞 上 铁 , 锐 意 改 革 谋 发 展 , 勇 当 千 里 马 

其tokenize后的结果为:

['腾', '飞', '上', '铁', ',', '锐', '意', '改', '革', '谋', '发', '展', ',', '勇', '当', '千', '里', '马']

第2步:建立词表

在介绍完tokenize的实现方法后,我们就可以正式通过torchtext.vocab中的Vocab方法来构建词典了,代码如下:

def build_vocab(tokenizer, filepath, min_freq=1, specials=None):
    if specials is None:
        specials = ['<unk>', '<pad>', '<bos>', '<eos>']
    counter = Counter()
    with open(filepath[0], encoding='utf8') as f:
        for string_ in f:
            counter.update(tokenizer(string_))
    with open(filepath[1], encoding='utf8') as f:
        for string_ in f:
            counter.update(tokenizer(string_))
    return Vocab(counter, specials=specials, min_freq=min_freq)

在上述代码中,第3行代码用来指定特殊的字符;第5-10行分别用来遍历in.txt文件和out.txt文件中的每一个样本(每行一个)并进行tokenize和计数,其中对于counter.update进行介绍可以参考[10];第8行则是返回最后得到词典。值得注意的是,由于在对联生成这一场景中编码器和解码器共用的是一个词表,所以这里同时对in.txtout.txt文件进行了遍历。

在完成上述过程后,我们将得到一个Vocab类的实例化对象,即:

{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, ',': 4, '风': 5, '春': 6, '一': 7, '人': 8, '月': 9, '山': 10, '心': 11, '花': 12, '天': 13, ...}

此时,我们就需要定义一个类,并在类的初始化过程中根据训练语料完成字典的构建,代码如下:

class LoadCoupletDataset():
    def __init__(self, train_file_paths=None, tokenizer=None,
                 batch_size=2, min_freq=1):
        # 根据训练预料建立字典,由于都是中文,所以共用一个即可
        self.tokenizer = tokenizer
        self.vocab = build_vocab(self.tokenizer, filepath=train_file_paths, min_freq=min_freq)
        self.specials = ['<unk>', '<pad>', '<bos>', '<eos>']
        self.PAD_IDX = self.vocab['<pad>']
        self.BOS_IDX = self.vocab['<bos>']
        self.EOS_IDX = self.vocab['<eos>']
        self.batch_size = batch_size

第3步:转换为Token序列

在得到构建的字典后,便可以通过如下函数来将训练集和测试集转换成Token序列:

    def data_process(self, filepaths):
        """
        将每一句话中的每一个词根据字典转换成索引的形式
        :param filepaths:
        :return:
        """
        raw_in_iter = iter(open(filepaths[0], encoding="utf8"))
        raw_out_iter = iter(open(filepaths[1], encoding="utf8"))
        data = []
        for (raw_in, raw_out) in zip(raw_in_iter, raw_out_iter):
            in_tensor_ = torch.tensor([self.vocab[token] for token in
                                       self.tokenizer(raw_in.rstrip("\n"))], dtype=torch.long)
            out_tensor_ = torch.tensor([self.vocab[token] for token in
                                        self.tokenizer(raw_out.rstrip("\n"))], dtype=torch.long)
            data.append((in_tensor_, out_tensor_))
        return data

在上述代码中,第11-4行分别用来将原始序列上联和目标序列下联转换为对应词表中的Token形式。在处理完成后,就会得到类似如下的结果:

[(tensor([   5,  549,  250, 1758,   11,  228,  651]), tensor([ 154, 1420,  310,  598,   29,  206,  164])), (tensor([  12,   29, 3218,  262,    6,  628,  419]), tensor([ 441,   62, 2049,   93,   66,  304,  111])), (tensor([1137,   40,   47,  286,  819,  364, 1383]), tensor([1803,   49,  586,  556,  126,   25, 1830])), (tensor([  7, 291, 138, 115, 216, 151,   9]), tensor([ 15, 311, 107,  57,  80,   5,  21]))]

其中左边的一列就是原始序列上联的Token形式,右边一列就是目标序列下联的Token形式,每一行构成一个样本。

第4步:padding处理

从上面的输出结果可以看到,无论是对于原始序列来说还是目标序列来说,在不同的样本中其对应长度都不尽相同。但是在将数据输入到相应模型时却需要保持同样的长度,因此在这里我们就需要对Token序列化后的样本进行padding处理。同时需要注意的是,一般在这种生成模型中,模型在训练过程中只需要保证同一个batch中所有的原始序列等长,所有的目标序列等长即可,也就是说不需要在整个数据集中所有样本都保证等长。

因此,在实际处理过程中无论是原始序列还是目标序列都会以每个batch中最长的样本为标准对其它样本进行padding,具体代码如下:

    def generate_batch(self, data_batch):
        in_batch, out_batch = [], []
        for (in_item, out_item) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
            in_batch.append(in_item)  # 编码器输入序列不需要加起止符
            # 在每个idx序列的首位加上 起始token 和 结束 token
            out = torch.cat([torch.tensor([self.BOS_IDX]), 
                             out_item, torch.tensor([self.EOS_IDX])], dim=0)
            out_batch.append(out)
        # 以最长的序列为标准进行填充
        in_batch = pad_sequence(in_batch, padding_value=self.PAD_IDX)  # [de_len,batch_size]
        out_batch = pad_sequence(out_batch, padding_value=self.PAD_IDX)  # [en_len,batch_size]
        return in_batch, out_batch

在上述代码中,第6-7行用来在目标序列的首尾加上特定的起止符;第10-11行则是分别对一个batch中的原始序列和目标序列以各自当中最长的样本为标准进行padding(这里的pad_sequence导入自torch.nn.utils.rnn)。

第5步:构造mask向量

在处理完成前面几个步骤后,进一步需要根据src_inputtgt_input来构造相关的mask向量,具体代码如下:

    def generate_square_subsequent_mask(self, sz, device):
        mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def create_mask(self, src, tgt, device='cpu'):
        src_seq_len = src.shape[0]
        tgt_seq_len = tgt.shape[0]
        tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device)  # [tgt_len,tgt_len]
        # Decoder的注意力Mask输入,用于掩盖当前position之后的position,所以这里是一个对称矩阵
        src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
        # Encoder的注意力Mask输入,这部分其实对于Encoder来说是没有用的,所以这里全是0
        src_padding_mask = (src == self.PAD_IDX).transpose(0, 1)
        # 用于mask掉Encoder的Token序列中的padding部分,[batch_size, src_len]
        tgt_padding_mask = (tgt == self.PAD_IDX).transpose(0, 1)
        # 用于mask掉Decoder的Token序列中的padding部分,batch_size, tgt_len
        return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

在上述代码中,第1-4行是用来生成一个形状为[sz,sz]的注意力掩码矩阵,用于在解码过程中掩盖当前position之后的position;第6-17行用来返回Transformer中各种情况下的mask矩阵,其中src_mask在这里并没有作用。

第6步:构造DataLoade与使用示例

经过前面5步的操作,整个数据集的构建就算是已经基本完成了,只需要再构造一个DataLoader迭代器即可,代码如下:

    def load_train_val_test_data(self, train_file_paths, test_file_paths):
        train_data = self.data_process(train_file_paths)
        test_data = self.data_process(test_file_paths)
        train_iter = DataLoader(train_data, batch_size=self.batch_size,
                                shuffle=True, collate_fn=self.generate_batch)
        test_iter = DataLoader(test_data, batch_size=self.batch_size,
                               shuffle=True, collate_fn=self.generate_batch)
        return train_iter, test_iter

在上述代码中,第2-3行便是分别用来将训练集和测试集转换为Token序列;第4-7行则是分别构造2个DataLoader,其中generate_batch将作为一个参数传入来对每个batch的样本进行处理。在完成类LoadCoupletDataset所有的编码过程后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    config = Config()
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
                                     batch_size=config.batch_size,
                                     tokenizer=my_tokenizer,
                                     min_freq=config.min_freq)
    train_iter, test_iter = data_loader.load_train_val_test_data(config.test_corpus_file_paths,
                                                                 config.test_corpus_file_paths)
    print(data_loader.PAD_IDX)
    for src, tgt in train_iter:
        tgt_input = tgt[:-1, :]
        tgt_out = tgt[1:, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = data_loader.create_mask(src, tgt_input)
        print(src)
        print(tgt)
        print("src shape:", src.shape)  # [in_tensor_len,batch_size]
        print("tgt shape:", tgt.shape)  # [in_tensor_len,batch_size]
        print("src input shape:", src.shape)
        print("src_padding_mask shape (batch_size, src_len): ", src_padding_mask.shape)
        print("tgt input shape:", tgt_input.shape)
        print("tgt_padding_mask shape: (batch_size, tgt_len) ", tgt_padding_mask.shape)
        print("tgt output shape:", tgt_out.shape)
        print("tgt_mask shape (tgt_len,tgt_len): ", tgt_mask.shape)

在介绍完数据集构建的整个过程后,下面就开始正式进入到翻译模型的构建中。如果对于这部分不是特别理解的话,建议先看第5节中的数据处理流程图进行理解。

7.2 对联生成模型

7.2.1 网络结构

总体来说,基于Transformer的对联生成模型的网络结构其实就是图7-1所展示的所有部分,只是在前面介绍Transformer网络结构时笔者并没有把Embedding部分的实现给加进去。这是因为对于不同的文本生成模型,其Embedding部分会不一样(例如在本场景中编码器和解码器共用一个TokenEmbedding即可,而在翻译模型中就需要两个),所以将两者进行了拆分。同时,待模型训练完成后,在inference过程中Encoder只需要执行一次,所以在此过程中也需要单独使用Transformer中的Encoder和Decoder。

首先,我们需要定义一个名为CoupletModel的类,其前向传播过程代码如下所示:

class CoupletModel(nn.Module):
    def __init__(self, vocab_size,
                 d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048,
                 dropout=0.1):
        super(CoupletModel, self).__init__()
        self.my_transformer = MyTransformer(d_model=d_model,
                                            nhead=nhead,
                                            num_encoder_layers=num_encoder_layers,
                                            num_decoder_layers=num_decoder_layers,
                                            dim_feedforward=dim_feedforward,
                                            dropout=dropout)
        self.pos_embedding = PositionalEncoding(d_model=d_model, dropout=dropout)
        self.token_embedding = TokenEmbedding(vocab_size, d_model)
        self.classification = nn.Linear(d_model, vocab_size)

    def forward(self, src=None, tgt=None, src_mask=None,
                tgt_mask=None, memory_mask=None, src_key_padding_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        """
        :param src: Encoder的输入 [src_len,batch_size]
        :param tgt: Decoder的输入 [tgt_len,batch_size]
        :param src_key_padding_mask: 用来Mask掉Encoder中不同序列的padding部分,[batch_size, src_len]
        :param tgt_key_padding_mask: 用来Mask掉Decoder中不同序列的padding部分 [batch_size, tgt_len]
         memory_key_padding_mask: 用来Mask掉Encoder输出的memory中不同序列的padding部分 [batch_size, src_len]
        :return:
        """
        src_embed = self.token_embedding(src)  # [src_len, batch_size, embed_dim]
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
        tgt_embed = self.token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
        outs = self.my_transformer(src=src_embed, tgt=tgt_embed, src_mask=src_mask,
                                   tgt_mask=tgt_mask, memory_mask=memory_mask,
                                   src_key_padding_mask=src_key_padding_mask,
                                   tgt_key_padding_mask=tgt_key_padding_mask,
                                   memory_key_padding_mask=memory_key_padding_mask)
        # [tgt_len,batch_size,embed_dim]
        logits = self.classification(outs)  # [tgt_len,batch_size,tgt_vocab_size]
        return logits

在上述代码中,第7-12行便是用来定义一个Transformer结构;第13-15分别用来定义Positional Embedding、Token Embedding和最后的分类器(需要注意的是这里是共用同一个Token Embedding);第28-38行便是用来执行整个前向传播过程,其中Transformer的整个前向传播过程在第4节中已经介绍过,在这里就不再赘述。

在定义完logits的前向传播过后,便可以通过如下形式进行使用:

if __name__ == '__main__':
    src_len = 7
    batch_size = 2
    dmodel = 32
    tgt_len = 8
    num_head = 4
    src = torch.tensor([[4, 3, 2, 6, 0, 0, 0],
                        [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1)  # 转换成 [src_len, batch_size]
    src_key_padding_mask = torch.tensor([[True, True, True, True, False, False, False],
                                         [True, True, True, True, True, False, False]])
    tgt = torch.tensor([[1, 3, 3, 5, 4, 3, 0, 0],
                        [1, 6, 8, 2, 9, 1, 0, 0]]).transpose(0, 1)
    tgt_key_padding_mask = torch.tensor([[True, True, True, True, True, True, False, False],
                                         [True, True, True, True, True, True, False, False]])
    trans_model = CoupletModel(vocab_size=10, d_model=dmodel,
                               nhead=num_head, num_encoder_layers=6,
                               num_decoder_layers=6, dim_feedforward=30,
                               dropout=0.1)
    tgt_mask = trans_model.my_transformer.generate_square_subsequent_mask(tgt_len)
    logits = trans_model(src, tgt=tgt, tgt_mask=tgt_mask,
                         src_key_padding_mask=src_key_padding_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=src_key_padding_mask)
    print(logits.shape)

接着,我们需要再定义一个EncoderDecoder在inference中进行使用,代码如下:

    def encoder(self, src):
        src_embed = self.token_embedding(src)  # [src_len, batch_size, embed_dim]
        src_embed = self.pos_embedding(src_embed)  # [src_len, batch_size, embed_dim]
        memory = self.my_transformer.encoder(src_embed)
        return memory # [src_len, batch_size, embed_dim]

    def decoder(self, tgt, memory, tgt_mask):
        tgt_embed = self.tgt_token_embedding(tgt)  # [tgt_len, batch_size, embed_dim]
        tgt_embed = self.pos_embedding(tgt_embed)  # [tgt_len, batch_size, embed_dim]
        outs = self.my_transformer.decoder(tgt_embed, memory=memory,
                                           tgt_mask=tgt_mask)  # [tgt_len,batch_size,embed_dim]
        return outs

在上述代码中,第1-5行用于在inference时对输入序列进行编码并得到memory(只需要执行一次);第7-11行用于根据memory和当前解码时刻的输入对输出进行预测,需要循环执行多次,这部分内容详见7.2.3模型预测部分。

7.2.2 模型训练

在定义完成整个对联生成模型的网络结构后下面就可以开始训练模型了。由于这部分代码较长,所以下面笔者依旧以分块的形式进行介绍:

第1步:载入数据集

def train_model(config):
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
                                     batch_size=config.batch_size,
                                     tokenizer=my_tokenizer,
                                     min_freq=config.min_freq)
    train_iter, test_iter = data_loader.load_train_val_test_data(
      config.train_corpus_file_paths,config.test_corpus_file_paths)

首先我们可以根据前面的介绍,通过类LoadCoupletDataset来载入数据集,其中config中定义了模型所涉及到的所有配置参数。

第2步:定义模型并初始化权重

    couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
                                 d_model=config.d_model,
                                 nhead=config.num_head,
                                 num_encoder_layers=config.num_encoder_layers,
                                 num_decoder_layers=config.num_decoder_layers,
                                 dim_feedforward=config.dim_feedforward,
                                 dropout=config.dropout)
    for p in couplet_model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

在载入数据后,便可以定义模型CoupletModel,并根据相关参数对其进行实例化;同时,可以对整个模型中的所有参数进行初始化操作。

第3步:定义损失学习率与优化器

        loss_fn = torch.nn.CrossEntropyLoss(ignore_index=data_loader.PAD_IDX)
    learning_rate = CustomSchedule(config.d_model)
    optimizer = torch.optim.Adam(couplet_model.parameters(),
                                 lr=0.,
                                 betas=(config.beta1, config.beta2), eps=config.epsilon)

在上述代码中,第1行是定义交叉熵损失函数,并同时指定需要忽略的索引ignore_index部分的损失。因为根据tgt_output可知,有些位置上的标签值其实是Padding后的结果,因此在计算损失的时候需要将这些位置给忽略掉。第2行代码则是论文中所提出来的动态学习率计算过程,具体在第5节中已经介绍过来,这里就不再赘述。

第4步:开始训练

    for epoch in range(config.epochs):
        losses = 0
        start_time = time.time()
        for idx, (src, tgt) in enumerate(train_iter):
            src = src.to(config.device)  # [src_len, batch_size]
            tgt = tgt.to(config.device)
            tgt_input = tgt[:-1, :]  # 解码部分的输入, [tgt_len,batch_size]
            src_mask, tgt_mask, src_padding_mask, tgt_padding_mask \
                = data_loader.create_mask(src, tgt_input, config.device)
            logits = couplet_model(
                src=src,  # Encoder的token序列输入,[src_len,batch_size]
                tgt=tgt_input,  # Decoder的token序列输入,[tgt_len,batch_size]
                src_mask=src_mask,  # Encoder的注意力Mask输入,这部分其实对于Encoder来说是没有用的
                tgt_mask=tgt_mask,
                # Decoder的注意力Mask输入,用于掩盖当前position之后的position [tgt_len,tgt_len]
                src_key_padding_mask=src_padding_mask,  # 用于mask掉Encoder的Token序列中的padding部分
                tgt_key_padding_mask=tgt_padding_mask,  # 用于mask掉Decoder的Token序列中的padding部分
                memory_key_padding_mask=src_padding_mask)  # 用于mask掉Encoder的Token序列中的padding部分
            # logits 输出shape为[tgt_len,batch_size,tgt_vocab_size]
            optimizer.zero_grad()
            tgt_out = tgt[1:, :]  # 解码部分的真实值  shape: [tgt_len,batch_size]
            loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
            # [tgt_len*batch_size, tgt_vocab_size] with [tgt_len*batch_size, ]
            loss.backward()
            ......

在上述代码中,第5-9行是用来得到模型各个部分的输入;第10-18行是计算模型整个前向传播的过程;第20-24行则是执行损失计算与反向传播;第26-28则是将每个step更新后的学习率送入到模型中并进行参数更新;第30行是用来计算模型预测的准确率。

以下便是模型训练过程中的输出:

-- INFO: Epoch: 0, Batch[29/3010], Train loss : 8.965, Train acc: 0.094
-- INFO: Epoch: 0, Batch[59/3010], Train loss : 8.618, Train acc: 0.098
-- INFO: Epoch: 0, Batch[89/3010], Train loss : 8.366, Train acc: 0.099
-- INFO: Epoch: 0, Batch[119/3010], Train loss : 8.137, Train acc: 0.109
...

7.2.3 模型预测

在介绍完模型的训练过程后接下来就来看模型的预测部分。生成模型的预测部分不像普通的分类任务只需要将网络最后的输出做argmax操作即可,生成模型在预测过程中往往需要按时刻一步步进行来进行。因此,下面我们这里定义一个couplet函数来执行这一过程,具体代码如下:

def couplet(model, src, data_loader, config):
    vocab = data_loader.vocab
    tokenizer = data_loader.tokenizer
    model.eval()
    tokens = [vocab.stoi[tok] for tok in tokenizer(src)]  # 构造一个样本
    num_tokens = len(tokens)
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1))  # 将src_len 作为第一个维度
    tgt_tokens = greedy_decode(model, src, max_len=num_tokens + 5,
                               start_symbol=data_loader.BOS_IDX, config=config,
                               data_loader=data_loader).flatten()  # 解码的预测结果
    return "".join([vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

在上述代码中,第5行是将待翻译的源序列进行序列化操作;第7-10行则是通过函数greedy_decode函数来对输入进行解码;第11行则是将最后解码后的结果由Token序列在转换成实际的目标语言。同时,greedy_decode函数的实现如下:

def greedy_decode(model, src, max_len, start_symbol, config, data_loader):
    src = src.to(config.device)
    memory = model.encoder(src)  # 对输入的Token序列进行解码翻译
    ys = torch.ones(1, 1).fill_(start_symbol). \
        type(torch.long).to(config.device)  # 解码的第一个输入,起始符号
    for i in range(max_len - 1):
        memory = memory.to(config.device)
        tgt_mask = (model.my_transformer.generate_square_subsequent_mask(ys.size(0))
                    .type(torch.bool)).to(config.device)  # 根据tgt_len产生一个注意力mask矩阵(对称的)
        out = model.decoder(ys, memory, tgt_mask)  # [tgt_len,tgt_vocab_size]
        out = out.transpose(0, 1)  # [tgt_vocab_size, tgt_len]
        prob = model.classification(out[:, -1])  # 只对对预测的下一个词进行分类
        _, next_word = torch.max(prob, dim=1)  # 选择概率最大者
        next_word = next_word.item()
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        # 将当前时刻解码的预测输出结果,同之前所有的结果堆叠作为输入再去预测下一个词。
        if next_word == data_loader.EOS_IDX:  # 如果当前时刻的预测输出为结束标志,则跳出循环结束预测。
            break
    return ys

在上述代码中,第3行是将源序列输入到Transformer的编码器中进行编码并得到Memory;第4-5行是初始化解码阶段输入的第1个时刻的,在这里也就是&lt;sos&gt;;第6-18行则是整个循环解码过程,在下一个时刻为EOS_IDX或者达到最大长度后停止;第8-9行是根据当前解码器输入的长度生成注意力掩码矩阵tgt_mask;第10行是根据memory以及当前时刻的输入对当前时刻的输出进行解码;第12-14行则是分类得到当前时刻的解码输出结果;第15行则是将当前时刻的解码输出结果头当前时刻之前所有的输入进行拼接,以此再对下一个时刻的输出进行预测。

最后,我们只需要调用如下函数便可以完成对原始输入上联的下联生成任务:

def do_couplet(srcs, config):
    data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
                                     batch_size=config.batch_size,
                                     tokenizer=my_tokenizer,
                                     min_freq=config.min_freq)
    couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
                                 d_model=config.d_model,
                                 nhead=config.num_head,
                                 num_encoder_layers=config.num_encoder_layers,
                                 num_decoder_layers=config.num_decoder_layers,
                                 dim_feedforward=config.dim_feedforward,
                                 dropout=config.dropout)
    couplet_model = couplet_model.to(config.device)
    loaded_paras = torch.load(config.model_save_dir + '/model.pkl')
    couplet_model.load_state_dict(loaded_paras)
    results = []
    for src in srcs:
        r = couplet(couplet_model, src, data_loader, config)
        results.append(r)
    return results

if __name__ == '__main__':
    srcs = ["晚风摇树树还挺",
            "忽忽几晨昏,离别间之,疾病间之,不及终年同静好",
            "风声、雨声、读书声,声声入耳",
            "上海自来水来自海上"]
    tgts = ["晨露润花花更红",
            "茕茕小儿女,孱羸若此,娇憨若此,更烦二老费精神",
            "家事、国事、天下事,事事关心",
            ""]
    config = Config()
    srcs = [" ".join(src) for src in srcs]
    results = do_couplet(srcs, config)
    for src, tgt, r in zip(srcs, tgts, results):
        print(f"上联:{''.join(src.split())}")
        print(f" AI:{r}")
        print(f"下联:{tgts}")
        print("=======")

在上述代码中,第6-15行是定义网络结构,以及恢复本地保存的网络权重;第16-20行则是开始执行多个上联的下联生成任务;第34-39行为生成示例,其输出结果为:

上联:晚风摇树树还挺
 AI: 朝露沾花花更红
下联:晨露润花花更红

上联:忽忽几晨昏,离别间之,疾病间之,不及终年同静好
 AI:茕茕小儿女,孱羸若此,娇憨若此,更烦二老费精神
下联:茕茕小儿女,孱羸若此,娇憨若此,更烦二老费精神

上联:风声、雨声、读书声,声声入耳
 AI:山色、水色、烟霞色,色色宜人
下联:家事、国事、天下事,事事关心

上联:上海自来水来自海上
 AI:中山落叶松叶落山中

上联:月来客栈,迎五湖好友,谈百态人生。
 AI:花发老舍,请四海嘉宾,喝一杯清茶。

以上完整代码可参见[14]。

8 总结

在本篇文章中,笔者首先详细地介绍了Transformer论文的动机以及自注意力机制的原理与多头注意力机制的作用;然后介绍了Transformer中的位置编码以及整个预测和训练过程中模型的编码解码过程;接着进一步介绍了Transformer的网络结构、自注意力机制的原理实现以及Transformer网络的实现;最后,笔者还通过三个实例(包括论文中的翻译模型、文本分类模型,以及对联生成模型)从代码的角度来介绍了整个Transformer网络的原理与使用示例。

本次内容就到此结束,感谢您的阅读!如果你觉得上述内容对你有所帮助,欢迎分享至一位你的朋友!若有任何疑问与建议,请添加笔者微信或加群进行交流。青山不改,绿水长流,我们月来客栈见!

引用

[1] Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need

[2] The Illustrated Transformer http://jalammar.github.io/illustrated-transformer/

[3] https://towardsdatascience.com/illustrated-guide-to-transformers-step-by-step-explanation-f74876522bc0

[4] https://pytorch.org/docs/stable/generated/torch.bmm.html?highlight=bmm#torch.bmm

[5] LANGUAGE TRANSLATION WITH TRANSFORMER https://pytorch.org/tutorials/beginner/translation_transformer.html

[6] The Annotated Transformer http://nlp.seas.harvard.edu/2018/04/03/attention.html

[7] SEQUENCE-TO-SEQUENCE MODELING WITH NN.TRANSFORMER AND TORCHTEXT https://pytorch.org/tutorials/beginner/transformer_tutorial.html

[8] Transformer model for language understanding https://tensorflow.google.cn/text/tutorials/transformer?hl=en#multi-head_attention

[9] https://github.com/multi30k/dataset

[10]你还在手动构造词表?试试torchtext.vocab

[11] https://github.com/moon-hotel/TransformerTranslation

[12] https://github.com/moon-hotel/TransformerClassification

[13] https://pytorch.org/tutorials/beginner/text_sentiment_ngrams_tutorial.html

[14] https://github.com/moon-hotel/TransformerCouplet

22010905029.jpg