谷歌BERT预训练源码解析（三）：训练过程

一时失言乱红尘 2021-10-19 20:44 308阅读 0赞

目录  
前言  
源码解析  
主函数  
自定义模型  
遮蔽词预测  
下一句预测  
规范化数据集  
前言  
本部分介绍BERT训练过程，BERT模型训练过程是在自己的TPU上进行的，这部分我没做过研究所以不做深入探讨。BERT针对两个任务同时训练。1.下一句预测。2.遮蔽词识别  
下面介绍BERT的预训练模型run\_pretraining.py是怎么训练的。

源码解析  
主函数  
训练过程主要用了estimator调度器。这个调度器支持自定义训练过程，将训练集传入之后自动训练。详情见注释

def main(\_):  
tf.logging.set\_verbosity(tf.logging.INFO)

if not FLAGS.do\_train and not FLAGS.do\_eval:  
raise ValueError("At least one of \`do\_train\` or \`do\_eval\` must be True.")

bert\_config = modeling.BertConfig.from\_json\_file(FLAGS.bert\_config\_file)

tf.gfile.MakeDirs(FLAGS.output\_dir)

input\_files = \[\]  
for input\_pattern in FLAGS.input\_file.split(","):  
input\_files.extend(tf.gfile.Glob(input\_pattern))

tf.logging.info("\*\*\* Input Files \*\*\*")  
for input\_file in input\_files:  
tf.logging.info(" %s" % input\_file)

tpu\_cluster\_resolver = None  
if FLAGS.use\_tpu and FLAGS.tpu\_name:  
tpu\_cluster\_resolver = tf.contrib.cluster\_resolver.TPUClusterResolver(  
FLAGS.tpu\_name, zone=FLAGS.tpu\_zone, project=FLAGS.gcp\_project)

is\_per\_host = tf.contrib.tpu.InputPipelineConfig.PER\_HOST\_V2  
run\_config = tf.contrib.tpu.RunConfig( \#训练参数  
cluster=tpu\_cluster\_resolver,  
master=FLAGS.master,  
model\_dir=FLAGS.output\_dir,  
save\_checkpoints\_steps=FLAGS.save\_checkpoints\_steps,  
tpu\_config=tf.contrib.tpu.TPUConfig(  
iterations\_per\_loop=FLAGS.iterations\_per\_loop,  
num\_shards=FLAGS.num\_tpu\_cores,  
per\_host\_input\_for\_training=is\_per\_host))

model\_fn = model\_fn\_builder( \#自定义模型，用于estimator训练  
bert\_config=bert\_config,  
init\_checkpoint=FLAGS.init\_checkpoint,  
learning\_rate=FLAGS.learning\_rate,  
num\_train\_steps=FLAGS.num\_train\_steps,  
num\_warmup\_steps=FLAGS.num\_warmup\_steps,  
use\_tpu=FLAGS.use\_tpu,  
use\_one\_hot\_embeddings=FLAGS.use\_tpu)

\# If TPU is not available, this will fall back to normal Estimator on CPU  
\# or GPU.  
estimator = tf.contrib.tpu.TPUEstimator( \#创建TPUEstimator  
use\_tpu=FLAGS.use\_tpu,  
model\_fn=model\_fn,  
config=run\_config,  
train\_batch\_size=FLAGS.train\_batch\_size,  
eval\_batch\_size=FLAGS.eval\_batch\_size)

if FLAGS.do\_train: \#训练过程  
tf.logging.info("\*\*\*\*\* Running training \*\*\*\*\*")  
tf.logging.info(" Batch size = %d", FLAGS.train\_batch\_size)  
train\_input\_fn = input\_fn\_builder( \#创建输入训练集  
input\_files=input\_files,  
max\_seq\_length=FLAGS.max\_seq\_length,  
max\_predictions\_per\_seq=FLAGS.max\_predictions\_per\_seq,  
is\_training=True)  
estimator.train(input\_fn=train\_input\_fn, max\_steps=FLAGS.num\_train\_steps)

if FLAGS.do\_eval: \#验证过程  
tf.logging.info("\*\*\*\*\* Running evaluation \*\*\*\*\*")  
tf.logging.info(" Batch size = %d", FLAGS.eval\_batch\_size)

eval\_input\_fn = input\_fn\_builder( \#创建验证集  
input\_files=input\_files,  
max\_seq\_length=FLAGS.max\_seq\_length,  
max\_predictions\_per\_seq=FLAGS.max\_predictions\_per\_seq,  
is\_training=False)

result = estimator.evaluate(  
input\_fn=eval\_input\_fn, steps=FLAGS.max\_eval\_steps)

output\_eval\_file = os.path.join(FLAGS.output\_dir, "eval\_results.txt")  
with tf.gfile.GFile(output\_eval\_file, "w") as writer:  
tf.logging.info("\*\*\*\*\* Eval results \*\*\*\*\*")  
for key in sorted(result.keys()):  
tf.logging.info(" %s = %s", key, str(result\[key\]))  
writer.write("%s = %s\\n" % (key, str(result\[key\])))  
1  
2  
3  
4  
5  
6  
7  
8  
9  
10  
11  
12  
13  
14  
15  
16  
17  
18  
19  
20  
21  
22  
23  
24  
25  
26  
27  
28  
29  
30  
31  
32  
33  
34  
35  
36  
37  
38  
39  
40  
41  
42  
43  
44  
45  
46  
47  
48  
49  
50  
51  
52  
53  
54  
55  
56  
57  
58  
59  
60  
61  
62  
63  
64  
65  
66  
67  
68  
69  
70  
71  
72  
73  
74  
75  
76  
77  
78  
79  
80  
81  
自定义模型  
首先获取数据内容，传入到上一篇定义的模型中。对下一句预测任务取出模型的\[CLS\]结果。对遮蔽词预测任务取出模型的最后结果。然后分别计算loss值，最后将loss值相加。详情见注释

def model\_fn\_builder(bert\_config, init\_checkpoint, learning\_rate,  
num\_train\_steps, num\_warmup\_steps, use\_tpu,  
use\_one\_hot\_embeddings):  
"""Returns \`model\_fn\` closure for TPUEstimator."""

def model\_fn(features, labels, mode, params): \# pylint: disable=unused-argument  
"""The \`model\_fn\` for TPUEstimator."""

tf.logging.info("\*\*\* Features \*\*\*")  
for name in sorted(features.keys()):  
tf.logging.info(" name = %s, shape = %s" % (name, features\[name\].shape))  
\#获取数据内容  
input\_ids = features\["input\_ids"\]  
input\_mask = features\["input\_mask"\]  
segment\_ids = features\["segment\_ids"\]  
masked\_lm\_positions = features\["masked\_lm\_positions"\]  
masked\_lm\_ids = features\["masked\_lm\_ids"\]  
masked\_lm\_weights = features\["masked\_lm\_weights"\]  
next\_sentence\_labels = features\["next\_sentence\_labels"\]

is\_training = (mode == tf.estimator.ModeKeys.TRAIN)  
传入到Bert模型中。  
model = modeling.BertModel(  
config=bert\_config,  
is\_training=is\_training,  
input\_ids=input\_ids,  
input\_mask=input\_mask,  
token\_type\_ids=segment\_ids,  
use\_one\_hot\_embeddings=use\_one\_hot\_embeddings)  
\#遮蔽预测的batch\_loss,平均loss，预测概率矩阵  
(masked\_lm\_loss,  
masked\_lm\_example\_loss, masked\_lm\_log\_probs) = get\_masked\_lm\_output(  
bert\_config, model.get\_sequence\_output(), model.get\_embedding\_table(),  
masked\_lm\_positions, masked\_lm\_ids, masked\_lm\_weights)  
\#下一句预测的batch\_loss,平均loss，预测概率矩阵  
(next\_sentence\_loss, next\_sentence\_example\_loss,  
next\_sentence\_log\_probs) = get\_next\_sentence\_output(  
bert\_config, model.get\_pooled\_output(), next\_sentence\_labels)  
\#loss相加  
total\_loss = masked\_lm\_loss + next\_sentence\_loss  
\#获取所有变量  
tvars = tf.trainable\_variables()

initialized\_variable\_names = \{\}  
scaffold\_fn = None  
\#如果有之前保存的模型  
if init\_checkpoint:  
(assignment\_map, initialized\_variable\_names  
) = modeling.get\_assignment\_map\_from\_checkpoint(tvars, init\_checkpoint)  
if use\_tpu:

def tpu\_scaffold():  
tf.train.init\_from\_checkpoint(init\_checkpoint, assignment\_map)  
return tf.train.Scaffold()

scaffold\_fn = tpu\_scaffold  
else:  
tf.train.init\_from\_checkpoint(init\_checkpoint, assignment\_map)

tf.logging.info("\*\*\*\* Trainable Variables \*\*\*\*")  
\#如果有之前保存的模型  
for var in tvars:  
init\_string = ""  
if var.name in initialized\_variable\_names:  
init\_string = ", \*INIT\_FROM\_CKPT\*"  
tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape,  
init\_string)  
  
output\_spec = None  
if mode == tf.estimator.ModeKeys.TRAIN:  
train\_op = optimization.create\_optimizer( \#自定义好的优化器  
total\_loss, learning\_rate, num\_train\_steps, num\_warmup\_steps, use\_tpu)

output\_spec = tf.contrib.tpu.TPUEstimatorSpec( \#Estimator要求返回一个EstimatorSpec对象  
mode=mode,  
loss=total\_loss,  
train\_op=train\_op,  
scaffold\_fn=scaffold\_fn)  
\#验证过程  
elif mode == tf.estimator.ModeKeys.EVAL:

def metric\_fn(masked\_lm\_example\_loss, masked\_lm\_log\_probs, masked\_lm\_ids,  
masked\_lm\_weights, next\_sentence\_example\_loss,  
next\_sentence\_log\_probs, next\_sentence\_labels):  
"""Computes the loss and accuracy of the model."""  
masked\_lm\_log\_probs = tf.reshape(masked\_lm\_log\_probs,  
\[-1, masked\_lm\_log\_probs.shape\[-1\]\]) \#概率矩阵转成\[batch\_size\*max\_pred\_pre\_seq,vocab\_size\]  
masked\_lm\_predictions = tf.argmax(  
masked\_lm\_log\_probs, axis=-1, output\_type=tf.int32) \#取最大值位置为输出  
masked\_lm\_example\_loss = tf.reshape(masked\_lm\_example\_loss, \[-1\]) \#每句loss列表 \[batch\_size\*max\_pred\_per\_seq\]  
masked\_lm\_ids = tf.reshape(masked\_lm\_ids, \[-1\])  
masked\_lm\_weights = tf.reshape(masked\_lm\_weights, \[-1\])  
masked\_lm\_accuracy = tf.metrics.accuracy( \#计算准确率  
labels=masked\_lm\_ids,  
predictions=masked\_lm\_predictions,  
weights=masked\_lm\_weights)  
masked\_lm\_mean\_loss = tf.metrics.mean( \#计算平均loss  
values=masked\_lm\_example\_loss, weights=masked\_lm\_weights)

next\_sentence\_log\_probs = tf.reshape(  
next\_sentence\_log\_probs, \[-1, next\_sentence\_log\_probs.shape\[-1\]\])  
next\_sentence\_predictions = tf.argmax( \#获取最大位置为输出  
next\_sentence\_log\_probs, axis=-1, output\_type=tf.int32)  
next\_sentence\_labels = tf.reshape(next\_sentence\_labels, \[-1\])  
next\_sentence\_accuracy = tf.metrics.accuracy( \#计算准确率  
labels=next\_sentence\_labels, predictions=next\_sentence\_predictions)  
next\_sentence\_mean\_loss = tf.metrics.mean( 计算平均loss  
values=next\_sentence\_example\_loss)

return \{  
"masked\_lm\_accuracy": masked\_lm\_accuracy,  
"masked\_lm\_loss": masked\_lm\_mean\_loss,  
"next\_sentence\_accuracy": next\_sentence\_accuracy,  
"next\_sentence\_loss": next\_sentence\_mean\_loss,  
\}

eval\_metrics = (metric\_fn, \[  
masked\_lm\_example\_loss, masked\_lm\_log\_probs, masked\_lm\_ids,  
masked\_lm\_weights, next\_sentence\_example\_loss,  
next\_sentence\_log\_probs, next\_sentence\_labels  
\])  
output\_spec = tf.contrib.tpu.TPUEstimatorSpec( \#Estimator要求返回一个EstimatorSpec对象  
mode=mode,  
loss=total\_loss,  
eval\_metrics=eval\_metrics,  
scaffold\_fn=scaffold\_fn)  
else:  
raise ValueError("Only TRAIN and EVAL modes are supported: %s" % (mode))

return output\_spec

return model\_fn  
1  
2  
3  
4  
5  
6  
7  
8  
9  
10  
11  
12  
13  
14  
15  
16  
17  
18  
19  
20  
21  
22  
23  
24  
25  
26  
27  
28  
29  
30  
31  
32  
33  
34  
35  
36  
37  
38  
39  
40  
41  
42  
43  
44  
45  
46  
47  
48  
49  
50  
51  
52  
53  
54  
55  
56  
57  
58  
59  
60  
61  
62  
63  
64  
65  
66  
67  
68  
69  
70  
71  
72  
73  
74  
75  
76  
77  
78  
79  
80  
81  
82  
83  
84  
85  
86  
87  
88  
89  
90  
91  
92  
93  
94  
95  
96  
97  
98  
99  
100  
101  
102  
103  
104  
105  
106  
107  
108  
109  
110  
111  
112  
113  
114  
115  
116  
117  
118  
119  
120  
121  
122  
123  
124  
125  
126  
127  
128  
129  
130  
131  
132  
遮蔽词预测  
输入BERT模型的最后一层encoder，输出遮蔽词预测任务的loss和概率矩阵。详情见注释

