Module keras.saving.utils_v1.export_utils
Utilities for creating SavedModels.
Expand source code
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# LINT.IfChange
"""Utilities for creating SavedModels."""
import tensorflow.compat.v2 as tf
import collections
import os
import time
from keras.saving.utils_v1 import export_output as export_output_lib
from keras.saving.utils_v1 import mode_keys
from keras.saving.utils_v1 import unexported_constants
from keras.saving.utils_v1.mode_keys import KerasModeKeys as ModeKeys
from tensorflow.python.platform import tf_logging as logging
# Mapping of the modes to appropriate MetaGraph tags in the SavedModel.
EXPORT_TAG_MAP = mode_keys.ModeKeyMap(**{
ModeKeys.PREDICT: [tf.saved_model.SERVING],
ModeKeys.TRAIN: [tf.saved_model.TRAINING],
ModeKeys.TEST: [unexported_constants.EVAL]})
# For every exported mode, a SignatureDef map should be created using the
# functions `export_outputs_for_mode` and `build_all_signature_defs`. By
# default, this map will contain a single Signature that defines the input
# tensors and output predictions, losses, and/or metrics (depending on the mode)
# The default keys used in the SignatureDef map are defined below.
SIGNATURE_KEY_MAP = mode_keys.ModeKeyMap(**{
ModeKeys.PREDICT: tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
ModeKeys.TRAIN: unexported_constants.DEFAULT_TRAIN_SIGNATURE_DEF_KEY,
ModeKeys.TEST: unexported_constants.DEFAULT_EVAL_SIGNATURE_DEF_KEY})
# Default names used in the SignatureDef input map, which maps strings to
# TensorInfo protos.
SINGLE_FEATURE_DEFAULT_NAME = 'feature'
SINGLE_RECEIVER_DEFAULT_NAME = 'input'
SINGLE_LABEL_DEFAULT_NAME = 'label'
### Below utilities are specific to SavedModel exports.
def build_all_signature_defs(receiver_tensors,
export_outputs,
receiver_tensors_alternatives=None,
serving_only=True):
"""Build `SignatureDef`s for all export outputs.
Args:
receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
input nodes where this receiver expects to be fed by default. Typically,
this is a single placeholder expecting serialized `tf.Example` protos.
export_outputs: a dict of ExportOutput instances, each of which has
an as_signature_def instance method that will be called to retrieve
the signature_def for all export output tensors.
receiver_tensors_alternatives: a dict of string to additional
groups of receiver tensors, each of which may be a `Tensor` or a dict of
string to `Tensor`. These named receiver tensor alternatives generate
additional serving signatures, which may be used to feed inputs at
different points within the input receiver subgraph. A typical usage is
to allow feeding raw feature `Tensor`s *downstream* of the
tf.io.parse_example() op. Defaults to None.
serving_only: boolean; if true, resulting signature defs will only include
valid serving signatures. If false, all requested signatures will be
returned.
Returns:
signature_def representing all passed args.
Raises:
ValueError: if export_outputs is not a dict
"""
if not isinstance(receiver_tensors, dict):
receiver_tensors = {SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
if export_outputs is None or not isinstance(export_outputs, dict):
raise ValueError('export_outputs must be a dict and not'
'{}'.format(type(export_outputs)))
signature_def_map = {}
excluded_signatures = {}
for output_key, export_output in export_outputs.items():
signature_name = '{}'.format(output_key or 'None')
try:
signature = export_output.as_signature_def(receiver_tensors)
signature_def_map[signature_name] = signature
except ValueError as e:
excluded_signatures[signature_name] = str(e)
if receiver_tensors_alternatives:
for receiver_name, receiver_tensors_alt in (
receiver_tensors_alternatives.items()):
if not isinstance(receiver_tensors_alt, dict):
receiver_tensors_alt = {
SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors_alt
}
for output_key, export_output in export_outputs.items():
signature_name = '{}:{}'.format(receiver_name or 'None', output_key or
'None')
try:
signature = export_output.as_signature_def(receiver_tensors_alt)
signature_def_map[signature_name] = signature
except ValueError as e:
excluded_signatures[signature_name] = str(e)
_log_signature_report(signature_def_map, excluded_signatures)
# The above calls to export_output_lib.as_signature_def should return only
# valid signatures; if there is a validity problem, they raise a ValueError,
# in which case we exclude that signature from signature_def_map above.
# The is_valid_signature check ensures that the signatures produced are
# valid for serving, and acts as an additional sanity check for export
# signatures produced for serving. We skip this check for training and eval
# signatures, which are not intended for serving.
if serving_only:
signature_def_map = {
k: v
for k, v in signature_def_map.items()
if tf.compat.v1.saved_model.is_valid_signature(v)
}
return signature_def_map
_FRIENDLY_METHOD_NAMES = {
tf.saved_model.CLASSIFY_METHOD_NAME: 'Classify',
tf.saved_model.REGRESS_METHOD_NAME: 'Regress',
tf.saved_model.PREDICT_METHOD_NAME: 'Predict',
unexported_constants.SUPERVISED_TRAIN_METHOD_NAME: 'Train',
unexported_constants.SUPERVISED_EVAL_METHOD_NAME: 'Eval',
}
def _log_signature_report(signature_def_map, excluded_signatures):
"""Log a report of which signatures were produced."""
sig_names_by_method_name = collections.defaultdict(list)
# We'll collect whatever method_names are present, but also we want to make
# sure to output a line for each of the three standard methods even if they
# have no signatures.
for method_name in _FRIENDLY_METHOD_NAMES:
sig_names_by_method_name[method_name] = []
for signature_name, sig in signature_def_map.items():
sig_names_by_method_name[sig.method_name].append(signature_name)
# TODO(b/67733540): consider printing the full signatures, not just names
for method_name, sig_names in sig_names_by_method_name.items():
if method_name in _FRIENDLY_METHOD_NAMES:
method_name = _FRIENDLY_METHOD_NAMES[method_name]
logging.info('Signatures INCLUDED in export for {}: {}'.format(
method_name, sig_names if sig_names else 'None'))
if excluded_signatures:
logging.info('Signatures EXCLUDED from export because they cannot be '
'be served via TensorFlow Serving APIs:')
for signature_name, message in excluded_signatures.items():
logging.info('\'{}\' : {}'.format(signature_name, message))
if not signature_def_map:
logging.warning('Export includes no signatures!')
elif (tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY not in
signature_def_map):
logging.warning('Export includes no default signature!')
# When we create a timestamped directory, there is a small chance that the
# directory already exists because another process is also creating these
# directories. In this case we just wait one second to get a new timestamp and
# try again. If this fails several times in a row, then something is seriously
# wrong.
MAX_DIRECTORY_CREATION_ATTEMPTS = 10
def get_timestamped_export_dir(export_dir_base):
"""Builds a path to a new subdirectory within the base directory.
Each export is written into a new subdirectory named using the
current time. This guarantees monotonically increasing version
numbers even across multiple runs of the pipeline.
The timestamp used is the number of seconds since epoch UTC.
Args:
export_dir_base: A string containing a directory to write the exported
graph and checkpoints.
Returns:
The full path of the new subdirectory (which is not actually created yet).
Raises:
RuntimeError: if repeated attempts fail to obtain a unique timestamped
directory name.
"""
attempts = 0
while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS:
timestamp = int(time.time())
result_dir = os.path.join(
tf.compat.as_bytes(export_dir_base), tf.compat.as_bytes(str(timestamp)))
if not tf.compat.v1.gfile.Exists(result_dir):
# Collisions are still possible (though extremely unlikely): this
# directory is not actually created yet, but it will be almost
# instantly on return from this function.
return result_dir
time.sleep(1)
attempts += 1
logging.warning(
'Directory {} already exists; retrying (attempt {}/{})'.format(
tf.compat.as_str(result_dir), attempts,
MAX_DIRECTORY_CREATION_ATTEMPTS))
raise RuntimeError('Failed to obtain a unique export directory name after '
'{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS))
def get_temp_export_dir(timestamped_export_dir):
"""Builds a directory name based on the argument but starting with 'temp-'.
This relies on the fact that TensorFlow Serving ignores subdirectories of
the base directory that can't be parsed as integers.
Args:
timestamped_export_dir: the name of the eventual export directory, e.g.
/foo/bar/<timestamp>
Returns:
A sister directory prefixed with 'temp-', e.g. /foo/bar/temp-<timestamp>.
"""
(dirname, basename) = os.path.split(timestamped_export_dir)
if isinstance(basename, bytes):
str_name = basename.decode('utf-8')
else:
str_name = str(basename)
temp_export_dir = os.path.join(
tf.compat.as_bytes(dirname),
tf.compat.as_bytes('temp-{}'.format(str_name)))
return temp_export_dir
def export_outputs_for_mode(
mode, serving_export_outputs=None, predictions=None, loss=None,
metrics=None):
"""Util function for constructing a `ExportOutput` dict given a mode.
The returned dict can be directly passed to `build_all_signature_defs` helper
function as the `export_outputs` argument, used for generating a SignatureDef
map.
Args:
mode: A `ModeKeys` specifying the mode.
serving_export_outputs: Describes the output signatures to be exported to
`SavedModel` and used during serving. Should be a dict or None.
predictions: A dict of Tensors or single Tensor representing model
predictions. This argument is only used if serving_export_outputs is not
set.
loss: A dict of Tensors or single Tensor representing calculated loss.
metrics: A dict of (metric_value, update_op) tuples, or a single tuple.
metric_value must be a Tensor, and update_op must be a Tensor or Op
Returns:
Dictionary mapping the a key to an `tf.estimator.export.ExportOutput` object
The key is the expected SignatureDef key for the mode.
Raises:
ValueError: if an appropriate ExportOutput cannot be found for the mode.
"""
if mode not in SIGNATURE_KEY_MAP:
raise ValueError(
'Export output type not found for mode: {}. Expected one of: {}.\n'
'One likely error is that V1 Estimator Modekeys were somehow passed to '
'this function. Please ensure that you are using the new ModeKeys.'
.format(mode, SIGNATURE_KEY_MAP.keys()))
signature_key = SIGNATURE_KEY_MAP[mode]
if mode_keys.is_predict(mode):
return get_export_outputs(serving_export_outputs, predictions)
elif mode_keys.is_train(mode):
return {signature_key: export_output_lib.TrainOutput(
loss=loss, predictions=predictions, metrics=metrics)}
else:
return {signature_key: export_output_lib.EvalOutput(
loss=loss, predictions=predictions, metrics=metrics)}
def get_export_outputs(export_outputs, predictions):
"""Validate export_outputs or create default export_outputs.
Args:
export_outputs: Describes the output signatures to be exported to
`SavedModel` and used during serving. Should be a dict or None.
predictions: Predictions `Tensor` or dict of `Tensor`.
Returns:
Valid export_outputs dict
Raises:
TypeError: if export_outputs is not a dict or its values are not
ExportOutput instances.
"""
if export_outputs is None:
default_output = export_output_lib.PredictOutput(predictions)
export_outputs = {
tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: default_output}
if not isinstance(export_outputs, dict):
raise TypeError('export_outputs must be dict, given: {}'.format(
export_outputs))
for v in export_outputs.values():
if not isinstance(v, export_output_lib.ExportOutput):
raise TypeError(
'Values in export_outputs must be ExportOutput objects. '
'Given: {}'.format(export_outputs))
_maybe_add_default_serving_output(export_outputs)
return export_outputs
def _maybe_add_default_serving_output(export_outputs):
"""Add a default serving output to the export_outputs if not present.
Args:
export_outputs: Describes the output signatures to be exported to
`SavedModel` and used during serving. Should be a dict.
Returns:
export_outputs dict with default serving signature added if necessary
Raises:
ValueError: if multiple export_outputs were provided without a default
serving key.
"""
if len(export_outputs) == 1:
(key, value), = export_outputs.items()
if key != tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
export_outputs[
tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = value
if len(export_outputs) > 1:
if (tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
not in export_outputs):
raise ValueError(
'Multiple export_outputs were provided, but none of them is '
'specified as the default. Do this by naming one of them with '
'signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY.')
return export_outputs
# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/export_utils.py)
Functions
def build_all_signature_defs(receiver_tensors, export_outputs, receiver_tensors_alternatives=None, serving_only=True)
-
Build
SignatureDef
s for all export outputs.Args
receiver_tensors
- a
Tensor
, or a dict of string toTensor
, specifying input nodes where this receiver expects to be fed by default. Typically, this is a single placeholder expecting serializedtf.Example
protos. export_outputs
- a dict of ExportOutput instances, each of which has an as_signature_def instance method that will be called to retrieve the signature_def for all export output tensors.
receiver_tensors_alternatives
- a dict of string to additional
groups of receiver tensors, each of which may be a
Tensor
or a dict of string toTensor
. These named receiver tensor alternatives generate additional serving signatures, which may be used to feed inputs at different points within the input receiver subgraph. A typical usage is to allow feeding raw featureTensor
s downstream of the tf.io.parse_example() op. Defaults to None. serving_only
- boolean; if true, resulting signature defs will only include valid serving signatures. If false, all requested signatures will be returned.
Returns
signature_def representing all passed args.
Raises
ValueError
- if export_outputs is not a dict
Expand source code
def build_all_signature_defs(receiver_tensors, export_outputs, receiver_tensors_alternatives=None, serving_only=True): """Build `SignatureDef`s for all export outputs. Args: receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying input nodes where this receiver expects to be fed by default. Typically, this is a single placeholder expecting serialized `tf.Example` protos. export_outputs: a dict of ExportOutput instances, each of which has an as_signature_def instance method that will be called to retrieve the signature_def for all export output tensors. receiver_tensors_alternatives: a dict of string to additional groups of receiver tensors, each of which may be a `Tensor` or a dict of string to `Tensor`. These named receiver tensor alternatives generate additional serving signatures, which may be used to feed inputs at different points within the input receiver subgraph. A typical usage is to allow feeding raw feature `Tensor`s *downstream* of the tf.io.parse_example() op. Defaults to None. serving_only: boolean; if true, resulting signature defs will only include valid serving signatures. If false, all requested signatures will be returned. Returns: signature_def representing all passed args. Raises: ValueError: if export_outputs is not a dict """ if not isinstance(receiver_tensors, dict): receiver_tensors = {SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors} if export_outputs is None or not isinstance(export_outputs, dict): raise ValueError('export_outputs must be a dict and not' '{}'.format(type(export_outputs))) signature_def_map = {} excluded_signatures = {} for output_key, export_output in export_outputs.items(): signature_name = '{}'.format(output_key or 'None') try: signature = export_output.as_signature_def(receiver_tensors) signature_def_map[signature_name] = signature except ValueError as e: excluded_signatures[signature_name] = str(e) if receiver_tensors_alternatives: for receiver_name, receiver_tensors_alt in ( receiver_tensors_alternatives.items()): if not isinstance(receiver_tensors_alt, dict): receiver_tensors_alt = { SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors_alt } for output_key, export_output in export_outputs.items(): signature_name = '{}:{}'.format(receiver_name or 'None', output_key or 'None') try: signature = export_output.as_signature_def(receiver_tensors_alt) signature_def_map[signature_name] = signature except ValueError as e: excluded_signatures[signature_name] = str(e) _log_signature_report(signature_def_map, excluded_signatures) # The above calls to export_output_lib.as_signature_def should return only # valid signatures; if there is a validity problem, they raise a ValueError, # in which case we exclude that signature from signature_def_map above. # The is_valid_signature check ensures that the signatures produced are # valid for serving, and acts as an additional sanity check for export # signatures produced for serving. We skip this check for training and eval # signatures, which are not intended for serving. if serving_only: signature_def_map = { k: v for k, v in signature_def_map.items() if tf.compat.v1.saved_model.is_valid_signature(v) } return signature_def_map
def export_outputs_for_mode(mode, serving_export_outputs=None, predictions=None, loss=None, metrics=None)
-
Util function for constructing a
ExportOutput
dict given a mode.The returned dict can be directly passed to
build_all_signature_defs()
helper function as theexport_outputs
argument, used for generating a SignatureDef map.Args
mode
- A
ModeKeys
specifying the mode. serving_export_outputs
- Describes the output signatures to be exported to
SavedModel
and used during serving. Should be a dict or None. predictions
- A dict of Tensors or single Tensor representing model predictions. This argument is only used if serving_export_outputs is not set.
loss
- A dict of Tensors or single Tensor representing calculated loss.
metrics
- A dict of (metric_value, update_op) tuples, or a single tuple. metric_value must be a Tensor, and update_op must be a Tensor or Op
Returns
Dictionary mapping the a key to an
tf.estimator.export.ExportOutput
object The key is the expected SignatureDef key for the mode.Raises
ValueError
- if an appropriate ExportOutput cannot be found for the mode.
Expand source code
def export_outputs_for_mode( mode, serving_export_outputs=None, predictions=None, loss=None, metrics=None): """Util function for constructing a `ExportOutput` dict given a mode. The returned dict can be directly passed to `build_all_signature_defs` helper function as the `export_outputs` argument, used for generating a SignatureDef map. Args: mode: A `ModeKeys` specifying the mode. serving_export_outputs: Describes the output signatures to be exported to `SavedModel` and used during serving. Should be a dict or None. predictions: A dict of Tensors or single Tensor representing model predictions. This argument is only used if serving_export_outputs is not set. loss: A dict of Tensors or single Tensor representing calculated loss. metrics: A dict of (metric_value, update_op) tuples, or a single tuple. metric_value must be a Tensor, and update_op must be a Tensor or Op Returns: Dictionary mapping the a key to an `tf.estimator.export.ExportOutput` object The key is the expected SignatureDef key for the mode. Raises: ValueError: if an appropriate ExportOutput cannot be found for the mode. """ if mode not in SIGNATURE_KEY_MAP: raise ValueError( 'Export output type not found for mode: {}. Expected one of: {}.\n' 'One likely error is that V1 Estimator Modekeys were somehow passed to ' 'this function. Please ensure that you are using the new ModeKeys.' .format(mode, SIGNATURE_KEY_MAP.keys())) signature_key = SIGNATURE_KEY_MAP[mode] if mode_keys.is_predict(mode): return get_export_outputs(serving_export_outputs, predictions) elif mode_keys.is_train(mode): return {signature_key: export_output_lib.TrainOutput( loss=loss, predictions=predictions, metrics=metrics)} else: return {signature_key: export_output_lib.EvalOutput( loss=loss, predictions=predictions, metrics=metrics)}
def get_export_outputs(export_outputs, predictions)
-
Validate export_outputs or create default export_outputs.
Args
export_outputs
- Describes the output signatures to be exported to
SavedModel
and used during serving. Should be a dict or None. predictions
- Predictions
Tensor
or dict ofTensor
.
Returns
Valid export_outputs dict
Raises
TypeError
- if export_outputs is not a dict or its values are not ExportOutput instances.
Expand source code
def get_export_outputs(export_outputs, predictions): """Validate export_outputs or create default export_outputs. Args: export_outputs: Describes the output signatures to be exported to `SavedModel` and used during serving. Should be a dict or None. predictions: Predictions `Tensor` or dict of `Tensor`. Returns: Valid export_outputs dict Raises: TypeError: if export_outputs is not a dict or its values are not ExportOutput instances. """ if export_outputs is None: default_output = export_output_lib.PredictOutput(predictions) export_outputs = { tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: default_output} if not isinstance(export_outputs, dict): raise TypeError('export_outputs must be dict, given: {}'.format( export_outputs)) for v in export_outputs.values(): if not isinstance(v, export_output_lib.ExportOutput): raise TypeError( 'Values in export_outputs must be ExportOutput objects. ' 'Given: {}'.format(export_outputs)) _maybe_add_default_serving_output(export_outputs) return export_outputs
def get_temp_export_dir(timestamped_export_dir)
-
Builds a directory name based on the argument but starting with 'temp-'.
This relies on the fact that TensorFlow Serving ignores subdirectories of the base directory that can't be parsed as integers.
Args
timestamped_export_dir
- the name of the eventual export directory, e.g.
/foo/bar/
Returns
A sister directory prefixed with 'temp-', e.g. /foo/bar/temp-
. Expand source code
def get_temp_export_dir(timestamped_export_dir): """Builds a directory name based on the argument but starting with 'temp-'. This relies on the fact that TensorFlow Serving ignores subdirectories of the base directory that can't be parsed as integers. Args: timestamped_export_dir: the name of the eventual export directory, e.g. /foo/bar/<timestamp> Returns: A sister directory prefixed with 'temp-', e.g. /foo/bar/temp-<timestamp>. """ (dirname, basename) = os.path.split(timestamped_export_dir) if isinstance(basename, bytes): str_name = basename.decode('utf-8') else: str_name = str(basename) temp_export_dir = os.path.join( tf.compat.as_bytes(dirname), tf.compat.as_bytes('temp-{}'.format(str_name))) return temp_export_dir
def get_timestamped_export_dir(export_dir_base)
-
Builds a path to a new subdirectory within the base directory.
Each export is written into a new subdirectory named using the current time. This guarantees monotonically increasing version numbers even across multiple runs of the pipeline. The timestamp used is the number of seconds since epoch UTC.
Args
export_dir_base
- A string containing a directory to write the exported graph and checkpoints.
Returns
The full path of the new subdirectory (which is not actually created yet).
Raises
RuntimeError
- if repeated attempts fail to obtain a unique timestamped directory name.
Expand source code
def get_timestamped_export_dir(export_dir_base): """Builds a path to a new subdirectory within the base directory. Each export is written into a new subdirectory named using the current time. This guarantees monotonically increasing version numbers even across multiple runs of the pipeline. The timestamp used is the number of seconds since epoch UTC. Args: export_dir_base: A string containing a directory to write the exported graph and checkpoints. Returns: The full path of the new subdirectory (which is not actually created yet). Raises: RuntimeError: if repeated attempts fail to obtain a unique timestamped directory name. """ attempts = 0 while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS: timestamp = int(time.time()) result_dir = os.path.join( tf.compat.as_bytes(export_dir_base), tf.compat.as_bytes(str(timestamp))) if not tf.compat.v1.gfile.Exists(result_dir): # Collisions are still possible (though extremely unlikely): this # directory is not actually created yet, but it will be almost # instantly on return from this function. return result_dir time.sleep(1) attempts += 1 logging.warning( 'Directory {} already exists; retrying (attempt {}/{})'.format( tf.compat.as_str(result_dir), attempts, MAX_DIRECTORY_CREATION_ATTEMPTS)) raise RuntimeError('Failed to obtain a unique export directory name after ' '{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS))