NLP系列——Transformer源码解析（TensorFlow版）

前言

  这篇博客是对transformer源码的解析，这个源码并非官方的，但是比官方代码更容易理解。
  采用TensorFlow框架，下面的解析过程只针对模型构建过程，其训练/测试等其他代码忽略。
  解读顺序按照model.py中函数顺序解读。
  文末会给出代码地址。

1. _init_()

  模型初始化，主要是初始化词向量矩阵

 def __init__(self, hp):
        self.hp = hp
        self.token2idx, self.idx2token = load_vocab(hp.vocab)
        self.embeddings = get_token_embeddings(self.hp.vocab_size, self.hp.d_model, zero_pad=True)

  hp是一个类，其变量是模型初始化的参数，例如学习率，词向量长度，词表路径等。
  load_vocab()用于加载词表，返回每个词对应的索引self.token2idx，以及每个索引对应的词self.idx2token。
  get_token_embeddings()用于构建词向量矩阵。第一个参数是词典大小；第二个参数是词向量长度；第三个参数为True表示词向量矩阵第一行全为0，因为索引为0的行表示padding的词向量，padding用于掩模。
  下面是get_token_embeddings()函数代码：

def get_token_embeddings(vocab_size, num_units, zero_pad=True):
    '''Constructs token embedding matrix.
    Note that the column of index 0's are set to zeros.
    vocab_size: scalar. V. （词典大小）
    num_units: embedding dimensionalty. E. （词向量长度：512）
    zero_pad: Boolean. If True, all the values of the first row (id = 0) should be constant zero
    To apply query/key masks easily, zero pad is turned on.

    Returns
    weight variable: (V, E)  返回词向量矩阵
    '''
    with tf.variable_scope("shared_weight_matrix"):
        # 初始化词向量矩阵
        embeddings = tf.get_variable('weight_mat',
                                     dtype=tf.float32,
                                     shape=(vocab_size, num_units),
                                     initializer=tf.contrib.layers.xavier_initializer())
        # 为了后续对 query/key 矩阵掩模操作，将词向量矩阵第一行设置为全0
        if zero_pad:
            embeddings = tf.concat((tf.zeros(shape=[1, num_units]), embeddings[1:, :]), 0)
    return embeddings

  代码中首先初始化一个（vocab_size, num_units）大小的矩阵embeddings，然后将第一行置零。

2. encode()

  transformer中的编码器部分。

def encode(self, xs, training=True):
    '''
    xs: 训练数据
    Returns
    memory: encoder outputs. (N, T1, d_model)
                            N: batch size;
                            T1: sentence length
                            d_model: 512, 词向量长度
    '''
    with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
        # xs: tuple of
        #               x: int32 tensor. (N, T1)
        #               x_seqlens: int32 tensor. (N,)  句子长度
        #               sents1: str tensor. (N,)
        x, seqlens, sents1 = xs

        # src_masks
        src_masks = tf.math.equal(x, 0)  # (N, T1)

        # embedding
        enc = tf.nn.embedding_lookup(self.embeddings, x)  # (N, T1, d_model)
        enc *= self.hp.d_model ** 0.5  # scale
		# 加上位置编码向量
        enc += positional_encoding(enc, self.hp.maxlen1)  
        enc = tf.layers.dropout(enc, self.hp.dropout_rate, training=training)

        # Blocks 编码器模块
        # num_blocks=6编码器中小模块数量，小模块指 multihead_attention + feed_forward
        for i in range(self.hp.num_blocks):
            with tf.variable_scope("num_blocks_{}".format(i), reuse=tf.AUTO_REUSE):
                # self-attention
                enc = multihead_attention(queries=enc,
                                          keys=enc,
                                          values=enc,
                                          key_masks=src_masks,
                                          num_heads=self.hp.num_heads,
                                          dropout_rate=self.hp.dropout_rate,
                                          training=training,
                                          causality=False)
                # feed forward
                enc = ff(enc, num_units=[self.hp.d_ff, self.hp.d_model])
    memory = enc
    return memory, sents1, src_masks

始终要记住的是N, T1, d_model三个参数的含义。
  N表示batch size；
  T1表示一个batch中最长句子的长度；
  d_model表示词向量的长度，默认参数是512

  xs表示一个batch中的训练数据；
  首先是查找表操作tf.nn.embedding_lookup()，得到训练数据中每个词的词向量，返回一个矩阵enc，enc的维度为[N, T1, d_model]；
  然后是对enc矩阵缩放，这个步骤貌似论文中没有提及，应该是作者自己加上去的；
  训练数据的词向量还要加上位置编码才能送入编码器中，也就是下面代码：

def positional_encoding(inputs,
                        maxlen,
                        masking=True,
                        scope="positional_encoding"):
    '''Sinusoidal Positional_Encoding. See 3.5
    inputs: 3d tensor. (N, T, E)
    maxlen: scalar. Must be >= T  一个batch中句子的最大长度
    masking: Boolean. If True, padding positions are set to zeros.
    scope: Optional scope for `variable_scope`.

    returns
    3d tensor that has the same shape as inputs.
    '''

    E <