机器翻译模型Transformer代码详细解析

谷歌一个月前发了一篇论文Attention is all you need,文中提出了一种新的架构叫做Transformer，用以来实现机器翻译。它抛弃了传统用CNN或者RNN的定式，取得了很好的效果，激起了工业界和学术界的广泛讨论。本人的另一篇博客也对改论文进行了一定的分析：对Attention is all you need 的理解。而在谷歌的论文发出不久，就有人用tensorflow实现了Transformer模型：A TensorFlow Implementation of the Transformer: Attention Is All You Need。这里我打算对该开源实现的代码进行细致的分析。
该实现相对原始论文有些许不同，比如为了方便使用了IWSLT 2016德英翻译的数据集，直接用的positional embedding，把learning rate一开始就调的很小等等，不过大同小异，主要模型没有区别。

该实现一共包括以下几个文件

• hyperparams.py 该文件包含所有需要用到的参数
• prepro.py 该文件生成源语言和目标语言的词汇文件。
• data_load.py 该文件包含所有关于加载数据以及批量化数据的函数。
• modules.py 该文件具体实现编码器和解码器网络
• train.py 训练模型的代码，定义了模型，损失函数以及训练和保存模型的过程
• eval.py 评估模型的效果

接下来针对每一个文件分别解析。
首先是hyperparams.py文件。
该实现所用到的所又的超参数都在这个文件里面。以下是该文件的所有代码：

class Hyperparams:
    '''Hyperparameters'''
    # data
    source_train = 'corpora/train.tags.de-en.de'
    target_train = 'corpora/train.tags.de-en.en'
    source_test = 'corpora/IWSLT16.TED.tst2014.de-en.de.xml'
    target_test = 'corpora/IWSLT16.TED.tst2014.de-en.en.xml'
    
    # training
    batch_size = 32 # alias = N
    lr = 0.0001 # learning rate. In paper, learning rate is adjusted to the global step.
    logdir = 'logdir' # log directory
    
    # model
    maxlen = 10 # Maximum number of words in a sentence. alias = T.
                # Feel free to increase this if you are ambitious.
    min_cnt = 20 # words whose occurred less than min_cnt are encoded as <UNK>.
    hidden_units = 512 # alias = C
    num_blocks = 6 # number of encoder/decoder blocks
    num_epochs = 20
    num_heads = 8
    dropout_rate = 0.1

可以看出该部分没有什么特别难以理解的，定义了一些要使用的超参数以便以后使用。首先是源语言以及目标语言的训练数据和测试数据的路径，其次设定了batch_size的大小以及初始学习速率还有日志的目录,batch_size 在后续代码中即所谓的N，参数中常会见到。最后定义了一些模型相关的参数，maxlen为一句话里最大词的长度为10个，在其他代码中就用的是T来表示，你也可以根据自己的喜好将这个参数调大;min_cnt被设置为20，该参数表示所有出现次数少于min_cnt次的都会被当作UNK来处理;hidden_units设置为512，隐藏节点的 个数，在代码中用C来表示。num_blocks和num_heads都是论文中提到的设定，epoch大小设置为20，此外还有dropout就不用多费口舌了。
以上就是该开源实现中超参数的设定，该部分到此为止，没有太多可以说的。

接下来看预处理的代码prepro.py ，该代码的作用是生成源语言和目标语言的词汇文件。
为了直观理解，首先看一下执行代码之后生成的词汇文件是啥样的，我这里截取了德语词汇文件的前几行：

<PAD>	1000000000
<UNK>	1000000000
<S>	1000000000
</S>	1000000000
die	85235
und	77082
der	56248
ist	51457

可以看出，文件把训练数据中出现的单词和其出现的次数做了统计，并且记录在生成的词汇文件中。第一列为单词，第二列为出现的次数。同时，设置了四个特殊的标记符号，把他们设定为出现次数很多放在文件的最前。
仍然是先贴代码。

from __future__ import print_function
from hyperparams import Hyperparams as hp
import tensorflow as tf
import numpy as np
import codecs
import os
import regex
from collections import Counter

