Refactor tfonnx.py (#1595)

* Refactor tfonnx

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

* Add Identity nodes even for subgraphs

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

* pylint

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

* Remove unnecessary subgraph check

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

* Fix dtype finding for multiple identity nodes in loops)

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

Author: TomWildenhain-Microsoft

tf2onnx/graph.py

@@ -503,35 +503,34 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
 
         self.reset_nodes(ops)
 
-        if not is_subgraph:
-            # add identity node after each output, in case it is renamed during conversion.
-            for o in self.outputs:
-                n = self.get_node_by_output_in_current_graph(o)
-                if n.is_graph_input():
-                    # Don't add identity if the node is also an input. We want to keep input names the same.
-                    continue
-                new_output_name = port_name(n.name + "_" + utils.make_name("raw_output_"))
-                n_shapes = n.output_shapes
-                n_dtypes = n.output_dtypes
-                body_graphs = n.graph.contained_graphs.pop(n.name, None)
-                self.remove_node(n.name)
-
-                new_outputs = [output if output != o else new_output_name for output in n.output]
-                # domain should be passed to new node
-                branches = {}
-                if body_graphs:
-                    for attr_name, body_graph in body_graphs.items():
-                        body_graph.parent_graph = self
-                        branches[attr_name] = body_graph
+        # add identity node after each output, in case it is renamed during conversion.
+        for o in self.outputs:
+            n = self.get_node_by_output_in_current_graph(o)
+            if n.is_graph_input():
+                # Don't add identity if the node is also an input. We want to keep input names the same.
+                continue
+            new_output_name = port_name(n.name + "_" + utils.make_name("raw_output_"))
+            n_shapes = n.output_shapes
+            n_dtypes = n.output_dtypes
+            body_graphs = n.graph.contained_graphs.pop(n.name, None)
+            self.remove_node(n.name)
+
+            new_outputs = [output if output != o else new_output_name for output in n.output]
+            # domain should be passed to new node
+            branches = {}
+            if body_graphs:
+                for attr_name, body_graph in body_graphs.items():
+                    body_graph.parent_graph = self
+                    branches[attr_name] = body_graph
 
-                _ = self.make_node(n.type, n.input, outputs=new_outputs, attr=n.attr, name=n.name,
-                                   skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes,
-                                   domain=n.domain, branches=branches)
+            _ = self.make_node(n.type, n.input, outputs=new_outputs, attr=n.attr, name=n.name,
+                               skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes,
+                               domain=n.domain, branches=branches)
 
-                self.replace_all_inputs(o, new_output_name, ops=self.get_nodes())
-                self.make_node("Identity", [new_output_name], outputs=[o], op_name_scope=n.name + "_" + "graph_outputs")
-                self.copy_shape(new_output_name, o)
-                self.copy_dtype(new_output_name, o)
+            self.replace_all_inputs(o, new_output_name, ops=self.get_nodes())
+            self.make_node("Identity", [new_output_name], outputs=[o], op_name_scope=n.name + "_" + "graph_outputs")
+            self.copy_shape(new_output_name, o)
+            self.copy_dtype(new_output_name, o)
 
     def create_new_graph_with_same_config(self):
         """Create a clean graph inheriting current graph's configuration."""
@@ -874,6 +873,23 @@ def is_const(self, output):
     def get_tensor_value(self, output, as_list=True):
         return self.get_node_by_output(output).get_tensor_value(as_list)
 
+    def rename_tensors(self, tensors_to_rename):
+        """Replace tensor names within nodes and graph inputs/outputs"""
+        def rename_list(l):
+            return [tensors_to_rename.get(t, t) for t in l]
+
+        def rename_keys(d):
+            return {tensors_to_rename.get(k, k): v for k, v in d.items()}
+
+        self._output_to_node_name = rename_keys(self._output_to_node_name)
+        self._output_to_consumers = rename_keys(self._output_to_consumers)
+        self._dtypes = rename_keys(self._dtypes)
+        self._output_shapes = rename_keys(self._output_shapes)
+        self.outputs = rename_list(self.outputs)
+        for node in self._nodes:
+            node._input = rename_list(node._input)
+            node._output = rename_list(node._output)
+
     def change_node_name(self, node, new_name):
         """Remove node in current graph."""
         utils.make_sure(new_name not in self._nodes_by_name, "node %s not unique ", new_name)
