Fixed processing tags in summarize_graph.py (#6726)

This commit is contained in:
Vladimir Gavrilov 2021-07-22 14:39:54 +03:00 committed by GitHub
parent 6dc771c2e7
commit abc6a336b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -70,9 +70,10 @@ if __name__ == "__main__": # pragma: no cover
if argv.input_model and argv.saved_model_dir:
print("[ ERROR ] Both keys were provided --input_model and --input_dir. Please, provide only one of them")
sys.exit(1)
graph_def, _ = load_tf_graph_def(graph_file_name=argv.input_model, is_binary=not argv.text,
checkpoint=argv.input_checkpoint,
model_dir=argv.saved_model_dir, saved_model_tags=argv.saved_model_tags)
tags = argv.saved_model_tags.split(",")
graph_def, _, _ = load_tf_graph_def(graph_file_name=argv.input_model, is_binary=not argv.text,
checkpoint=argv.input_checkpoint,
model_dir=argv.saved_model_dir, saved_model_tags=tags)
summary = summarize_graph(graph_def)
print("{} input(s) detected:".format(len(summary['inputs'])))
for input in summary['inputs']: