bert pretraining源码刨析

发表于   |   分类于

Bert Pretraining代码刨析

本篇文章结合实际例子对bert pretraining源码进行刨析，更清晰的展现bert的整个过程。pretraining主要包括三个文件，create_pretraining_data.py, run_pretraining.py, modeling.py。

create_pretraining_data

构造pretraining的训练数据，我们使用的数据(chg_data)是周杰伦的四句歌词

1
2
3
4
5
素胚勾勒出青花笔锋浓转淡转淡
瓶身描绘的牡丹一如你初妆初妆

冉冉檀香透过窗心事我了然了然
宣纸上走笔至此搁一半一半

构造训练数据的代码如下：

1
2
3
4
5
6
7
8
9
10
CUDA_VISIBLE_DEVICES=3 python create_pretraining_data.py \
                   --input_file=./chg_data \
                   --output_file=./tmp/tf_examples.tfrecord \
                   --vocab_file=$BERT_BASE_DIR/vocab.txt \
                   --do_lower_case=True \
                   --max_seq_length=128 \
                   --max_predictions_per_seq=20 \
                   --masked_lm_prob=0.15 \
                   --random_seed=12345 \
                   --dupe_factor=5

下面对create_pretraining_data的代码进行一下解析。首先，create_training_instances函数读取chg_data数据，根据空行，将数据分成两个文档，如下所示。

1
[[['素', '胚', '勾', '勒', '出', '青', '花', '笔', '锋', '浓', '转', '淡'], ['瓶', '身', '描', '绘', '的', '牡', '丹', '一', '如', '你', '初', '妆']], [['冉', '冉',     '檀', '香', '透', '过', '窗', '心', '事', '我', '了', '然'], ['宣', '纸', '上', '走', '笔', '至', '此', '搁', '一', '半']]]

create_training_instances中的关键代码如下，功能是将数据进行拆分，构造sentence／next_sentence对

1
2
3
4
5
6
7
# 构造多组first_sen & next_sen对，并对数据进行mask
for _ in range(dupe_factor):
  for document_index in range(len(all_documents)):
    instances.extend(
        create_instances_from_document(
            all_documents, document_index, max_seq_length, short_seq_prob,
            masked_lm_prob, max_predictions_per_seq, vocab_words, rng))

最终生成的数据instance如下，并将数据写入tf_examples.tfrecord，用于之后的pertaining。当然这是命中80%概率进行mask的情况。

1
2
3
4
5
tokens : [CLS] [MASK] 胚 勾 勒 出 青 花 [MASK] 锋 浓 转 淡 [SEP] 瓶 身 描 绘 的 牡 丹 一 如 [MASK] 初 [MASK] [SEP]
segment_ids :  0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1
is_random_next : False
masked_lm_positions: 1 8 23 25
masked_lm_labels: 素 笔 你 妆

tf_record格式的数据如下，需要注意的是next_sentence_labels中0表示是next_sentence, 1表示不是next_sentence。

1
2
3
4
5
6
7
8
tokens: [CLS] [MASK] 胚 勾 勒 出 青 花 [MASK] 锋 浓 转 淡 [SEP] 瓶 身 描 绘 的 牡 丹 一 如 [MASK] 初 [MASK] [SEP]
input_ids: 101 103 5524 1256 1239 1139 7471 5709 103 7226 3849 6760 3909 102 4486 6716 2989 5313 4638 4285 710 671 1963 103 1159 103 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
intput_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
segment_ids: 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
masked_lm_positions: 1 8 23 25 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
masked_lm_ids: 5162 5011 872 1966 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
masked_lm_weights: 1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
next_sentence_labels: 0

run_pretraining

主要包括两个loss，一个是get_masked_lm_output用于预测mask后的词，一个是get_next_sentence_output用来预测next_sen的概率。最终的loss是这两个loss的和，用来作为优化目标。

modeling

该模块完整实现了transformer的encoder部分的代码，实现了multi-head attention结构。

最终的输出包括get_sequence_output返回所有token的隐层表示，get_pooled_output返回[CLS]对应的隐层表示。

下面这段代码要表达的是，[CLS]位置不是直接使用encode的输出，而是又增加了一个hidden layer，作为最终[CLS]的输出。

1
2
3
4
5
6
7
8
9
with tf.variable_scope("pooler"):
# We "pool" the model by simply taking the hidden state corresponding
# to the first token. We assume that this has been pre-trained
    first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1)
    self.pooled_output = tf.layers.dense(
            first_token_tensor,
            config.hidden_size,
            activation=tf.tanh,
            kernel_initializer=create_initializer(config.initializer_range))

总结

看论文的同时还要去看下源代码来加深理解，更需要我们跑一些case来验证我们的想法。

------ 本文结束 ------
k