@@ -1232,15 +1248,23 @@ def follow_inputs(self, node, num, space=""):
             return []
         return val
 
-    def dump_node_statistics(self):
+    def dump_node_statistics(self, include_attrs=False, include_subgraphs=True):
+        """Return a counter of op types (and optionally attribute names) within the graph"""
         op_cnt = collections.Counter()
+        attr_cnt = collections.Counter()
         for n in self.get_nodes():
             op_cnt[n.type] += 1
+            for k in n.attr.keys():
+                attr_cnt[k] += 1
             body_graphs = n.get_body_graphs()
-            if body_graphs:
+            if body_graphs and include_subgraphs:
                 for b_g in body_graphs.values():
-                    op_cnt += b_g.dump_node_statistics()
+                    g_op_cnt, g_attr_cnt = b_g.dump_node_statistics(include_attrs=True, include_subgraphs=True)
+                    op_cnt += g_op_cnt
+                    attr_cnt += g_attr_cnt
 
+        if include_attrs:
+            return op_cnt, attr_cnt
         return op_cnt
 
     def remove_input(self, node, to_be_removed, input_index=None):

tf2onnx/onnx_opset/controlflow.py

@@ -623,12 +623,13 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
 
     g.outputs = [cond_outputs[0]] + g.outputs[2:] + scan_outputs
 
-    # FIXME: onnx does not have a variant type so we try to fish for the dtype in a prior TensorListSetItem.
+    # onnx does not have a variant type so we try to fish for the dtype in a prior TensorListSetItem.
     for o in g.outputs:
         if g.get_dtype(o) == onnx_pb.TensorProto.UNDEFINED:
-            node = g.get_node_by_output(o)
-            if node.type in ["Identity"]:
-                g.set_dtype(o, node.inputs[0].output_dtypes[0])
+            curr_o = o
+            while g.get_node_by_output(curr_o).type == "Identity":
+                curr_o = g.get_node_by_output(curr_o).input[0]
+            g.copy_dtype(curr_o, o)
 
     for node in g.ragged_variant_list_reads:
         # Requires opset 11

tf2onnx/tfonnx.py

@@ -33,7 +33,7 @@
 # pylint: disable=useless-return,broad-except,logging-not-lazy,unused-argument,missing-docstring
 # pylint: disable=unused-variable
 
-def fold_constants_using_tf(g, outputs_to_values, outputs_to_dtypes):
+def fold_constants_using_tf(g, outputs_to_values):
     ops = list(g.get_nodes())
     # pylint: disable=too-many-nested-blocks
     keep_looking = True
@@ -409,14 +409,13 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     del verbose
 
     opset = utils.find_opset(opset)
-    if not is_subgraph:
-        logger.info("Using tensorflow=%s, onnx=%s, tf2onnx=%s/%s",
-                    get_tf_version(), utils.get_onnx_version(), tf2onnx.__version__, tf2onnx.version.git_version[:6])
-        logger.info("Using opset <onnx, %s>", opset)
-        if opset > schemas.get_max_supported_opset_version():
-            logger.warning("Currently installed onnx package %s is too low to support opset %s, "
-                           "please upgrade onnx package to avoid potential conversion issue.",
-                           utils.get_onnx_version(), opset)
+    logger.info("Using tensorflow=%s, onnx=%s, tf2onnx=%s/%s",
+                get_tf_version(), utils.get_onnx_version(), tf2onnx.__version__, tf2onnx.version.git_version[:6])
+    logger.info("Using opset <onnx, %s>", opset)
+    if opset > schemas.get_max_supported_opset_version():
+        logger.warning("Currently installed onnx package %s is too low to support opset %s, "
+                       "please upgrade onnx package to avoid potential conversion issue.",
+                       utils.get_onnx_version(), opset)
 
     if shape_override is None:
         shape_override = {}
