Merge pull request #456 from lucienwang1009/remove_useless_inputs

remove useless inputs

Author: guschmue

tests/run_pretrained_models.py

@@ -222,18 +222,20 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
             dir_name = os.path.dirname(self.local)
         print("\tdownloaded", model_path)
 
-        inputs = list(self.input_names.keys())
+        input_names = list(self.input_names.keys())
         outputs = self.output_names
         if self.model_type in ["checkpoint"]:
-            graph_def, inputs, outputs = loader.from_checkpoint(model_path, inputs, outputs)
+            graph_def, input_names, outputs = loader.from_checkpoint(model_path, input_names, outputs)
         elif self.model_type in ["saved_model"]:
-            graph_def, inputs, outputs = loader.from_saved_model(model_path, inputs, outputs)
+            graph_def, input_names, outputs = loader.from_saved_model(model_path, input_names, outputs)
         else:
-            graph_def, inputs, outputs = loader.from_graphdef(model_path, inputs, outputs)
+            graph_def, input_names, outputs = loader.from_graphdef(model_path, input_names, outputs)
 
         # create the input data
         inputs = {}
         for k, v in self.input_names.items():
+            if k not in input_names:
+                continue
             if isinstance(v, six.text_type) and v.startswith("np."):
                 inputs[k] = eval(v)  # pylint: disable=eval-used
             else:
@@ -312,10 +314,12 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                     print("\tResults: OK")
                 return True
             except Exception as ex:
-                print("\tResults: ", ex)
+                tb = traceback.format_exc()
+                print("\tResults", ex, tb)
 
         except Exception as ex:
-            print("\trun_onnx", "FAIL", ex)
+            tb = traceback.format_exc()
+            print("\trun_onnx", "FAIL", ex, tb)
 
         return False
 
@@ -399,7 +403,8 @@ def main():
                              fold_const=args.fold_const)
         except Exception as ex:
             ret = None
-            print(ex)
+            tb = traceback.format_exc()
+            print(ex, tb)
         finally:
             if not args.debug:
                 utils.delete_directory(TEMP_DIR)

tests/run_pretrained_models.yaml

@@ -32,10 +32,30 @@ saved_model_with_redundant_inputs:
   model_type: saved_model
   input_get: get_ramp
   inputs:
+    "X:0": [1, 10]
     "Placeholder:0": [1, 10]
   outputs:
     - Add:0
 
+graphdef_with_redundant_inputs:
+  model: tests/models/regression/graphdef/frozen.pb
+  input_get: get_ramp
+  inputs:
+    "X:0": [1, 10]
+    "Placeholder:0": [1, 10]
+  outputs:
+    - Add:0
+
+checkpoint_with_redundant_inputs:
+  model: tests/models/regression/checkpoint/model.meta
+  model_type: checkpoint
+  input_get: get_ramp
+  inputs:
+    "X:0": [1]
+    "Placeholder:0": [1, 10]
+  outputs:
+    - pred:0
+
 benchtf-fc:
   model: tests/models/fc-layers/frozen.pb
   input_get: get_ramp

tf2onnx/loader.py

@@ -38,6 +38,20 @@ def freeze_session(sess, keep_var_names=None, output_names=None, clear_devices=T
         return frozen_graph
 
 
+def remove_redundant_inputs(frozen_graph, input_names):
+    """Remove redundant inputs not in frozen graph."""
+    frozen_inputs = []
+    # get inputs in frozen graph
+    for n in frozen_graph.node:
+        for inp in input_names:
+            if utils.node_name(inp) == n.name:
+                frozen_inputs.append(inp)
+    deleted_inputs = list(set(input_names) - set(frozen_inputs))
+    if deleted_inputs:
+        log.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
+    return frozen_inputs
+
+
 def from_graphdef(model_path, input_names, output_names):
     """Load tensorflow graph from graphdef."""
     # make sure we start with clean default graph
@@ -48,6 +62,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)
+    input_names = remove_redundant_inputs(frozen_graph, input_names)
     # clean up
     tf.reset_default_graph()
     return frozen_graph, input_names, output_names
@@ -63,6 +78,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)
+    input_names = remove_redundant_inputs(frozen_graph, input_names)
     # clean up
     tf.reset_default_graph()
     return frozen_graph, input_names, output_names
@@ -93,15 +109,9 @@ def from_saved_model(model_path, input_names, output_names):
             for _, output_tensor in sorted(outputs_tensor_info.items()):
                 outputs[output_tensor.name] = sess.graph.get_tensor_by_name(output_tensor.name)
         frozen_graph = freeze_session(sess, output_names=list(outputs.keys()))
-        frozen_inputs = []
-        # get inputs in frozen graph
-        for n in frozen_graph.node:
-            for inp, _ in inputs.items():
-                if utils.node_name(inp) == n.name:
-                    frozen_inputs.append(inp)
-        deleted_inputs = list(set(inputs.keys()) - set(frozen_inputs))
-        if deleted_inputs:
-            log.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
+    if input_names is None:
+        input_names = inputs.keys()
+    input_names = remove_redundant_inputs(frozen_graph, input_names)
     # clean up
     tf.reset_default_graph()
-    return frozen_graph, frozen_inputs, outputs.keys()
+    return frozen_graph, input_names, outputs.keys()