TensorFlow|基于Transformer的自然语言推理（SNLI）

在经历了看论文，看源码，看Bert源码之后，整理思路，实现了一下Transformer，并搭建了一个小型的Transformer做了一下SNLI任务。

同时吸取以前的教训，这次好好的写了注释

1.Transofrmer

原理不再重述，其他博客中讲的很好，

比如：https://jalammar.github.io/illustrated-transformer/

和他的翻译版：https://blog.csdn.net/qq_41664845/article/details/84969266

直接进入代码

1.1.激活函数

Transformer原文中使用的都是Relu，但Bert包括之后的工作，大多采用的是Gelu（高斯误差线性单元），效果更好（只是参考了论文中的数据对比，还并未亲自实验对比）。

抱着举贤不举亲的原则，就算平时使用的大多Relu，在此也将默认的激活函数设为Gelu。

关于Gelu的原论文：https://arxiv.org/abs/1606.08415

Gelu：

def gelu(inputs):
    """
    gelu: https://arxiv.org/abs/1606.08415
    :param inputs: [Tensor]
    :return: [Tensor] outputs after activation
    """
    cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2 / np.pi) * (inputs + 0.044715 * tf.pow(inputs, 3))))
    return inputs * cdf

获得激活函数的方法（设置默认gelu）：

def get_activation(activation_name):
    """
    get activate function
    :param activation_name: [Tensor]
    :return: [Function] activation function
    """
    if activation_name is None:
        return gelu
    else:
        act = activation_name.lower()
        if act == "relu":
            return tf.nn.relu
        elif act == "gelu":
            return gelu
        elif act == "tanh":
            return tf.tanh
        else:
            raise ValueError("Unsupported activation: %s" % act)

1.2.嵌入（embedding）

Transformer除了词嵌入，还做了位置嵌入（Positional Encoding），来使每个单词携带位置信息，否则可以想象它只是一个复杂一些的，通过训练获得每个单词权重的词袋模型了。

同时为了完成SNLI这类需要最终输出shape一致的任务，采用了Bert的想法，对每个输入的起始加入[CLS]token，使用该token的最终输出做预测，而这样做的话，需要加入segment embedding来更好的区分两个不同的句子（参考Bert）

1.2.1.词嵌入（Word Embedding）

这里可以通过随机初始化嵌入矩阵，也可以通过载入其他任务（比如Glove，Fast text）产生的词嵌入矩阵来完成这部分，只需要在restore的时候声明一下即可。paper中提到需要对embedding做scale，这里照做。

def get_embedding(inputs, vocab_size, channels, scale=True, scope="embedding", reuse=None):
    """
    embedding
    :param inputs: [Tensor] Tensor with first dimension of "batch_size"
    :param vocab_size: [Int] Vocabulary size
    :param channels: [Int] Embedding size
    :param scale: [Boolean] If True, the output will be multiplied by sqrt num_units
    :param scope: [String] name of "variable_scope"
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs of embedding of sentence with shape of "batch_size * length * channels"
    """
    with tf.variable_scope(scope, reuse=reuse):
        lookup_table = tf.get_variable('lookup_table',
                                       dtype=tf.float32,
                                       shape=[vocab_size, channels],
                                       initializer=tf.contrib.layers.xavier_initializer())
        lookup_table = tf.concat((tf.zeros(shape=[1, channels], dtype=tf.float32),
                                  lookup_table[1:, :]), 0)

        outputs = tf.nn.embedding_lookup(lookup_table, inputs)

        if scale:
            outputs = outputs * math.sqrt(channels)

    return outputs

1.2.2.位置嵌入（Position Embedding）

获得和inputs经过word embedding之后相同shape的位置嵌入，没有使用word embedding之后的作为输入，是考虑这样可以为之后的mask提供便利

def get_positional_encoding(inputs, channels, scale=False, scope="positional_embedding", reuse=None):
    """
    positional encoding
    :param inputs: [Tensor] with dimension of "batch_size * max_length"
    :param channels: [Int] Embedding size
    :param scale: [Boolean] If True, the output will be multiplied by sqrt num_units
    :param scope: [String] name of "variable_scope"
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs after positional encoding
    """
    batch_size = tf.shape(inputs)[0]
    max_length = tf.shape(inputs)[1]
    with tf.variable_scope(scope, reuse=reuse):
        position_ind = tf.tile(tf.expand_dims(tf.range(tf.to_int32(1), tf.add(max_length, 1)), 0), [batch_size, 1])

        # Convert to a tensor
        lookup_table = tf.convert_to_tensor(get_timing_signal_1d(max_length, channels))

        lookup_table = tf.concat((tf.zeros(shape=[1, channels]),
                                  lookup_table[:, :]), 0)
        position_inputs = tf.where(tf.equal(inputs, 0), tf.zeros_like(inputs), position_ind)

        outputs = tf.nn.embedding_lookup(lookup_table, position_inputs)

        if scale:
            outputs = outputs * math.sqrt(channels)

    return tf.cast(outputs, tf.float32)

