def _show_tag_sets(saved_model_dir):
"""Prints the tag-sets stored in SavedModel directory.
Prints all the tag-sets for MetaGraphs stored in SavedModel directory.
Args:
saved_model_dir: Directory containing the SavedModel to inspect.
"""
tag_sets = saved_model_utils.get_saved_model_tag_sets(saved_model_dir)
print('The given SavedModel contains the following tag-sets:')
for tag_set in sorted(tag_sets):
print(', '.join(sorted(tag_set)))
def _show_tag_sets(saved_model_dir):
"""Prints the tag-sets stored in SavedModel directory.
Prints all the tag-sets for MetaGraphs stored in SavedModel directory.
Args:
saved_model_dir: Directory containing the SavedModel to inspect.
"""
tag_sets = saved_model_utils.get_saved_model_tag_sets(saved_model_dir)
print('The given SavedModel contains the following tag-sets:')
for tag_set in sorted(tag_sets):
print(', '.join(sorted(tag_set)))
def testGetSavedModelTagSets(self):
saved_model_dir = os.path.join(test.get_temp_dir(), "test_tags")
builder = saved_model_builder.SavedModelBuilder(saved_model_dir)
# Force test to run in graph mode since SavedModelBuilder.save requires a
# session to work.
with ops.Graph().as_default():
# Graph with a single variable. SavedModel invoked to:
# - add with weights.
# - a single tag (from predefined constants).
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess,
[tag_constants.TRAINING])
# Graph that updates the single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
# - a single tag (from predefined constants).
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
builder.add_meta_graph([tag_constants.SERVING])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple predefined tags.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 44)
builder.add_meta_graph(
[tag_constants.SERVING, tag_constants.GPU])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple predefined tags for serving on TPU.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 44)
builder.add_meta_graph(
[tag_constants.SERVING, tag_constants.TPU])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple custom tags.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 45)
builder.add_meta_graph(["foo", "bar"])
# Save the SavedModel to disk.
builder.save()
actual_tags = saved_model_utils.get_saved_model_tag_sets(
saved_model_dir)
expected_tags = [["train"], ["serve"], ["serve", "gpu"],
["serve", "tpu"], ["foo", "bar"]]
self.assertEqual(expected_tags, actual_tags)
def testGetSavedModelTagSets(self):
saved_model_dir = os.path.join(test.get_temp_dir(), "test_tags")
builder = saved_model_builder.SavedModelBuilder(saved_model_dir)
# Graph with a single variable. SavedModel invoked to:
# - add with weights.
# - a single tag (from predefined constants).
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess, [tag_constants.TRAINING])
# Graph that updates the single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
# - a single tag (from predefined constants).
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
builder.add_meta_graph([tag_constants.SERVING])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple predefined tags.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 44)
builder.add_meta_graph([tag_constants.SERVING, tag_constants.GPU])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple predefined tags for serving on TPU.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 44)
builder.add_meta_graph([tag_constants.SERVING, tag_constants.TPU])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple custom tags.
with self.session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 45)
builder.add_meta_graph(["foo", "bar"])
# Save the SavedModel to disk.
builder.save()
actual_tags = saved_model_utils.get_saved_model_tag_sets(saved_model_dir)
expected_tags = [["train"], ["serve"], ["serve", "gpu"], ["serve", "tpu"],
["foo", "bar"]]
self.assertEqual(expected_tags, actual_tags)
def verify_outputs(args, onnx_model):
tag_sets = saved_model_utils.get_saved_model_tag_sets(args.saved_model)
for tag_set in tag_sets:
tag_set = ','.join(tag_set)
meta_graph_def = saved_model_utils.get_meta_graph_def(
args.saved_model, tag_set)
signature_def_map = meta_graph_def.signature_def
for signature_def_key in signature_def_map.keys():
outputs_tensor_info = signature_def_map[signature_def_key].outputs
for output_key, output_tensor in outputs_tensor_info.items():
rename_output(onnx_model, output_key, output_tensor)
print("Inputs in model: {}".format(", ".join([
"'{}'".format(o.name) for o in onnx_model.graph.input
if not has_initializer(onnx_model, o.name)
])))
print("Outputs in model: {}".format(", ".join(
["'{}'".format(o.name) for o in onnx_model.graph.output])))
def _show_all(saved_model_dir):
"""Prints tag-set, SignatureDef and Inputs/Outputs information in SavedModel.
Prints all tag-set, SignatureDef and Inputs/Outputs information stored in
SavedModel directory.
Args:
saved_model_dir: Directory containing the SavedModel to inspect.
"""
tag_sets = saved_model_utils.get_saved_model_tag_sets(saved_model_dir)
for tag_set in sorted(tag_sets):
print("\nMetaGraphDef with tag-set: '%s' "
"contains the following SignatureDefs:" % ', '.join(tag_set))
tag_set = ','.join(tag_set)
signature_def_map = get_signature_def_map(saved_model_dir, tag_set)
for signature_def_key in sorted(signature_def_map.keys()):
print('\nsignature_def[\'' + signature_def_key + '\']:')
_show_inputs_outputs(saved_model_dir, tag_set, signature_def_key,
indent=1)
def run_main(unused_args):
input_model_dir = FLAGS.input_saved_model_dir
output_model_dir = FLAGS.output_saved_model_dir
sig_key = FLAGS.signature_key
inp_tags = FLAGS.saved_model_tags
if FLAGS.saved_model_tags == "":
tag_set = []
else:
tag_set = [tag for tag in inp_tags.split(",")]
avail_tags = saved_model_utils.get_saved_model_tag_sets(input_model_dir)
found = False
for tag in tag_set:
if [tag] in avail_tags:
found = True
else:
found = False
break
if not found:
print ("Supplied tags", tag_set, "is not in available tag set,\
please use one or more of these", avail_tags, "Using --saved_model_tags")
exit(1)
sig_def = saved_model_utils.get_meta_graph_def(input_model_dir, inp_tags)
pretrained_model = load.load(input_model_dir, tag_set)
if sig_key not in list(pretrained_model.signatures.keys()):
print (sig_key, "is not in ", list(pretrained_model.signatures.keys()),
"provide one of those using --signature_key")
exit(1)
infer = pretrained_model.signatures[sig_key]
frozen_func = convert_to_constants.convert_variables_to_constants_v2(infer,lower_control_flow=True)
frozen_func.graph.structured_outputs = nest.pack_sequence_as(
infer.graph.structured_outputs,
frozen_func.graph.structured_outputs)
souts = frozen_func.graph.structured_outputs
inputs = frozen_func.inputs
input_nodes = [(tensor.name.split(":"))[0] for tensor in inputs]
output_nodes = [(souts[name].name.split(":"))[0] for name in souts]
gdef = frozen_func.graph.as_graph_def()
opt_graph = optimize_for_inference_lib.optimize_for_inference(gdef, input_nodes, output_nodes,
[tensor.dtype.as_datatype_enum for tensor in inputs] )
with session.Session() as sess:
graph = importer.import_graph_def(opt_graph,name="")
signature_inputs = {(tensor.name.split(":"))[0]: model_utils.build_tensor_info(tensor)
for tensor in inputs}
signature_outputs = {name: model_utils.build_tensor_info(souts[name])
for name in souts}
signature_def = signature_def_utils.build_signature_def(
signature_inputs, signature_outputs,
signature_constants.PREDICT_METHOD_NAME)
signature_def_map = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
}
builder = saved_model_builder.SavedModelBuilder(output_model_dir)
builder.add_meta_graph_and_variables(sess, tags=[tag_constants.SERVING],
signature_def_map=signature_def_map)
builder.save()
