Allow for cutting and renaming of IO. Use structured outputs by default (#1355)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/run_pretrained_models.py

@@ -296,7 +296,7 @@ def run_tensorflow(self, sess, inputs):
         return result
 
     def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None,
-                const_node_values=None, initialized_tables=None, tflite_path=None):
+                const_node_values=None, initialized_tables=None, tflite_path=None, tensors_to_rename=None):
         """Convert graph to tensorflow."""
         if extra_opset is None:
             extra_opset = []
@@ -306,7 +306,8 @@ def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, i
                                 extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
                                 input_names=input_names, output_names=self.output_names,
                                 const_node_values=const_node_values, initialized_tables=initialized_tables,
-                                tflite_path=tflite_path, dequantize=self.dequantize)
+                                tflite_path=tflite_path, dequantize=self.dequantize,
+                                tensors_to_rename=tensors_to_rename)
 
     def run_caffe2(self, name, model_proto, inputs):
         """Run test again caffe2 backend."""
@@ -320,7 +321,7 @@ def run_caffe2(self, name, model_proto, inputs):
             self.onnx_runtime = time.time() - start
         return results
 
-    def run_onnxruntime(self, name, model_proto, inputs, external_tensor_storage=None):
+    def run_onnxruntime(self, name, model_proto, inputs, outputs, external_tensor_storage=None):
         """Run test against onnxruntime backend."""
         import onnxruntime as rt
         model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True,
@@ -334,11 +335,11 @@ def run_onnxruntime(self, name, model_proto, inputs, external_tensor_storage=Non
             m = rt.InferenceSession(model_path, opt)
         else:
             m = rt.InferenceSession(model_path)
-        results = m.run(self.output_names, inputs)
+        results = m.run(outputs, inputs)
         if self.perf:
             start = time.time()
             for _ in range(PERFITER):
-                _ = m.run(self.output_names, inputs)
+                _ = m.run(outputs, inputs)
             self.onnx_runtime = time.time() - start
         return results
 
@@ -371,19 +372,20 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
         initialized_tables = {}
         outputs = self.output_names
         tflite_path = None
+        to_rename = None
         if self.model_type in ["checkpoint"]:
             graph_def, input_names, outputs = tf_loader.from_checkpoint(model_path, input_names, outputs)
         elif self.model_type in ["saved_model"]:
-            loaded = tf_loader.from_saved_model(model_path, input_names, outputs, self.tag, self.signatures,
+            loaded = tf_loader.from_saved_model(model_path, None, None, self.tag, self.signatures,
                                                 self.concrete_function, self.large_model,
                                                 return_concrete_func=not self.run_tf_frozen,
-                                                return_initialized_tables=True)
+                                                return_initialized_tables=True, return_tensors_to_rename=True)
             if not self.run_tf_frozen:
                 # Must maintain ref to imported since concrete_func uses weak refs
                 # pylint: disable=unused-variable
-                graph_def, input_names, outputs, concrete_func, imported, initialized_tables = loaded
+                graph_def, input_names, outputs, concrete_func, imported, initialized_tables, to_rename = loaded
             else:
-                graph_def, input_names, outputs, initialized_tables = loaded
+                graph_def, input_names, outputs, initialized_tables, to_rename = loaded
         elif self.model_type in ["keras"]:
             graph_def, input_names, outputs = tf_loader.from_keras(model_path, input_names, outputs)
         elif self.model_type in ["tflite"]:
@@ -434,10 +436,8 @@ def run_tflite():
             # If there is only a single output a dict might not be returned
             if isinstance(tf_results_d, tf.Tensor):
                 tf_results = [tf_results_d]
-            elif self.structured_outputs is None:
-                tf_results = list(tf_results_d.values())
             else:
-                tf_results = [tf_results_d[output] for output in self.structured_outputs]
+                tf_results = [tf_results_d[k] for k in sorted(tf_results_d.keys())]
             tf_results = [tf_res.numpy() for tf_res in tf_results]
             if self.perf:
                 logger.info("Running TF perf")
@@ -507,7 +507,8 @@ def run_tflite():
                 onnx_graph = self.to_onnx(tf_graph, opset=opset, extra_opset=extra_opset,
                                           shape_override=shape_override, input_names=inputs.keys(),
                                           const_node_values=const_node_values,
-                                          initialized_tables=initialized_tables, tflite_path=tflite_path)
+                                          initialized_tables=initialized_tables, tflite_path=tflite_path,
+                                          tensors_to_rename=to_rename)
                 onnx_graph = optimizer.optimize_graph(onnx_graph)
                 print("ONNX", onnx_graph.dump_node_statistics())
                 external_tensor_storage = ExternalTensorStorage() if self.large_model else None
@@ -532,7 +533,11 @@ def run_tflite():
             if backend == "caffe2":
                 onnx_results = self.run_caffe2(name, model_proto, inputs)
             elif backend == "onnxruntime":
-                onnx_results = self.run_onnxruntime(name, model_proto, inputs, external_tensor_storage)
+                if to_rename is None:
+                    struc_outputs = self.output_names
+                else:
+                    struc_outputs = [to_rename.get(k, k) for k in self.output_names]
+                onnx_results = self.run_onnxruntime(name, model_proto, inputs, struc_outputs, external_tensor_storage)
             else:
                 raise ValueError("unknown backend")
             logger.info("Run_ONNX OK")

tests/test_convert.py

@@ -88,6 +88,34 @@ def test_convert_checkpoint(self):
                                        '--output',
                                        'converted_checkpoint.onnx']))
 
+    def test_convert_graphdef_cut_input(self):
+        """ convert graphdef, change input to start from Mul:0 and rename it """
+        self.assertTrue(run_test_case(['',
+                                       '--input',
+                                       'tests/models/regression/graphdef/frozen.pb',
+                                       '--inputs',
+                                       'Mul:0',
+                                       '--rename-inputs',
+                                       'new_input',
+                                       '--outputs',
+                                       'pred:0',
+                                       '--output',
+                                       'converted_graphdef_cut_input.onnx']))
+
+    def test_convert_graphdef_cut_output(self):
+        """ convert graphdef, change output to Mul:0 and rename it """
+        self.assertTrue(run_test_case(['',
+                                       '--input',
+                                       'tests/models/regression/graphdef/frozen.pb',
+                                       '--inputs',
+                                       'X:0',
+                                       '--rename-outputs',
+                                       'new_output',
+                                       '--outputs',
+                                       'Mul:0',
+                                       '--output',
+                                       'converted_graphdef_cut_output.onnx']))
+
 
 if __name__ == '__main__':
     unittest.main()

tf2onnx/convert.py

@@ -58,12 +58,14 @@ def get_args():
     parser.add_argument("--tflite", help="input from tflite model")
     parser.add_argument("--large_model", help="use the large model format (for models > 2GB)", action="store_true")
     parser.add_argument("--output", help="output model file")
-    parser.add_argument("--inputs", help="model input_names")
-    parser.add_argument("--outputs", help="model output_names")
+    parser.add_argument("--inputs", help="model input_names (optional for saved_model, keras, and tflite)")
+    parser.add_argument("--outputs", help="model output_names (optional for saved_model, keras, and tflite)")
     parser.add_argument("--ignore_default", help="comma-separated list of names of PlaceholderWithDefault "
                                                  "ops to change into Placeholder ops")
     parser.add_argument("--use_default", help="comma-separated list of names of PlaceholderWithDefault ops to "
                                               "change into Identity ops using their default value")
+    parser.add_argument("--rename-inputs", help="input names to use in final model (optional)")
+    parser.add_argument("--rename-outputs", help="output names to use in final model (optional)")
     parser.add_argument("--opset", type=int, default=None, help="opset version to use for onnx domain")
     parser.add_argument("--dequantize", help="Remove quantization from model. Only supported for tflite currently.",
                         action="store_true")