def get\_masked\_lm\_output(bert\_config, input\_tensor, output\_weights, positions,  
label\_ids, label\_weights):  
\#这里的input\_tensor是模型中传回的最后一层结果 \[batch\_size,seq\_length,hidden\_size\]。  
\#output\_weights是词向量表 \[vocab\_size,embedding\_size\]  
"""Get loss and log probs for the masked LM."""  
\#获取positions位置的所有encoder（即要预测的那些位置的encoder）  
input\_tensor = gather\_indexes(input\_tensor, positions) \#\[batch\_size\*max\_pred\_pre\_seq,hidden\_size\]

with tf.variable\_scope("cls/predictions"):  
\# We apply one more non-linear transformation before the output layer.  
\# This matrix is not used after pre-training.  
with tf.variable\_scope("transform"):  
input\_tensor = tf.layers.dense( \#传入一个全连接层 输出shape \[batch\_size\*max\_pred\_pre\_seq,hidden\_size\]  
input\_tensor,  
units=bert\_config.hidden\_size,  
activation=modeling.get\_activation(bert\_config.hidden\_act),  
kernel\_initializer=modeling.create\_initializer(  
bert\_config.initializer\_range))  
input\_tensor = modeling.layer\_norm(input\_tensor)

\# The output weights are the same as the input embeddings, but there is  
\# an output-only bias for each token.  
output\_bias = tf.get\_variable(  
"output\_bias",  
shape=\[bert\_config.vocab\_size\],  
initializer=tf.zeros\_initializer())  
logits = tf.matmul(input\_tensor, output\_weights, transpose\_b=True) \#\[batch\_size\*max\_pred\_pre\_seq,vocab\_size\]  
logits = tf.nn.bias\_add(logits, output\_bias) \#加bias  
log\_probs = tf.nn.log\_softmax(logits, axis=-1) \#\[batch\_size\*max\_pred\_pre\_seq,vocab\_size\]

label\_ids = tf.reshape(label\_ids, \[-1\]) \#\[batch\_size\*max\_pred\_per\_seq\]  
label\_weights = tf.reshape(label\_weights, \[-1\])

one\_hot\_labels = tf.one\_hot( \#\[batch\_size\*max\_pred\_per\_seq,vocab\_size\]  
label\_ids, depth=bert\_config.vocab\_size, dtype=tf.float32) \#label id转one hot

\# The \`positions\` tensor might be zero-padded (if the sequence is too  
\# short to have the maximum number of predictions). The \`label\_weights\`  
\# tensor has a value of 1.0 for every real prediction and 0.0 for the  
\# padding predictions.  
per\_example\_loss = -tf.reduce\_sum(log\_probs \* one\_hot\_labels, axis=\[-1\]) \#\[batch\_size\*max\_pred\_per\_seq\]  
numerator = tf.reduce\_sum(label\_weights \* per\_example\_loss) \#\[1\] 一个batch的loss  
denominator = tf.reduce\_sum(label\_weights) + 1e-5  
loss = numerator / denominator \#平均loss

return (loss, per\_example\_loss, log\_probs)  
1  
2  
3  
4  
5  
6  
7  
8  
9  
10  
11  
12  
13  
14  
15  
16  
17  
18  
19  
20  
21  
22  
23  
24  
25  
26  
27  
28  
29  
30  
31  
32  
33  
34  
35  
36  
37  
38  
39  
40  
41  
42  
43  
44  
45  
46  
下一句预测  
输入BERT模型CLS的encoder，输出下一句预测任务的loss和概率矩阵,详情见注释

def get\_next\_sentence\_output(bert\_config, input\_tensor, labels):  
\#input\_tensor shape \[batch\_size,hidden\_size\]  
"""Get loss and log probs for the next sentence prediction."""

\# Simple binary classification. Note that 0 is "next sentence" and 1 is  
\# "random sentence". This weight matrix is not used after pre-training.  
with tf.variable\_scope("cls/seq\_relationship"):  
output\_weights = tf.get\_variable(  
"output\_weights",  
shape=\[2, bert\_config.hidden\_size\],  
initializer=modeling.create\_initializer(bert\_config.initializer\_range))  
output\_bias = tf.get\_variable(  
"output\_bias", shape=\[2\], initializer=tf.zeros\_initializer()) \#\[batch\_size,hidden\_size\]  
  