def make_vocab(fpath, fname):
    '''Constructs vocabulary.
    
    Args:
      fpath: A string. Input file path.
      fname: A string. Output file name.
    
    Writes vocabulary line by line to `preprocessed/fname`
    '''  
    text = codecs.open(fpath, 'r', 'utf-8').read()
    text = regex.sub("[^\s\p{Latin}']", "", text)
    words = text.split()
    word2cnt = Counter(words)
    if not os.path.exists('preprocessed'): os.mkdir('preprocessed')
    with codecs.open('preprocessed/{}'.format(fname), 'w', 'utf-8') as fout:
        fout.write("{}\t1000000000\n{}\t1000000000\n{}\t1000000000\n{}\t1000000000\n".format("<PAD>", "<UNK>", "<S>", "</S>"))
        for word, cnt in word2cnt.most_common(len(word2cnt)):
            fout.write(u"{}\t{}\n".format(word, cnt))

if __name__ == '__main__':
    make_vocab(hp.source_train, "de.vocab.tsv")
    make_vocab(hp.target_train, "en.vocab.tsv")
    print("Done")

代码中make_vocab函数就是生成词汇文件的函数。该函数一共有两个参数，fpath表示输入文件的路径，具体而言就是训练数据，而另一个参数fname即要输出的词汇文件名。
该函数一行一行地将词汇写入到’preprocessed/fname’中。
可以注意到一开始使用codecs中的open函数来打开并读取文件的。那么这个和我们平常使用的open函数有什么区别呢？基本上在处理语言的时候都要在unicode这种编码上边搞，可以看到codecs.open的时候直接将文件爱呢转换为内部unicode，其中第三个参数就是源文件的编码格式。关于codecs具体可以参考python模块之codecs
。
读取文件之后用正则表达式对读入的数据进行了处理，sub函数用于替换字符串中的匹配项，一共有三个参数，将第三个参数所代表的字符串中等所有满足第一个参数示例的形式的字符都用第二个参数来代替。
接下来将读取的文本按照空白分割成words之后放入Counter进行计数，计数的结果类似于一个字典，key为词，value为出现的次数。然后创建爱呢保存预处理文件的目录。同样利用codecs李的open函数创建一个要输出的文件，首先将四个准备好的特殊词写入文件在开始的四行。然后利用most_common函数依词出现的频率将训练集中出现的词和其对应的计数一行一行写入文件。
分别用德语和英语文件作为参数运行该函数即可得到词汇文件。

接下来分析第三个文件data_load.py，该文件包含所有关于加载数据以及批量化数据的函数。还是先上代码。

from __future__ import print_function
from hyperparams import Hyperparams as hp
import tensorflow as tf
import numpy as np
import codecs
import regex

def load_de_vocab():
    vocab = [line.split()[0] for line in codecs.open('preprocessed/de.vocab.tsv', 'r', 'utf-8').read().splitlines() if int(line.split()[1])>=hp.min_cnt]
    word2idx = {
   word: idx for idx, word in enumerate(vocab)}
    idx2word = {
   idx: word for idx, word in enumerate(vocab)}
    return word2idx, idx2word

这里一部分一部分代码进行分析。
首先是load_de_vocab()函数。该函数的目的是给德语的每个词分配一个id并返回两个字典，一个是根据词找id，一个是根据id找词。函数直接利用codecs的open来读取之前在预处理的时候生成的词汇文件。注意这里读每行的时候去掉了那些出现次数少于hp.min_cnt(根据设定为20)的词汇。读完之后有一个词汇列表。然后便利该列表的枚举enumerate(vocab)生成词和其对应id的两个字典。
接下来是load_en_vocab()函数的代码：

def load_en_vocab():
    vocab = [line.split()[0] for line in codecs.open('preprocessed/en.vocab.tsv', 'r', 'utf-8').read().splitlines() if int(line.split()[1])>=hp.min_cnt]
    word2idx = {
   word: idx for idx, word in enumerate(vocab)}
    idx2word = {
   idx: word for idx, word in enumerate(vocab)}
    return word2idx, idx2word

该函数和之前的生成德语word/id字典的函数一样，只不过生成的是英语的word/id字典，方法都一样，不用多说。

接下来是creat_data函数。

