def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
with ops.name_scope('head'):
logits = head_lib._check_logits(logits, self.logits_dimension) # pylint:disable=protected-access
# Predict.
pred_keys = prediction_keys.PredictionKeys
with ops.name_scope(None, 'predictions', (logits, )):
probabilities = math_ops.sigmoid(logits,
name=pred_keys.PROBABILITIES)
predictions = {
pred_keys.LOGITS: logits,
pred_keys.PROBABILITIES: probabilities,
}
if mode == model_fn.ModeKeys.PREDICT:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
'':
export_output.ClassificationOutput(
scores=probabilities)
})
# Eval.
unweighted_loss, processed_labels = self.create_loss(
features=features, mode=mode, logits=logits, labels=labels)
# Averages loss over classes.
per_example_loss = math_ops.reduce_mean(unweighted_loss,
axis=-1,
keep_dims=True)
weights = head_lib._weights(features, self._weight_column) # pylint:disable=protected-access
training_loss = losses.compute_weighted_loss(
per_example_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=processed_labels,
probabilities=probabilities,
weights=weights,
per_example_loss=per_example_loss))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(
head_lib._summary_key(self._head_name,
metric_keys.MetricKeys.LOSS), # pylint:disable=protected-access
training_loss)
summary.scalar(
head_lib._summary_key( # pylint:disable=protected-access
self._head_name, metric_keys.MetricKeys.LOSS_MEAN),
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def testTrainOpMissing(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Missing train_op'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant(1.))
def testLossMissing(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Missing loss'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
train_op=control_flow_ops.no_op())
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
with ops.name_scope(self._name, 'head'):
logits = head_lib._check_logits(logits, self.logits_dimension) # pylint:disable=protected-access
# Predict.
pred_keys = prediction_keys.PredictionKeys
with ops.name_scope(None, 'predictions', (logits, )):
probabilities = math_ops.sigmoid(logits,
name=pred_keys.PROBABILITIES)
predictions = {
pred_keys.LOGITS: logits,
pred_keys.PROBABILITIES: probabilities,
}
if mode == model_fn.ModeKeys.PREDICT:
classifier_output = head_lib._classification_output( # pylint:disable=protected-access
scores=probabilities,
n_classes=self._n_classes,
label_vocabulary=self._label_vocabulary)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
_DEFAULT_SERVING_KEY:
classifier_output,
head_lib._CLASSIFY_SERVING_KEY:
classifier_output, # pylint:disable=protected-access
head_lib._PREDICT_SERVING_KEY: ( # pylint:disable=protected-access
export_output.PredictOutput(predictions))
})
# Eval.
unweighted_loss, processed_labels = self.create_loss(
features=features, mode=mode, logits=logits, labels=labels)
weights = head_lib._weights(features, self._weight_column) # pylint:disable=protected-access
training_loss = losses.compute_weighted_loss(
unweighted_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=processed_labels,
probabilities=probabilities,
weights=weights,
unweighted_loss=unweighted_loss))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(
head_lib._summary_key(self._name, metric_keys.MetricKeys.LOSS), # pylint:disable=protected-access
training_loss)
summary.scalar(
head_lib._summary_key( # pylint:disable=protected-access
self._name, metric_keys.MetricKeys.LOSS_MEAN),
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def testPredictionsTensor(self):
"""Tests that no error is raised when predictions is Tensor (not dict)."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT,
predictions=constant_op.constant(1.))
def testRequiredArgumentsSet(self):
"""Tests that no errors are raised when all required arguments are set."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant(1.),
train_op=control_flow_ops.no_op())
def create_estimator_spec(
self, features, mode, logits, labels=None, train_op_fn=None,
regularization_losses=None):
"""Returns an `EstimatorSpec`.
Args:
features: Input `dict` of `Tensor` or `SparseTensor` objects.
mode: Estimator's `ModeKeys`.
logits: logits `Tensor` with shape `[D0, D1, ... DN, n_classes]`.
For many applications, the shape is `[batch_size, n_classes]`.
labels: Labels with shape matching `logits`. Can be multi-hot `Tensor`
with shape `[D0, D1, ... DN, n_classes]` or `SparseTensor` with
`dense_shape` `[D0, D1, ... DN, ?]`. `labels` is required argument when
`mode` equals `TRAIN` or `EVAL`.
train_op_fn: Function that takes a scalar loss `Tensor` and returns
`train_op`. Required in TRAIN mode.
regularization_losses: A list of additional scalar losses to be added to
the training loss, such as regularization losses. These losses are
usually expressed as a batch average, so for best results users need to
set `loss_reduction=SUM_OVER_BATCH_SIZE` or
`loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to
avoid scaling errors.
Returns:
`EstimatorSpec`.
Raises:
ValueError: If `train_op_fn` is `None` in TRAIN mode.
"""
with ops.name_scope(self._name, 'head'):
logits = head_lib._check_logits_final_dim(logits, self.logits_dimension) # pylint:disable=protected-access
# Predict.
pred_keys = prediction_keys.PredictionKeys
with ops.name_scope(None, 'predictions', (logits,)):
probabilities = math_ops.sigmoid(logits, name=pred_keys.PROBABILITIES)
predictions = {
pred_keys.LOGITS: logits,
pred_keys.PROBABILITIES: probabilities,
}
if mode == model_fn.ModeKeys.PREDICT:
classifier_output = head_lib._classification_output( # pylint:disable=protected-access
scores=probabilities, n_classes=self._n_classes,
label_vocabulary=self._label_vocabulary)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
_DEFAULT_SERVING_KEY: classifier_output,
head_lib._CLASSIFY_SERVING_KEY: classifier_output, # pylint:disable=protected-access
head_lib._PREDICT_SERVING_KEY: ( # pylint:disable=protected-access
export_output.PredictOutput(predictions))
})
(training_loss, unreduced_loss, weights,
processed_labels) = self.create_loss(
features=features, mode=mode, logits=logits, labels=labels)
if regularization_losses:
regularization_loss = math_ops.add_n(regularization_losses)
regularized_training_loss = math_ops.add_n(
[training_loss, regularization_loss])
else:
regularization_loss = None
regularized_training_loss = training_loss
# Eval.
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=regularized_training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=processed_labels,
probabilities=probabilities,
weights=weights,
unreduced_loss=unreduced_loss,
regularization_loss=regularization_loss))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
# Only summarize mean_loss for SUM reduction to preserve backwards
# compatibility. Otherwise skip it to avoid unnecessary computation.
if self._loss_reduction == losses.Reduction.SUM:
example_weight_sum = math_ops.reduce_sum(
weights * array_ops.ones_like(unreduced_loss))
mean_loss = training_loss / example_weight_sum
else:
mean_loss = None
with ops.name_scope(''):
keys = metric_keys.MetricKeys
summary.scalar(
head_lib._summary_key(self._name, keys.LOSS), # pylint:disable=protected-access
regularized_training_loss)
if mean_loss is not None:
summary.scalar(
head_lib._summary_key(self._name, keys.LOSS_MEAN), # pylint:disable=protected-access
mean_loss)
if regularization_loss is not None:
summary.scalar(
head_lib._summary_key(self._name, keys.LOSS_REGULARIZATION), # pylint:disable=protected-access
regularization_loss)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=regularized_training_loss,
train_op=train_op_fn(regularized_training_loss))
def testRequiredArgumentsSet(self):
"""Tests that no errors are raised when all required arguments are set."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions={'loss': constant_op.constant(1.)})
def _model_fn(features, labels, mode, params, model):
"""Model defination for the SSD model based on ResNet-50.
Args:
features: the input image tensor with shape [batch_size, height, width, 3].
The height and width are fixed and equal.
labels: the input labels in a dictionary. The labels include class targets
and box targets which are dense label maps. The labels are generated from
get_input_fn function in data/dataloader.py
mode: the mode of TPUEstimator including TRAIN, EVAL, and PREDICT.
params: the dictionary defines hyperparameters of model. The default
settings are in default_hparams function in this file.
model: the SSD model outputs class logits and box regression outputs.
Returns:
spec: the EstimatorSpec or TPUEstimatorSpec to run training, evaluation,
or prediction.
"""
if mode == tf.estimator.ModeKeys.PREDICT:
labels = features
features = labels.pop('image')
features -= tf.constant(constants.NORMALIZATION_MEAN,
shape=[1, 1, 3],
dtype=features.dtype)
COEF_STD = 1.0 / tf.constant(
constants.NORMALIZATION_STD, shape=[1, 1, 3], dtype=features.dtype)
features *= COEF_STD
def _model_outputs():
return model(features,
params,
is_training_bn=(mode == tf.estimator.ModeKeys.TRAIN))
if params['dtype'] == 'bf16':
with tf.compat.v1.tpu.bfloat16_scope():
cls_outputs, box_outputs = _model_outputs()
levels = cls_outputs.keys()
for level in levels:
cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32)
box_outputs[level] = tf.cast(box_outputs[level], tf.float32)
else:
cls_outputs, box_outputs = _model_outputs()
levels = cls_outputs.keys()
# First check if it is in PREDICT mode.
if mode == tf.estimator.ModeKeys.PREDICT:
flattened_cls, flattened_box = concat_outputs(cls_outputs, box_outputs,
True)
ssd_box_coder = faster_rcnn_box_coder.FasterRcnnBoxCoder(
scale_factors=constants.BOX_CODER_SCALES)
anchors = box_list.BoxList(
tf.convert_to_tensor(dataloader.DefaultBoxes()('ltrb')))
decoded_boxes = box_coder.batch_decode(encoded_boxes=flattened_box,
box_coder=ssd_box_coder,
anchors=anchors)
pred_scores = tf.nn.softmax(flattened_cls, axis=2)
pred_scores, indices = select_top_k_scores(
pred_scores, constants.MAX_NUM_EVAL_BOXES)
predictions = dict(
labels,
indices=indices,
pred_scores=pred_scores,
pred_box=decoded_boxes,
)
if params['visualize_dataloader']:
# this is for inference visualization.
predictions['image'] = features
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
# Load pretrained model from checkpoint.
if params['resnet_checkpoint'] and mode == tf.estimator.ModeKeys.TRAIN:
def scaffold_fn():
"""Loads pretrained model through scaffold function."""
tf.train.init_from_checkpoint(
params['resnet_checkpoint'], {
'/': 'resnet%s/' % constants.RESNET_DEPTH,
})
return tf.train.Scaffold()
else:
scaffold_fn = None
# Set up training loss and learning rate.
update_learning_rate_schedule_parameters(params)
global_step = tf.train.get_or_create_global_step()
learning_rate = learning_rate_schedule(params, global_step)
# cls_loss and box_loss are for logging. only total_loss is optimized.
loss, cls_loss, box_loss = detection_loss(cls_outputs, box_outputs, labels)
total_loss = loss + params['weight_decay'] * tf.add_n(
[tf.nn.l2_loss(v) for v in tf.trainable_variables()])
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=constants.MOMENTUM)
if params['distributed_optimizer']:
optimizer = params['distributed_optimizer'](optimizer)
# Batch norm requires update_ops to be added as a train_op dependency.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
train_op = tf.group(optimizer.minimize(total_loss, global_step),
update_ops)
save_summary_steps = params['save_summary_steps']
if save_summary_steps is not None:
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar("class_loss", cls_loss)
tf.summary.scalar("box_loss", box_loss)
return model_fn_lib.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op,
training_hooks=None,
scaffold=scaffold_fn())
if mode == tf.estimator.ModeKeys.EVAL:
raise NotImplementedError
def estimator_spec(self,
mode,
default_serving_output_alternative_key=None):
"""Creates an equivalent `EstimatorSpec`.
Args:
mode: One of `ModeKeys`. Specifies if this training, evaluation or
prediction.
default_serving_output_alternative_key: Required for multiple heads. If
you have multiple entries in `output_alternatives` dict (comparable to
multiple heads), `EstimatorSpec` requires a default head that will be
used if a Servo request does not explicitly mention which head to infer
on. Pass the key of the output alternative here that you want to
designate as default. A separate ExportOutpout for this default head
wil be added to the export_outputs dict with the special key
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, unless there is
already an enry in output_alternatives with this special key.
Returns:
Instance of `EstimatorSpec` that is equivalent to this `ModelFnOps`
Raises:
ValueError: If problem type is unknown.
"""
def _scores(output_tensors):
scores = output_tensors.get(prediction_key.PredictionKey.SCORES)
if scores is None:
scores = output_tensors.get(
prediction_key.PredictionKey.PROBABILITIES)
return scores
def _classes(output_tensors): # pylint: disable=missing-docstring
classes = output_tensors.get(prediction_key.PredictionKey.CLASSES)
if classes is None:
logging.warning(
'classes is None, Servo inference will not have class ids.'
)
return None
elif classes.dtype != dtypes.string:
# Servo classification can only serve string classes
logging.warning(
'classes is not string, Servo inference will not have class ids.'
)
return None
return classes
def _export_output(problem_type, predictions): # pylint: disable=missing-docstring
if problem_type == constants.ProblemType.LINEAR_REGRESSION:
return core_export_lib.RegressionOutput(_scores(predictions))
if (problem_type == constants.ProblemType.CLASSIFICATION or
problem_type == constants.ProblemType.LOGISTIC_REGRESSION):
return core_export_lib.ClassificationOutput(
scores=_scores(predictions), classes=_classes(predictions))
if problem_type == constants.ProblemType.UNSPECIFIED:
return core_export_lib.PredictOutput(predictions)
raise ValueError('Unknown problem_type=%s' % problem_type)
# Converts output_alternatives
export_outputs_dict = None
if self.output_alternatives:
output_alternatives = self.output_alternatives
# Adds default output_alternative if needed.
if (len(output_alternatives) > 1
and signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
not in output_alternatives):
output_alternatives = output_alternatives.copy()
output_alternatives[
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
output_alternatives[
default_serving_output_alternative_key])
export_outputs_dict = {
key: _export_output(*val)
for key, val in output_alternatives.items()
}
def _get_eval_metric_ops():
"""Returns self.eval_metric_ops without loss metric."""
result = {}
for key, value in six.iteritems(self.eval_metric_ops):
if key != metric_key.MetricKey.LOSS:
result[key] = value
return result
return core_model_fn_lib.EstimatorSpec(
mode=mode,
predictions=self.predictions,
loss=self.loss,
train_op=self.train_op,
eval_metric_ops=_get_eval_metric_ops(),
export_outputs=export_outputs_dict,
training_chief_hooks=self.training_chief_hooks,
training_hooks=self.training_hooks,
scaffold=self.scaffold)
def my_model_fn(features, mode):
loss = math_ops.reduce_sum(features)
train_op = array_ops.identity(loss)
return model_fn_lib.EstimatorSpec(mode,
loss=loss,
train_op=train_op)
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
# Predict.
with ops.name_scope(self._name, 'head'):
logits = _check_logits(logits, self._logits_dimension)
predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits}
if mode == model_fn.ModeKeys.PREDICT:
regression_output = export_output.RegressionOutput(
value=logits)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
_DEFAULT_SERVING_KEY:
regression_output,
_REGRESS_SERVING_KEY:
regression_output,
_PREDICT_SERVING_KEY:
export_output.PredictOutput(predictions)
})
# Eval.
unweighted_loss, _ = self.create_loss(features=features,
mode=mode,
logits=logits,
labels=labels)
weights = _weights(features, self._weight_column)
training_loss = losses.compute_weighted_loss(
unweighted_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
# Estimator already adds a metric for loss.
eval_metric_ops = {
metric_keys.MetricKeys.LOSS_MEAN:
metrics_lib.mean(unweighted_loss, weights=weights)
}
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=eval_metric_ops)
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(
_summary_key(self._name, metric_keys.MetricKeys.LOSS),
training_loss)
summary.scalar(
_summary_key(self._name, metric_keys.MetricKeys.LOSS_MEAN),
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
# Predict.
with ops.name_scope(self._name, 'head'):
with ops.name_scope(None, 'predictions', (logits, )):
pred_keys = prediction_keys.PredictionKeys
logits = _check_logits(logits, self.logits_dimension)
logistic = math_ops.sigmoid(logits, name=pred_keys.LOGISTIC)
two_class_logits = array_ops.concat(
(array_ops.zeros_like(logits), logits),
1,
name='two_class_logits')
probabilities = nn.softmax(two_class_logits,
name=pred_keys.PROBABILITIES)
class_ids = array_ops.reshape(math_ops.argmax(two_class_logits,
axis=1), (-1, 1),
name='classes')
if self._label_vocabulary:
table = lookup_ops.index_to_string_table_from_tensor(
vocabulary_list=self._label_vocabulary,
name='class_string_lookup')
classes = table.lookup(class_ids)
else:
classes = string_ops.as_string(class_ids,
name='str_classes')
predictions = {
pred_keys.LOGITS: logits,
pred_keys.LOGISTIC: logistic,
pred_keys.PROBABILITIES: probabilities,
pred_keys.CLASS_IDS: class_ids,
pred_keys.CLASSES: classes,
}
if mode == model_fn.ModeKeys.PREDICT:
classifier_output = _classification_output(
scores=probabilities,
n_classes=2,
label_vocabulary=self._label_vocabulary)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
_DEFAULT_SERVING_KEY:
classifier_output,
_CLASSIFY_SERVING_KEY:
classifier_output,
_REGRESS_SERVING_KEY:
export_output.RegressionOutput(value=logistic),
_PREDICT_SERVING_KEY:
export_output.PredictOutput(predictions)
})
# Eval.
unweighted_loss, processed_labels = self.create_loss(
features=features, mode=mode, logits=logits, labels=labels)
weights = _weights(features, self._weight_column)
training_loss = losses.compute_weighted_loss(
unweighted_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=processed_labels,
logits=logits,
logistic=logistic,
class_ids=class_ids,
unweighted_loss=unweighted_loss,
weights=weights))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(
_summary_key(self._name, metric_keys.MetricKeys.LOSS),
training_loss)
summary.scalar(
_summary_key(self._name, metric_keys.MetricKeys.LOSS_MEAN),
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def testLossMissing(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Missing loss'):
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions={'loss': constant_op.constant(1.)})
def model_fn(features, labels, mode):
"""model_fn for keras Estimator."""
model = _clone_and_build_model(mode, keras_model, custom_objects, features,
labels)
model_output_names = []
# We need to make sure that the output names of the last layer in the model
# is the same for each of the cloned models. This is required for mirrored
# strategy when we call regroup.
if distribute_lib.has_distribution_strategy():
for name in model.output_names:
name = re.compile(r'_\d$').sub('', name)
model_output_names.append(name)
else:
model_output_names = model.output_names
# Get inputs to EstimatorSpec
predictions = dict(zip(model_output_names, model.outputs))
loss = None
train_op = None
eval_metric_ops = None
# Set loss and metric only during train and evaluate.
if mode is not model_fn_lib.ModeKeys.PREDICT:
if mode is model_fn_lib.ModeKeys.TRAIN:
model._make_train_function() # pylint: disable=protected-access
else:
model._make_test_function() # pylint: disable=protected-access
loss = model.total_loss
if model.metrics:
# TODO(fchollet): support stateful metrics
eval_metric_ops = {}
# When each metric maps to an output
if isinstance(model.metrics, dict):
for i, output_name in enumerate(model.metrics.keys()):
metric_name = model.metrics[output_name]
if callable(metric_name):
metric_name = metric_name.__name__
# When some outputs use the same metric
if list(model.metrics.values()).count(metric_name) > 1:
metric_name += '_' + output_name
eval_metric_ops[metric_name] = metrics_module.mean(
model.metrics_tensors[i - len(model.metrics)])
else:
for i, metric_name in enumerate(model.metrics):
if callable(metric_name):
metric_name = metric_name.__name__
eval_metric_ops[metric_name] = metrics_module.mean(
model.metrics_tensors[i])
# Set train_op only during train.
if mode is model_fn_lib.ModeKeys.TRAIN:
train_op = model.train_function.updates_op
if not model._is_graph_network:
# Reset model state to original state,
# to avoid `model_fn` being destructive for the initial model argument.
_in_place_subclassed_model_state_restoration(keras_model)
return model_fn_lib.EstimatorSpec(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
export_outputs={
_DEFAULT_SERVING_KEY:
export_lib.export_output.PredictOutput(predictions)
})
def testPredictionsMissingIsOkay(self):
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.EVAL,
loss=constant_op.constant(1.))
def eval_model_fn_no_eval_metric_ops(features, labels, mode, params):
del features, labels, params
return model_fn_lib.EstimatorSpec(
mode=mode, loss=constant_op.constant(_EXPECTED_LOSS))
def testPredictionsMissing(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Missing predictions'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT)
def _test_train_model_fn(features, labels, mode, params):
"""A dummy model_fn for testing purpose."""
del features, labels, params
loss = constant_op.constant(_EXPECTED_LOSS)
return model_fn_lib.EstimatorSpec(
mode=mode, loss=loss, train_op=array_ops.identity(loss))
def testPredictionsNumber(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(
TypeError, r'predictions\[number\] must be Tensor'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT,
predictions={'number': 1.})
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
with ops.name_scope('head'):
logits = _check_logits(logits, self.logits_dimension)
# Predict.
pred_keys = prediction_keys.PredictionKeys
with ops.name_scope(None, 'predictions', (logits, )):
# class_ids's shape is [batch_size]
class_ids = math_ops.argmax(logits,
1,
name=pred_keys.CLASS_IDS)
class_ids = array_ops.expand_dims(class_ids, axis=(1, ))
if self._label_vocabulary:
table = lookup_ops.index_to_string_table_from_tensor(
vocabulary_list=self._label_vocabulary,
name='class_string_lookup')
classes = table.lookup(class_ids)
else:
classes = string_ops.as_string(class_ids,
name='str_classes')
probabilities = nn.softmax(logits,
name=pred_keys.PROBABILITIES)
predictions = {
pred_keys.LOGITS: logits,
pred_keys.PROBABILITIES: probabilities,
# Expand to [batch_size, 1]
pred_keys.CLASS_IDS: class_ids,
pred_keys.CLASSES: classes,
}
if mode == model_fn.ModeKeys.PREDICT:
batch_size = array_ops.shape(probabilities)[0]
export_class_list = self._label_vocabulary
if not export_class_list:
export_class_list = string_ops.as_string(
math_ops.range(self._n_classes))
export_output_classes = array_ops.tile(
input=array_ops.expand_dims(input=export_class_list,
axis=0),
multiples=[batch_size, 1])
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
'':
export_output.ClassificationOutput(
scores=probabilities,
# `ClassificationOutput` requires string classes.
classes=export_output_classes)
})
# Eval.
label_ids = self._label_ids(
_check_labels(_maybe_expand_dim(labels), 1))
unweighted_loss = losses.sparse_softmax_cross_entropy(
labels=label_ids,
logits=logits,
reduction=losses.Reduction.NONE)
# Restore the squeezed dim, so unweighted_loss matches the weights shape.
unweighted_loss = array_ops.expand_dims(unweighted_loss,
axis=(1, ))
weights = _weights(features, self._weight_column)
training_loss = losses.compute_weighted_loss(
unweighted_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=label_ids,
probabilities=probabilities,
logits=logits,
class_ids=class_ids,
unweighted_loss=unweighted_loss,
weights=weights))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(metric_keys.MetricKeys.LOSS, training_loss)
summary.scalar(
metric_keys.MetricKeys.LOSS_MEAN,
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def testLoss1DTensor(self):
"""Tests that no errors are raised when loss is 1D tensor."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1.]),
train_op=control_flow_ops.no_op())
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=None):
"""See `Head`."""
# Predict.
with ops.name_scope('head'):
with ops.name_scope(None, 'predictions', (logits, )):
pred_keys = prediction_keys.PredictionKeys
logits = _check_logits(logits, self.logits_dimension)
logistic = math_ops.sigmoid(logits, name=pred_keys.LOGISTIC)
two_class_logits = array_ops.concat(
(array_ops.zeros_like(logits), logits),
1,
name='two_class_logits')
scores = nn.softmax(two_class_logits,
name=pred_keys.PROBABILITIES)
class_ids = array_ops.reshape(math_ops.argmax(two_class_logits,
axis=1), (-1, 1),
name='classes')
if self._label_vocabulary:
table = lookup_ops.index_to_string_table_from_tensor(
vocabulary_list=self._label_vocabulary,
name='class_string_lookup')
classes = table.lookup(class_ids)
else:
classes = string_ops.as_string(class_ids,
name='str_classes')
predictions = {
pred_keys.LOGITS: logits,
pred_keys.LOGISTIC: logistic,
pred_keys.PROBABILITIES: scores,
pred_keys.CLASS_IDS: class_ids,
pred_keys.CLASSES: classes,
}
if mode == model_fn.ModeKeys.PREDICT:
batch_size = array_ops.shape(logistic)[0]
export_class_list = self._label_vocabulary
if not export_class_list:
export_class_list = string_ops.as_string([0, 1])
export_output_classes = array_ops.tile(
input=array_ops.expand_dims(input=export_class_list,
axis=0),
multiples=[batch_size, 1])
classifier_output = export_output.ClassificationOutput(
scores=scores,
# `ClassificationOutput` requires string classes.
classes=export_output_classes)
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
'':
classifier_output, # to be same as other heads.
'classification':
classifier_output, # to be called by name.
_DEFAULT_SERVING_KEY:
classifier_output, # default
'regression':
export_output.RegressionOutput(value=logistic)
})
# Eval.
labels = _check_labels(_maybe_expand_dim(labels),
self.logits_dimension)
if self._label_vocabulary is not None:
labels = lookup_ops.index_table_from_tensor(
vocabulary_list=tuple(self._label_vocabulary),
name='class_id_lookup').lookup(labels)
labels = math_ops.to_float(labels)
labels = _assert_range(labels, 2)
unweighted_loss = nn.sigmoid_cross_entropy_with_logits(
labels=labels, logits=logits, name='loss')
weights = _weights(features, self._weight_column)
training_loss = losses.compute_weighted_loss(
unweighted_loss,
weights=weights,
reduction=losses.Reduction.SUM)
if mode == model_fn.ModeKeys.EVAL:
return model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=training_loss,
eval_metric_ops=self._eval_metric_ops(
labels=labels,
logits=logits,
logistic=logistic,
scores=scores,
class_ids=class_ids,
unweighted_loss=unweighted_loss,
weights=weights))
# Train.
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
with ops.name_scope(''):
summary.scalar(metric_keys.MetricKeys.LOSS, training_loss)
summary.scalar(
metric_keys.MetricKeys.LOSS_MEAN,
losses.compute_weighted_loss(unweighted_loss,
weights=weights,
reduction=losses.Reduction.MEAN))
return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=training_loss,
train_op=train_op_fn(training_loss))
def testLossNotScalar(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Loss must be scalar'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1., 2.]),
train_op=control_flow_ops.no_op())
def create_estimator_spec(self,
features,
mode,
logits,
labels=None,
train_op_fn=tfgan_train.gan_train_ops):
"""Returns `EstimatorSpec` that a model_fn can return.
See `Head` for more details.
Args:
features: Must be `None`.
mode: Estimator's `ModeKeys`.
logits: A GANModel tuple.
labels: Must be `None`.
train_op_fn: Function that takes a GANModel, GANLoss, generator optimizer,
and discriminator optimizer, and returns a `GANTrainOps` tuple. For
example, this function can come from TFGAN's `train.py` library, or can
be custom.
Returns:
`EstimatorSpec`.
Raises:
ValueError: If `features` isn't `None`.
ValueError: If `train_op_fn` isn't provided in train mode.
"""
_validate_logits_and_labels(logits, labels)
if features is not None:
raise ValueError(
'`features` should be `None`. Instead, found: %s' % features)
gan_model = logits # rename variable for clarity
with ops.name_scope('GANHead'):
if mode == model_fn_lib.ModeKeys.PREDICT:
return model_fn_lib.EstimatorSpec(
mode=model_fn_lib.ModeKeys.PREDICT,
predictions=gan_model.generated_data,
export_outputs={
'predict':
export_output.PredictOutput(gan_model.generated_data)
})
elif mode == model_fn_lib.ModeKeys.EVAL:
gan_loss = self.create_loss(features=None,
mode=mode,
logits=gan_model,
labels=None)
scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss
with ops.name_scope(
None, 'metrics',
[gan_loss.generator_loss, gan_loss.discriminator_loss]):
eval_metric_ops = {
_summary_key(self._name, 'generator_loss'):
metrics_lib.mean(gan_loss.generator_loss),
_summary_key(self._name, 'discriminator_loss'):
metrics_lib.mean(gan_loss.discriminator_loss)
}
if self._get_eval_metric_ops_fn is not None:
custom_eval_metric_ops = self._get_eval_metric_ops_fn(
gan_model)
if not isinstance(custom_eval_metric_ops, dict):
raise TypeError(
'get_eval_metric_ops_fn must return a dict, '
'received: {}'.format(custom_eval_metric_ops))
eval_metric_ops.update(custom_eval_metric_ops)
return model_fn_lib.EstimatorSpec(
mode=model_fn_lib.ModeKeys.EVAL,
predictions=gan_model.generated_data,
loss=scalar_loss,
eval_metric_ops=eval_metric_ops)
elif mode == model_fn_lib.ModeKeys.TRAIN:
if train_op_fn is None:
raise ValueError('train_op_fn can not be None.')
gan_loss = self.create_loss(None, mode, gan_model, None)
scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss
train_ops = train_op_fn(gan_model, gan_loss,
self._generator_optimizer,
self._discriminator_optimizer)
training_hooks = self._get_hooks_fn(train_ops)
return model_fn_lib.EstimatorSpec(
loss=scalar_loss,
mode=model_fn_lib.ModeKeys.TRAIN,
train_op=train_ops.global_step_inc_op,
training_hooks=training_hooks)
else:
raise ValueError('Mode not recognized: %s' % mode)