@@ -87,7 +89,7 @@ def get_args():
         # for backward compativility
         args.graphdef = args.input
     if args.graphdef or args.checkpoint:
-        if not args.input and not args.outputs:
+        if not args.inputs or not args.outputs:
             parser.error("graphdef and checkpoint models need to provide inputs and outputs")
     if not any([args.graphdef, args.checkpoint, args.saved_model, args.keras, args.tflite]):
         parser.print_help()
@@ -100,6 +102,10 @@ def get_args():
         args.ignore_default = args.ignore_default.split(",")
     if args.use_default:
         args.use_default = args.use_default.split(",")
+    if args.rename_outputs:
+        args.rename_outputs = args.rename_outputs.split(",")
+    if args.rename_inputs:
+        args.rename_inputs = args.rename_inputs.split(",")
     if args.inputs_as_nchw:
         args.inputs_as_nchw = args.inputs_as_nchw.split(",")
     if args.target:
@@ -135,6 +141,7 @@ def main():
     tflite_path = None
     custom_ops = {}
     initialized_tables = None
+    tensors_to_rename = {}
     if args.custom_ops:
         using_tf_opset = False
         for op in args.custom_ops.split(","):
@@ -162,9 +169,9 @@ def main():
         graph_def, inputs, outputs = tf_loader.from_checkpoint(args.checkpoint, args.inputs, args.outputs)
         model_path = args.checkpoint
     if args.saved_model:
-        graph_def, inputs, outputs, initialized_tables = tf_loader.from_saved_model(
-            args.saved_model, args.inputs, args.outputs, args.tag,
-            args.signature_def, args.concrete_function, args.large_model, return_initialized_tables=True)
+        graph_def, inputs, outputs, initialized_tables, tensors_to_rename = tf_loader.from_saved_model(
+            args.saved_model, args.inputs, args.outputs, args.tag, args.signature_def, args.concrete_function,
+            args.large_model, return_initialized_tables=True, return_tensors_to_rename=True)
         model_path = args.saved_model
     if args.keras:
         graph_def, inputs, outputs = tf_loader.from_keras(
@@ -181,6 +188,11 @@ def main():
         logger.info("inputs: %s", inputs)
         logger.info("outputs: %s", outputs)
 
+    if args.rename_inputs:
+        tensors_to_rename.update(zip(inputs, args.rename_inputs))
+    if args.rename_outputs:
+        tensors_to_rename.update(zip(outputs, args.rename_outputs))
+
     tf_graph = None
     const_node_values = None
     if graph_def is not None:
@@ -206,6 +218,7 @@ def main():
                              ignore_default=args.ignore_default,
                              use_default=args.use_default,
                              const_node_values=const_node_values,
+                             tensors_to_rename=tensors_to_rename,
                              initialized_tables=initialized_tables,
                              tflite_path=tflite_path,
                              dequantize=args.dequantize)
@@ -218,6 +231,8 @@ def main():
     # write onnx graph
     logger.info("")
     logger.info("Successfully converted TensorFlow model %s to ONNX", model_path)
+    logger.info("Model inputs: %s", onnx_graph.input_names)
+    logger.info("Model outputs: %s", onnx_graph.outputs)
     if args.output:
         if args.large_model:
             utils.save_onnx_zip(args.output, model_proto, tensor_storage)

tf2onnx/shape_inference.py

@@ -100,7 +100,7 @@ def infer_shape_for_op(op):
             op.outputs[0].set_shape(new_shape)
             logger.debug("set placeholder op [%s] with new shape %s", op.outputs[0].name, new_shape)
             return True
-        logger.warning("Shape of placeholder '%s' is unknown, treated it as a scalar. Please use the --input flag "
+        logger.warning("Shape of placeholder '%s' is unknown, treated it as a scalar. Please use the --inputs flag "
                        "and append the shape to the input name if this input is not a scalar.", op.name)
         op.outputs[0].set_shape([])
         return True