Pylint

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1849,7 +1849,7 @@ def func(m):
         feed_dict = {_INPUT: m_val}
         if "input" in g.input_names:
             # TFLite inputs don't have port numbers
-            feed_dict = { k.split(":")[0]: v for k, v in feed_dict.items() }
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
         results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(8)))
@@ -1867,7 +1867,7 @@ def func(n, m):
         feed_dict = {_INPUT: n_val, _INPUT1: m_val}
         if "input" in g.input_names:
             # TFLite inputs don't have port numbers
-            feed_dict = { k.split(":")[0]: v for k, v in feed_dict.items() }
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
         results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(2, 10)))
@@ -1887,7 +1887,7 @@ def func(n, m, s):
         feed_dict = {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val}
         if "input" in g.input_names:
             # TFLite inputs don't have port numbers
-            feed_dict = { k.split(":")[0]: v for k, v in feed_dict.items() }
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
         results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(2, 10)))

tf2onnx/tf_utils.py

@@ -319,10 +319,12 @@ def replace_placeholders_with_tables(graph_def, placeholder_to_table_info):
             n.attr['value_dtype'].type = val_dtype
 
 def read_tf_node_def_attrs(node_def, input_dtypes, input_shapes):
-    from tf2onnx.tf_loader import tf_session, tf_placeholder
+    """Given a tf node def, returns a dict of attribute names to values"""
+    from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
     del node_def.input[:]
     node_def.name = "node"
 
+    # read_tf_node_attrs uses some tf methods that require the node to be loaded into a valid TF graph
     g = tf.Graph()
     with g.as_default():
         for i, (dtype, shape) in enumerate(zip(input_dtypes, input_shapes)):
@@ -340,6 +342,7 @@ def read_tf_node_def_attrs(node_def, input_dtypes, input_shapes):
 
 
 def read_tf_node_attrs(node):
+    """Given a tf Node, returns a dict of attribute names to values"""
     attr = {}
     attr_cnt = collections.Counter()
 

tf2onnx/tflite_utils.py

@@ -18,7 +18,7 @@
 from tf2onnx.tflite.TensorType import TensorType as TFLiteTensorType
 from tf2onnx.tflite.Model import Model
 from tf2onnx.flexbuffers import read_flexbuffer
-from tf2onnx.tf_utils import read_tf_node_attrs, read_tf_node_def_attrs
+from tf2onnx.tf_utils import read_tf_node_def_attrs
 from tf2onnx import utils
 
 logger = logging.getLogger(__name__)
@@ -126,6 +126,7 @@ def get_options_class(name):
 def read_tflite_model(tflite_path):
     """
     Given the path to a tflite model, returns tuple (tflite_graphs, opcodes_map, model)
+    Graphs are topologically sorted and the main graph is last
     Pass these to parse_tflite_graph
     """
     with open(tflite_path, 'rb') as f:
@@ -160,14 +161,16 @@ def read_tflite_model(tflite_path):
                 tensor_shapes[name] = details["shape"].tolist()
     except Exception as e:    # pylint: disable=broad-except
         logger.warning("Error loading model into tflite interpreter: %s", e)
+    tflite_graphs = get_model_subgraphs(model)
     return tflite_graphs, opcodes_map, model, tensor_shapes
 
 
-def get_subgraph_dependencies(tflite_g, opcodes_map, model):
-    """Returns a list of subgraph names referenced by the provided graph"""
+def get_subgraph_dependencies(model, graph_idx):
+    """Returns a list of subgraph indices referenced by the indicated graph"""
     dependencies = []
-    for i in range(tflite_g.OperatorsLength()):
-        op = tflite_g.Operators(i)
+    g = model.Subgraphs(graph_idx)
+    for i in range(g.OperatorsLength()):
+        op = g.Operators(i)
         options_type_name = lookup_enum(op.BuiltinOptionsType(), 'BuiltinOptions')
         option_class = get_options_class(options_type_name)
         if option_class is not None:
@@ -176,21 +179,20 @@ def get_subgraph_dependencies(tflite_g, opcodes_map, model):
             for attr in FUNCTION_ATTRS:
                 if hasattr(options, attr):
                     value = getattr(options, attr)()
-                    dependencies.append(model.Subgraphs(value).Name().decode())
+                    dependencies.append(value)
     return dependencies
 
 
-def topsort_tfl_subgraphs(tflite_graphs, opcodes_map, model):
+def get_model_subgraphs(model):
     """Returns topologically sorted subgraphs of a model. Guarantees main graph is placed at the end."""
-    main_g = tflite_graphs[0].Name().decode()
+    main_g = 0
     dependencies = {}
-    name_to_graph = {}
-    for g in tflite_graphs:
-        name = g.Name().decode()
-        name_to_graph[name] = g
-        ds = get_subgraph_dependencies(g, opcodes_map, model)
-        utils.make_sure(main_g not in ds, "Main graph %s is a dependency of subgraph %s", main_g, name)
-        dependencies[name] = ds
+    idx_to_graph = {}
+    for i in range(model.SubgraphsLength()):
+        idx_to_graph[i] = model.Subgraphs(i)
+        ds = get_subgraph_dependencies(model, i)
+        utils.make_sure(main_g not in ds, "Main graph %s is a dependency of subgraph %s", main_g, i)
+        dependencies[i] = ds
 
     ordered = []
     visited = set()
@@ -206,10 +208,10 @@ def visit(g):
         ordered.append(g)
         visiting.remove(g)
 
-    for g in reversed(tflite_graphs):
-        visit(g.Name().decode())
+    for g in reversed(range(model.SubgraphsLength())):
+        visit(g)
 
-    return [name_to_graph[n] for n in ordered]
+    return [idx_to_graph[i] for i in ordered]
 
 
 def get_quantization_attr(quant_params):

tf2onnx/tfonnx.py

@@ -26,7 +26,7 @@
 from tf2onnx.shape_inference import infer_shape
 from tf2onnx.tf_loader import is_function, resolve_functions, set_function
 from tf2onnx.tf_utils import tensorflow_to_onnx, get_tf_version, compute_const_folding_using_tf
-from tf2onnx.tflite_utils import read_tflite_model, parse_tflite_graph, topsort_tfl_subgraphs
+from tf2onnx.tflite_utils import read_tflite_model, parse_tflite_graph
 
 from . import constants, logging, schemas, utils, handler
 
@@ -462,7 +462,6 @@ def rename_tensors_in_nodes(onnx_nodes):
 
     if tflite_path is not None:
         tflite_graphs, opcodes, model, tensor_shapes = read_tflite_model(tflite_path)
-        tflite_graphs = topsort_tfl_subgraphs(tflite_graphs, opcodes, model)
         main_g = None
         inputs_as_nchw = rename_tensors_in_list(inputs_as_nchw)
         for i, tfl_graph in enumerate(tflite_graphs):