@@ -440,34 +439,17 @@ def check_io(input_names, output_names, output_shapes):
                              non_exists)
                 raise ValueError("Inputs/Outputs Not Found")
 
-    def rename_tensors_in_dict(d):
-        if tensors_to_rename is None:
-            return d
-        return {tensors_to_rename.get(k, k): v for k, v in d.items()}
-
-    def rename_tensors_in_list(tensors):
-        if tensors_to_rename is None or tensors is None:
-            return tensors
-        return [tensors_to_rename.get(t, t) for t in tensors]
-
-    def rename_tensors_in_nodes(onnx_nodes):
-        if tensors_to_rename is None:
-            return
-        for n in onnx_nodes:
-            n.input[:] = rename_tensors_in_list(n.input)
-            n.output[:] = rename_tensors_in_list(n.output)
-
     if tflite_path is not None:
         tflite_graphs, opcodes, model, tensor_shapes = read_tflite_model(tflite_path)
         main_g = None
-        inputs_as_nchw = rename_tensors_in_list(inputs_as_nchw)
+        subgraphs = []
         for i, tfl_graph in enumerate(tflite_graphs):
             is_main_g = i == len(tflite_graphs) - 1
             prefix = '' if is_main_g else tfl_graph.Name().decode() + '_'
             tensor_shapes_from_interpreter = None
             if is_main_g:
                 tensor_shapes_from_interpreter = tensor_shapes
-            onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, f_inputs, f_outputs, graph_name = \
+            onnx_nodes, _, _, output_shapes, dtypes, f_inputs, f_outputs, graph_name = \
                 parse_tflite_graph(tfl_graph, opcodes, model, prefix, tensor_shapes_from_interpreter)
             g_inputs = f_inputs
             g_outputs = f_outputs
@@ -478,63 +460,73 @@ def rename_tensors_in_nodes(onnx_nodes):
                     g_inputs = input_names
                 if output_names is not None:
                     g_outputs = output_names
-            rename_tensors_in_nodes(onnx_nodes)
-            g_inputs = rename_tensors_in_list(g_inputs)
-            g_outputs = rename_tensors_in_list(g_outputs)
-            output_shapes = rename_tensors_in_dict(output_shapes)
-            dtypes = rename_tensors_in_dict(dtypes)
-            g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset, g_inputs, g_outputs, is_subgraph)
-            fg = process_parsed_graph(g, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
-                                      g_outputs, {}, {}, {}, op_cnt, attr_cnt, is_tflite=True, dequantize=dequantize)
-            fg.graph_name = graph_name
+            g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset, g_inputs, g_outputs,
+                      not is_main_g, graph_name)
             if is_main_g:
-                main_g = fg
+                main_g = g
             else:
-                set_function(graph_name, fg)
+                subgraphs.append(g)
+
+        g = process_graphs(main_g, subgraphs, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter,
+                           target, {}, tensors_to_rename, is_tflite=True, dequantize=dequantize)
+        return g
+
+    # make tf2onnx internal subgraphs from the tensorflow subgraphs
+    ordered_func = resolve_functions(tf_graph)
+    subgraphs = []
+    for func in ordered_func:
+        f_inputs_names = [t.name for t in func.inputs]
+        f_output_names = [t.name for t in func.outputs]
+
+        outputs_to_values, _ = compute_const_folding_using_tf(func, const_node_values, output_names)
+
+        onnx_nodes, _, _, output_shapes, dtypes, _ = \
+            tensorflow_to_onnx(func, shape_override, const_node_values, ignore_default, use_default)
 
-        return main_g
+        fg = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset, f_inputs_names, f_output_names,
+                   is_subgraph=True, graph_name=func.name)
+        fold_constants_using_tf(fg, outputs_to_values)
+        subgraphs.append(fg)
 
     is_func = is_function(tf_graph)
     if not is_func:
         tf_graph = infer_shape(tf_graph, shape_override)
 
-    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
+    outputs_to_values, _ = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
 
