def model_fn(features, labels, mode, params):
if mode == tf.estimator.ModeKeys.PREDICT: # predict mode
_, _, logits = function_builder.get_race_loss(
FLAGS, features, False)
predictions_ = tf.argmax(logits, axis=-1,
output_type=tf.int32) # 返回最大值的索引
predictions = {
'predictions': predictions_,
}
predict_spec = tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions)
return predict_spec
def model_fn(features, labels, mode, params):
#### Training or Evaluation
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
total_loss, per_example_loss, logits = function_builder.get_race_loss(
FLAGS, features, is_training)
#### Check model parameters
num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
logger.info('#params: {}'.format(num_params))
#### load pretrained models
scaffold_fn = model_utils.init_from_checkpoint(FLAGS)
#### Evaluation mode
if mode == tf.estimator.ModeKeys.EVAL:
assert FLAGS.num_hosts == 1
def metric_fn(per_example_loss, label_ids, logits, is_real_example):
predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
eval_input_dict = {
'labels': label_ids,
'predictions': predictions,
'weights': is_real_example
}
accuracy = tf.metrics.accuracy(**eval_input_dict)
loss = tf.metrics.mean(values=per_example_loss,
weights=is_real_example)
return {
'eval_accuracy': accuracy,
'eval_loss': loss}
is_real_example = tf.cast(features["is_real_example"],
dtype=tf.float32)
#### Constucting evaluation TPUEstimatorSpec with new cache.
label_ids = tf.reshape(features['label_ids'], [-1])
metric_args = [per_example_loss, label_ids, logits, is_real_example]
if FLAGS.use_tpu:
eval_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=(metric_fn, metric_args),
scaffold_fn=scaffold_fn)
else:
eval_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
eval_metric_ops=metric_fn(*metric_args))
return eval_spec
#### Configuring the optimizer
train_op, learning_rate, _ = model_utils.get_train_op(FLAGS, total_loss)
monitor_dict = {}
monitor_dict["lr"] = learning_rate
#### Constucting training TPUEstimatorSpec with new cache.
if FLAGS.use_tpu:
#### Creating host calls
host_call = None
train_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, loss=total_loss, train_op=train_op,
host_call=host_call,
scaffold_fn=scaffold_fn)
else:
train_spec = tf.estimator.EstimatorSpec(
mode=mode, loss=total_loss, train_op=train_op)
return train_spec
