def testCreationOfAssetsKeyCollectionIsDeterministic(self):
tmp_asset_dir = os.path.join(self.get_temp_dir(), "assets")
tf.gfile.MakeDirs(tmp_asset_dir)
filenames = [
os.path.join(tmp_asset_dir, "file%d.txt" % n) for n in range(10)
]
for filename in filenames:
_write_string_to_file(filename, "I am file %s" % filename)
with tf.Graph().as_default() as graph:
assets = [tf.constant(f, name=os.path.basename(f)) for f in filenames]
for asset in assets:
graph.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, asset)
saved_model_lib.add_signature("default", {}, {"default": assets[0]})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
saved_model_proto = copy.deepcopy(handler._proto)
export_dir = os.path.join(self.get_temp_dir(), "assets_key_test")
saved_model_lib._make_assets_key_collection(saved_model_proto, export_dir)
meta_graph = list(saved_model_proto.meta_graphs)[0]
asset_tensor_names = []
for asset_any_proto in meta_graph.collection_def[
tf.saved_model.constants.ASSETS_KEY].any_list.value:
asset_proto = meta_graph_pb2.AssetFileDef()
asset_any_proto.Unpack(asset_proto)
asset_tensor_names.append(asset_proto.tensor_info.name)
self.assertEqual(asset_tensor_names, sorted(asset_tensor_names))
def testWithMultipleAssetsWithSameBasename(self):
tmp_asset_dir = os.path.join(self.get_temp_dir(), "asset")
file_a = os.path.join(tmp_asset_dir, "a", "hello.txt")
file_b = os.path.join(tmp_asset_dir, "b", "hello.txt")
tf_v1.gfile.MakeDirs(os.path.dirname(file_a))
tf_v1.gfile.MakeDirs(os.path.dirname(file_b))
_write_string_to_file(file_a, "hello A")
_write_string_to_file(file_b, "hello B")
with tf.Graph().as_default() as graph:
asset_a = tf.constant(file_a, name="file_a")
asset_b = tf.constant(file_b, name="file_b")
graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS, asset_a)
graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS, asset_b)
saved_model_lib.add_signature("default", {}, {"default": asset_a})
export_dir = os.path.join(self.get_temp_dir(), "exported")
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
handler.export(export_dir)
tf_v1.gfile.DeleteRecursively(tmp_asset_dir)
loaded_handler = saved_model_lib.load(export_dir)
with _instantiate_meta_graph(loaded_handler).as_default():
with tf_v1.Session() as sess:
self.assertEqual(_read_file_to_string(sess.run("file_a:0")),
"hello A")
self.assertEqual(_read_file_to_string(sess.run("file_b:0")),
"hello B")
def add_signature(name=None, inputs=None, outputs=None):
"""Adds a signature to the module definition.
NOTE: This must be called within a `module_fn` that is defining a Module.
Args:
name: Signature name as a string. If omitted, it is interpreted as 'default'
and is the signature used when `Module.__call__` `signature` is not
specified.
inputs: A dict from input name to Tensor or SparseTensor to feed when
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
outputs: A dict from output name to Tensor or SparseTensor to return from
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
Raises:
ValueError: if the arguments are invalid.
"""
if not name:
name = "default"
if inputs is None:
inputs = {}
if outputs is None:
outputs = {}
if not isinstance(inputs, dict):
inputs = {"default": inputs}
if not isinstance(outputs, dict):
outputs = {"default": outputs}
err = find_signature_input_colocation_error(name, inputs)
if err: raise ValueError(err)
saved_model_lib.add_signature(name, inputs, outputs)
def add_signature(name=None, inputs=None, outputs=None):
"""Adds a signature to the module definition.
NOTE: This must be called within a `module_fn` that is defining a Module.
Args:
name: Signature name as a string. If omitted, it is interpreted as 'default'
and is the signature used when `Module.__call__` `signature` is not
specified.
inputs: A dict from input name to Tensor or SparseTensor to feed when
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
outputs: A dict from output name to Tensor or SparseTensor to return from
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
Raises:
ValueError: if the arguments are invalid.
"""
if not name:
name = "default"
if inputs is None:
inputs = {}
if outputs is None:
outputs = {}
if not isinstance(inputs, dict):
inputs = {"default": inputs}
if not isinstance(outputs, dict):
outputs = {"default": outputs}
err = find_signature_input_colocation_error(name, inputs)
if err: raise ValueError(err)
saved_model_lib.add_signature(name, inputs, outputs)
def testAssets(self):
original_asset_file = os.path.join(self.get_temp_dir(), "hello.txt")
_write_string_to_file(original_asset_file, "hello world")
with tf.Graph().as_default() as graph:
asset_tensor = tf.constant(original_asset_file, name="file")
graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS,
asset_tensor)
saved_model_lib.add_signature("default", {},
{"default": asset_tensor})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
export_dir = os.path.join(self.get_temp_dir(), "exported")
handler.export(export_dir)
# Check that asset file got written to the expected place:
exported_asset_file = os.path.join(export_dir, "assets", "hello.txt")
self.assertTrue(tf_v1.gfile.Exists(exported_asset_file))
loaded_handler = saved_model_lib.load(export_dir)
with _instantiate_meta_graph(loaded_handler).as_default():
with tf_v1.Session() as sess:
self.assertEqual(sess.run("file:0"),
tf.compat.as_bytes(exported_asset_file))
def testCreationOfAssetsKeyCollectionIsDeterministic(self):
tmp_asset_dir = os.path.join(self.get_temp_dir(), "assets")
tf_v1.gfile.MakeDirs(tmp_asset_dir)
filenames = [
os.path.join(tmp_asset_dir, "file%d.txt" % n) for n in range(10)
]
for filename in filenames:
_write_string_to_file(filename, "I am file %s" % filename)
with tf.Graph().as_default() as graph:
assets = [
tf.constant(f, name=os.path.basename(f)) for f in filenames
]
for asset in assets:
graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS, asset)
saved_model_lib.add_signature("default", {},
{"default": assets[0]})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
saved_model_proto = copy.deepcopy(handler._proto)
export_dir = os.path.join(self.get_temp_dir(), "assets_key_test")
saved_model_lib._make_assets_key_collection(saved_model_proto,
export_dir)
meta_graph = list(saved_model_proto.meta_graphs)[0]
asset_tensor_names = []
for asset_any_proto in meta_graph.collection_def[
tf_v1.saved_model.constants.ASSETS_KEY].any_list.value:
asset_proto = meta_graph_pb2.AssetFileDef()
asset_any_proto.Unpack(asset_proto)
asset_tensor_names.append(asset_proto.tensor_info.name)
self.assertEqual(asset_tensor_names, sorted(asset_tensor_names))
def testWithMultipleAssetsWithSameBasename(self):
tmp_asset_dir = os.path.join(self.get_temp_dir(), "asset")
file_a = os.path.join(tmp_asset_dir, "a", "hello.txt")
file_b = os.path.join(tmp_asset_dir, "b", "hello.txt")
tf.gfile.MakeDirs(os.path.dirname(file_a))
tf.gfile.MakeDirs(os.path.dirname(file_b))
_write_string_to_file(file_a, "hello A")
_write_string_to_file(file_b, "hello B")
with tf.Graph().as_default() as graph:
asset_a = tf.constant(file_a, name="file_a")
asset_b = tf.constant(file_b, name="file_b")
graph.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, asset_a)
graph.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, asset_b)
saved_model_lib.add_signature("default", {}, {"default": asset_a})
export_dir = os.path.join(self.get_temp_dir(), "exported")
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
handler.export(export_dir)
tf.gfile.DeleteRecursively(tmp_asset_dir)
loaded_handler = saved_model_lib.load(export_dir)
with _instantiate_meta_graph(loaded_handler).as_default():
with tf.Session() as sess:
self.assertEqual(_read_file_to_string(sess.run("file_a:0")), "hello A")
self.assertEqual(_read_file_to_string(sess.run("file_b:0")), "hello B")
def testSignatureImplementationIsInvisible(self):
with tf.Graph().as_default() as graph:
saved_model_lib.add_signature("test", {}, {})
self.assertEqual(graph.get_all_collection_keys(), [])
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
meta_graph, = handler.meta_graphs
self.assertEqual(len(meta_graph.collection_def), 0)
self.assertEqual(len(meta_graph.signature_def), 1)
def testSignatureImplementationIsInvisible(self):
with tf.Graph().as_default() as graph:
saved_model_lib.add_signature("test", {}, {})
self.assertEqual(graph.get_all_collection_keys(), [])
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
meta_graph, = handler.meta_graphs
self.assertEqual(len(meta_graph.collection_def), 0)
self.assertEqual(len(meta_graph.signature_def), 1)
def testEmptyCollectionsDoNotShowUpInMetaGraphDef(self):
with tf.Graph().as_default() as graph:
tf.Variable("name")
self.assertEqual(len(graph.get_all_collection_keys()), 2)
for collection_key in graph.get_all_collection_keys():
del graph.get_collection_ref(collection_key)[:]
saved_model_lib.add_signature("default", {}, {"default": tf.constant(1)})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
meta_graph, = handler.meta_graphs
self.assertEqual(len(meta_graph.collection_def), 0)
def testTags(self):
with tf.Graph().as_default() as graph:
saved_model_lib.add_signature("default", {}, {"default": tf.constant(1)})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph, ["tag1"])
handler.add_graph_copy(graph, ["tag1", "tag2"])
self.assertAllEqual(sorted(handler.get_tags()),
sorted([set(["tag1"]), set(["tag1", "tag2"])]))
self.assertIsNotNone(handler.get_meta_graph_copy(["tag1"]))
self.assertIsNotNone(handler.get_meta_graph_copy(["tag2", "tag1"]))
with self.assertRaises(KeyError):
handler.get_meta_graph_copy(["tag2"])
def testEmptyCollectionsDoNotShowUpInMetaGraphDef(self):
with tf.Graph().as_default() as graph:
tf.Variable("name")
self.assertEqual(len(graph.get_all_collection_keys()), 2)
for collection_key in graph.get_all_collection_keys():
del graph.get_collection_ref(collection_key)[:]
saved_model_lib.add_signature("default", {}, {"default": tf.constant(1)})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
meta_graph, = handler.meta_graphs
self.assertEqual(len(meta_graph.collection_def), 0)
def testTags(self):
with tf.Graph().as_default() as graph:
saved_model_lib.add_signature("default", {}, {"default": tf.constant(1)})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph, ["tag1"])
handler.add_graph_copy(graph, ["tag1", "tag2"])
self.assertAllEqual(sorted(handler.get_tags()),
sorted([set(["tag1"]), set(["tag1", "tag2"])]))
self.assertTrue(handler.get_meta_graph_copy(["tag1"]) is not None)
self.assertTrue(handler.get_meta_graph_copy(["tag2", "tag1"]) is not None)
with self.assertRaises(KeyError):
handler.get_meta_graph_copy(["tag2"])
def testSignatures(self):
with tf.Graph().as_default() as graph:
input_a = tf.constant(2)
input_b = tf.constant(3)
mul = input_a * input_b
saved_model_lib.add_signature("six", {}, {"out": mul})
saved_model_lib.add_signature("mul2", {"in": input_b}, {"out": mul})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
signatures = handler.get_meta_graph_copy().signature_def
self.assertEqual(set(signatures.keys()), set(["six", "mul2"]))
self.assertAllEqual(list(signatures["six"].inputs.keys()), [])
self.assertAllEqual(list(signatures["six"].outputs.keys()), ["out"])
self.assertAllEqual(list(signatures["mul2"].inputs.keys()), ["in"])
self.assertAllEqual(list(signatures["mul2"].outputs.keys()), ["out"])
def testSignatures(self):
with tf.Graph().as_default() as graph:
input_a = tf.constant(2)
input_b = tf.constant(3)
mul = input_a * input_b
saved_model_lib.add_signature("six", {}, {"out": mul})
saved_model_lib.add_signature("mul2", {"in": input_b}, {"out": mul})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
signatures = handler.get_meta_graph_copy().signature_def
self.assertEqual(set(signatures.keys()), set(["six", "mul2"]))
self.assertAllEqual(list(signatures["six"].inputs.keys()), [])
self.assertAllEqual(list(signatures["six"].outputs.keys()), ["out"])
self.assertAllEqual(list(signatures["mul2"].inputs.keys()), ["in"])
self.assertAllEqual(list(signatures["mul2"].outputs.keys()), ["out"])
def add_signature(name=None, inputs=None, outputs=None):
"""Adds a signature to the module definition.
Warning: Deprecated. This belongs to the hub.Module API and TF1 Hub format.
For TF2, switch to plain SavedModels.
NOTE: This must be called within a `module_fn` that is defining a hub.Module.
THIS FUNCTION IS DEPRECATED.
Args:
name: Signature name as a string. If omitted, it is interpreted as 'default'
and is the signature used when `Module.__call__` `signature` is not
specified.
inputs: A dict from input name to Tensor or SparseTensor to feed when
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
outputs: A dict from output name to Tensor or SparseTensor to return from
applying the signature. If a single tensor is passed, it is interpreted
as a dict with a single 'default' entry.
Raises:
ValueError: if the arguments are invalid.
"""
if not name:
name = "default"
if inputs is None:
inputs = {}
if outputs is None:
outputs = {}
if not isinstance(inputs, dict):
inputs = {"default": inputs}
if not isinstance(outputs, dict):
outputs = {"default": outputs}
message = find_signature_inputs_from_multivalued_ops(inputs)
if message: logging.error(message)
message = find_signature_input_colocation_error(name, inputs)
if message: raise ValueError(message)
saved_model_lib.add_signature(name, inputs, outputs)
def testAssets(self):
original_asset_file = os.path.join(self.get_temp_dir(), "hello.txt")
_write_string_to_file(original_asset_file, "hello world")
with tf.Graph().as_default() as graph:
asset_tensor = tf.constant(original_asset_file, name="file")
graph.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, asset_tensor)
saved_model_lib.add_signature("default", {}, {"default": asset_tensor})
handler = saved_model_lib.SavedModelHandler()
handler.add_graph_copy(graph)
export_dir = os.path.join(self.get_temp_dir(), "exported")
handler.export(export_dir)
# Check that asset file got written to the expected place:
exported_asset_file = os.path.join(export_dir, "assets", "hello.txt")
self.assertTrue(tf.gfile.Exists(exported_asset_file))
loaded_handler = saved_model_lib.load(export_dir)
with _instantiate_meta_graph(loaded_handler).as_default():
with tf.Session() as sess:
self.assertEqual(sess.run("file:0"),
tf.compat.as_bytes(exported_asset_file))
