-
Notifications
You must be signed in to change notification settings - Fork 0
/
component_bert_sequence_label_train.py
288 lines (253 loc) · 12.5 KB
/
component_bert_sequence_label_train.py
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
#!/usr/bin/python
# coding:utf8
"""
@author: Cong Yu
@time: 2019-06-05 10:42
"""
import tensorflow as tf
import numpy as np
import pandas as pd
import collections
import modeling
import optimization
from sklearn.externals import joblib
import os
from sklearn.metrics import classification_report
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'
config = {
"in_1": "./out/seq_train.tf_record", # 第一个输入为 训练文件
"in_2": "./out/seq_dev.tf_record", # 第二个输入为 验证文件
"bert_config": "./bert/bert_config.json", # bert模型配置文件
"init_checkpoint": "./bert/bert_model.ckpt", # 预训练bert模型
# "init_checkpoint": "./bin/bert.ckpt-114000", # 预训练bert模型
"train_examples_len": 20864,
"dev_examples_len": 710,
"num_labels": 10,
"train_batch_size": 32,
"dev_batch_size": 32,
"num_train_epochs": 2,
"eval_per_step": 100,
"learning_rate": 5e-5,
"warmup_proportion": 0.1,
"max_seq_len": 128, # 输入文本片段的最大 char级别 长度
"out": "./bin_seq/", # 保存模型路径
"out_1": "./bin_seq_1/" # 保存模型路径
}
def load_bert_config(path):
"""
bert 模型配置文件
"""
return modeling.BertConfig.from_json_file(path)
def create_model(bert_config, is_training, input_ids, input_mask, segment_ids, labels, keep_prob, num_labels,
use_one_hot_embeddings):
"""Creates a classification model."""
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)
# In the demo, we are doing a simple classification task on the entire
# segment.
#
# If you want to use the token-level output, use model.get_sequence_output()
# instead.
# output_layer = model.get_pooled_output()
output_layer = model.get_sequence_output()
hidden_size = output_layer.shape[-1].value
print(output_layer.shape)
output_weight = tf.get_variable(
"output_weights", [num_labels, hidden_size],
initializer=tf.truncated_normal_initializer(stddev=0.02)
)
output_bias = tf.get_variable(
"output_bias", [num_labels], initializer=tf.zeros_initializer()
)
with tf.variable_scope("loss"):
if is_training:
output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)
output_layer = tf.reshape(output_layer, [-1, hidden_size])
logits = tf.matmul(output_layer, output_weight, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits, [-1, config["max_seq_len"], num_labels])
print(labels.shape)
print(logits.shape)
predict = tf.argmax(tf.nn.softmax(logits, axis=-1), 2, name="tag")
print(predict.shape)
# input_mask = tf.squeeze(input_mask, axis=1)
print("input_mask", input_mask.shape)
sum_loss = tf.reduce_sum(input_mask)
print("sum_loss", sum_loss.shape)
loss_s = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
print("loss_s", loss_s.shape)
loss_s = tf.squeeze(loss_s)
# 不计算 mask 的loss
loss = tf.reduce_sum(loss_s * tf.cast(input_mask, tf.float32)) / tf.cast(sum_loss, tf.float32)
equals = tf.reduce_sum(
tf.cast(tf.equal(tf.cast(predict, tf.int64), labels), tf.float32) * tf.cast(input_mask, tf.float32))
print(equals.shape)
acc = equals / tf.cast(sum_loss, tf.float32)
return (loss, acc, logits, predict)
def get_input_data(input_file, seq_length, batch_size):
def parser(record):
name_to_features = {
"input_ids": tf.FixedLenFeature([seq_length], tf.int64),
"input_mask": tf.FixedLenFeature([seq_length], tf.int64),
"segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
"label_ids": tf.FixedLenFeature([seq_length], tf.int64),
}
example = tf.parse_single_example(record, features=name_to_features)
input_ids = example["input_ids"]
input_mask = example["input_mask"]
segment_ids = example["segment_ids"]
labels = example["label_ids"]
return input_ids, input_mask, segment_ids, labels
dataset = tf.data.TFRecordDataset(input_file)
# 数据类别集中,需要较大的buffer_size,才能有效打乱,或者再 数据处理的过程中进行打乱
dataset = dataset.map(parser).repeat().batch(batch_size).shuffle(buffer_size=3000)
iterator = dataset.make_one_shot_iterator()
input_ids, input_mask, segment_ids, labels = iterator.get_next()
return input_ids, input_mask, segment_ids, labels
def main():
print("print start load the params...")
tf.logging.set_verbosity(tf.logging.INFO)
tf.gfile.MakeDirs(config["out"])
train_examples_len = config["train_examples_len"]
dev_examples_len = config["dev_examples_len"]
learning_rate = config["learning_rate"]
eval_per_step = config["eval_per_step"]
num_labels = config["num_labels"]
print(num_labels)
num_train_steps = int(train_examples_len / config["train_batch_size"] * config["num_train_epochs"])
print("num_train_steps:", num_train_steps)
num_dev_steps = int(dev_examples_len / config["dev_batch_size"])
num_warmup_steps = int(num_train_steps * config["warmup_proportion"])
use_one_hot_embeddings = False
is_training = True
use_tpu = False
seq_len = config["max_seq_len"]
init_checkpoint = config["init_checkpoint"]
print("print start compile the bert model...")
# 定义输入输出
input_ids = tf.placeholder(tf.int64, shape=[None, seq_len], name='input_ids')
input_mask = tf.placeholder(tf.int64, shape=[None, seq_len], name='input_mask')
segment_ids = tf.placeholder(tf.int64, shape=[None, seq_len], name='segment_ids')
labels = tf.placeholder(tf.int64, shape=[None, seq_len], name='labels')
keep_prob = tf.placeholder(tf.float32, name='keep_prob') # , name='is_training'
bert_config_ = load_bert_config(config["bert_config"])
(total_loss, acc, logits, probabilities) = create_model(bert_config_, is_training, input_ids,
input_mask, segment_ids, labels, keep_prob,
num_labels, use_one_hot_embeddings)
train_op = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps, False)
print("print start train the bert model(multi label)...")
batch_size = config["train_batch_size"]
input_ids2, input_mask2, segment_ids2, labels2 = get_input_data(config["in_1"], seq_len, batch_size)
dev_batch_size = config["dev_batch_size"]
init_global = tf.global_variables_initializer()
saver = tf.train.Saver(tf.global_variables(), max_to_keep=3) # 保存最后top3模型
with tf.Session() as sess:
sess.run(init_global)
tvars = tf.trainable_variables()
initialized_variable_names = {}
print("start load the pretrain model")
scaffold_fn = None
if init_checkpoint:
tvars = tf.trainable_variables()
print("trainable_variables", len(tvars))
(assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars,
init_checkpoint)
print("initialized_variable_names:", len(initialized_variable_names))
saver_ = tf.train.Saver([v for v in tvars if v.name in initialized_variable_names])
saver_.restore(sess, init_checkpoint)
tvars = tf.global_variables()
not_initialized_vars = [v for v in tvars if v.name not in initialized_variable_names]
tf.logging.info('--all size %s; not initialized size %s' % (len(tvars), len(not_initialized_vars)))
if len(not_initialized_vars):
sess.run(tf.variables_initializer(not_initialized_vars))
for v in not_initialized_vars:
tf.logging.info('--not initialized: %s, shape = %s' % (v.name, v.shape))
else:
sess.run(tf.global_variables_initializer())
# if init_checkpoint:
# saver.restore(sess, init_checkpoint)
# print("checkpoint restored from %s" % init_checkpoint)
print("********* bert_multi_sequence_train start *********")
# tf.summary.FileWriter("output/",sess.graph)
def train_step(ids, mask, segment, y, step):
feed = {input_ids: ids,
input_mask: mask,
segment_ids: segment,
labels: y,
keep_prob: 0.9}
_, out_loss, acc_, p_ = sess.run([train_op, total_loss, acc, probabilities], feed_dict=feed)
print("step :{},loss :{}, acc :{}".format(step, out_loss, acc_))
return out_loss, p_, y
def dev_step(ids, mask, segment, y):
feed = {input_ids: ids,
input_mask: mask,
segment_ids: segment,
labels: y,
keep_prob: 1.0
}
out_loss, acc_, p_ = sess.run([total_loss, acc, probabilities], feed_dict=feed)
print("loss :{}, acc :{}".format(out_loss, acc_))
return out_loss, p_, y
min_total_loss_dev = 999999
for i in range(num_train_steps):
# batch 数据
i += 1
ids_train, mask_train, segment_train, y_train = sess.run([input_ids2, input_mask2, segment_ids2, labels2])
train_step(ids_train, mask_train, segment_train, y_train, i)
if i % eval_per_step == 0:
total_loss_dev = 0
dev_input_ids2, dev_input_mask2, dev_segment_ids2, dev_labels2 = get_input_data(config["in_2"], seq_len,
dev_batch_size)
total_pre_dev = []
total_true_dev = []
for j in range(num_dev_steps): # 一个 epoch 的 轮数
ids_dev, mask_dev, segment_dev, y_dev = sess.run(
[dev_input_ids2, dev_input_mask2, dev_segment_ids2, dev_labels2])
out_loss, pre, y = dev_step(ids_dev, mask_dev, segment_dev, y_dev)
total_loss_dev += out_loss
total_pre_dev.extend(pre)
total_true_dev.extend(y_dev)
#
print("dev result report:")
# print(classification_report(total_true_dev, total_pre_dev))
if total_loss_dev < min_total_loss_dev:
print("save model:\t%f\t>%f" % (min_total_loss_dev, total_loss_dev))
min_total_loss_dev = total_loss_dev
saver.save(sess, config["out"] + 'bert.ckpt', global_step=i)
sess.close()
# remove dropout
print("remove dropout in predict")
tf.reset_default_graph()
is_training = False
input_ids = tf.placeholder(tf.int64, shape=[None, seq_len], name='input_ids')
input_mask = tf.placeholder(tf.int64, shape=[None, seq_len], name='input_mask')
segment_ids = tf.placeholder(tf.int64, shape=[None, seq_len], name='segment_ids')
labels = tf.placeholder(tf.int64, shape=[None, seq_len], name='labels')
keep_prob = tf.placeholder(tf.float32, name='keep_prob') # , name='is_training'
bert_config_ = load_bert_config(config["bert_config"])
(total_loss, _, logits, probabilities) = create_model(bert_config_, is_training, input_ids,
input_mask, segment_ids, labels, keep_prob,
num_labels, use_one_hot_embeddings)
init_global = tf.global_variables_initializer()
saver = tf.train.Saver(tf.global_variables(), max_to_keep=1) # 保存最后top3模型
try:
checkpoint = tf.train.get_checkpoint_state(config["out"])
input_checkpoint = checkpoint.model_checkpoint_path
print("[INFO] input_checkpoint:", input_checkpoint)
except Exception as e:
input_checkpoint = config["out"]
print("[INFO] Model folder", config["out"], repr(e))
with tf.Session() as sess:
sess.run(init_global)
saver.restore(sess, input_checkpoint)
saver.save(sess, config["out_1"] + 'bert.ckpt')
sess.close()
if __name__ == "__main__":
print("********* component_bert_sequence_label_train start *********")
main()