def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True.")
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
tf.gfile.MakeDirs(FLAGS.output_dir)
processor = DataProcessor()
label_list = processor.get_labels(FLAGS.data_dir)
tokenizer = tokenization.FullTokenizer(
vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
tpu_cluster_resolver = None
run_config = tf.estimator.RunConfig(
model_dir=FLAGS.output_dir,
save_checkpoints_steps=FLAGS.save_checkpoints_steps)
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir)
num_train_steps = int(
len(train_examples) / FLAGS.batch_size * FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
model_fn = model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_one_hot_embeddings=False)
estimator = tf.estimator.Estimator(model_fn=model_fn, config=run_config,
params={'batch_size': FLAGS.batch_size})
if FLAGS.do_train:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
file_based_convert_examples_to_features(
train_examples, label_list, FLAGS.max_seq_length, tokenizer, train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", FLAGS.batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
file_based_convert_examples_to_features(
eval_examples, label_list, FLAGS.max_seq_length, tokenizer, eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d", len(eval_examples))
tf.logging.info(" Batch size = %d", FLAGS.batch_size)
# This tells the estimator to run through the entire set.
eval_steps = None
eval_drop_remainder = False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder)
result = estimator.evaluate(input_fn=eval_input_fn, steps=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])))
if FLAGS.do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
file_based_convert_examples_to_features(predict_examples, label_list,
FLAGS.max_seq_length, tokenizer,
predict_file)
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d", len(predict_examples))
tf.logging.info(" Batch size = %d", FLAGS.batch_size)
predict_drop_remainder = False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir, "test_results.tsv")
example_index = 0
with tf.gfile.GFile(output_predict_file, "w") as writer:
tf.logging.info("***** Predict results *****")
for prediction in result:
example = predict_examples[example_index]
example_index += 1
output_line = example.text_a + "\t"
output_line += "\t".join(
str(class_probability) for class_probability in prediction) + "\n"
writer.write(output_line)
if FLAGS.do_train:
# Save model
feature_columns = [
tf.feature_column.numeric_column(key="input_ids", shape=(FLAGS.max_seq_length,), dtype=tf.int64),
tf.feature_column.numeric_column(key="input_mask", shape=(FLAGS.max_seq_length,), dtype=tf.int64),
tf.feature_column.numeric_column(key="segment_ids", shape=(FLAGS.max_seq_length,), dtype=tf.int64),
tf.feature_column.numeric_column(key="label_ids", shape=(1,), dtype=tf.int64)]
feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
inputs = {'input_ids': tf.placeholder(shape=(None, FLAGS.max_seq_length), dtype=tf.int64, name='serving_input_ids'),
'input_mask': tf.placeholder(shape=(None, FLAGS.max_seq_length), dtype=tf.int64, name='serving_input_mask'),
'segment_ids': tf.placeholder(shape=(None, FLAGS.max_seq_length), dtype=tf.int64, name='serving_segment_ids'),
'label_ids': tf.placeholder(shape=(None, 1), dtype=tf.int64, name='serving_label_ids')}
serving_input_receiver_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(inputs)
# 构建接收函数，并导出模型。
estimator.export_savedmodel(os.path.join(FLAGS.output_dir, "model"), serving_input_receiver_fn,
strip_default_attrs=True, )
save_label(label_list, os.path.join(FLAGS.output_dir, "model"))
save_vocab(os.path.join(FLAGS.output_dir, "model"))
save_information(os.path.join(FLAGS.output_dir, "model"))
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
processors = {
"text class": TextClassProcessor,
}
tokenization.validate_case_matches_checkpoint(FLAGS.do_lower_case,
FLAGS.init_checkpoint)
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True."
)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
tf.gfile.MakeDirs(FLAGS.output_dir)
task_name = FLAGS.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]() #sim
label_list = processor.get_labels()
tokenizer = tokenization.FullTokenizer(
vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case
) #vocab_file=chinese_L-12_H-768_A-12/vocab.txt "do_lower_case", True,
tpu_cluster_resolver = None
if FLAGS.use_tpu and FLAGS.tpu_name: # "tpu_name", None,
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, #None
master=FLAGS.master, #"master", None,
model_dir=FLAGS.output_dir, #output_dir=tmp_dev/sim_model
save_checkpoints_steps=FLAGS.
save_checkpoints_steps, #"save_checkpoints_steps", 1000,
keep_checkpoint_max=1,
tpu_config=tf.contrib.tpu.TPUConfig(
iterations_per_loop=FLAGS.
iterations_per_loop, #"iterations_per_loop", 1000,
num_shards=FLAGS.num_tpu_cores, #"num_tpu_cores", 8,
per_host_input_for_training=is_per_host))
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir)
num_train_steps = int(
len(train_examples) / FLAGS.train_batch_size *
FLAGS.num_train_epochs
) #train_batch_size=32 num_train_epochs=3.0
num_warmup_steps = int(
num_train_steps *
FLAGS.warmup_proportion) #"warmup_proportion", 0.1,
model_fn = model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list), #2
init_checkpoint=FLAGS.
init_checkpoint, #init_checkpoint=chinese_L-12_H-768_A-12/bert_model.ckpt
learning_rate=FLAGS.learning_rate, #learning_rate=5e-5
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_tpu=FLAGS.use_tpu, #"use_tpu", False,
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(
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,
predict_batch_size=FLAGS.predict_batch_size)
if FLAGS.do_train:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
train_file = train_file.replace('\\', '/')
file_based_convert_examples_to_features(train_examples, label_list,
FLAGS.max_seq_length,
tokenizer, train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
num_actual_eval_examples = len(eval_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on. These do NOT count towards the metric (all tf.metrics
# support a per-instance weight, and these get a weight of 0.0).
while len(eval_examples) % FLAGS.eval_batch_size != 0:
eval_examples.append(PaddingInputExample())
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
eval_file = eval_file.replace('\\', '/')
file_based_convert_examples_to_features(eval_examples, label_list,
FLAGS.max_seq_length,
tokenizer, eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(eval_examples), num_actual_eval_examples,
len(eval_examples) - num_actual_eval_examples)
tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size)
# This tells the estimator to run through the entire set.
eval_steps = None
# However, if running eval on the TPU, you will need to specify the
# number of steps.
if FLAGS.use_tpu:
assert len(eval_examples) % FLAGS.eval_batch_size == 0
eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size)
eval_drop_remainder = True if FLAGS.use_tpu else False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder)
result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt")
output_eval_file = output_eval_file.replace('\\', '/')
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])))
if FLAGS.do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
num_actual_predict_examples = len(predict_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on.
while len(predict_examples) % FLAGS.predict_batch_size != 0:
predict_examples.append(PaddingInputExample())
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
predict_file = predict_file.replace('\\', '/')
file_based_convert_examples_to_features(predict_examples, label_list,
FLAGS.max_seq_length,
tokenizer, predict_file)
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(predict_examples), num_actual_predict_examples,
len(predict_examples) - num_actual_predict_examples)
tf.logging.info(" Batch size = %d", FLAGS.predict_batch_size)
predict_drop_remainder = True if FLAGS.use_tpu else False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir,
"model1997_test_results.tsv")
output_predict_file = output_predict_file.replace('\\', '/')
with tf.gfile.GFile(output_predict_file, "w") as writer:
num_written_lines = 0
tf.logging.info("***** Predict results *****")
for (i, prediction) in enumerate(result):
probabilities = prediction["probabilities"]
if i >= num_actual_predict_examples:
break
output_line = "\t".join(
str(class_probability)
for class_probability in probabilities) + "\n"
writer.write(output_line)
num_written_lines += 1
print(num_written_lines)
assert num_written_lines == num_actual_predict_examples
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
validate_flags_or_throw(bert_config)
tf.gfile.MakeDirs(FLAGS.output_dir)
tokenizer = tokenization.FullTokenizer(
vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
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))
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = read_squad_examples(
input_file=FLAGS.train_file, is_training=True)
num_train_steps = int(
len(train_examples) / FLAGS.train_batch_size * FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
# Pre-shuffle the input to avoid having to make a very large shuffle
# buffer in in the `input_fn`.
rng = random.Random(12345)
rng.shuffle(train_examples)
model_fn = model_fn_builder(
bert_config=bert_config,
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=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(
use_tpu=FLAGS.use_tpu,
model_fn=model_fn,
config=run_config,
train_batch_size=FLAGS.train_batch_size,
predict_batch_size=FLAGS.predict_batch_size)
if FLAGS.do_train:
# We write to a temporary file to avoid storing very large constant tensors
# in memory.
train_writer = FeatureWriter(
filename=os.path.join(FLAGS.output_dir, "train.tf_record"),
is_training=True)
convert_examples_to_features(
examples=train_examples,
tokenizer=tokenizer,
max_seq_length=FLAGS.max_seq_length,
doc_stride=FLAGS.doc_stride,
max_query_length=FLAGS.max_query_length,
is_training=True,
output_fn=train_writer.process_feature)
train_writer.close()
tf.logging.info("***** Running training *****")
tf.logging.info(" Num orig examples = %d", len(train_examples))
tf.logging.info(" Num split examples = %d", train_writer.num_features)
tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
del train_examples
train_input_fn = input_fn_builder(
input_file=train_writer.filename,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_predict:
eval_examples = read_squad_examples(
input_file=FLAGS.predict_file, is_training=False)
eval_writer = FeatureWriter(
filename=os.path.join(FLAGS.output_dir, "eval.tf_record"),
is_training=False)
eval_features = []
def append_feature(feature):
eval_features.append(feature)
eval_writer.process_feature(feature)
convert_examples_to_features(
examples=eval_examples,
tokenizer=tokenizer,
max_seq_length=FLAGS.max_seq_length,
doc_stride=FLAGS.doc_stride,
max_query_length=FLAGS.max_query_length,
is_training=False,
output_fn=append_feature)
eval_writer.close()
tf.logging.info("***** Running predictions *****")
tf.logging.info(" Num orig examples = %d", len(eval_examples))
tf.logging.info(" Num split examples = %d", len(eval_features))
tf.logging.info(" Batch size = %d", FLAGS.predict_batch_size)
all_results = []
predict_input_fn = input_fn_builder(
input_file=eval_writer.filename,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=False)
# If running eval on the TPU, you will need to specify the number of
# steps.
all_results = []
for result in estimator.predict(
predict_input_fn, yield_single_examples=True):
if len(all_results) % 1000 == 0:
tf.logging.info("Processing example: %d" % (len(all_results)))
unique_id = int(result["unique_ids"])
start_logits = [float(x) for x in result["start_logits"].flat]
end_logits = [float(x) for x in result["end_logits"].flat]
all_results.append(
RawResult(
unique_id=unique_id,
start_logits=start_logits,
end_logits=end_logits))
output_prediction_file = os.path.join(FLAGS.output_dir, "predictions.json")
output_nbest_file = os.path.join(FLAGS.output_dir, "nbest_predictions.json")
output_null_log_odds_file = os.path.join(FLAGS.output_dir, "null_odds.json")
write_predictions(eval_examples, eval_features, all_results,
FLAGS.n_best_size, FLAGS.max_answer_length,
FLAGS.do_lower_case, output_prediction_file,
output_nbest_file, output_null_log_odds_file)
def sentEmb(S):
embedding_size = 768
dim_cls = 256
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
vocabulary_size = len(tokenizer.vocab)
tf.reset_default_graph()
input_ids = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='input_ids')
input_mask = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='input_mask')
segment_ids = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='segment_ids')
word_inputs = tf.placeholder(tf.int32, shape=[None])
word_labels = tf.placeholder(tf.int32, shape=[None, 1])
model = modeling.BertModel(config=bert_config,
is_training=FLAGS.do_train,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=False)
sequence_output = model.get_sequence_output()
#first_token_tensor = tf.squeeze(sequence_output[:, 0:1, :], axis=1)
first_token_tensor = tf.reduce_mean(sequence_output, axis=1)
out1 = tf.layers.dense(first_token_tensor,
units=dim_cls,
activation=tf.nn.relu,
name='fine_s2v_cls')
embeddings = model.get_embedding_table()
embed = tf.nn.embedding_lookup(embeddings, word_inputs)
feature = tf.concat([out1, embed], axis=-1)
# a = 0.01
# feature = a*embed+(1-a)*first_token_tensor
nce_weights = tf.Variable(tf.truncated_normal(
[vocabulary_size, dim_cls + embedding_size],
stddev=1.0 / math.sqrt(dim_cls + embedding_size)),
name='fine_s2v_cce_w')
nce_biases = tf.Variable(tf.zeros([vocabulary_size]),
dtype=tf.float32,
name='fine_s2v_cce_b')
nce_weights0 = nce_weights[:, :dim_cls]
nce_weights1 = nce_weights[:, dim_cls:]
# Compute the average NCE loss for the batch.
# tf.nce_loss automatically draws a new sample of the negative labels each
# time we evaluate the loss.
logits = tf.matmul(feature, nce_weights, transpose_b=True) + nce_biases
pro_predict = tf.nn.softmax(logits, axis=-1)
label_predict = tf.argmax(pro_predict, axis=-1)
logits_emb = tf.matmul(embed, nce_weights1, transpose_b=True) + nce_biases
pro_predict_emb = tf.nn.softmax(logits_emb, axis=-1)
label_predict_emb = tf.argmax(pro_predict_emb, axis=-1)
logits_cls = tf.matmul(out1, nce_weights0, transpose_b=True) + nce_biases
pro_predict_cls = tf.nn.softmax(logits_cls, axis=-1)
label_predict_cls = tf.argmax(pro_predict_cls, axis=-1)
output = {'sent2vec': out1}
saver = tf.train.Saver(max_to_keep=None)
sess = tf.Session()
module_file = tf.train.latest_checkpoint(FLAGS.output_dir)
saver.restore(sess, module_file)
# X_input = []
# X_mask = []
# X_segmet = []
# for i in range(len(S)):
# text_a = S[i]
# example = InputExample(guid='guid', text_a=text_a, label='0')
# feature = convert_single_example(10, example, ['0','1'], FLAGS.max_seq_length, tokenizer)
# X_input.append(feature.input_ids)
# X_segmet.append(feature.segment_ids)
# X_mask.append(feature.input_mask)
T = []
for i in range(len(S)):
if i % 100 == 0:
print(i, len(S))
text_a = S[i]
example = InputExample(guid='guid', text_a=text_a, label='0')
feature = convert_single_example(10, example, ['0', '1'],
FLAGS.max_seq_length, tokenizer)
feed_dict = {
input_ids: [feature.input_ids],
segment_ids: [feature.segment_ids],
input_mask: [feature.input_mask]
}
y = {
key: sess.run(output[key], feed_dict=feed_dict)[0]
for key in output
}
T.append([i, S[i], y])
return T
def main(_):
embedding_size = 768
num_sampled = 128
learning_rate = 1e-3
step_printloss = 100
step_savemodel = 1000
dim_cls = 256
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
vocabulary_size = len(tokenizer.vocab)
tf.reset_default_graph()
input_ids = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='input_ids')
input_mask = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='input_mask')
segment_ids = tf.placeholder(tf.int32,
shape=[None, FLAGS.max_seq_length],
name='segment_ids')
word_inputs = tf.placeholder(tf.int32, shape=[None])
word_labels = tf.placeholder(tf.int32, shape=[None, 1])
model = modeling.BertModel(config=bert_config,
is_training=FLAGS.do_train,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=False)
sequence_output = model.get_sequence_output()
#first_token_tensor = tf.squeeze(sequence_output[:, 0:1, :], axis=1)
first_token_tensor = tf.reduce_mean(sequence_output, axis=1)
out1 = tf.layers.dense(first_token_tensor,
units=dim_cls,
activation=tf.nn.relu,
name='fine_s2v_cls')
embeddings = model.get_embedding_table()
embed = tf.nn.embedding_lookup(embeddings, word_inputs)
feature = tf.concat([out1, embed], axis=-1)
#a = 0.01
#feature = a*embed+(1-a)*first_token_tensor
nce_weights = tf.Variable(tf.truncated_normal(
[vocabulary_size, dim_cls + embedding_size],
stddev=1.0 / math.sqrt(dim_cls + embedding_size)),
name='fine_s2v_cce_w')
nce_biases = tf.Variable(tf.zeros([vocabulary_size]),
dtype=tf.float32,
name='fine_s2v_cce_b')
nce_weights0 = nce_weights[:, :dim_cls]
nce_weights1 = nce_weights[:, dim_cls:]
# Compute the average NCE loss for the batch.
# tf.nce_loss automatically draws a new sample of the negative labels each
# time we evaluate the loss.
loss = tf.reduce_mean(
tf.nn.nce_loss(weights=nce_weights,
biases=nce_biases,
inputs=feature,
labels=word_labels,
num_sampled=num_sampled,
num_classes=vocabulary_size))
logits = tf.matmul(feature, nce_weights, transpose_b=True) + nce_biases
pro_predict = tf.nn.softmax(logits, axis=-1)
label_predict = tf.argmax(pro_predict, axis=-1)
logits_emb = tf.matmul(embed, nce_weights1, transpose_b=True) + nce_biases
pro_predict_emb = tf.nn.softmax(logits_emb, axis=-1)
label_predict_emb = tf.argmax(pro_predict_emb, axis=-1)
logits_cls = tf.matmul(out1, nce_weights0, transpose_b=True) + nce_biases
pro_predict_cls = tf.nn.softmax(logits_cls, axis=-1)
label_predict_cls = tf.argmax(pro_predict_cls, axis=-1)
update_var_list = [] # 该list中的变量参与参数更新
tvars = tf.trainable_variables()
for tvar in tvars:
if "bert" not in tvar.name:
update_var_list.append(tvar)
train_op = tf.train.AdamOptimizer(learning_rate).minimize(
loss, var_list=update_var_list)
saver = tf.train.Saver(max_to_keep=10)
global_step = tf.train.get_or_create_global_step()
model_train_op = tf.group(train_op, [tf.assign_add(global_step, 1)])
sess = tf.Session()
sess.run(tf.global_variables_initializer())
if FLAGS.init_checkpoint:
try:
tvars = tf.trainable_variables()
(assignment_map,
initialized_variable_names) = get_assignment_map_from_checkpoint(
tvars, FLAGS.init_checkpoint)
tf.train.init_from_checkpoint(FLAGS.init_checkpoint,
assignment_map)
except:
saver.restore(sess, FLAGS.init_checkpoint)
iter = iterData(FLAGS.data_dir,
tokenizer,
batch_size=FLAGS.train_batch_size,
epochs=int(FLAGS.num_train_epochs))
data = next(iter)
while data != '__STOP__':
epoch, batch_input_ids, batch_input_mask, batch_segment_ids, batch_word_inputs, batch_word_labels = data
batch_word_labels = np.array(batch_word_labels)
batch_word_labels = np.reshape(batch_word_labels,
[len(batch_word_labels), 1])
feed_dict = {
input_ids: batch_input_ids,
input_mask: batch_input_mask,
segment_ids: batch_segment_ids,
word_inputs: batch_word_inputs,
word_labels: batch_word_labels
}
batch_loss, step, _ = sess.run([loss, global_step, model_train_op],
feed_dict=feed_dict)
if step % step_printloss == 0:
print('train total steps %d (epoch %d) with loss=%0.4f' %
(step, epoch, batch_loss))
if step % step_savemodel == 0:
saver.save(sess,
os.path.join(FLAGS.output_dir, 'model.ckpt'),
global_step=global_step)
# import random
# random.shuffle(batch_input_ids)
# feed_dict = {input_ids: batch_input_ids, input_mask: batch_input_mask, segment_ids: batch_segment_ids,
# word_inputs: batch_word_inputs, word_labels: batch_word_labels}
label_predict_, label_predict_emb_, label_predict_cls_ = sess.run(
[label_predict, label_predict_emb, label_predict_cls],
feed_dict=feed_dict)
acc = sum([
label_predict_[j] == batch_word_labels[j][0]
for j in range(len(label_predict_))
]) / (0.01 + len(label_predict_))
acc_emb = sum([
label_predict_emb_[j] == batch_word_labels[j][0]
for j in range(len(label_predict_))
]) / (0.01 + len(label_predict_))
acc_cls = sum([
label_predict_cls_[j] == batch_word_labels[j][0]
for j in range(len(label_predict_))
]) / (0.01 + len(label_predict_))
print('acc:%0.4f,%0.4f,%0.4f' % (acc, acc_emb, acc_cls))
#embed_,first_token_tensor_,feature_ = sess.run([embed,first_token_tensor,feature],feed_dict = feed_dict)
data = next(iter)
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
processors = {
"cola": ColaProcessor,
"mnli": MnliProcessor,
"mrpc": MrpcProcessor,
"xnli": XnliProcessor,
"sentence_pair": SentencePairClassificationProcessor,
"sentiment_analysis": SentimentAnalysisFineGrainProcessor,
}
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True."
)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
tf.gfile.MakeDirs(FLAGS.output_dir)
task_name = FLAGS.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels(FLAGS.data_dir)
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
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))
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir)
num_train_steps = int(
len(train_examples) / FLAGS.train_batch_size *
FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
model_fn = model_fn_builder(bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=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(
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,
predict_batch_size=FLAGS.predict_batch_size)
if FLAGS.do_train:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
file_based_convert_examples_to_features(train_examples, label_list,
FLAGS.max_seq_length,
tokenizer, train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
file_based_convert_examples_to_features(eval_examples, label_list,
FLAGS.max_seq_length,
tokenizer, eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d", len(eval_examples))
tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size)
# This tells the estimator to run through the entire set.
eval_steps = None
# However, if running eval on the TPU, you will need to specify the
# number of steps.
if FLAGS.use_tpu:
# Eval will be slightly WRONG on the TPU because it will truncate
# the last batch.
eval_steps = int(len(eval_examples) / FLAGS.eval_batch_size)
eval_drop_remainder = True if FLAGS.use_tpu else False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder)
result = estimator.evaluate(input_fn=eval_input_fn, steps=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])))
if FLAGS.do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
file_based_convert_examples_to_features(predict_examples, label_list,
FLAGS.max_seq_length,
tokenizer, predict_file)
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d", len(predict_examples))
tf.logging.info(" Batch size = %d", FLAGS.predict_batch_size)
if FLAGS.use_tpu:
# Warning: According to tpu_estimator.py Prediction on TPU is an
# experimental feature and hence not supported here
raise ValueError("Prediction in TPU not supported")
predict_drop_remainder = True if FLAGS.use_tpu else False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir,
"test_results.tsv")
with tf.gfile.GFile(output_predict_file, "w") as writer:
tf.logging.info("***** Predict results *****")
for prediction in result:
output_line = "\t".join(
str(class_probability)
for class_probability in prediction) + "\n"
writer.write(output_line)
def main(_):
import time
start_time = time.time()
if FLAGS.use_perseus and FLAGS.use_horovod:
raise ValeuError(
"Could not set use_perseus and use_horovod at the same time.")
if FLAGS.use_perseus:
import perseus.tensorflow.horovod as hvd
hvd.init()
if FLAGS.use_horovod:
import horovod.tensorflow as hvd
hvd.init()
if FLAGS.use_perseus or FLAGS.use_horovod:
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())
if FLAGS.use_xla:
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
tf.logging.set_verbosity(tf.logging.INFO)
processors = {
"cola": ColaProcessor,
"mnli": MnliProcessor,
"mrpc": MrpcProcessor,
"news": NewsProcessor,
"xnli": XnliProcessor,
}
tokenization.validate_case_matches_checkpoint(FLAGS.do_lower_case,
FLAGS.init_checkpoint)
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True."
)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
if not tf.gfile.Exists(FLAGS.output_dir):
tf.gfile.MakeDirs(FLAGS.output_dir)
task_name = FLAGS.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels()
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
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
save_checkpoints_steps = FLAGS.save_checkpoints_steps
learning_rate = FLAGS.learning_rate
if FLAGS.use_perseus or FLAGS.use_horovod:
model_dir = FLAGS.output_dir if hvd.rank() == 0 else None
save_checkpoints_steps = save_checkpoints_steps // hvd.size()
else:
model_dir = FLAGS.output_dir
run_config = tf.contrib.tpu.RunConfig(
cluster=tpu_cluster_resolver,
master=FLAGS.master,
model_dir=model_dir,
session_config=config,
save_checkpoints_steps=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))
train_examples = None
num_train_steps = None
num_warmup_steps = None
num_workers = 1
worker_index = 0
if FLAGS.use_perseus or FLAGS.use_horovod:
bcast_hook = [hvd.BroadcastGlobalVariablesHook(0)]
num_workers = hvd.size()
worker_index = hvd.rank()
else:
bcast_hook = []
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir)
num_train_steps = int(
len(train_examples) / FLAGS.train_batch_size *
FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
if FLAGS.use_perseus or FLAGS.use_horovod:
num_train_steps = num_train_steps // hvd.size()
num_warmup_steps = num_warmup_steps // hvd.size()
num_workers = hvd.size()
worker_index = hvd.rank()
model_fn = model_fn_builder(bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_tpu=FLAGS.use_tpu,
use_perseus=FLAGS.use_perseus,
use_horovod=FLAGS.use_horovod,
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(
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,
predict_batch_size=FLAGS.predict_batch_size)
if FLAGS.do_train:
if FLAGS.use_perseus or FLAGS.use_horovod:
train_file = os.path.join(
FLAGS.output_dir, "train-" + str(hvd.rank()) + ".tf_record")
else:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
file_based_convert_examples_to_features(train_examples, label_list,
FLAGS.max_seq_length,
tokenizer, train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True,
num_workers=num_workers,
worker_index=worker_index)
estimator.train(input_fn=train_input_fn,
max_steps=num_train_steps,
hooks=bcast_hook)
do_eval = FLAGS.do_eval
if FLAGS.use_perseus or FLAGS.use_horovod:
if hvd.rank() == 0:
do_eval = FLAGS.do_eval
else:
do_eval = False
if do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
num_actual_eval_examples = len(eval_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on. These do NOT count towards the metric (all tf.metrics
# support a per-instance weight, and these get a weight of 0.0).
while len(eval_examples) % FLAGS.eval_batch_size != 0:
eval_examples.append(PaddingInputExample())
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
file_based_convert_examples_to_features(eval_examples, label_list,
FLAGS.max_seq_length,
tokenizer, eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(eval_examples), num_actual_eval_examples,
len(eval_examples) - num_actual_eval_examples)
tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size)
# This tells the estimator to run through the entire set.
eval_steps = None
# However, if running eval on the TPU, you will need to specify the
# number of steps.
if FLAGS.use_tpu:
assert len(eval_examples) % FLAGS.eval_batch_size == 0
eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size)
eval_drop_remainder = True if FLAGS.use_tpu else False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder,
num_workers=1,
worker_index=0)
result = estimator.evaluate(input_fn=eval_input_fn, steps=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])))
do_predict = FLAGS.do_predict
if FLAGS.use_perseus or FLAGS.use_horovod:
if hvd.rank() == 0:
do_predict = FLAGS.do_predict
else:
do_predict = False
if do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
num_actual_predict_examples = len(predict_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on.
while len(predict_examples) % FLAGS.predict_batch_size != 0:
predict_examples.append(PaddingInputExample())
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
file_based_convert_examples_to_features(predict_examples, label_list,
FLAGS.max_seq_length,
tokenizer, predict_file)
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(predict_examples), num_actual_predict_examples,
len(predict_examples) - num_actual_predict_examples)
tf.logging.info(" Batch size = %d", FLAGS.predict_batch_size)
predict_drop_remainder = True if FLAGS.use_tpu else False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder,
num_workers=1,
worker_index=0)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir,
"test_results.tsv")
with tf.gfile.GFile(output_predict_file, "w") as writer:
num_written_lines = 0
tf.logging.info("***** Predict results *****")
for (i, prediction) in enumerate(result):
probabilities = prediction["probabilities"]
if i >= num_actual_predict_examples:
break
output_line = "\t".join(
str(class_probability)
for class_probability in probabilities) + "\n"
writer.write(output_line)
num_written_lines += 1
assert num_written_lines == num_actual_predict_examples
if FLAGS.do_export:
estimator._export_to_tpu = False
estimator.export_savedmodel(FLAGS.export_dir, serving_input_fn)
eclapse_time = time.time() - start_time
print(f'overall time is {eclapse_time} s')
def main(train_cfg='config/pretrain.json',
model_cfg='config/bert_base.json',
data_file='../tbc/books_large_all.txt',
model_file=None,
data_parallel=True,
vocab='../uncased_L-12_H-768_A-12/vocab.txt',
save_dir='../exp/bert/pretrain',
log_dir='../exp/bert/pretrain/runs',
max_len=512,
max_pred=20,
mask_prob=0.15):
cfg = Config(**json.load(open(config, "r")))
cfg_optim = train.Config(**json.load(open(cfg.cfg_optim, "r")))
cfg_model = models.Config(**json.load(open(cfg.cfg_model, "r")))
set_seeds(cfg.seed)
tokenizer = tokenization.FullTokenizer(vocab_file=vocab,
do_lower_case=True)
tokenize = lambda x: tokenizer.tokenize(tokenizer.convert_to_unicode(x))
pipeline = [
Preprocess4Pretrain(max_pred, mask_prob, list(tokenizer.vocab.keys()),
tokenizer.convert_tokens_to_ids, max_len)
]
data_iter = SentPairDataLoader(data_file,
cfg.batch_size,
tokenize,
max_len,
pipeline=pipeline)
model = BertModel4Pretrain(model_cfg)
criterion1 = nn.CrossEntropyLoss(reduction='none')
criterion2 = nn.CrossEntropyLoss()
optimizer = optim.optim4GPU(cfg, model)
trainer = train.Trainer(cfg, model, data_iter, optimizer, save_dir,
get_device())
writer = SummaryWriter(log_dir=log_dir) # for tensorboardX
def get_loss(model, batch, global_step): # make sure loss is tensor
input_ids, segment_ids, input_mask, masked_ids, masked_pos, masked_weights, is_next = batch
logits_lm, logits_clsf = model(input_ids, segment_ids, input_mask,
masked_pos)
loss_lm = criterion1(logits_lm.transpose(1, 2),
masked_ids) # for masked LM
loss_lm = (loss_lm * masked_weights.float()).mean()
loss_clsf = criterion2(logits_clsf,
is_next) # for sentence classification
writer.add_scalars(
'data/scalar_group', {
'loss_lm': loss_lm.item(),
'loss_clsf': loss_clsf.item(),
'loss_total': (loss_lm + loss_clsf).item(),
'lr': optimizer.get_lr()[0],
}, global_step)
return loss_lm + loss_clsf
trainer.train(get_loss, model_file, None, data_parallel)
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
processors = {
"sentence_pair": SentencePairClassificationProcessor,
"lcqmc_pair": LCQMCPairClassificationProcessor,
"lcqmc": LCQMCPairClassificationProcessor
}
tokenization.validate_case_matches_checkpoint(FLAGS.do_lower_case, FLAGS.init_checkpoint)
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError("At least one of `do_train`, `do_eval` or `do_predict' must be True.")
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
tf.gfile.MakeDirs(FLAGS.output_dir)
task_name = FLAGS.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels()
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
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
# Cloud TPU: Invalid TPU configuration, ensure ClusterResolver is passed to tpu.
print("###tpu_cluster_resolver:", tpu_cluster_resolver)
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))
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir) # TODO
print("###length of total train_examples:", len(train_examples))
num_train_steps = int(len(train_examples) / FLAGS.train_batch_size * FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
model_fn = model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=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(
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,
predict_batch_size=FLAGS.predict_batch_size)
if FLAGS.do_train:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
train_file_exists = os.path.exists(train_file)
print("###train_file_exists:", train_file_exists, " ;train_file:", train_file)
if not train_file_exists: # if tf_record file not exist, convert from raw text file. # TODO
file_based_convert_examples_to_features(
train_examples, label_list, FLAGS.max_seq_length, tokenizer, train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
num_actual_eval_examples = len(eval_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on. These do NOT count towards the metric (all tf.metrics
# support a per-instance weight, and these get a weight of 0.0).
while len(eval_examples) % FLAGS.eval_batch_size != 0:
eval_examples.append(PaddingInputExample())
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
file_based_convert_examples_to_features(eval_examples, label_list, FLAGS.max_seq_length, tokenizer, eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(eval_examples), num_actual_eval_examples,
len(eval_examples) - num_actual_eval_examples)
tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size)
# This tells the estimator to run through the entire set.
eval_steps = None
# However, if running eval on the TPU, you will need to specify the
# number of steps.
if FLAGS.use_tpu:
assert len(eval_examples) % FLAGS.eval_batch_size == 0
eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size)
eval_drop_remainder = True if FLAGS.use_tpu else False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder)
#######################################################################################################################
# evaluate all checkpoints; you can use the checkpoint with the best dev accuarcy
steps_and_files = []
filenames = tf.gfile.ListDirectory(FLAGS.output_dir)
for filename in filenames:
if filename.endswith(".index"):
ckpt_name = filename[:-6]
cur_filename = os.path.join(FLAGS.output_dir, ckpt_name)
global_step = int(cur_filename.split("-")[-1])
tf.logging.info("Add {} to eval list.".format(cur_filename))
steps_and_files.append([global_step, cur_filename])
steps_and_files = sorted(steps_and_files, key=lambda x: x[0])
output_eval_file = os.path.join(FLAGS.data_dir, "eval_results_albert_zh.txt")
print("output_eval_file:", output_eval_file)
tf.logging.info("output_eval_file:" + output_eval_file)
with tf.gfile.GFile(output_eval_file, "w") as writer:
for global_step, filename in sorted(steps_and_files, key=lambda x: x[0]):
result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps, checkpoint_path=filename)
tf.logging.info("***** Eval results %s *****" % (filename))
writer.write("***** Eval results %s *****\n" % (filename))
for key in sorted(result.keys()):
tf.logging.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
#######################################################################################################################
# result = estimator.evaluate(input_fn=eval_input_fn, steps=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])))
if FLAGS.do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
num_actual_predict_examples = len(predict_examples)
if FLAGS.use_tpu:
# TPU requires a fixed batch size for all batches, therefore the number
# of examples must be a multiple of the batch size, or else examples
# will get dropped. So we pad with fake examples which are ignored
# later on.
while len(predict_examples) % FLAGS.predict_batch_size != 0:
predict_examples.append(PaddingInputExample())
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
file_based_convert_examples_to_features(
predict_examples, label_list, FLAGS.max_seq_length, tokenizer,predict_file)
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d (%d actual, %d padding)",
len(predict_examples), num_actual_predict_examples,
len(predict_examples) - num_actual_predict_examples)
tf.logging.info(" Batch size = %d", FLAGS.predict_batch_size)
predict_drop_remainder = True if FLAGS.use_tpu else False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir, "test_results.tsv")
with tf.gfile.GFile(output_predict_file, "w") as writer:
num_written_lines = 0
tf.logging.info("***** Predict results *****")
for (i, prediction) in enumerate(result):
probabilities = prediction["probabilities"]
if i >= num_actual_predict_examples:
break
output_line = "\t".join(
str(class_probability)
for class_probability in probabilities) + "\n"
writer.write(output_line)
num_written_lines += 1
assert num_written_lines == num_actual_predict_examples
def main():
tf.logging.set_verbosity(tf.logging.INFO)
processors = {"ner": NerProcessor_from_file, "kashgari": kashgariProcessor}
if not do_train and not do_eval and not do_predict:
raise ValueError(
"At least one of `do_train` or `do_eval` or `do_predict` must be True."
)
bert_config = modeling.BertConfig.from_json_file(bert_config_file)
if max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(max_seq_length, bert_config.max_position_embeddings))
global task_name, output_dir
task_name = task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
global num_labels
label_list = processor.get_labels()
num_labels = len(label_list) + 1
tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file,
do_lower_case=do_lower_case)
tpu_cluster_resolver = None
if use_tpu and False:
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
tpu_name, zone=tpu_zone, project=gcp_project)
is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
run_config = tf.contrib.tpu.RunConfig(
cluster=tpu_cluster_resolver,
master=master,
model_dir=output_dir,
save_checkpoints_steps=save_checkpoints_steps,
tpu_config=tf.contrib.tpu.TPUConfig(
iterations_per_loop=iterations_per_loop,
num_shards=num_tpu_cores,
per_host_input_for_training=is_per_host))
train_examples = None
num_train_steps = None
num_warmup_steps = None
global traindata_dir
if do_train:
train_examples = processor.get_train_examples(traindata_dir)
num_train_steps = int(
len(train_examples) / train_batch_size * num_train_epochs)
num_warmup_steps = int(num_train_steps * warmup_proportion)
model_fn = model_fn_builder(bert_config=bert_config,
num_labels=num_labels + 1,
init_checkpoint=init_checkpoint,
learning_rate=learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_tpu=use_tpu,
use_one_hot_embeddings=use_tpu)
estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=use_tpu,
model_fn=model_fn,
config=run_config,
train_batch_size=train_batch_size,
eval_batch_size=eval_batch_size,
predict_batch_size=predict_batch_size)
if do_train:
train_file = os.path.join(output_dir, "train.tf_record")
filed_based_convert_examples_to_features(train_examples, label_list,
max_seq_length, tokenizer,
train_file)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", train_batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(input_file=train_file,
seq_length=max_seq_length,
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if do_eval:
eval_examples = processor.get_dev_examples(traindata_dir)
eval_file = os.path.join(output_dir, "eval.tf_record")
filed_based_convert_examples_to_features(eval_examples, label_list,
max_seq_length, tokenizer,
eval_file)
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d", len(eval_examples))
tf.logging.info(" Batch size = %d", eval_batch_size)
eval_steps = None
if use_tpu:
eval_steps = int(len(eval_examples) / eval_batch_size)
eval_drop_remainder = True if use_tpu else False
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=max_seq_length,
is_training=False,
drop_remainder=eval_drop_remainder)
result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
output_eval_file = os.path.join(output_dir, "eval_results.txt")
with open(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])))
if do_predict:
token_path = os.path.join(output_dir, "token_test.txt")
with open('./output/label2id.pkl', 'rb') as rf:
label2id = pickle.load(rf)
id2label = {value: key for key, value in label2id.items()}
if os.path.exists(token_path):
os.remove(token_path)
predict_examples = processor.get_test_examples(traindata_dir)
predict_file = os.path.join(output_dir, "predict.tf_record")
filed_based_convert_examples_to_features(predict_examples,
label_list,
max_seq_length,
tokenizer,
predict_file,
mode="test")
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d", len(predict_examples))
tf.logging.info(" Batch size = %d", predict_batch_size)
if use_tpu:
# Warning: According to tpu_estimator.py Prediction on TPU is an
# experimental feature and hence not supported here
raise ValueError("Prediction in TPU not supported")
predict_drop_remainder = True if use_tpu else False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(output_dir, "label_test.csv")
