目录
- 前言
- 1. \__init\__()
- 2. encode()
- 2. decode()
- 3. linear projection
- 4. 代码地址
前言
这篇博客是对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 = 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()用法解释：
# 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函数，e0=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。
