Allow for cutting and renaming of IO. Use structured outputs by default

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 = outputs
+                else:
+                    struc_outputs = [to_rename.get(k, k) for k in outputs]
+                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

tf2onnx/tf_loader.py

@@ -170,10 +170,8 @@ 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)
+    node_names = set(n.name for n in frozen_graph.node)
+    frozen_inputs = [inp for inp in input_names if utils.node_name(inp) in node_names]
     deleted_inputs = list(set(input_names) - set(frozen_inputs))
     if deleted_inputs:
         logger.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
@@ -372,8 +370,8 @@ def _from_saved_model_v2(model_path, input_names, output_names, tag, signature_d
         utils.make_sure(hasattr(imported, "__call__"), err_no_call)
         utils.make_sure(concrete_function_index < len(imported.__call__.concrete_functions),
                         err_index, concrete_function_index, len(imported.__call__.concrete_functions) - 1)
-        sig = imported.__call__.concrete_functions[concrete_function_index].structured_input_signature[0]
-        concrete_func = imported.__call__.get_concrete_function(*sig)
+        args, kwargs = imported.__call__.concrete_functions[concrete_function_index].structured_input_signature
+        concrete_func = imported.__call__.get_concrete_function(*args, **kwargs)
     elif signature_def:
         utils.make_sure(signature_def[0] in valid_sigs, err_sig_nomatch, signature_def[0])
         concrete_func = imported.signatures[signature_def[0]]
@@ -382,14 +380,28 @@ def _from_saved_model_v2(model_path, input_names, output_names, tag, signature_d
         logger.warning(wrn_sig_1, valid_sigs[0])
         concrete_func = imported.signatures[valid_sigs[0]]
 
-    inputs = [tensor.name for tensor in concrete_func.inputs if tensor.dtype != tf.dtypes.resource]
-    outputs = [tensor.name for tensor in concrete_func.outputs if tensor.dtype != tf.dtypes.resource]
-
-    # filter by user specified inputs/outputs
-    if input_names:
-        inputs = list(set(input_names) & set(inputs))
-    if output_names:
-        outputs = list(set(output_names) & set(outputs))
+    tensors_to_rename = {}
+    if input_names is None:
+        args, kwargs = concrete_func.structured_input_signature
+        structured_inputs = [t.name for t in args if isinstance(t, tf.TensorSpec)] + sorted(kwargs.keys())
+        structured_inputs = set(inp + ":0" for inp in structured_inputs)
+        inputs = [tensor.name for tensor in concrete_func.inputs if tensor.dtype != tf.dtypes.resource]
+        if any(inp in structured_inputs for inp in inputs):
+            inputs = [inp for inp in inputs if inp in structured_inputs]
+    else:
+        inputs = input_names
+
+    if output_names is None:
+        outputs = [tensor.name for tensor in concrete_func.outputs if tensor.dtype != tf.dtypes.resource]
+        if isinstance(concrete_func.structured_outputs, dict):
+            structured_outputs = sorted(concrete_func.structured_outputs.keys())
+            tensors_to_rename.update(zip(outputs, structured_outputs))
+            logger.info("Output names: %r", structured_outputs)
+        else:
+            logger.info("Output names: %r", outputs)
+    else:
+        outputs = output_names
+        logger.warning("Outputs not left as None; will use provided names not structured output names.")
 
     # Avoid errors due to bug in TF freezing
     removed_resource_to_placeholder, graph_captures_copy, func_captures_copy = \
@@ -425,23 +437,25 @@ def _from_saved_model_v2(model_path, input_names, output_names, tag, signature_d
 
     replace_placeholders_with_tables(frozen_graph, placeholder_to_table_info)
 
-    return frozen_graph, inputs, outputs, concrete_func, imported, initialized_tables
+    return frozen_graph, inputs, outputs, concrete_func, imported, initialized_tables, tensors_to_rename
 
 def from_saved_model(model_path, input_names, output_names, tag=None,
                      signatures=None, concrete_function=None, large_model=False,
-                     return_concrete_func=False, return_initialized_tables=False):
+                     return_concrete_func=False, return_initialized_tables=False, return_tensors_to_rename=False):
     """Load tensorflow graph from saved_model."""
     if signatures is None:
         signatures = []
     tf_reset_default_graph()
     if is_tf2():
-        frozen_graph, input_names, output_names, concrete_func, imported, initialized_tables = \
+        frozen_graph, input_names, output_names, concrete_func, imported, initialized_tables, tensors_to_rename = \
             _from_saved_model_v2(model_path, input_names, output_names, tag, signatures, concrete_function, large_model)
         result = [frozen_graph, input_names, output_names]
         if return_concrete_func:
             result += [concrete_func, imported]
         if return_initialized_tables:
             result += [initialized_tables]
+        if return_tensors_to_rename:
+            result += [tensors_to_rename]
     else:
         with tf_session() as sess:
             frozen_graph, input_names, output_names = \
@@ -451,6 +465,8 @@ def from_saved_model(model_path, input_names, output_names, tag=None,
                 result += [None, None]
             if return_initialized_tables:
                 result += [{}]
+            if return_tensors_to_rename:
+                result += [{}]
 
     tf_reset_default_graph()
     return result