def create_data(source_sents, target_sents): 
    de2idx, idx2de = load_de_vocab()
    en2idx, idx2en = load_en_vocab()
    
    # Index
    x_list, y_list, Sources, Targets = [], [], [], []
    for source_sent, target_sent in zip(source_sents, target_sents):
        x = [de2idx.get(word, 1) for word in (source_sent + u" </S>").split()] # 1: OOV, </S>: End of Text
        y = [en2idx.get(word, 1) for word in (target_sent + u" </S>").split()] 
        if max(len(x), len(y)) <=hp.maxlen:
            x_list.append(np.array(x))
            y_list.append(np.array(y))
            Sources.append(source_sent)
            Targets.append(target_sent)
    
    # Pad      
    X = np.zeros([len(x_list), hp.maxlen], np.int32)
    Y = np.zeros([len(y_list), hp.maxlen], np.int32)
    for i, (x, y) in enumerate(zip(x_list, y_list)):
        X[i] = np.lib.pad(x, [0, hp.maxlen-len(x)], 'constant', constant_values=(0, 0))
        Y[i] = np.lib.pad(y, [0, hp.maxlen-len(y)], 'constant', constant_values=(0, 0))
    
    return X, Y, Sources, Targets

该函数一共有两个参数，source_sents和target_sents。可以理解为源语言和目标语言的句子列表。每个列表中的一个元素就是一个句子。
首先利用之前定义的两个函数生成双语语言的word/id字典。
同时遍历这两个参数指示的句子列表。一次遍历一个句子对，在该次遍历中，给每个句子末尾后加一个文本结束符$</s>$ 用以表示句子末尾。加上该结束符的句子又被遍历每个词，同时利用双语word/id字典读取word对应的id加入一个新列表中，若该word不再字典中则id用1代替（即UNK的id）。如此则生辰概率两个用一串id表示的双语句子的列表。然后判断这两个句子的长度是否都没超过设定的句子最大长度hp.maxlen,如果没超过，则将这两个双语句子id列表加入模型要用的双语句子id列表x_list,y_list中，同时将满足最大句子长度的原始句子（用word表示的）也加入到句子列表Sources以及Targets中。
函数后半部分为Pad操作。关于numpy中的pad操作可以参考numpy–prod和pad运算。这里说该函数的pad运算，由于x和y都是一维的，所有只有前后两个方向可以pad，所以pad函数的第二个参数是一个含有两个元素的列表，第一个元素为0说明给x或者y前面什么也不pad，即pad上0个数，第二个元素为hp.maxlen-len(x)以及hp.maxlen-len(x)代表给x和y后面pad上x和y初始元素个数和句子最大长度差的那么多数值，至于pad成什么数值，后面的constant_values给出了，即pad上去的id值为0，这也是我们词汇表中PAD的id。经过pad的操作可以保证用id表示的句子列表都是等长的。
最终返回等长的句子id数组X，Y，以及原始句子李标Sources以及Targets。X和Y的shape都为[len(x_list),hp.maxlen]。其中len(x_list)为句子的总个数，hp.maxlen为设定的最大句子长度。

接下来有一个函数为load_train_data(),还是上代码：

def load_train_data():
    de_sents = [regex.sub("[^\s\p{Latin}']", "", line) for line in codecs.open(hp.source_train, 'r', 'utf-8').read().split("\n") if line and line[0] != "<"]
    en_sents = [regex.sub("[^\s\p{Latin}']", "", line) for line in codecs.open(hp.target_train, 'r', 'utf-8').read().split("\n") if line and line[0] != "<"]
    
    X, Y, Sources, Targets = create_data(de_sents, en_sents)
    return X, Y

顾名思义，该函数的作用是加载训练数据，加载的方式很简单，就是加载刚才create_data返回的等长句子id数组。load_train_data的作用只不过是给create_data提供了de_sents和en_sents两个参数而已。
而de_sents和en_sents这两个句子列表同样是通过codecs里的open读取训练数据生成的。读取之后按照换行符\n分隔开每一句，在这些句子中选择那些那些行开头符号不是‘<’的句子(句首为<是数据描述的行，并非真实数据的部分)。在这些分离好的句子中同样用正则表达式进行处理。

接下来是load_test_data()函数。

def load_test_data():
    def _refine(line):
        line = regex.sub("<[^>]+>", "", line)
        line = regex.sub("[^\s\p{Latin}']", "", line