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

這篇博客是對transformer源碼的解析，這個源碼並非官方的，但是比官方代碼更容易理解。
採用TensorFlow框架，下面的解析過程只針對模型構建過程，其訓練/測試等其他代碼忽略。
解讀順序按照model.py中函數順序解讀。
文末會給出代碼地址。文章結構如下：

__init__()
encode()
decode()
代碼地址

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 = inputs.get_shape().as_list()[-1]  # static （d_model)
    N, T = tf.shape(inputs)[0], tf.shape(inputs)[1]  # dynamic
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # position indices
        position_ind = tf.tile(tf.expand_dims(tf.range(T), 0), [N, 1])  # (N, T)

        # First part of the PE function: sin and cos argument
        position_enc = np.array([
            [pos / np.power(10000, (i - i % 2) / E) for i in range(E)]
            for pos in range(maxlen)])

        # Second part, apply the cosine to even columns and sin to odds.
        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1
        position_enc = tf.convert_to_tensor(position_enc, tf.float32)  # (maxlen, E)

        # lookup
        outputs = tf.nn.embedding_lookup(position_enc, position_ind)

        # masks
        if masking:
            # tf.where()用法解釋：https://blog.csdn.net/a_a_ron/article/details/79048446
            # tf.where(input, a, b) 將a中對應input中true的位置的元素不變，其餘元素替換成b中對應位置的元素
            # 下面操作中將inputs中值不爲0的位置的數替換爲對應位置outputs上的數，inputs中值爲0的位置則保留
            # 顯然這種掩模是必須的，因爲inputs中值爲0的是padding結果，這些位置自然不用參與計算，所以這些位置
            # 也就不應該有位置編碼值
            outputs = tf.where(tf.equal(inputs, 0), inputs, outputs)

        return tf.to_float(outputs)

位置編碼用的是三角函數，最後是掩模操作，transformer中有兩種掩模，一種是padding mask，另一種是sequence mask，後者只在解碼器中用到。
什麼是 padding mask 呢？因爲每個批次輸入序列長度是不一樣的也就是說，我們要對輸入序列進行對齊。如果輸入的序列太長，則是截取左邊的內容，把多餘的直接捨棄。如果序列太短，通常的做法是給在較短的序列後面填充 0，但是這些填充的位置，其實是沒什麼意義的，所以我們的attention機制不應該把注意力放在這些位置上。填充0在進行 softmax 的時候就會產生問題, 回顧 softmax函數，e⁰=1是有值的，這樣softmax中用0填充的位置實際上是參與了運算，等於是讓無意義的位置參與計算了，所以我們需要進行一些處理，讓這些無效位置不參與運算。具體的做法是，把這些位置的值加上一個非常大的負數(-2^32+1)，這樣的話，經過 softmax，這些位置的概率就會接近0。
上面代碼中只是將原來是padding的位置保留其值，後續將這些位置值變爲很大的負數是在softmax中。
mask原理代碼中解釋已非常清楚。
加上位置編碼後，這裏還有個隨機失活操作，這個論文中也沒有提到，也許是加上隨機失活效果更好。

接下來就是編碼器最重要部分。編碼器中有很多block，每個block由multihead_attention + feed_forward組成。self.hp.num_blocks指定block數量。
首先來看multihead_attention，也就是自注意力機制和多頭機制。

def multihead_attention(queries, keys, values,
                        key_masks,
                        num_heads=8,
                        dropout_rate=0,
                        training=True,
                        causality=False,
                        scope="multihead_attention"):
    '''Applies multihead attention. See 3.2.2
    queries: A 3d tensor with shape of [N, T_q, d_model].
    keys: A 3d tensor with shape of [N, T_k, d_model].
    values: A 3d tensor with shape of [N, T_k, d_model].
    key_masks: A 2d tensor with shape of [N, key_seqlen]
    num_heads: An int. Number of heads.
    dropout_rate: A floating point number.
    training: Boolean. Controller of mechanism for dropout.
    causality: Boolean. If true, units that reference the future are masked.
    scope: Optional scope for `variable_scope`.
        
    Returns
      A 3d tensor with shape of (N, T_q, C)  
    '''
    d_model = queries.get_shape().as_list()[-1]  # 詞向量長度
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # Linear projections，分別乘以Q_kernel,K_kernel,V_kernel矩陣，得到Q,K,V矩陣
        # 在tf.layers.dense()定義中可發現`kernel` is a weights matrix created by the layer
        # 所以Q_kernel,K_kernel,V_kernel矩陣是隱式給出的，並且這裏activation=None，默認爲linear activation
        Q = tf.layers.dense(queries, d_model, use_bias=True)  # (N, T_q, d_model)
        K = tf.layers.dense(keys, d_model, use_bias=True)  # (N, T_k, d_model)
        V = tf.layers.dense(values, d_model, use_bias=True)  # (N, T_k, d_model)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, d_model/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, d_model/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, d_model/h)

        # Attention
        outputs = scaled_dot_product_attention(Q_, K_, V_, key_masks, causality, dropout_rate, training)

        # Restore shape，多頭矩陣合併
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)  # (N, T_q, d_model)

        # Residual connection
        outputs += queries

        # Normalize
        outputs = ln(outputs)

    return outputs

參數 queries, keys, values 就是前面得到的詞向量矩陣，這三個矩陣相同。
爲了計算自注意力，將 queries, keys, values 分別乘以Q’， K’， V’ 矩陣，得到Q，K，V三個不同矩陣，Q’， K’， V’ 並沒有顯式的初始化，而是在 tf.layers.dense() 隱式完成。
然後將 Q，K，V 的每個矩陣分成多頭，這個劃分注意是針對最後一個維度劃分的，也就是詞向量長度，例如詞向量長度爲512，多頭爲8，那麼劃分後每個詞向量長度爲64，共8份，看懂代碼中對維度的註釋就明白了。

Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, d_model/h)
K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, d_model/h)
V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, d_model/h)

多頭劃分後就可以計算每個句子的自注意力，也就是下面操作，

outputs = scaled_dot_product_attention(Q_, K_, V_, key_masks, causality, dropout_rate, training)

具體如下，步驟很簡單，顯示Q_和K_矩陣相乘計算注意力分數，爲了梯度穩定，有個縮放操作，然後是softmax歸一化。然後是掩模操作，最後將注意力矩陣和V_矩陣相乘。
每個步驟代碼中都有註釋，比較易懂。

def scaled_dot_product_attention(Q, K, V, key_masks,
                                 causality=False, dropout_rate=0.,
                                 training=True,
                                 scope="scaled_dot_product_attention"):
    '''See 3.2.1.
    Q: Packed queries. 3d tensor. [N, T_q, d_k].
    K: Packed keys. 3d tensor. [N, T_k, d_k].
    V: Packed values. 3d tensor. [N, T_k, d_v].
    key_masks: A 2d tensor with shape of [N, key_seqlen]
    causality: If True, applies masking for future blinding
    dropout_rate: A floating point number of [0, 1].
    training: boolean for controlling droput
    scope: Optional scope for `variable_scope`.
    '''
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        d_k = Q.get_shape().as_list()[-1]

        # dot product，Q與轉置後的K相乘，得到注意力矩陣
        outputs = tf.matmul(Q, tf.transpose(K, [0, 2, 1]))  # (N, T_q, T_k)

        # scale，將注意力矩陣變成標準正態分佈，使得softmax歸一化後結果更穩定
        outputs /= d_k ** 0.5

        # key masking，和position_encoding中masking作用相同，將句子中padding位置的注意力清零，
        # 因爲這些位置並沒有字，所以也不存在注意力，但上面計算注意力時，即使該位置爲0也會有值
        outputs = mask(outputs, key_masks=key_masks, type="key")

        # causality or future blinding masking
        if causality:
            outputs = mask(outputs, type="future")  # outputs矩陣上三角全置爲-2^32+1

        # softmax
        outputs = tf.nn.softmax(outputs)
        attention = tf.transpose(outputs, [0, 2, 1])  # 爲什麼注意力矩陣需要轉置？
        tf.summary.image("attention", tf.expand_dims(attention[:1], -1))

        # # query masking
        # outputs = mask(outputs, Q, K, type="query")

        # dropout
        outputs = tf.layers.dropout(outputs, rate=dropout_rate, training=training)

        # weighted sum (context vectors)
        outputs = tf.matmul(outputs, V)  # (N, T_q, d_v)

    return outputs

掩模操作mask()代碼如下，type=key表示padding mask， type=future表示sentence mask。

def mask(inputs, key_masks=None, type=None):
    """Masks paddings on keys or queries to inputs
    inputs: 3d tensor. (h*N, T_q, T_k)
    key_masks: 3d tensor. (N, 1, T_k)
    type: string. "key" | "future"

    e.g.,
    >> inputs = tf.zeros([2, 2, 3], dtype=tf.float32)
    >> key_masks = tf.constant([[0., 0., 1.],
                                [0., 1., 1.]])
    >> mask(inputs, key_masks=key_masks, type="key")
    array([[[ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09],
        [ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09]],

       [[ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09],
        [ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09]],

       [[ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09],
        [ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09]],

       [[ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09],
        [ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09]]], dtype=float32)
    """
    padding_num = -2 ** 32 + 1
    if type in ("k", "key", "keys"):
        key_masks = tf.to_float(key_masks)
        key_masks = tf.tile(key_masks, [tf.shape(inputs)[0] // tf.shape(key_masks)[0], 1])  # (h*N, seqlen)
        key_masks = tf.expand_dims(key_masks, 1)  # (h*N, 1, seqlen)
        outputs = inputs + key_masks * padding_num
    elif type in ("f", "future", "right"):
        diag_vals = tf.ones_like(inputs[0, :, :])  # (T_q, T_k) 全1矩陣
        tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (T_q, T_k) 矩陣上三角置零
        future_masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(inputs)[0], 1, 1])  # (N, T_q, T_k)
        paddings = tf.ones_like(future_masks) * padding_num   # (N, T_q, T_k) 全padding_num矩陣
        outputs = tf.where(tf.equal(future_masks, 0), paddings, inputs)  # inputs矩陣中上三角用paddings中的值代替
    else:
        print("Check if you entered type correctly!")

    return outputs

自注意力計算到這裏就完成了，接下來是block中的feed forward。

# feed forward
enc = ff(enc, num_units=[self.hp.d_ff, self.hp.d_model])

詳細代碼如下：

def ff(inputs, num_units, scope="position_wise_feedforward"):
    '''position-wise feed forward net. See 3.3
    
    inputs: A 3d tensor with shape of [N, T, C].
    num_units: A list of two integers.
                num_units[0]=d_ff: 隱藏層大小（2048）
                num_units[1]=d_model: 詞向量長度（512）
    scope: Optional scope for `variable_scope`.

    Returns:
      A 3d tensor with the same shape and dtype as inputs
    '''
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # Inner layer
        outputs = tf.layers.dense(inputs, num_units[0], activation=tf.nn.relu)

        # Outer layer
        outputs = tf.layers.dense(outputs, num_units[1])

        # Residual connection
        outputs += inputs

        # Normalize
        outputs = ln(outputs)

    return outputs

這裏面的inner layer和 outer layer可能不太好理解，可以認爲就是兩個全連接矩陣，num_units[0]=d_ff表示隱藏層大小（2048），num_units[1]=d_model表示詞向量長度（512），也可以看做是一個1*1卷積核做卷積操作。然後是殘差連接，最後是LN標準化。
編碼器中的一個block就講完了，然後是循環self.hp.num_blocks，每個block的輸出作爲下一個block的輸入，最後一個block的輸出就是整個編碼器的輸出。

2. decode()

Transformer中解碼器部分。

def decode(self, ys, memory, src_masks, training=True):
    '''
    memory: encoder outputs. (N, T1, d_model)
    src_masks: (N, T1)

    Returns
    logits: (N, T2, V). float32.
    y_hat: (N, T2). int32
    y: (N, T2). int32
    sents2: (N,). string.
    '''
    with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
        decoder_inputs, y, seqlens, sents2 = ys

        # tgt_masks
        tgt_masks = tf.math.equal(decoder_inputs, 0)  # (N, T2)

        # embedding， encoder 和 decoder 共用一個 embeddings
        dec = tf.nn.embedding_lookup(self.embeddings, decoder_inputs)  # (N, T2, d_model)
        dec *= self.hp.d_model ** 0.5  # scale

        dec += positional_encoding(dec, self.hp.maxlen2)
        dec = tf.layers.dropout(dec, self.hp.dropout_rate, training=training)

        # Blocks
        for i in range(self.hp.num_blocks):
            with tf.variable_scope("num_blocks_{}".format(i), reuse=tf.AUTO_REUSE):
                # Masked self-attention (Note that causality is True at this time)
                dec = multihead_attention(queries=dec,
                                          keys=dec,
                                          values=dec,
                                          key_masks=tgt_masks,
                                          num_heads=self.hp.num_heads,
                                          dropout_rate=self.hp.dropout_rate,
                                          training=training,
                                          causality=True,
                                          scope="self_attention")

                # Vanilla attention
                dec = multihead_attention(queries=dec,
                                          keys=memory,
                                          values=memory,
                                          key_masks=src_masks,
                                          num_heads=self.hp.num_heads,
                                          dropout_rate=self.hp.dropout_rate,
                                          training=training,
                                          causality=False,
                                          scope="vanilla_attention")
                # Feed Forward
                dec = ff(dec, num_units=[self.hp.d_ff, self.hp.d_model])

    # Final linear projection (embedding weights are shared)
    weights = tf.transpose(self.embeddings)  # (d_model, vocab_size)
    logits = tf.einsum('ntd,dk->ntk', dec, weights)  # (N, T2, vocab_size)，矩陣相乘，消除d_model維度
    y_hat = tf.to_int32(tf.argmax(logits, axis=-1))

    return logits, y_hat, y, sents2

解碼器也是由多個block組成，每個block由 multihead_attention + multihead_attention + feed_forward組成，其中第一個 multihead_attention是自注意力計算，第二個 multihead_attention是注意力計算，有很大的區別，仔細看給他們的參數就會發現。
解碼器相比編碼器就多了一個部分，即encode-decode-attention。
在訓練階段，解碼器的輸入是直接從訓練數據給出的，而不是將解碼器輸出結果循環的送入解碼器輸入。所以上面代碼中只有一個循環，就是解碼器中block個數。
首先是計算自注意力，注意解碼器中的自注意力計算和編碼器中自注意力計算方式稍有不同。主要是多了一個sequence mask操作（causality=True時才需要這個步驟），這個是爲了讓當前詞看不都其後面的詞。
sequence mask就是mask()函數參數type=future時的操作，具體如下：

diag_vals = tf.ones_like(inputs[0, :, :])  # (T_q, T_k) 全1矩陣
tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (T_q, T_k) 矩陣上三角置零
future_masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(inputs)[0], 1, 1])  # (N, T_q, T_k)
paddings = tf.ones_like(future_masks) * padding_num   # (N, T_q, T_k) 全padding_num矩陣
outputs = tf.where(tf.equal(future_masks, 0), paddings, inputs)  # inputs矩陣中上三角用paddings中的值代替

具體步驟看代碼註釋即可看懂。其他步驟和編碼器中相同。
完成自注意力計算後，下一步就是encode-decode-attention計算，這是注意力計算，而不是自注意力計算。代碼如下：

# encode-decode attention
dec = multihead_attention(queries=dec,
                           keys=memory,
                           values=memory,
                           key_masks=src_masks,
                           num_heads=self.hp.num_heads,
                           dropout_rate=self.hp.dropout_rate,
                           training=training,
                           causality=False,
                           scope="vanilla_attention")

不管是注意力還是自注意力的計算，都是調用multihead_attention函數。這裏encode-decode-attention的計算中 K,V 矩陣是編碼器的最終輸出，也就是參數中的
keys=memory, values=memory，而 Q 矩陣是解碼器中自注意力的輸出dec。
最後是feed_forward層，和編碼器中相同。

3. linear projection

transformer中最後部分是線性映射。
解碼器最後輸出浮點向量，如何將它轉成詞？這是最後的線性層和softmax層的主要工作。
線性層是個簡單的全連接層，將解碼器的最後輸出映射到一個非常大的logits向量上。假設模型已知有1萬個單詞（輸出的詞表）從訓練集中學習得到。那麼，logits向量就有1萬維，每個值表示是某個詞的可能傾向值。
softmax層將這些分數轉換成概率值（都是正值，且加和爲1），最高值對應的維度上的詞就是這一步的輸出單詞。
代碼如下：

# Final linear projection (embedding weights are shared)
weights = tf.transpose(self.embeddings)  # (d_model, vocab_size)
logits = tf.einsum('ntd,dk->ntk', dec, weights)  # (N, T2, vocab_size)，矩陣相乘，消除d_model維度
y_hat = tf.to_int32(tf.argmax(logits, axis=-1))

4. 代碼地址

github
主要是model.py和modules.py。

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

1. _init_()

2. encode()

2. decode()

3. linear projection

4. 代碼地址

劍指offer第二版（Python3）--面試題64 : 求1+2+3+...+n

劍指offer第二版（Python3）--面試題56 :數組中數字出現的次數

劍指offer第二版（Python3）--面試題55 :二叉樹的深度、平衡二叉樹

劍指offer 第二版（Python3）--面試題12：矩陣中的路徑

劍指offer 第二版（Python3）--面試題8：二叉樹的下一個節點

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結