Refactor tfonnx

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

Author: TomWildenhain-Microsoft

tf2onnx/graph.py

@@ -874,6 +874,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 +1249,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/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,74 @@ 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)
-
-        return main_g
-
-    is_func = is_function(tf_graph)
-    if not is_func:
-        tf_graph = infer_shape(tf_graph, shape_override)
+                subgraphs.append(g)
 
-    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
+        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
 
-    onnx_nodes, op_cnt, attr_cnt, 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)
+        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]
-            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)
+
+            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)
+
+            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, _ = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
+
+    onnx_nodes, _, _, output_shapes, dtypes, _ = \
+        tensorflow_to_onnx(tf_graph, shape_override, const_node_values, ignore_default, use_default)
 
     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
 
-    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)
+
+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 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 +580,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 +617,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 = \