diff options
Diffstat (limited to 'tensorflow/python/tools/freeze_graph.py')
| -rw-r--r-- | tensorflow/python/tools/freeze_graph.py | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/tensorflow/python/tools/freeze_graph.py b/tensorflow/python/tools/freeze_graph.py index bd046a7fd0..8069950925 100644 --- a/tensorflow/python/tools/freeze_graph.py +++ b/tensorflow/python/tools/freeze_graph.py @@ -115,10 +115,15 @@ def freeze_graph_with_def_protos( output_node_names.split(","), variable_names_blacklist=variable_names_blacklist) - with gfile.GFile(output_graph, "wb") as f: - f.write(output_graph_def.SerializeToString()) + # Write GraphDef to file if output path has been given. + if output_graph: + with gfile.GFile(output_graph, "wb") as f: + f.write(output_graph_def.SerializeToString()) + print("%d ops in the final graph." % len(output_graph_def.node)) + return output_graph_def + def _parse_input_graph_proto(input_graph, input_binary): """Parser input tensorflow graph into GraphDef proto.""" |