review feedback

Author: guschmue

README.md

@@ -41,7 +41,7 @@ For pytorch/caffe2, follow the instructions here:
 We tested with pytorch/caffe2 and onnxruntime and unit tests are passing for those.
 
 ## Supported Tensorflow and Python Versions
-We tested with tensorflow 1.5-1.13 and anaconda **3.5,3.6**.
+We tested with tensorflow 1.5-1.12 and anaconda **3.5,3.6**.
 
 # Installation
 ## From Pypi

tests/models/regression/checkpoint/checkpoint

@@ -1,2 +1,2 @@
-model_checkpoint_path: "C:\\src\\tensorflow-onnx\\tests\\models\\regression\\checkpoint\\model"
-all_model_checkpoint_paths: "C:\\src\\tensorflow-onnx\\tests\\models\\regression\\checkpoint\\model"
+model_checkpoint_path: "model"
+all_model_checkpoint_paths: "model"

tf2onnx/convert.py

@@ -108,7 +108,7 @@ def main():
 
     # todo: consider to enable const folding by default?
     graph_def = tf_optimize(args.inputs, args.outputs, graph_def, args.fold_const)
-    tf.reset_default_graph()
+
     with tf.Graph().as_default() as tf_graph:
         tf.import_graph_def(graph_def, name='')
     with tf.Session(graph=tf_graph):

tf2onnx/loader.py

@@ -17,7 +17,7 @@
 
 def freeze_session(sess, keep_var_names=None, output_names=None, clear_devices=True):
     """Freezes the state of a session into a pruned computation graph."""
-    output_names = [i.replace(":0", "") for i in output_names]
+    output_names = [i.split(':')[:-1][0] for i in output_names]
     graph = sess.graph
     with graph.as_default():
         freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
@@ -42,7 +42,7 @@ def from_graphdef(model_path, input_names, output_names):
             graph_def.ParseFromString(f.read())
             tf.import_graph_def(graph_def, name='')
             frozen_graph = freeze_session(sess, output_names=output_names)
-    # clean up after us
+    # clean up
     tf.reset_default_graph()
     return frozen_graph, input_names, output_names
 
@@ -57,7 +57,7 @@ def from_checkpoint(model_path, input_names, output_names):
         # restore from model_path minus the ".meta"
         saver.restore(sess, model_path[:-5])
         frozen_graph = freeze_session(sess, output_names=output_names)
-    # clean up after us
+    # clean up
     tf.reset_default_graph()
     return frozen_graph, input_names, output_names
 
@@ -86,12 +86,7 @@ def from_saved_model(model_path, input_names, output_names):
             outputs_tensor_info = get_signature_def(meta_graph_def, k).outputs
             for _, output_tensor in sorted(outputs_tensor_info.items()):
                 outputs[output_tensor.name] = sess.graph.get_tensor_by_name(output_tensor.name)
-        # freeze uses the node name derived from output:0 so only pass in output:0;
-        # it will provide all outputs of that node.
-        for o in list(outputs.keys()):
-            if not o.endswith(":0"):
-                del outputs[o]
         frozen_graph = freeze_session(sess, output_names=list(outputs.keys()))
-    # clean up after us
+    # clean up
     tf.reset_default_graph()
     return frozen_graph, inputs.keys(), outputs.keys()