-    onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, _ = \
+    onnx_nodes, _, _, output_shapes, dtypes, _ = \
         tensorflow_to_onnx(tf_graph, shape_override, const_node_values, ignore_default, use_default)
-    if not is_subgraph:
-        # make tf2onnx internal subgraphs from the tensorflow subgraphs
-        ordered_func = resolve_functions(tf_graph)
-        for func in ordered_func:
-            f_inputs_names = [t.name for t in func.inputs]
-            f_output_names = [t.name for t in func.outputs]
-            fg = process_tf_graph(func, continue_on_error, False, target, opset,
-                                  custom_op_handlers, custom_rewriter,
-                                  extra_opset, shape_override, inputs_as_nchw,
-                                  f_inputs_names, f_output_names, is_subgraph=True,
-                                  const_node_values=const_node_values, tensors_to_rename=tensors_to_rename,
-                                  initialized_tables=initialized_tables)
-            fg.graph_name = func.name
-            set_function(func.name, fg)
 
     check_io(input_names, output_names, output_shapes)
+    main_g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset, input_names, output_names,
+                   is_subgraph)
+    fold_constants_using_tf(main_g, outputs_to_values)
+    g = process_graphs(main_g, subgraphs, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter,
+                       target, initialized_tables, tensors_to_rename)
+    return g
+
+
+def process_graphs(main_g, subgraphs, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
+                   initialized_tables, tensors_to_rename, is_tflite=False, dequantize=False):
 
-    if not is_subgraph:
-        rename_tensors_in_nodes(onnx_nodes)
-        input_names = rename_tensors_in_list(input_names)
-        output_names = rename_tensors_in_list(output_names)
-        output_shapes = rename_tensors_in_dict(output_shapes)
-        dtypes = rename_tensors_in_dict(dtypes)
-        inputs_as_nchw = rename_tensors_in_list(inputs_as_nchw)
-    g = Graph(onnx_nodes, output_shapes, dtypes, target, opset, extra_opset, input_names, output_names, is_subgraph)
-    g = process_parsed_graph(g, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
-                             output_names, initialized_tables, outputs_to_values, outputs_to_dtypes, op_cnt, attr_cnt)
+    if tensors_to_rename is not None:
+        main_g.rename_tensors(tensors_to_rename)
+        inputs_as_nchw = [tensors_to_rename.get(t, t) for t in inputs_as_nchw]
+
+    for g in subgraphs:
+        fg = process_parsed_graph(g, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
+                                  initialized_tables, is_tflite, dequantize)
+        set_function(fg.graph_name, fg)
+    g = process_parsed_graph(main_g, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
+                             initialized_tables, is_tflite,
+                             dequantize)
     return g
 
 
 def process_parsed_graph(g, custom_op_handlers, inputs_as_nchw, continue_on_error, custom_rewriter, target,
-                         output_names, initialized_tables, outputs_to_values, outputs_to_dtypes, op_cnt, attr_cnt,
-                         is_tflite=False, dequantize=False):
+                         initialized_tables, is_tflite=False, dequantize=False):
+
+    op_cnt, attr_cnt = g.dump_node_statistics(include_attrs=True, include_subgraphs=False)
 
     if is_tflite:
         tfl_rewriters = []
@@ -587,8 +579,6 @@ def compat_handler(ctx, node, **kwargs):
     if inputs_as_nchw:
         transpose_inputs(g, inputs_as_nchw)
 
-    fold_constants_using_tf(g, outputs_to_values, outputs_to_dtypes)
-
     # pre-processing graph rewrites
     # bi-directional re-writer should be placed after single directional re-writer
     rewriters = [
@@ -626,7 +616,7 @@ def compat_handler(ctx, node, **kwargs):
     run_rewriters(g, rewriters, continue_on_error)
 
     # some nodes may already copied into inner Graph, so remove them from main Graph.
-    g.delete_unused_nodes(output_names)
+    g.delete_unused_nodes(g.outputs)
     topological_sort(g, continue_on_error)
 
     mapped_op, unmapped_op, exceptions = \