通过get_timing_signal_1d()方法获得 [ 句子长度 * embedding维度 ]的矩阵

def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4, start_index=0):
    """
    positional encoding的方法
    :param length: [Int] max_length size
    :param channels: [Int] Embedding size
    :param min_timescale: [Float]
    :param max_timescale: [Float]
    :param start_index: [Int] index of first position
    :return: [Tensor] positional encoding of shape "length * channels"
    """
    position = tf.to_float(tf.range(start_index, length))
    num_timescales = channels // 2
    log_timescale_increment = (math.log(float(min_timescale) / float(max_timescale)) /
                               (tf.to_float(num_timescales) - 1))
    inv_timescales = min_timescale * tf.exp(tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)

    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
    return signal

1.2.3.Segment Embedding

该嵌入仅仅是为了让模型能够更好的区分输入的两个句子，其实通过[SEP]这个token以及能够区分两个句子了，但是对于模型来说显然还不够，在不加入segment embedding的情况下，模型的表现不太良好。

对于[PAD]这个token，所有的embedding（seg、pos）都设为了全零向量，以便后面attention的时候加入mask

def get_seg_embedding(inputs, channels, scale=True, scope="seg_embedding", reuse=None):
    """
    segment embedding
    :param inputs: [Tensor] with first dimension of "batch_size" like [1 1 1 2 2 2 2 0 0 0 ...]
    :param channels: [Int] Embedding size
    :param scale: [Boolean] If True, the output will be multiplied by sqrt num_units
    :param scope: [String] name of "variable_scope"
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs of embedding of sentence with shape of "batch_size * length * channels"
    """
    with tf.variable_scope(scope, reuse=reuse):
        lookup_table = tf.get_variable('lookup_table',
                                       dtype=tf.float32,
                                       shape=[3, channels],
                                       initializer=tf.contrib.layers.xavier_initializer())
        lookup_table = tf.concat((tf.zeros(shape=[1, channels], dtype=tf.float32),
                                  lookup_table[1:, :]), 0)

        outputs = tf.nn.embedding_lookup(lookup_table, inputs)
        if scale:
            outputs = outputs * math.sqrt(channels)

    return outputs

1.3.Self-Attention和Encoder-Decoder Attention

到这里，输入的处理就算完成了，到了重头戏Attention机制

两个输入的tensor总觉的一行用英语讲不清楚，就写在这里吧，from tensor对于两个Attention都是一致的就是输入，to tensor对于self-attention来说也是一致的，但对于encoder-decoder attention来说是最后一层encoder的输出，用来捕捉decoder和encoder之间的attention关系。

因为前面做了处理，所有的[PAD]这个token的embedding都是全零，所以对这个维度求绝对值后reduce sum之后，零就是[PAD]这个token，这样就不用再额外的添加一个mask ids作为输入了。

按照paper中的描述

def multi_head_attention(from_tensor: tf.Tensor,  to_tensor: tf.Tensor, channels=None, num_units=None, num_heads=8,
                         dropout_rate=0, is_training=True, attention_mask_flag=False, scope="multihead_attention",
                         activation=None, reuse=None):
    """
    multihead attention
    :param from_tensor: [Tensor]
    :param to_tensor: [Tensor] 
    :param channels: [Int] channel of last dimension of output
    :param num_units: [Int] channel size of matrix Q, K, V
    :param num_heads: [Int] head number of attention
    :param dropout_rate: [Float] dropout rate when 0 means no dropout
    :param is_training: [Boolean] whether it is training, If true, use dropout
    :param attention_mask_flag: [Boolean] If true, units that reference the future are masked
    :param scope: [String] name of "variable_scope"
    :param activation: [String] name of activate function
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs after multihead self attention with shape of "batch_size * max_length * (channels*num_heads)"
    """
    with tf.variable_scope(scope, reuse=reuse):
        if channels is None:
            channels = from_tensor.get_shape().as_list()[-1]
        if num_units is None:
            num_units = channels//num_heads
        activation_fn = get_activation(activation)
        # shape [batch_size, max_length, channels*num_heads]
        query_layer = tf.layers.dense(from_tensor, num_units * num_heads, activation=activation_fn)
        key_layer = tf.layers.dense(to_tensor, num_units * num_heads, activation=activation_fn)
        value_layer = tf.layers.dense(to_tensor, num_units * num_heads, activation=activation_fn)

        # shape [batch_size*num_heads, max_length, channels]
        query_layer_ = tf.concat(tf.split(query_layer, num_heads, axis=2), axis=0)
        key_layer_ = tf.concat(tf.split(key_layer, num_heads, axis=2), axis=0)
        value_layer_ = tf.concat(tf.split(value_layer, num_heads, axis=2), axis=0)

        # shape = [batch_size*num_heads, max_length, max_length]
        attention_scores = tf.matmul(query_layer_, tf.transpose(key_layer_, [0, 2, 1]))
        # Scale
        attention_scores = tf.multiply(attention_scores, 1.0 / tf.sqrt(float(channels)))
        # attention masks
        attention_masks = tf.sign(tf.abs(tf.reduce_sum(to_tensor, axis=-1)))
        attention_masks = tf.tile(attention_masks, [num_heads, 1])
        attention_masks = tf.tile(tf.expand_dims(attention_masks, axis=1), [1, tf.shape(from_tensor)[1], 1])
        neg_inf_matrix = tf.multiply(tf.ones_like(attention_scores), (-math.pow(2, 32) + 1))
        attention_scores = tf.where(tf.equal(attention_masks, 0), neg_inf_matrix, attention_scores)

        if attention_mask_flag:
            diag_vals = tf.ones_like(attention_scores[0, :, :])
            tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()

            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(attention_scores)[0], 1, 1])
            neg_inf_matrix = tf.multiply(tf.ones_like(masks), (-math.pow(2, 32) + 1))
            attention_scores = tf.where(tf.equal(masks, 0), neg_inf_matrix, attention_scores)

        # attention probability
        attention_probs = tf.nn.softmax(attention_scores)

        # query mask
        query_masks = tf.sign(tf.abs(tf.reduce_sum(from_tensor, axis=-1)))
        query_masks = tf.tile(query_masks, [num_heads, 1])
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(to_tensor)[1]])

        attention_probs *= query_masks

        # dropout
        attention_probs = tf.layers.dropout(attention_probs, rate=dropout_rate,
                                            training=tf.convert_to_tensor(is_training))
        outputs = tf.matmul(attention_probs, value_layer_)
        # shape [batch_size, max_length, channels*num_heads]
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)

        # reshape to from tensor
        outputs = tf.layers.dense(outputs, channels, activation=activation_fn)
        # Residual connection
        outputs += from_tensor
        # group normalization
        outputs = group_norm(outputs)
    return outputs

1.4.Feed Ward

论文中的Position-wise Feed-Forward Networks，论文中第二层的激活函数为线性激活函数，将第二层的activation function参数改为None才是原论文的做法，这里出于一些实验的原因没有照做

def feed_forward(inputs, channels, hidden_dims=None, scope="multihead_attention", activation=None, reuse=None):
    """
    :param inputs: [Tensor] with first dimension of "batch_size"
    :param channels: [Int] Embedding size
    :param hidden_dims: [List] hidden dimensions
    :param scope: [String] name of "variable_scope"
    :param activation: [String] name of activate function
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs after feed forward with shape of "batch_size * max_length * channels"
    """
    if hidden_dims is None:
        hidden_dims = 2*channels
    with tf.variable_scope(scope, reuse=reuse):
        activation_fn = get_activation(activation)

        params = {"inputs": inputs, "num_outputs": hidden_dims, "activation_fn": activation_fn}
        outputs = tf.contrib.layers.fully_connected(**params)

        params = {"inputs": outputs, "num_outputs": channels, "activation_fn": activation_fn}  # activation_fn可以改为None
        outputs = tf.contrib.layers.fully_connected(**params)
        outputs += inputs
        outputs = group_norm(outputs)
    return outputs

1.5.Layer Normalization

对了，还有layer normalization。

