def construct(): # pylint: disable=invalid-name
"""Function for constructing a EvalSavedModel."""
start_time = datetime.datetime.now()
result = load.EvalSavedModel(eval_saved_model_path)
if add_metrics_callbacks:
features_dict, predictions_dict, labels_dict = (
result.get_features_predictions_labels_dicts())
features_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
features_dict)
predictions_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
predictions_dict)
labels_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(labels_dict)
with result.graph_as_default():
metric_ops = {}
for add_metrics_callback in add_metrics_callbacks:
new_metric_ops = add_metrics_callback(features_dict, predictions_dict,
labels_dict)
overlap = set(new_metric_ops.keys()) & set(metric_ops.keys())
if overlap:
raise ValueError('metric keys should not conflict, but an '
'earlier callback already added the metrics '
'named %s' % overlap)
metric_ops.update(new_metric_ops)
result.register_additional_metric_ops(metric_ops)
end_time = datetime.datetime.now()
model_load_seconds_distribution.update(
int((end_time - start_time).total_seconds()))
return result
def _LegacyEvalInputReceiver( # pylint: disable=invalid-name
features: types.TensorTypeMaybeDict,
labels: Optional[types.TensorTypeMaybeDict],
receiver_tensors: Dict[Text, types.TensorType],
input_refs: Optional[types.TensorType] = None) -> EvalInputReceiverType:
"""Returns a legacy eval_input_receiver_fn.
This is for testing purposes only.
Args:
features: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the features to be passed to the model.
labels: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the labels to be passed to the model. If your
model is an unsupervised model whose `model_fn` does not accept a `labels`
argument, you may pass None instead.
receiver_tensors: A dict of string to `Tensor` containing exactly key named
'examples', which maps to the single input node that the receiver expects
to be fed by default. Typically this is a placeholder expecting serialized
`tf.Example` protos.
input_refs: Optional. A 1-D integer `Tensor` that is batch-aligned with
`features` and `labels` which is an index into
receiver_tensors['examples'] indicating where this slice of features /
labels came from. If not provided, defaults to range(0,
len(receiver_tensors['examples'])).
Raises:
ValueError: receiver_tensors did not contain exactly one key named
"examples".
"""
# Force list representation for Python 3 compatibility.
if list(receiver_tensors.keys()) != ['examples']:
raise ValueError('receiver_tensors must contain exactly one key named '
'examples.')
# Workaround for TensorFlow issue #17568. Note that we pass the
# identity-wrapped features and labels to model_fn, but we have to feed
# the non-identity wrapped Tensors during evaluation.
#
# Also note that we can't wrap predictions, so metrics that have control
# dependencies on predictions will cause the predictions to be recomputed
# during their evaluation.
wrapped_features = util.wrap_tensor_or_dict_of_tensors_in_identity(features)
if labels is not None:
wrapped_labels = util.wrap_tensor_or_dict_of_tensors_in_identity(labels)
receiver = export_lib.SupervisedInputReceiver(
features=wrapped_features,
labels=wrapped_labels,
receiver_tensors=receiver_tensors)
else:
receiver = export_lib.UnsupervisedInputReceiver(
features=wrapped_features, receiver_tensors=receiver_tensors)
if input_refs is None:
input_refs = tf.range(tf.size(input=receiver_tensors['examples']))
# Note that in the collection we store the unwrapped versions, because
# we want to feed the unwrapped versions.
_add_tfma_collections(features, labels, input_refs)
util.add_build_data_collection()
return receiver
def register_add_metric_callbacks(
self,
add_metrics_callbacks: List[types.AddMetricsCallbackType]) -> None:
"""Register additional metric callbacks.
Runs the given list of callbacks for adding additional metrics to the graph.
For more details about add_metrics_callbacks, see the docstring for
EvalSharedModel.add_metrics_callbacks in types.py.
Args:
add_metrics_callbacks: A list of metric callbacks to add to the metrics
graph.
Raises:
ValueError: There was a metric name conflict: a callback tried to add a
metric whose name conflicted with a metric that was added by an earlier
callback.
"""
with self._graph.as_default():
features_dict, predictions_dict, labels_dict = (
self.get_features_predictions_labels_dicts())
features_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
features_dict)
predictions_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
predictions_dict)
labels_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
labels_dict)
metric_ops = {}
for add_metrics_callback in add_metrics_callbacks:
new_metric_ops = add_metrics_callback(features_dict,
predictions_dict,
labels_dict)
overlap = set(new_metric_ops.keys()) & set(metric_ops.keys())
if overlap:
raise ValueError(
'metric keys should not conflict, but an '
'earlier callback already added the metrics '
'named %s' % overlap)
metric_ops.update(new_metric_ops)
self.register_additional_metric_ops(metric_ops)
def __new__(cls, features, labels, receiver_tensors):
# Workaround for TensorFlow issue #17568. Note that we pass the
# identity-wrapped features and labels to model_fn, but we have to feed
# the non-identity wrapped Tensors during evaluation.
#
# Also note that we can't wrap predictions, so metrics that have control
# dependencies on predictions will cause the predictions to be recomputed
# during their evaluation.
wrapped_features = util.wrap_tensor_or_dict_of_tensors_in_identity(
features)
wrapped_labels = util.wrap_tensor_or_dict_of_tensors_in_identity(
labels)
receiver = super(EvalInputReceiver,
cls).__new__(cls,
features=wrapped_features,
labels=wrapped_labels,
receiver_tensors=receiver_tensors)
# Note that in the collection we store the unwrapped versions, because
# we want to feed the unwrapped versions.
_add_features_labels_version_collections(features, labels)
return receiver
def EvalInputReceiver( # pylint: disable=invalid-name
features: types.TensorTypeMaybeDict,
labels: Optional[types.TensorTypeMaybeDict],
receiver_tensors: types.TensorTypeMaybeDict,
input_refs: Optional[types.TensorType] = None,
iterator_initializer: Optional[Text] = None) -> EvalInputReceiverType:
"""Returns an appropriate receiver for eval_input_receiver_fn.
This is a wrapper around TensorFlow's InputReceiver that adds additional
entries and prefixes to the input tensors so that features and labels can be
discovered at evaluation time. It also wraps the features and labels tensors
in identity to workaround TensorFlow issue #17568.
The resulting signature_def.inputs will have the following form:
inputs/<input> - placeholders that are used for input processing (i.e
receiver_tensors). If receiver_tensors is a tensor and
not a dict, then this will just be named 'inputs'.
input_refs - reference to input_refs tensor (see below).
features/<feature> - references to tensors passed in features. If features
is a tensor and not a dict, then this will just be
named 'features'.
labels/<label> - references to tensors passed in labels. If labels is
a tensor and not a dict, then this will just be named
'labels'.
Args:
features: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the features to be passed to the model.
labels: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the labels to be passed to the model. If your
model is an unsupervised model whose `model_fn` does not accept a `labels`
argument, you may pass None instead.
receiver_tensors: A dict of string to `Tensor` containing exactly key named
'examples', which maps to the single input node that the receiver expects
to be fed by default. Typically this is a placeholder expecting serialized
`tf.Example` protos.
input_refs: Optional (unless iterator_initializer used). A 1-D integer
`Tensor` that is batch-aligned with `features` and `labels` which is an
index into receiver_tensors['examples'] indicating where this slice of
features / labels came from. If not provided, defaults to range(0,
len(receiver_tensors['examples'])).
iterator_initializer: Optional name of tf.compat.v1.data.Iterator
initializer used when the inputs are fed using an iterator. This is
intended to be used by models that cannot handle a single large input due
to memory resource constraints. For example, a model that takes a
tf.train.SequenceExample record as input but only processes smaller
batches of examples within the overall sequence at a time. The caller is
responsible for setting the input_refs appropriately (i.e. all examples
belonging to the same tf.train.Sequence should have the same input_ref).
Raises:
ValueError: receiver_tensors did not contain exactly one key named
"examples" or iterator_initializer used without input_refs.
"""
if list(receiver_tensors.keys()) != ['examples']:
raise ValueError('receiver_tensors must contain exactly one key named '
'examples.')
if input_refs is None:
if iterator_initializer is not None:
raise ValueError('input_refs is required if iterator_initializer is used')
input_refs = tf.range(tf.size(input=list(receiver_tensors.values())[0]))
updated_receiver_tensors = {}
def add_tensors(prefix, tensor_or_dict):
if isinstance(tensor_or_dict, dict):
for key in tensor_or_dict:
updated_receiver_tensors[prefix + '/' + key] = tensor_or_dict[key]
else:
updated_receiver_tensors[prefix] = tensor_or_dict
add_tensors(constants.SIGNATURE_DEF_INPUTS_PREFIX, receiver_tensors)
add_tensors(constants.FEATURES_NAME, features)
if labels is not None:
add_tensors(constants.LABELS_NAME, labels)
updated_receiver_tensors[constants.SIGNATURE_DEF_INPUT_REFS_KEY] = (
input_refs)
if iterator_initializer:
updated_receiver_tensors[
constants.SIGNATURE_DEF_ITERATOR_INITIALIZER_KEY] = (
tf.constant(iterator_initializer))
updated_receiver_tensors[constants.SIGNATURE_DEF_TFMA_VERSION_KEY] = (
tf.constant(version.VERSION_STRING))
# TODO(b/119308261): Remove once all evaluator binaries have been updated.
_add_tfma_collections(features, labels, input_refs)
util.add_build_data_collection()
# Workaround for TensorFlow issue #17568. Note that we pass the
# identity-wrapped features and labels to model_fn, but we have to feed
# the non-identity wrapped Tensors during evaluation.
#
# Also note that we can't wrap predictions, so metrics that have control
# dependencies on predictions will cause the predictions to be recomputed
# during their evaluation.
wrapped_features = util.wrap_tensor_or_dict_of_tensors_in_identity(features)
if labels is not None:
wrapped_labels = util.wrap_tensor_or_dict_of_tensors_in_identity(labels)
return export_lib.SupervisedInputReceiver(
features=wrapped_features,
labels=wrapped_labels,
receiver_tensors=updated_receiver_tensors)
else:
return export_lib.UnsupervisedInputReceiver(
features=wrapped_features, receiver_tensors=updated_receiver_tensors)
def EvalInputReceiver( # pylint: disable=invalid-name
features,
labels,
receiver_tensors,
input_refs=None):
"""Returns an appropriate receiver for eval_input_receiver_fn.
This is a wrapper around TensorFlow's InputReceiver that adds additional
entries and prefixes to the input tensors so that features and labels can be
discovered at evaluation time. It also wraps the features and labels tensors
in identity to workaround TensorFlow issue #17568.
The resulting signature_def.inputs will have the following form:
inputs/<input> - placeholders that are used for input processing (i.e
receiver_tensors). If receiver_tensors is a tensor and
not a dict, then this will just be named 'inputs'.
input_refs - reference to input_refs tensor (see below).
features/<feature> - references to tensors passed in features. If features
is a tensor and not a dict, then this will just be
named 'features'.
labels/<label> - references to tensors passed in labels. If labels is
a tensor and not a dict, then this will just be named
'labels'.
Args:
features: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the features to be passed to the model.
labels: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
`SparseTensor`, specifying the labels to be passed to the model. If your
model is an unsupervised model whose `model_fn` does not accept a `labels`
argument, you may pass None instead.
receiver_tensors: A dict of string to `Tensor` containing exactly key named
'examples', which maps to the single input node that the receiver expects
to be fed by default. Typically this is a placeholder expecting serialized
`tf.Example` protos.
input_refs: Optional. A 1-D integer `Tensor` that is batch-aligned with
`features` and `labels` which is an index into
receiver_tensors['examples'] indicating where this slice of features /
labels came from. If not provided, defaults to range(0,
len(receiver_tensors['examples'])).
Raises:
ValueError: receiver_tensors did not contain exactly one key named
"examples".
"""
if list(receiver_tensors.keys()) != ['examples']:
raise ValueError('receiver_tensors must contain exactly one key named '
'examples.')
if input_refs is None:
input_refs = tf.range(tf.size(list(receiver_tensors.values())[0]))
updated_receiver_tensors = {}
def add_tensors(prefix, tensor_or_dict):
if isinstance(tensor_or_dict, dict):
for key in tensor_or_dict:
updated_receiver_tensors[prefix + '/' +
key] = tensor_or_dict[key]
else:
updated_receiver_tensors[prefix] = tensor_or_dict
add_tensors(constants.SIGNATURE_DEF_INPUTS_PREFIX, receiver_tensors)
add_tensors(constants.SIGNATURE_DEF_FEATURES_PREFIX, features)
if labels is not None:
add_tensors(constants.SIGNATURE_DEF_LABELS_PREFIX, labels)
updated_receiver_tensors[constants.SIGNATURE_DEF_INPUT_REFS_KEY] = (
input_refs)
updated_receiver_tensors[constants.SIGNATURE_DEF_TFMA_VERSION_KEY] = (
tf.constant(version.VERSION_STRING))
_add_tfma_collections(features, labels, input_refs)
# Workaround for TensorFlow issue #17568. Note that we pass the
# identity-wrapped features and labels to model_fn, but we have to feed
# the non-identity wrapped Tensors during evaluation.
#
# Also note that we can't wrap predictions, so metrics that have control
# dependencies on predictions will cause the predictions to be recomputed
# during their evaluation.
wrapped_features = util.wrap_tensor_or_dict_of_tensors_in_identity(
features)
if labels is not None:
wrapped_labels = util.wrap_tensor_or_dict_of_tensors_in_identity(
labels)
return export_lib.SupervisedInputReceiver(
features=wrapped_features,
labels=wrapped_labels,
receiver_tensors=updated_receiver_tensors)
else:
return export_lib.UnsupervisedInputReceiver(
features=wrapped_features,
receiver_tensors=updated_receiver_tensors)
def export_eval_savedmodel(
estimator,
export_dir_base,
eval_input_receiver_fn,
checkpoint_path = None):
"""Export a EvalSavedModel for the given estimator.
Args:
estimator: Estimator to export the graph for.
export_dir_base: Base path for export. Graph will be exported into a
subdirectory of this base path.
eval_input_receiver_fn: Eval input receiver function.
checkpoint_path: Path to a specific checkpoint to export. If set to None,
exports the latest checkpoint.
Returns:
Path to the directory where the eval graph was exported.
Raises:
ValueError: Could not find a checkpoint to export.
"""
with tf.Graph().as_default() as g:
eval_input_receiver = eval_input_receiver_fn()
tf.train.create_global_step(g)
tf.set_random_seed(estimator.config.tf_random_seed)
# Workaround for TensorFlow issue #17568. Note that we pass the
# identity-wrapped features and labels to model_fn, but we have to feed
# the non-identity wrapped Tensors during evaluation.
#
# Also note that we can't wrap predictions, so metrics that have control
# dependencies on predictions will cause the predictions to be recomputed
# during their evaluation.
wrapped_features = util.wrap_tensor_or_dict_of_tensors_in_identity(
eval_input_receiver.features)
wrapped_labels = util.wrap_tensor_or_dict_of_tensors_in_identity(
eval_input_receiver.labels)
if isinstance(estimator, tf.estimator.Estimator):
# This is a core estimator
estimator_spec = estimator.model_fn(
features=wrapped_features,
labels=wrapped_labels,
mode=tf.estimator.ModeKeys.EVAL,
config=estimator.config)
else:
# This is a contrib estimator
model_fn_ops = estimator._call_model_fn( # pylint: disable=protected-access
features=wrapped_features,
labels=wrapped_labels,
mode=tf.estimator.ModeKeys.EVAL)
estimator_spec = lambda x: None
estimator_spec.predictions = model_fn_ops.predictions
estimator_spec.eval_metric_ops = model_fn_ops.eval_metric_ops
estimator_spec.scaffold = model_fn_ops.scaffold
# Save metric using eval_metric_ops.
for user_metric_key, (value_op, update_op) in (
estimator_spec.eval_metric_ops.items()):
tf.add_to_collection('%s/%s' % (encoding.METRICS_COLLECTION,
encoding.KEY_SUFFIX),
encoding.encode_key(user_metric_key))
tf.add_to_collection('%s/%s' % (encoding.METRICS_COLLECTION,
encoding.VALUE_OP_SUFFIX),
encoding.encode_tensor_node(value_op))
tf.add_to_collection('%s/%s' % (encoding.METRICS_COLLECTION,
encoding.UPDATE_OP_SUFFIX),
encoding.encode_tensor_node(update_op))
# Save all prediction nodes.
# Predictions can either be a Tensor, or a dict of Tensors.
predictions = estimator_spec.predictions
if not isinstance(predictions, dict):
predictions = {encoding.DEFAULT_PREDICTIONS_DICT_KEY: predictions}
for prediction_key, prediction_node in predictions.items():
_encode_and_add_to_node_collection(encoding.PREDICTIONS_COLLECTION,
prediction_key, prediction_node)
############################################################
## Features, label (and weight) graph
# Placeholder for input example to label graph.
tf.add_to_collection(encoding.INPUT_EXAMPLE_COLLECTION,
encoding.encode_tensor_node(
eval_input_receiver.receiver_tensors['examples']))
# Save all label nodes.
# Labels can either be a Tensor, or a dict of Tensors.
labels = eval_input_receiver.labels
if not isinstance(labels, dict):
labels = {encoding.DEFAULT_LABELS_DICT_KEY: labels}
for label_key, label_node in labels.items():
_encode_and_add_to_node_collection(encoding.LABELS_COLLECTION, label_key,
label_node)
# Save features.
for feature_name, feature_node in eval_input_receiver.features.items():
_encode_and_add_to_node_collection(encoding.FEATURES_COLLECTION,
feature_name, feature_node)
############################################################
## Export as normal
if not checkpoint_path:
checkpoint_path = tf.train.latest_checkpoint(estimator.model_dir)
if not checkpoint_path:
raise ValueError(
'Could not find trained model at %s.' % estimator.model_dir)
export_dir = _get_timestamped_export_dir(export_dir_base)
temp_export_dir = _get_temp_export_dir(export_dir)
if estimator.config.session_config is None:
session_config = config_pb2.ConfigProto(allow_soft_placement=True)
else:
session_config = estimator.config.session_config
with tf.Session(config=session_config) as session:
if estimator_spec.scaffold and estimator_spec.scaffold.saver:
saver_for_restore = estimator_spec.scaffold.saver
else:
saver_for_restore = tf.train.Saver(sharded=True)
saver_for_restore.restore(session, checkpoint_path)
if estimator_spec.scaffold and estimator_spec.scaffold.local_init_op:
local_init_op = estimator_spec.scaffold.local_init_op
else:
local_init_op = tf.train.Scaffold._default_local_init_op()
# pylint: enable=protected-access
# Perform the export
builder = tf.saved_model.builder.SavedModelBuilder(temp_export_dir)
builder.add_meta_graph_and_variables(
session,
[tf.saved_model.tag_constants.SERVING],
# Don't export any signatures, since this graph is not actually
# meant for serving.
signature_def_map=None,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=local_init_op)
builder.save(False)
gfile.Rename(temp_export_dir, export_dir)
return export_dir
