def __init__(self, args):
self._shard_name_template = None if args.shard_files else ''
self._output_format = args.output_format
self._output_dir = args.output_dir
# Get the BQ schema if csv.
if self._output_format == 'csv':
from tensorflow.python.saved_model import tag_constants
from tensorflow.contrib.session_bundle import bundle_shim
from tensorflow.core.framework import types_pb2
session, meta_graph = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
args.trained_model_dir, tags=[tag_constants.SERVING])
signature = meta_graph.signature_def['serving_default']
self._schema = []
for friendly_name in sorted(signature.outputs):
tensor_info_proto = signature.outputs[friendly_name]
# TODO(brandondutra): Could dtype be DT_INVALID?
# Consider getting the dtype from the graph via
# session.graph.get_tensor_by_name(tensor_info_proto.name).dtype)
dtype = tensor_info_proto.dtype
if dtype == types_pb2.DT_FLOAT or dtype == types_pb2.DT_DOUBLE:
bq_type = 'FLOAT'
elif dtype == types_pb2.DT_INT32 or dtype == types_pb2.DT_INT64:
bq_type = 'INTEGER'
else:
bq_type = 'STRING'
self._schema.append({'mode': 'NULLABLE',
'name': friendly_name,
'type': bq_type})
session.close()
def __init__(self, args):
self._shard_name_template = None if args.shard_files else ''
self._output_format = args.output_format
self._output_dir = args.output_dir
# Get the BQ schema if csv.
if self._output_format == 'csv':
from tensorflow.python.saved_model import tag_constants
from tensorflow.contrib.session_bundle import bundle_shim
from tensorflow.core.framework import types_pb2
session, meta_graph = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
args.trained_model_dir, tags=[tag_constants.SERVING])
signature = meta_graph.signature_def['serving_default']
self._schema = []
for friendly_name in sorted(signature.outputs):
tensor_info_proto = signature.outputs[friendly_name]
# TODO(brandondutra): Could dtype be DT_INVALID?
# Consider getting the dtype from the graph via
# session.graph.get_tensor_by_name(tensor_info_proto.name).dtype)
dtype = tensor_info_proto.dtype
if dtype == types_pb2.DT_FLOAT or dtype == types_pb2.DT_DOUBLE:
bq_type = 'FLOAT'
elif dtype == types_pb2.DT_INT32 or dtype == types_pb2.DT_INT64:
bq_type = 'INTEGER'
else:
bq_type = 'STRING'
self._schema.append({'mode': 'NULLABLE',
'name': friendly_name,
'type': bq_type})
session.close()
def load_model(model_path, config=None, **unused_kwargs):
"""Loads the model at the specified path.
Args:
model_path: the path to either session_bundle or SavedModel
config: tf.ConfigProto containing session configuration options.
unused_kwargs: kwargs for compatiblity purpose.
Returns:
A pair of (Session, SignatureDef) objects.
Raises:
PredictionError: if the model could not be loaded.
"""
# Ideally, we could just always use bundle_shim to load legacy and
# regular graphs. However, bundle_shim and supporting functions are
# only available on recent versions of TF (~0.12). It's not even
# really possible to detect whether or not we're going to be able
# to use these functions, so in true Python function, it's better
# to ask forgiveness than permission...we try to import bundle_shim,
# which may fail, then we try to use bundle_shim, which may also fail
# for legacy graphs. In other failure case, we back off to our older
# custom session_bundle implementation.
try:
from tensorflow.contrib.session_bundle import bundle_shim # pylint: disable=g-import-not-at-top
from tensorflow.python.saved_model import tag_constants # pylint: disable=g-import-not-at-top
# We expect that the customer will export saved model and use
# tag_constants.SERVING for serving graph. This assumption also extends to
# model server.
session, meta_graph = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
model_path, tags=[tag_constants.SERVING], config=config))
except Exception: # pylint: disable=broad-except
session, meta_graph = session_bundle.load_session_bundle_from_path(
model_path, config=config)
if session is None:
raise PredictionError(PredictionError.FAILED_TO_LOAD_MODEL,
"Could not load model from %s" % model_path)
# Before the SavedModel spec came into existence the inputs and outputs
# of a model were specified using TensorFlow collections. Check if this model
# uses that spec.
graph = session.graph
collection_keys = graph.get_all_collection_keys()
if INPUTS_KEY in collection_keys and OUTPUTS_KEY in collection_keys:
signature = _get_legacy_signature(graph)
else:
# Otherwise, use (possibly upgraded from session_bundle) SavedModel.
signature = _get_signature_from_meta_graph(graph, meta_graph)
return session, signature
def _load_tf_model(model_dir):
from tensorflow.python.saved_model import tag_constants
from tensorflow.contrib.session_bundle import bundle_shim
model_dir = os.path.join(model_dir, 'model')
session, meta_graph = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
model_dir, tags=[tag_constants.SERVING])
signature = meta_graph.signature_def['serving_default']
inputs = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.inputs.items()}
outputs = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.outputs.items()}
return session, inputs, outputs
def start_bundle(self, element=None):
from tensorflow.python.saved_model import tag_constants
from tensorflow.contrib.session_bundle import bundle_shim
self._session, meta_graph = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
self._trained_model_dir, tags=[tag_constants.SERVING])
signature = meta_graph.signature_def['serving_default']
# get the mappings between aliases and tensor names
# for both inputs and outputs
self._input_alias_map = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.inputs.items()}
self._output_alias_map = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.outputs.items()}
self._aliases, self._tensor_names = zip(*self._output_alias_map.items())
def start_bundle(self, element=None):
from tensorflow.python.saved_model import tag_constants
from tensorflow.contrib.session_bundle import bundle_shim
self._session, meta_graph = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
self._trained_model_dir, tags=[tag_constants.SERVING])
signature = meta_graph.signature_def['serving_default']
# get the mappings between aliases and tensor names
# for both inputs and outputs
self._input_alias_map = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.inputs.items()}
self._output_alias_map = {friendly_name: tensor_info_proto.name
for (friendly_name, tensor_info_proto) in signature.outputs.items()}
self._aliases, self._tensor_names = zip(*self._output_alias_map.items())
def load_graph(output_graph_path):
#sess, meta_graph_def = session_bundle.load_session_bundle_from_path(output_graph_path)
sess, meta_graph_def = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
output_graph_path)
with sess.as_default():
collection_def = meta_graph_def.collection_def
signatures_any = collection_def[
constants.SIGNATURES_KEY].any_list.value
signatures = manifest_pb2.Signatures()
signatures_any[0].Unpack(signatures)
default_signature = signatures.default_signature
input_name = default_signature.classification_signature.input.tensor_name
output_name = default_signature.classification_signature.scores.tensor_name
classes = default_signature.classification_signature.classes.tensor_name
classes = sess.run(sess.graph.get_tensor_by_name(classes))
return (sess, input_name, output_name, classes)
def testSavedModelBasic(self):
base_path = test.test_src_dir_path(SAVED_MODEL_PATH)
ops.reset_default_graph()
sess, meta_graph_def = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path,
tags=[tag_constants.SERVING],
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})))
self.assertTrue(sess)
# Check basic signature def property.
signature_def = meta_graph_def.signature_def
self.assertEqual(signature_def["regress_x_to_y"].method_name,
signature_constants.REGRESS_METHOD_NAME)
with sess.as_default():
output1 = sess.run(["filename_tensor:0"])
self.assertEqual([compat.as_bytes("foo.txt")], output1)
def load_graph(self, output_graph_path):
#sess, meta_graph_def = session_bundle.load_session_bundle_from_path(
# output_graph_path)
sess, meta_graph_def = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(output_graph_path)
#sess, meta_graph_def = session_bundle.load_session_bundle_from_path(
# output_graph_path)
with sess.as_default():
collection_def = meta_graph_def.collection_def
signatures_any = collection_def[
constants.SIGNATURES_KEY].any_list.value
signatures = manifest_pb2.Signatures()
signatures_any[0].Unpack(signatures)
default_signature = signatures.default_signature
input_name = default_signature.classification_signature.input.tensor_name
output_name = default_signature.classification_signature.scores.tensor_name
classes = default_signature.classification_signature.classes.tensor_name
classes = sess.run(sess.graph.get_tensor_by_name(classes))
return (sess, input_name, output_name,classes)
def testSavedModelBasic(self):
base_path = tf.test.test_src_dir_path(SAVED_MODEL_PATH)
tf.reset_default_graph()
sess, meta_graph_def = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path,
tags=[tag_constants.SERVING],
target="",
config=tf.ConfigProto(device_count={"CPU": 2})))
self.assertTrue(sess)
# Check basic signature def property.
signature_def = meta_graph_def.signature_def
self.assertEqual(len(signature_def), 2)
self.assertEqual(
signature_def[signature_constants.REGRESS_METHOD_NAME].method_name,
signature_constants.REGRESS_METHOD_NAME)
signature = signature_def["tensorflow/serving/regress"]
asset_path = os.path.join(base_path, saved_model_constants.ASSETS_DIRECTORY)
with sess.as_default():
output1 = sess.run(["filename_tensor:0"])
self.assertEqual(["foo.txt"], output1)
def testLegacyBasic(self):
base_path = tf.test.test_src_dir_path(SESSION_BUNDLE_PATH)
tf.reset_default_graph()
sess, meta_graph_def = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path,
tags=[""],
target="",
config=tf.ConfigProto(device_count={"CPU": 2})))
self.assertTrue(sess)
asset_path = os.path.join(base_path, constants.ASSETS_DIRECTORY)
with sess.as_default():
path1, path2 = sess.run(["filename1:0", "filename2:0"])
self.assertEqual(
compat.as_bytes(os.path.join(asset_path, "hello1.txt")), path1)
self.assertEqual(
compat.as_bytes(os.path.join(asset_path, "hello2.txt")), path2)
collection_def = meta_graph_def.collection_def
signatures_any = collection_def[constants.SIGNATURES_KEY].any_list.value
self.assertEqual(len(signatures_any), 1)
def testLegacyBasic(self):
base_path = test.test_src_dir_path(SESSION_BUNDLE_PATH)
ops.reset_default_graph()
sess, meta_graph_def = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path,
tags=[""],
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})))
self.assertTrue(sess)
asset_path = os.path.join(base_path, constants.ASSETS_DIRECTORY)
with sess.as_default():
path1, path2 = sess.run(["filename1:0", "filename2:0"])
self.assertEqual(
compat.as_bytes(os.path.join(asset_path, "hello1.txt")), path1)
self.assertEqual(
compat.as_bytes(os.path.join(asset_path, "hello2.txt")), path2)
collection_def = meta_graph_def.collection_def
signatures_any = collection_def[constants.SIGNATURES_KEY].any_list.value
self.assertEqual(len(signatures_any), 1)
def testSavedModelBasic(self):
base_path = tf.test.test_src_dir_path(SAVED_MODEL_PATH)
tf.reset_default_graph()
sess, meta_graph_def = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path,
tags=[tag_constants.SERVING],
target="",
config=tf.ConfigProto(device_count={"CPU": 2})))
self.assertTrue(sess)
# Check basic signature def property.
signature_def = meta_graph_def.signature_def
self.assertEqual(len(signature_def), 2)
self.assertEqual(
signature_def[signature_constants.REGRESS_METHOD_NAME].method_name,
signature_constants.REGRESS_METHOD_NAME)
signature = signature_def["tensorflow/serving/regress"]
asset_path = os.path.join(base_path,
saved_model_constants.ASSETS_DIRECTORY)
with sess.as_default():
output1 = sess.run(["filename_tensor:0"])
self.assertEqual(["foo.txt"], output1)
def testBadPath(self):
base_path = tf.test.test_src_dir_path("/no/such/a/dir")
tf.reset_default_graph()
with self.assertRaises(RuntimeError) as cm:
_, _ = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path)
def testBadPath(self):
base_path = tf.test.test_src_dir_path("/no/such/a/dir")
tf.reset_default_graph()
with self.assertRaises(RuntimeError) as cm:
_, _ = bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
base_path)
def make_tensor_func_from_saved_model(model_dir,
tags,
signature_name=None,
input_keys_in_signature=None,
output_keys_in_signature=None):
"""Create a tensor-in-tensor-out function as a transform used in tft.map.
When tft.map is called with this function as first parameter, the second
parameter (input columns) should match the `input_keys_in_signature`
in their orders.
Args:
model_dir: A path containing a saved model.
tags: The tags specifying which metagraph to load from the saved model.
signature_name: Specify signature of the loaded model. The default value
None can be used if there is only one signature in the MetaGraphDef.
input_keys_in_signature: A list of strings which should match the inputs
in the signature of the saved model. The purpose of this parameter is to
specify the order of the input columns passed to tft.map when called
with the returned tensor_fn. The default value None can be used if there
is only one input.
output_keys_in_signature: A list of strings which should be a subset of
the outputs in the signature of the saved model. The returned tensor_fn
will return the corresponding tensors, in the same order. The default
value None can be used if there is only one output from signature.
Returns:
A tensor-in-tensor-out function which can be used in tft.map.
Raises:
ValueError: If
`signature_name` is None but the saved model contains multiple signature, or
`input_keys_in_signature` do not match the signature inputs, or
`output_keys_in_signature` is not a subset of the signature outputs, or
the metagraph from saved model contains TABLE_INITIALIZERS operations.
"""
# Load model, get graph, inputs and outputs.
loaded_graph = tf.Graph()
with loaded_graph.as_default():
session, meta_graph = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
model_dir, tags=tags))
if signature_name:
signature = meta_graph.signature_def[signature_name]
elif len(meta_graph.signature_def) > 1:
raise ValueError(
'The saved model contains multiple signatures "%s". Specify a '
'signature_name.' % ','.join(meta_graph.signature_def.keys()))
else:
signature = meta_graph.signature_def.values()[0]
inputs = {
key: tensor_info_proto.name
for (key, tensor_info_proto) in signature.inputs.items()
}
outputs = {
key: tensor_info_proto.name
for (key, tensor_info_proto) in signature.outputs.items()
}
# Get input tensor names.
if input_keys_in_signature is not None:
if set(input_keys_in_signature) != set(inputs.keys()):
raise ValueError(
'keys in input logical names do not match inputs of saved model. '
+
'Model signature has "%s" but input logical names has "%s".' %
(','.join(inputs.keys()), ','.join(input_keys_in_signature)))
input_tensor_names = [inputs[key] for key in input_keys_in_signature]
else:
if len(inputs) > 1:
raise ValueError(
'The signature from saved model contains multiple inputs "%s". '
'Input logical names are required.' % ','.join(inputs.keys()))
else:
input_tensor_names = [inputs.values()[0]]
# Get output tensor names.
if output_keys_in_signature:
if not set(output_keys_in_signature) <= set(outputs.keys()):
raise ValueError(
'output names are not a subset of outputs of saved model. ' +
'output names has "%s" but model signature has "%s".' %
(','.join(output_keys_in_signature), ','.join(outputs.keys())))
output_tensor_names = [
outputs[key] for key in output_keys_in_signature
]
else:
if len(outputs) > 1:
raise ValueError(
'The signature from saved model contains multiple outputs "%s". '
'Output names are required.' % ','.join(outputs.keys()))
output_tensor_names = [outputs.values()[0]]
if tf.get_collection(tf.GraphKeys.TABLE_INITIALIZERS):
raise ValueError(
'Models with table init ops in metagraph are not supported.')
# Convert_variables_to_constants() requires op name.
output_op_names = [
loaded_graph.get_tensor_by_name(x).op.name for x in output_tensor_names
]
constant_graph_def = tf.graph_util.convert_variables_to_constants(
session, loaded_graph.as_graph_def(), output_op_names)
def tensor_fn(*si):
input_name_to_tensor_map = dict(zip(input_tensor_names, si))
output_tensors = tf.import_graph_def(
constant_graph_def,
input_map=input_name_to_tensor_map,
return_elements=output_tensor_names)
return output_tensors[0]
return tensor_fn
def apply_saved_model(model_dir,
inputs,
tags,
signature_name=None,
output_keys_in_signature=None):
"""Applies a SavedModel to some `Tensor`s.
Applies a SavedModel to `inputs`. The SavedModel is specified with
`model_dir`, `tags` and `signature_name`. Note that the SavedModel will be
converted to an all-constants graph.
Args:
model_dir: A path containing a SavedModel.
inputs: A dict whose keys are the names from the input signature and whose
values are `Tensor`s. If there is only one input in the model's input
signature then `inputs` can be a single `Tensor`.
tags: The tags specifying which metagraph to load from the SavedModel.
signature_name: Specify signature of the loaded model. The default value
None can be used if there is only one signature in the MetaGraphDef.
output_keys_in_signature: A list of strings which should be a subset of
the outputs in the signature of the SavedModel. The returned `Tensor`s
will correspond to specified output `Tensor`s, in the same order. The
default value None can be used if there is only one output from
signature.
Returns:
A `Tensor` or list of `Tensor`s representing the application of the
SavedModel.
Raises:
ValueError: if
`inputs` is invalid type, or
`signature_name` is None but the SavedModel contains multiple signature, or
`inputs` do not match the signature inputs, or
`output_keys_in_signature` is not a subset of the signature outputs.
"""
# Load model, get graph, inputs and outputs.
loaded_graph = tf.Graph()
loaded_initializer_op_names = []
with loaded_graph.as_default():
session, meta_graph = (
bundle_shim.load_session_bundle_or_saved_model_bundle_from_path(
model_dir, tags=tags))
loaded_initializer_op_names = [
op.name
for op in tf.get_collection(tf.GraphKeys.TABLE_INITIALIZERS)
]
if signature_name:
signature = meta_graph.signature_def[signature_name]
elif len(meta_graph.signature_def) > 1:
raise ValueError(
'The SavedModel contains multiple signatures (%r) but signature_name '
'was not specified.' % (meta_graph.signature_def.keys(), ))
else:
signature = next(six.itervalues(meta_graph.signature_def))
# Generate mapping from tensors in the graph to the input tensors.
if isinstance(inputs, dict):
if set(signature.inputs.keys()) != set(inputs.keys()):
raise ValueError(
'The keys in `inputs` (%r) do not match inputs of the SavedModel '
'(%r).' % (inputs.keys(), signature.inputs.keys()))
input_name_to_tensor_map = {
signature.inputs[key].name: inputs[key]
for key in inputs.keys()
}
elif len(signature.inputs) != 1:
raise ValueError(
'The SavedModel does not have exactly one input (had inputs %r) but '
'`inputs` was not a dict.' % (signature.inputs.keys(), ))
else:
input_name_to_tensor_map = {
next(six.itervalues(signature.inputs)).name: inputs
}
# Get output tensor names.
if output_keys_in_signature:
if not set(output_keys_in_signature) <= set(signature.outputs.keys()):
raise ValueError(
'output_keys_in_signature (%r) is not a subset of outputs of the '
'SavedModel (%r).' %
(output_keys_in_signature, signature.outputs.keys()))
output_tensor_names = [
signature.outputs[key].name for key in output_keys_in_signature
]
output_single_tensor = False
elif len(signature.outputs) != 1:
raise ValueError(
'The SavedModel does not have exactly one output (had outputs %r) but '
'output_keys_in_signature was not specified.' %
(signature.outputs.keys(), ))
else:
output_tensor_names = [next(six.itervalues(signature.outputs)).name]
output_single_tensor = True
# Convert_variables_to_constants() requires op name.
output_op_names = [
loaded_graph.get_tensor_by_name(tensor_name).op.name
for tensor_name in output_tensor_names
]
constant_graph_def = tf.graph_util.convert_variables_to_constants(
session, loaded_graph.as_graph_def(),
output_op_names + loaded_initializer_op_names)
returned_elements = tf.import_graph_def(
constant_graph_def,
input_map=input_name_to_tensor_map,
return_elements=output_tensor_names + loaded_initializer_op_names)
returned_output_tensors = returned_elements[:len(output_tensor_names)]
returned_initializer_ops = returned_elements[len(output_tensor_names):]
for initializer_op in returned_initializer_ops:
tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, initializer_op)
if output_single_tensor:
assert len(output_tensor_names) == 1
return returned_output_tensors[0]
else:
return returned_output_tensors