def group_norm(inputs: tf.Tensor, epsilon=1e-8, scope="layer_normalization", reuse=None):
    """
    layer normalization
    :param inputs: [Tensor] with first dimension of "batch_size"
    :param epsilon: [Float] a number for preventing ZeroDivision
    :param scope: [String] name of "variable_scope"
    :param reuse: [Boolean] tf parameter reuse
    :return: [Tensor] outputs after normalized
    """
    with tf.variable_scope(scope, reuse=reuse):
        inputs_shape = inputs.get_shape()
        params_shape = inputs_shape[-1:]
        mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)
        beta = tf.Variable(tf.zeros(params_shape))
        gamma = tf.Variable(tf.ones(params_shape))
        normalized = (inputs - mean) * tf.rsqrt(variance + epsilon)
        outputs = gamma * normalized + beta
    return outputs

2.Transformer for SNLI

基本工作都做好了，接下来使用之前写好的代码来搭建一个6层的Transformer

先定义一些模型的细节配置：

class ConfigModel(object):
    vocab_size_en = len(word_dict_en)
    channels = 400
    learning_rate = 0.0005
    layer_num = 6
    is_training = True
    is_transfer_learning = False
    restore_embedding = False
    shuffle_pool_size = 2560
    dropout_rate = 0.1
    num_heads = 8
    batch_size = 64
    max_length = 100
    num_tags = 3

然后搭建模型，整体按照Bert的思路搭建，最后取 [CLS] token的输出：

class TransformerSNLICls():
    def __init__(self, inputs, segs, label, config):
        self.inputs = tf.to_int32(inputs)  # batch_size*max_length
        self.segs = tf.to_int32(segs)  # 标识属于第几个句子 ([1 1 2 2 2 0 0 0 ...])
        self.target = tf.to_int32(label)
        self.vocab_size_en = config.vocab_size_en
        self.channels = config.channels
        self.num_heads = config.num_heads
        self.dropout_rate = config.dropout_rate
        self.is_training = config.is_training
        self.num_layer = config.layer_num
        self.learning_rate = config.learning_rate
        # {'_PAD': 0, '_BEGIN': 1, '_EOS': 2, '_CLS': 3, '_SEP': 4, '_MASK': 5}
        self.inputs = tf.concat((tf.ones_like(self.inputs[:, :1])*3, self.inputs), axis=-1)
        self.segs = tf.concat((tf.ones_like(self.segs[:, :1]), self.segs), axis=-1)

        with tf.variable_scope("encoder"):
            self.encode = get_embedding(self.inputs, self.vocab_size_en, self.channels, scope="en_embed")
            self.encode += get_positional_encoding(self.inputs, self.channels, scope="en_pe")
            self.encode += get_seg_embedding(self.segs, self.channels, scope="en_se")
            self.encode = tf.layers.dropout(self.encode, rate=self.dropout_rate,
                                            training=tf.convert_to_tensor(self.is_training))
            for i in range(self.num_layer):
                with tf.variable_scope("encoder_layer_{}".format(i)):
                    self.encode = multi_head_attention(self.encode, self.encode, self.channels,
                                                       num_heads=self.num_heads,
                                                       dropout_rate=self.dropout_rate,
                                                       is_training=self.is_training,
                                                       attention_mask_flag=False)
                    self.encode = feed_forward(self.encode, self.channels)
            self.encode_cls = tf.reshape(self.encode[:, :1, :], [-1, self.channels])

        self.output = tf.layers.dense(self.encode_cls, config.num_tags)
        self.preds = tf.to_int32(tf.argmax(self.output, axis=-1))
        self.acc = tf.reduce_mean(tf.to_float(tf.equal(self.preds, self.target)))
        if self.is_training:
            self.loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.output, labels=self.target))
            self.global_step = tf.Variable(0, name='global_step', trainable=False)
            self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate,
                                                    beta1=0.9, beta2=0.98, epsilon=1e-8)
            self.grads = self.optimizer.compute_gradients(self.loss)
            if config.is_transfer_learning:
                var_list = tf.trainable_variables()
                layer_name_list = ["encoder_layer_" + str(i) for i in range(4)]
                var_list_ = [v for v in var_list if v.name.split("/")[1] in layer_name_list]
                var_list_ += [v for v in var_list if "lookup_table" in v.name]
                for index, grad in enumerate(self.grads):
                    if grad[1] in var_list_:
                        self.grads[index] = (grad[0]*0.2, grad[1])
            self.train_op = self.optimizer.apply_gradients(self.grads, global_step=self.global_step)

到这里模型就准备好了，接下来对数据进行一些处理

3.SNLI to TF Record

将所有数据进行处理并生成segment ids：

def get_data(snli_name, max_length=config.max_length//2, word_dict=word_dict_en):
    sentence_1 = list()
    sentence_2 = list()
    label = list()
    texts = list()
    seg_ids = list()
    with open(os.path.join(data_path, snli_name), 'r') as f:
        for item in jsonlines.Reader(f):
            try:
                label.append(label_to_num_dict[item["gold_label"]])
            except KeyError:
                continue
            sentence_1.append(normalize_text(item["sentence1"]))
            sentence_2.append(normalize_text(item["sentence2"]))

    en_data_num_1 = text_to_numbers(sentence_1, word_dict_en, max_length=max_length)
    en_data_num_2 = text_to_numbers(sentence_2, word_dict_en, max_length=max_length)

    for i_ in range(len(en_data_num_1)):
        texts.append(en_data_num_1[i_] + [word_dict["_SEP"]] + en_data_num_2[i_] + [word_dict["_SEP"]])
        seg_ids.append((len(en_data_num_1[i_])+1)*[1]+(len(en_data_num_2[i_])+1)*[2])
    return texts, label, seg_ids

并转换为tf record格式，以便tensorflow 更高效的读取：

def write_binary(record_name, texts_, label_, seg_ids_):
    writer = tf.python_io.TFRecordWriter(record_name)
    for it, text in tqdm(enumerate(texts_)):
        example = tf.train.Example(
            features=tf.train.Features(
                feature={
                    "text_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=text)),
                    "seg_ids": tf.train.Feature(int64_list=tf.train.Int64List(value=seg_ids_[it])),
                    "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label_[it]])),
                }
            )
        )
        serialized = example.SerializeToString()
        writer.write(serialized)
    writer.close()

4.训练

加载模型，加载数据：

if __name__ == '__main__':
    with tf.Session() as sess:
        data_set_train = get_dataset(train_snli_name_tf)
        data_set_train = data_set_train.shuffle(config.shuffle_pool_size).repeat(). \
            padded_batch(config.batch_size, padded_shapes=([config.max_length], [config.max_length], []))
        data_set_train_iter = data_set_train.make_one_shot_iterator()
        train_handle = sess.run(data_set_train_iter.string_handle())

        data_set_test = get_dataset(os.path.join(test_snli_name_tf))
        if test_total_acc:
            data_set_test = data_set_test.shuffle(config.shuffle_pool_size). \
                padded_batch(config.batch_size, padded_shapes=([config.max_length], [config.max_length], []))
        else:
            data_set_test = data_set_test.shuffle(config.shuffle_pool_size).repeat(). \
                padded_batch(config.batch_size, padded_shapes=([config.max_length], [config.max_length], []))
        data_set_test_iter = data_set_test.make_one_shot_iterator()
        test_handle = sess.run(data_set_test_iter.string_handle())

        handle = tf.placeholder(tf.string, shape=[])
        iterator = tf.data.Iterator.from_string_handle(handle, data_set_train.output_types,
                                                       data_set_train.output_shapes)

        inputs, segs, target = iterator.get_next()

        tsl = TransformerSNLICls(inputs, segs, target, config)
        sess.run(tf.global_variables_initializer())
        saver = tf.train.Saver(max_to_keep=1)

开始训练：

        print("starting training")
        for i in range(12000):
            train_feed = {handle: train_handle}
            sess.run(tsl.train_op, train_feed)
            if (i+1) % 100 == 0:
                pred, acc, loss = sess.run([tsl.preds, tsl.acc, tsl.loss], train_feed)
                print("Generation train {} : acc: {}  loss: {} ".format(i, acc, loss))
            if (i+1) % 200 == 0:
                tpred, tacc, tloss = sess.run([tsl.preds, tsl.acc, tsl.loss], {handle: test_handle})
                print("Generation test {} : acc: {}  loss: {} ".format(i, tacc, tloss))
            if (i+1) % 2000 == 0:
                print("Generation train {} model saved ".format(i))
                saver.save(sess, os.path.join(model_save_path, model_name.format(model_choose)))
        saver.save(sess, os.path.join(model_save_path, model_name.format(model_choose)))

最后，初步在整个测试集上达到78.7%的准确度