logits = tf.matmul(input\_tensor, output\_weights, transpose\_b=True) \#\[batch\_size,2\]  
logits = tf.nn.bias\_add(logits, output\_bias) \#\[batch\_size,2\]  
log\_probs = tf.nn.log\_softmax(logits, axis=-1)  
labels = tf.reshape(labels, \[-1\])  
one\_hot\_labels = tf.one\_hot(labels, depth=2, dtype=tf.float32) \#\[batch\_size,2\]  
per\_example\_loss = -tf.reduce\_sum(one\_hot\_labels \* log\_probs, axis=-1) \#\[batch\_size\]  
loss = tf.reduce\_mean(per\_example\_loss) \#\[1\]  
return (loss, per\_example\_loss, log\_probs)  
1  
2  
3  
4  
5  
6  
7  
8  
9  
10  
11  
12  
13  
14  
15  
16  
17  
18  
19  
20  
21  
22  
规范化数据集  
Estimator要求模型的输入为特定格式(from\_tensor\_slices)，所以要对数据进行类封装

def input\_fn\_builder(input\_files,  
max\_seq\_length,  
max\_predictions\_per\_seq,  
is\_training,  
num\_cpu\_threads=4):  
"""Creates an \`input\_fn\` closure to be passed to TPUEstimator."""

def input\_fn(params):  
"""The actual input function."""  
batch\_size = params\["batch\_size"\]

name\_to\_features = \{  
"input\_ids":  
tf.FixedLenFeature(\[max\_seq\_length\], tf.int64),  
"input\_mask":  
tf.FixedLenFeature(\[max\_seq\_length\], tf.int64),  
"segment\_ids":  
tf.FixedLenFeature(\[max\_seq\_length\], tf.int64),  
"masked\_lm\_positions":  
tf.FixedLenFeature(\[max\_predictions\_per\_seq\], tf.int64),  
"masked\_lm\_ids":  
tf.FixedLenFeature(\[max\_predictions\_per\_seq\], tf.int64),  
"masked\_lm\_weights":  
tf.FixedLenFeature(\[max\_predictions\_per\_seq\], tf.float32),  
"next\_sentence\_labels":  
tf.FixedLenFeature(\[1\], tf.int64),  
\}

\# For training, we want a lot of parallel reading and shuffling.  
\# For eval, we want no shuffling and parallel reading doesn't matter.  
if is\_training:  
d = tf.data.Dataset.from\_tensor\_slices(tf.constant(input\_files))  
d = d.repeat() \#重复  
d = d.shuffle(buffer\_size=len(input\_files)) \#打乱

\# \`cycle\_length\` is the number of parallel files that get read.  
cycle\_length = min(num\_cpu\_threads, len(input\_files))

\# \`sloppy\` mode means that the interleaving is not exact. This adds  
\# even more randomness to the training pipeline.  
d = d.apply(  
tf.contrib.data.parallel\_interleave( \#生成嵌套数据集，并且输出其元素隔行交错  
tf.data.TFRecordDataset,  
sloppy=is\_training,  
cycle\_length=cycle\_length))  
d = d.shuffle(buffer\_size=100)  
else:  
d = tf.data.TFRecordDataset(input\_files)  
\# Since we evaluate for a fixed number of steps we don't want to encounter  
\# out-of-range exceptions.  
d = d.repeat()

\# We must \`drop\_remainder\` on training because the TPU requires fixed  
\# size dimensions. For eval, we assume we are evaluating on the CPU or GPU  
\# and we \*don't\* want to drop the remainder, otherwise we wont cover  
\# every sample.  
d = d.apply(  
tf.contrib.data.map\_and\_batch( \#结构转换  
lambda record: \_decode\_record(record, name\_to\_features),  
batch\_size=batch\_size,  
num\_parallel\_batches=num\_cpu\_threads,  
drop\_remainder=True))  
return d

return input\_fn  
\---------------------  
作者：保持一份率性  
来源：CSDN  
原文：https://blog.csdn.net/weixin\_39470744/article/details/84619903  
版权声明：本文为博主原创文章，转载请附上博文链接！

![1432127-20190729165838647-485280105.jpg][]

转载于:https://www.cnblogs.com/jfdwd/p/11264955.html

[1432127-20190729165838647-485280105.jpg]: /images/20211019/3760b3e6e0824547a25f348859a79fcb.png