def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
# pylint: disable=no-member
"""The `model_fn` for Estimator."""
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"]
hard_label_ids = features["hard_label_ids"]
soft_label_ids = features["soft_label_ids"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss,
probabilities) = create_model(bert_config, is_training, input_ids,
input_mask, segment_ids, hard_label_ids,
soft_label_ids, use_one_hot_embeddings)
tvars = tf.trainable_variables()
initialized_variable_names = {}
if init_checkpoint:
(assignment_map, initialized_variable_names
) = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
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)
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps,
False, None, FLAGS.use_fp16)
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op)
elif mode == tf.estimator.ModeKeys.EVAL:
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss)
else:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions={"probabilities": probabilities})
return output_spec
def model_fn(features, labels, mode, params):
# pylint: disable=unused-argument
input_ids = features['input_ids']
label_ids = features['label_ids']
if mode == tf.estimator.ModeKeys.TRAIN:
self.hparams.dropout_rate = FLAGS.dropout_rate
else:
self.hparams.dropout_rate = 0.
total_loss, probabilities = model.classifier_model(
hparams=self.hparams,
input_tokens=input_ids,
labels=label_ids,
past=None,
reuse=tf.AUTO_REUSE)
tvars = tf.trainable_variables()
init_checkpoint = os.path.join('data', 'models', FLAGS.model_name)
(assignment_map, initialized_variable_names
) = model.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
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)
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps, False, None, FLAGS.use_fp16)
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op)
elif mode == tf.estimator.ModeKeys.EVAL:
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss)
else:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions={"probabilities": probabilities})
return output_spec
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"]
label_ids = features["label_ids"]
is_real_example = None
if "is_real_example" in features:
is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
else:
is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss, per_example_loss, logits, probabilities) = create_model(
bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
num_labels, use_one_hot_embeddings, fp16, weight_list)
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)
is_multi_gpu = use_gpu and int(num_gpu_cores) >= 2
if mode == tf.estimator.ModeKeys.TRAIN:
if is_multi_gpu:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps, fp16=fp16)
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold=scaffold_fn)
else:
train_op = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu, fp16=fp16)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, label_ids, logits, is_real_example):
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
accuracy = tf.metrics.accuracy(
labels=label_ids, predictions=predictions, weights=is_real_example)
# add more metrics
pr, pr_op = tf.metrics.precision(
labels=label_ids, predictions=predictions, weights=is_real_example)
re, re_op = tf.metrics.recall(
labels=label_ids, predictions=predictions, weights=is_real_example)
f1 = (2 * pr * re) / (pr + re) # f1-score for binary classification
loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
return {
"eval_accuracy": accuracy,
"eval_precision": (pr, pr_op),
"eval_recall": (re, re_op),
"eval_f1score": (f1, tf.identity(f1)),
"eval_loss": loss
}
eval_metrics = (metric_fn,
[per_example_loss, label_ids, logits, is_real_example])
# eval on single-gpu only
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
else:
if is_multi_gpu:
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={"probabilities": probabilities})
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions={"probabilities": probabilities},
scaffold_fn=scaffold_fn)
return output_spec
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"]
label_ids = features["label_ids"]
is_real_example = None
if "is_real_example" in features:
is_real_example = tf.cast(features["is_real_example"],
dtype=tf.float32)
else:
is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss, per_example_loss, logits, probabilities) = create_model(
bert_config, is_training, input_ids, input_mask, segment_ids,
label_ids, num_labels, use_one_hot_embeddings, weight_list)
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
is_multi_gpu = use_gpu and int(num_gpu_cores) >= 2
if mode == tf.estimator.ModeKeys.TRAIN:
if is_multi_gpu:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps)
tensors_to_log = {'train loss': total_loss}
logging_hook = tf.train.LoggingTensorHook(
tensors=tensors_to_log, every_n_iter=1)
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
training_hooks=[logging_hook],
scaffold=scaffold_fn)
else:
# output loss during training https://github.com/google-research/bert/issues/70
#train_op = optimization.create_optimizer(
# total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)
train_op, new_global_step = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps, use_tpu)
tensors_to_log = {
'train loss': total_loss,
'global step': new_global_step
}
logging_hook = tf.train.LoggingTensorHook(
tensors=tensors_to_log, every_n_iter=1)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
training_hooks=[logging_hook],
scaffold_fn=scaffold_fn)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, label_ids, logits,
is_real_example):
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
accuracy = tf.metrics.accuracy(labels=label_ids,
predictions=predictions,
weights=is_real_example)
loss = tf.metrics.mean(values=per_example_loss,
weights=is_real_example)
mean_per_class_accuracy = tf.metrics.mean_per_class_accuracy(
labels=label_ids,
predictions=predictions,
weights=is_real_example)
#precision = tf.metrics.F1Score(labels=label_ids, predictions=predictions, weights=is_real_example)
#recall = tf.metrics.recall(labels=label_ids, predictions=predictions, weights=is_real_example)
return {
"eval_accuracy": accuracy,
"eval_loss": loss,
"mean_per_class_accuracy": mean_per_class_accuracy
}
eval_metrics = (metric_fn, [
per_example_loss, label_ids, logits, is_real_example
])
#eval on single-gpu
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
else:
if is_multi_gpu:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions={"probabilities": probabilities})
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions={"probabilities": probabilities},
scaffold_fn=scaffold_fn)
return output_spec
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']
y = features['y']
is_training = mode == tf.estimator.ModeKeys.TRAIN
model = multilanguagebert_transformer.Model(
is_training = is_training,
input_ids = input_ids,
input_mask = input_mask,
token_type_ids = segment_ids,
Y = y,
)
o = model.get_sequence_output()
Y_seq_len = tf.count_nonzero(y, 1, dtype = tf.int32)
masks = tf.sequence_mask(Y_seq_len, tf.shape(y)[1], dtype = tf.float32)
logits = tf.expand_dims(tf.expand_dims(o, axis = 2), axis = 2)
feature = tf.expand_dims(tf.expand_dims(y, axis = 2), axis = 2)
loss_num, loss_denom = model_t2t._loss_single(
logits, 'targets', feature, weights = masks
)
total_loss = loss_num / loss_denom
# total_loss = tf.contrib.seq2seq.sequence_loss(
# logits = o, targets = y, weights = masks
# )
y_t = tf.argmax(o, axis = 2)
y_t = tf.cast(y_t, tf.int32)
prediction = tf.boolean_mask(y_t, masks)
mask_label = tf.boolean_mask(y, masks)
correct_pred = tf.equal(prediction, mask_label)
correct_index = tf.cast(correct_pred, tf.float32)
total_accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
tvars = tf.trainable_variables()
initialized_variable_names = {}
scaffold_fn = None
if init_checkpoint:
(
assignment_map,
initialized_variable_names,
) = get_assignment_map_from_checkpoint(tvars, init_checkpoint)
print(initialized_variable_names)
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:
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps, num_warmup_steps
)
else:
train_op = optimization.create_optimizer(
total_loss,
learning_rate,
num_train_steps,
num_warmup_steps,
use_tpu,
)
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode = mode,
loss = total_loss,
train_op = train_op,
scaffold = scaffold_fn,
)
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode = mode,
loss = total_loss,
train_op = train_op,
scaffold_fn = scaffold_fn,
)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(loss, accuracy):
return {'total_loss': loss, 'total_accuracy': accuracy}
eval_metrics = (metric_fn, [total_loss, total_accuracy])
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode = mode,
loss = total_loss,
eval_metrics = eval_metrics,
scaffold = scaffold_fn,
)
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
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
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']
# Note: We keep this feature name `next_sentence_labels` to be compatible
# with the original data created by lanzhzh@. However, in the ALBERT case
# it does represent sentence_order_labels.
sentence_order_labels = features['next_sentence_labels']
is_training = mode == tf.estimator.ModeKeys.TRAIN
model = modeling.AlbertModel(
config = albert_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,
)
(
masked_lm_loss,
masked_lm_example_loss,
masked_lm_log_probs,
) = get_masked_lm_output(
albert_config,
model.get_sequence_output(),
model.get_embedding_table(),
masked_lm_positions,
masked_lm_ids,
masked_lm_weights,
)
(
sentence_order_loss,
sentence_order_example_loss,
sentence_order_log_probs,
) = get_sentence_order_output(
albert_config, model.get_pooled_output(), sentence_order_labels
)
total_loss = masked_lm_loss + sentence_order_loss
tvars = tf.trainable_variables()
initialized_variable_names = {}
scaffold_fn = None
if init_checkpoint:
tf.logging.info(
'number of hidden group %d to initialize',
albert_config.num_hidden_groups,
)
num_of_initialize_group = 1
if FLAGS.init_from_group0:
num_of_initialize_group = albert_config.num_hidden_groups
if albert_config.net_structure_type > 0:
num_of_initialize_group = albert_config.num_hidden_layers
(
assignment_map,
initialized_variable_names,
) = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint, num_of_initialize_group
)
if use_tpu:
def tpu_scaffold():
for gid in range(num_of_initialize_group):
tf.logging.info('initialize the %dth layer', gid)
tf.logging.info(assignment_map[gid])
tf.train.init_from_checkpoint(
init_checkpoint, assignment_map[gid]
)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
for gid in range(num_of_initialize_group):
tf.logging.info('initialize the %dth layer', gid)
tf.logging.info(assignment_map[gid])
tf.train.init_from_checkpoint(
init_checkpoint, assignment_map[gid]
)
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
)
output_spec = tf.estimator.EstimatorSpec(
mode = mode,
loss = total_loss,
train_op = train_op,
scaffold = scaffold_fn,
)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(*args):
"""Computes the loss and accuracy of the model."""
(
masked_lm_example_loss,
masked_lm_log_probs,
masked_lm_ids,
masked_lm_weights,
sentence_order_example_loss,
sentence_order_log_probs,
sentence_order_labels,
) = args[:7]
masked_lm_log_probs = tf.reshape(
masked_lm_log_probs, [-1, masked_lm_log_probs.shape[-1]]
)
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]
)
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(
values = masked_lm_example_loss, weights = masked_lm_weights
)
metrics = {
'masked_lm_accuracy': masked_lm_accuracy,
'masked_lm_loss': masked_lm_mean_loss,
}
sentence_order_log_probs = tf.reshape(
sentence_order_log_probs,
[-1, sentence_order_log_probs.shape[-1]],
)
sentence_order_predictions = tf.argmax(
sentence_order_log_probs, axis = -1, output_type = tf.int32
)
sentence_order_labels = tf.reshape(sentence_order_labels, [-1])
sentence_order_accuracy = tf.metrics.accuracy(
labels = sentence_order_labels,
predictions = sentence_order_predictions,
)
sentence_order_mean_loss = tf.metrics.mean(
values = sentence_order_example_loss
)
metrics.update(
{
'sentence_order_accuracy': sentence_order_accuracy,
'sentence_order_loss': sentence_order_mean_loss,
}
)
return metrics
metric_values = [
masked_lm_example_loss,
masked_lm_log_probs,
masked_lm_ids,
masked_lm_weights,
sentence_order_example_loss,
sentence_order_log_probs,
sentence_order_labels,
]
eval_metrics = (metric_fn, metric_values)
output_spec = contrib_tpu.TPUEstimatorSpec(
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
def model_fn(features, labels, mode, params, config=None): # pylint: disable=unused-argument
"""
The `model_fn` for TPUEstimator.
模型有训练，验证和测试三种阶段，而且对于不同模式，对数据有不同的处理方式。例如在训练阶段，我们需要将数据喂给模型，
模型基于输入数据给出预测值，然后我们在通过预测值和真实值计算出loss，最后用loss更新网络参数，
而在评估阶段，我们则不需要反向传播更新网络参数，换句话说，mdoel_fn需要对三种模式设置三套代码。
Args:
features: dict of Tensor, This is batch_features from input_fn,`Tensor` or dict of `Tensor` (depends on data passed to `fit`
labels: This is batch_labels from input_fn. features, labels是从输入函数input_fn中返回的特征和标签batch
mode: An instance of tf.estimator.ModeKeys
params: Additional configuration for hyper-parameters. 是一个字典，它可以传入许多参数用来构建网络或者定义训练方式等
Return:
tf.estimator.EstimatorSpec
"""
print("features={}".format(features))
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"]
label_ids = features["label_ids"]
if "is_real_example" in features:
# 类型强制转换为tf.float32
is_real_example = tf.cast(features["is_real_example"],
dtype=tf.float32)
else:
# 创建一个将所有元素都设置为1的张量Tensor.
is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
# 根据mode判断是否为训练模式
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
# 基于特征数据创建模型，并计算loss等
print("create_model:\ninput_ids={}".format(input_ids.shape))
print("label_ids={}".format(label_ids.shape))
(total_loss, per_example_loss, logits,
probabilities) = create_model(bert_config, is_training, input_ids,
None, None, label_ids, num_labels,
use_one_hot_embeddings, None)
print("total_loss={}".format(total_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)
# 训练模式
if mode == tf.estimator.ModeKeys.TRAIN:
if FLAGS.num_gpu_cores > 1:
train_op = custom_optimization.create_optimizer(
total_loss,
learning_rate,
num_train_steps,
num_warmup_steps,
fp16=fp16)
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op,
scaffold=scaffold_fn)
else:
train_op = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps, use_tpu)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
# 评估模式
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, label_ids, logits,
is_real_example):
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
accuracy = tf.metrics.accuracy(labels=label_ids,
predictions=predictions,
weights=is_real_example)
loss = tf.metrics.mean(values=per_example_loss,
weights=is_real_example)
# add more metrics
pr, pr_op = tf.metrics.precision(labels=label_ids,
predictions=predictions,
weights=is_real_example)
re, re_op = tf.metrics.recall(labels=label_ids,
predictions=predictions,
weights=is_real_example)
# if FLAGS.classifier_mode == "multi-class":
# # multi-class
# # pr, pr_op = tf_metrics.precision(label_ids, predictions, num_labels, average="macro")
# # re, re_op = tf_metrics.recall(label_ids, predictions, num_labels, average="macro")
# f1 = tf_metrics.f1(label_ids, predictions, num_labels, average="macro")
# else:
# # binary classifier
# f1 = tf.contrib.metrics.f1_score(label_ids, predictions)
# # f1, f1_op = (2 * pr * re) / (pr + re) # f1-score for binary classification
# 返回结果：dict: {key: value(tuple: (metric_tensor, update_op)) }
return {
"eval_accuracy": accuracy,
"eval_loss": loss,
"eval_precision": (pr, pr_op),
"eval_recall": (re, re_op),
# "eval_f1": f1,
}
eval_metrics = (metric_fn, [
per_example_loss, label_ids, logits, is_real_example
])
if FLAGS.num_gpu_cores > 1:
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
eval_metric_ops=metric_fn(per_example_loss, label_ids,
logits, is_real_example),
scaffold=scaffold_fn,
)
else:
# eval on single-gpu only
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
else:
# tf.estimator.ModeKeys.PREDICT 预测模式
# 基于logits计算最大的概率所在索引的label
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
if FLAGS.num_gpu_cores > 1:
# 多GPUs
output_spec = tf.estimator.EstimatorSpec(mode=mode,
predictions={
"probabilities":
probabilities,
"predictions":
predictions
})
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions={
"probabilities": probabilities,
"predictions": predictions
},
scaffold_fn=scaffold_fn)
return output_spec
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
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,
)
(
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(),
model.get_embedding_table_2(),
masked_lm_positions,
masked_lm_ids,
masked_lm_weights,
)
(
next_sentence_loss,
next_sentence_example_loss,
next_sentence_log_probs,
) = get_next_sentence_output(bert_config, model.get_pooled_output(),
next_sentence_labels)
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:
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps)
# train_op = optimization.create_optimizer(
# total_loss,
# learning_rate,
# num_train_steps,
# num_warmup_steps,
# use_tpu,
# )
else:
train_op = optimization.create_optimizer(
total_loss,
learning_rate,
num_train_steps,
num_warmup_steps,
use_tpu,
)
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold=scaffold_fn,
)
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
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]])
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])
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(
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(
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,
],
)
if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold=scaffold_fn,
)
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
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
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
"""
features： This is batch_features from input_fn
labels： This is batch_labels from input_fn
mode： An instance of tf.estimator.ModeKeys
params： Additional configuration
"""
tf.logging.info("*** Features ***")
for name in sorted(features.keys()):
tf.logging.info(" name = %s, shape = %s" %
(name, features[name].shape))
input_sequence = features["input_sequence"]
input_mask = features["input_mask"]
segment_ids = features["segment_ids"]
edit_sequence = features["edit_sequence"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss, per_example_loss,
logits, probabilities) = gec_create_model(
bert_config, is_training, input_sequence, input_mask, segment_ids,
edit_sequence, use_one_hot_embeddings, mode, copy_weight,
use_bert_more, insert_ids, multitoken_insert_ids,
subtract_replaced_from_replacement)
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:
if FLAGS.use_tpu and FLAGS.tpu_name:
# TPU train
train_op = optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps, use_tpu)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
else:
# GPUs or CPU train
train_op = custom_optimization.create_optimizer(
total_loss, learning_rate, num_train_steps,
num_warmup_steps)
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, edit_sequence, logits):
predictions = tf.argmax(
logits[:, :, 3:], axis=-1, output_type=tf.int32) + 3
mask = tf.equal(edit_sequence, 0)
mask = tf.logical_or(mask, tf.equal(edit_sequence, 1))
mask = tf.logical_or(mask, tf.equal(edit_sequence, 2))
mask = tf.logical_or(mask, tf.equal(edit_sequence, 3))
mask = tf.to_float(tf.logical_not(mask))
accuracy = tf.metrics.accuracy(edit_sequence, predictions,
mask)
loss = tf.metrics.mean(per_example_loss)
result_dict = {}
result_dict["eval_accuracy"] = accuracy
result_dict["eval_loss"] = loss
return {
"eval_accuracy": accuracy,
"eval_loss": loss,
}
eval_metrics = (metric_fn,
[per_example_loss, edit_sequence, logits])
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
else:
# first three edit ids unk, sos, eos are dummy. We do not consider them in predictions
predictions = tf.argmax(
logits[:, :, 3:], axis=-1, output_type=tf.int32) + 3
if FLAGS.use_tpu and FLAGS.tpu_name:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions={
"predictions": predictions,
"logits": logits
},
scaffold_fn=scaffold_fn)
else:
# multiple GPUs
output_spec = tf.estimator.EstimatorSpec(mode=mode,
predictions={
"predictions":
predictions,
"logits": logits
})
return output_spec
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"]
label_ids = features["label_ids"]
seq_len = features["seq_len"]
output_mask = features["output_mask"]
output_mask_float = tf.to_float(output_mask)
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss, per_example_loss,
predictions) = create_model(bert_config, is_training, input_ids,
input_mask, segment_ids, label_ids,
output_mask_float, num_labels,
use_one_hot_embeddings, fp16)
tvars = tf.trainable_variables()
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)
if mode == tf.estimator.ModeKeys.TRAIN:
if use_gpu and int(num_gpu_cores) >= 2:
train_op = custom_optimization.create_optimizer(
total_loss,
learning_rate,
num_train_steps,
num_warmup_steps,
fp16=fp16)
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op,
scaffold=scaffold_fn)
else:
train_op = optimization.create_optimizer(total_loss,
learning_rate,
num_train_steps,
num_warmup_steps,
use_tpu,
fp16=fp16)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, label_ids, predictions,
output_mask):
accuracy = tf.metrics.accuracy(label_ids, predictions,
output_mask)
loss = tf.metrics.mean(per_example_loss)
return {
"eval_accuracy": accuracy,
"eval_loss": loss,
}
eval_metrics = (metric_fn, [
per_example_loss, label_ids, predictions, output_mask
])
if use_gpu and int(num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
eval_metric_ops=eval_metrics[0](*eval_metrics[1]))
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
else:
predictions = {"predictions": predictions, "seq_len": seq_len}
if use_gpu and int(num_gpu_cores) >= 2:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions=predictions)
else:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions=predictions,
scaffold_fn=scaffold_fn)
return output_spec
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"]
label_ids = features["label_ids"]
is_real_example = None
if "is_real_example" in features:
is_real_example = tf.cast(features["is_real_example"],
dtype=tf.float32)
else:
is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
(total_loss, per_example_loss, logits, probabilities,
predict) = create_model(bert_config, is_training, input_ids,
input_mask, segment_ids, label_ids,
num_labels, use_one_hot_embeddings, fp16)
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)
# summary
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
accuracy = tf.metrics.accuracy(labels=label_ids,
predictions=predictions)
tf.summary.scalar('accuracy', accuracy[1])
export_outputs = {
'predict_output':
tf.estimator.export.PredictOutput({
'predict': predict,
'probabilities': probabilities
})
}
if mode == tf.estimator.ModeKeys.TRAIN:
# 单机多卡
if use_gpu and int(num_gpu_cores) >= 2:
train_op = custom_optimization.create_optimizer(
total_loss,
params['learning_rate'],
params['num_train_steps'],
params['num_warmup_steps'],
fp16=fp16)
else:
train_op = optimization.create_optimizer(
total_loss, params['learning_rate'],
params['num_train_steps'], params['num_warmup_steps'],
False)
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
export_outputs=export_outputs,
scaffold=scaffold_fn)
if mode == tf.estimator.ModeKeys.EVAL:
eval_loss = tf.metrics.mean(values=per_example_loss)
eval_metrics = {"eval_accuracy": accuracy, "eval_loss": eval_loss}
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
eval_metric_ops=eval_metrics,
export_outputs=export_outputs,
scaffold=scaffold_fn)
if mode == tf.estimator.ModeKeys.PREDICT:
predicted_classes = tf.argmax(logits, 1)
predictions = {
'class_ids': predicted_classes[:, tf.newaxis],
'probabilities': probabilities,
'logits': logits
}
output_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs=export_outputs,
scaffold=scaffold_fn)
return output_spec
"""