def model_fn(features, labels, mode, params):
"""doc."""
#### Training or Evaluation
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
assert is_training
#### Retrieve `mems` from `params["cache"]`
mems = {}
idx = 0
if FLAGS.mem_len > 0:
mems["mems"] = params["cache"]
#### Get loss from inputs
total_loss, new_mems, monitor_dict = function_builder.get_loss(
FLAGS, features, labels, mems, is_training)
#### Turn `new_mems` into `new_cache`
new_cache = []
if FLAGS.mem_len > 0:
new_cache += new_mems["mems"]
#### Check model parameters
num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
tf.logging.info("#params: {}".format(num_params))
#### Configuring the optimizer
train_op, learning_rate, gnorm = model_utils.get_train_op(
FLAGS, total_loss)
monitor_dict["lr"] = learning_rate
monitor_dict["gnorm"] = gnorm
#### Customized initial checkpoint
scaffold_fn = model_utils.init_from_checkpoint(FLAGS, global_vars=True)
#### Creating host calls
host_call = function_builder.construct_scalar_host_call(
monitor_dict=monitor_dict,
model_dir=FLAGS.model_dir,
prefix="train/",
reduce_fn=tf.reduce_mean)
#### Constucting training TPUEstimatorSpec with new cache.
train_spec = tf.contrib.tpu.TPUEstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op,
host_call=host_call,
scaffold_fn=scaffold_fn)
train_spec.cache = new_cache
return train_spec
def model_fn(features, labels, mode, params):
#### Training or Evaluation
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
# Get loss from inputs
if FLAGS.is_regression:
(total_loss, per_example_loss,
logits) = function_builder.get_regression_loss(
FLAGS, features, is_training)
else:
(total_loss, per_example_loss, logits,
p) = function_builder.get_classification_loss(
FLAGS, features, n_class, is_training)
# Check model parameters
num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
tf.logging.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)
"""
change
"""
conf_mat = get_metrics_ops(label_ids, predictions, 3,
is_real_example)
loss = tf.metrics.mean(values=per_example_loss,
weights=is_real_example)
return {
'eval_accuracy': accuracy,
'eval_loss': loss,
"conf_mat": conf_mat
}
def regression_metric_fn(per_example_loss, label_ids, logits,
is_real_example):
loss = tf.metrics.mean(values=per_example_loss,
weights=is_real_example)
pearsonr = tf.contrib.metrics.streaming_pearson_correlation(
logits, label_ids, weights=is_real_example)
return {'eval_loss': loss, 'eval_pearsonr': pearsonr}
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])
if FLAGS.is_regression:
metric_fn = regression_metric_fn
else:
metric_fn = metric_fn
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
elif mode == tf.estimator.ModeKeys.PREDICT:
label_ids = tf.reshape(features["label_ids"], [-1])
predictions = {
"logits": logits,
"labels": label_ids,
"is_real": features["is_real_example"],
"p": p,
}
if FLAGS.use_tpu:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions=predictions,
scaffold_fn=scaffold_fn)
else:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions=predictions)
return output_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
if not FLAGS.is_regression:
label_ids = tf.reshape(features['label_ids'], [-1])
predictions = tf.argmax(logits,
axis=-1,
output_type=label_ids.dtype)
is_correct = tf.equal(predictions, label_ids)
accuracy = tf.reduce_mean(tf.cast(is_correct, tf.float32))
monitor_dict["accuracy"] = accuracy
host_call = function_builder.construct_scalar_host_call(
monitor_dict=monitor_dict,
model_dir=FLAGS.model_dir,
prefix="train/",
reduce_fn=tf.reduce_mean)
else:
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
def model_fn(features, labels, mode, params):
#### Training or Evaluation
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
#### Get loss from inputs
outputs = function_builder.get_qa_outputs(FLAGS, features, is_training)
#### Check model parameters
num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
tf.logging.info('#params: {}'.format(num_params))
scaffold_fn = None
#### Evaluation mode
if mode == tf.estimator.ModeKeys.PREDICT:
if FLAGS.init_checkpoint:
tf.logging.info(
"init_checkpoint not being used in predict mode.")
predictions = {
"unique_ids": features["unique_ids"],
"start_top_index": outputs["start_top_index"],
"start_top_log_probs": outputs["start_top_log_probs"],
"end_top_index": outputs["end_top_index"],
"end_top_log_probs": outputs["end_top_log_probs"],
"cls_logits": outputs["cls_logits"]
}
if FLAGS.use_tpu:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions=predictions,
scaffold_fn=scaffold_fn)
else:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions=predictions)
return output_spec
### Compute loss
seq_length = tf.shape(features["input_ids"])[1]
def compute_loss(log_probs, positions):
one_hot_positions = tf.one_hot(positions,
depth=seq_length,
dtype=tf.float32)
loss = -tf.reduce_sum(one_hot_positions * log_probs, axis=-1)
loss = tf.reduce_mean(loss)
return loss
start_loss = compute_loss(outputs["start_log_probs"],
features["start_positions"])
end_loss = compute_loss(outputs["end_log_probs"],
features["end_positions"])
total_loss = (start_loss + end_loss) * 0.5
cls_logits = outputs["cls_logits"]
is_impossible = tf.reshape(features["is_impossible"], [-1])
regression_loss = tf.nn.sigmoid_cross_entropy_with_logits(
labels=is_impossible, logits=cls_logits)
regression_loss = tf.reduce_mean(regression_loss)
# note(zhiliny): by default multiply the loss by 0.5 so that the scale is
# comparable to start_loss and end_loss
total_loss += regression_loss * 0.5
#### Configuring the optimizer
train_op, learning_rate, _ = model_utils.get_train_op(
FLAGS, total_loss)
monitor_dict = {}
monitor_dict["lr"] = learning_rate
#### load pretrained models
scaffold_fn = model_utils.init_from_checkpoint(FLAGS)
#### Constucting training TPUEstimatorSpec with new cache.
if FLAGS.use_tpu:
host_call = function_builder.construct_scalar_host_call(
monitor_dict=monitor_dict,
model_dir=FLAGS.model_dir,
prefix="train/",
reduce_fn=tf.reduce_mean)
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