Created fold_constants_using_tf

Author: TomWildenhain-Microsoft

tf2onnx/tf_utils.py

@@ -144,7 +144,7 @@ def compute_const_folding_using_tf(g, const_node_values):
     """Find nodes with constant inputs and compute their values using TF"""
     if const_node_values is None:
         const_node_values = {}
-    from tf2onnx.tf_loader import tf_session, tf_reset_default_graph
+    from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
 
     ops = g.get_operations()
     outputs_to_values = {}
@@ -167,28 +167,34 @@ def compute_const_folding_using_tf(g, const_node_values):
             # Find ops with constant inputs and compute their values
             input_names = [i.name for i in node.inputs]
             output_names = [i.name for i in node.outputs]
-            can_fold = len(input_names) > 0 and all(inp in outputs_to_values for inp in input_names)
+            can_fold = node.type not in ['Enter']
+            can_fold = can_fold and len(input_names) > 0 and all(inp in outputs_to_values for inp in input_names)
             # We can only fold nodes with a single output
             can_fold = can_fold and len(output_names) == 1 and output_names[0] not in outputs_to_values
             # Skip if value already computed, used, and discarded
             can_fold = can_fold and output_names[0] not in unneeded_outputs
             if can_fold:
-                g = tf.Graph()
-                with g.as_default():
+                # Make a mini graph containing just the node to fold
+                g2 = tf.Graph()
+                with g2.as_default():
                     for inp in input_names:
-                        tf.compat.v1.placeholder(outputs_to_dtypes[inp], name=inp.split(':')[0])
-                mini_graph_def = g.as_graph_def()
-                mini_graph_def.node.append(node.node_def)
-                tf_reset_default_graph()
-                feed_dict = {}
-                for inp in input_names:
-                    feed_dict[inp] = outputs_to_values[inp]
-                with tf_session() as sess:
-                    tf.import_graph_def(mini_graph_def, name='')
-                    results = sess.run(output_names, feed_dict=feed_dict)
-                outputs_to_values[output_names[0]] = results[0]
-                outputs_to_dtypes[output_names[0]] = node.outputs[0].dtype
-                progress = True
+                        tf_placeholder(outputs_to_dtypes[inp], name=inp.split(':')[0])
+                    mini_graph_def = g2.as_graph_def()
+                    mini_graph_def.node.append(node.node_def)
+                g3 = tf.Graph()
+                with g3.as_default():
+                    feed_dict = {}
+                    for inp in input_names:
+                        feed_dict[inp] = outputs_to_values[inp]
+                    try:
+                        with tf_session() as sess:
+                            tf.import_graph_def(mini_graph_def, name='')
+                            results = sess.run(output_names, feed_dict=feed_dict)
+                        outputs_to_values[output_names[0]] = results[0]
+                        outputs_to_dtypes[output_names[0]] = node.outputs[0].dtype
+                        progress = True
+                    except Exception:  # pylint: disable=broad-except
+                        logger.debug("Could not fold node %s", node.name)
         unneeded_outputs.update(outputs_to_values.keys())
         for node in ops:
             # Mark values we need to keep
@@ -205,6 +211,12 @@ def compute_const_folding_using_tf(g, const_node_values):
                 del outputs_to_values[node]
                 del outputs_to_dtypes[node]
 
+    for node in ops:
+        # We don't need the constants any more
+        if node.type in ["Const", "ConstV2"] and node.outputs[0].name in outputs_to_values:
+            del outputs_to_values[node.outputs[0].name]
+            del outputs_to_dtypes[node.outputs[0].name]
+
     logger.info("Computed %d values for constant folding", len(outputs_to_values))
     return outputs_to_values, outputs_to_dtypes
 

tf2onnx/tfonnx.py

@@ -23,7 +23,7 @@
 from tf2onnx.rewriter import *  # pylint: disable=wildcard-import
 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
+from tf2onnx.tf_utils import tensorflow_to_onnx, get_tf_version, compute_const_folding_using_tf
 
 from . import constants, logging, schemas, utils, handler
 
@@ -33,6 +33,35 @@
 # 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):
+    ops = g.get_nodes()
+    # pylint: disable=too-many-nested-blocks
+    keep_looking = True
+    while keep_looking:
+        keep_looking = False
+        for idx, op in enumerate(ops):
+            if op.output and op.output[0] in outputs_to_values:
+                logger.info("folding node using tf type=%s, name=%s" % (op.type, op.name))
+                val = outputs_to_values[op.output[0]]
+
+                new_node_name = utils.make_name(op.name)
+                new_output_name = new_node_name
+                old_output_name = op.output[0]
+                old_node_name = op.name
+                logger.debug("create const node [%s] replacing [%s]", new_node_name, old_node_name)
+                ops[idx] = g.make_const(new_node_name, val)
+
+                logger.debug("replace old output [%s] with new output [%s]", old_output_name, new_output_name)
+                # need to re-write the consumers input name to use the const name
+                consumers = g.find_output_consumers(old_output_name)
+                if consumers:
+                    for consumer in consumers:
+                        g.replace_input(consumer, old_output_name, new_output_name)
+
+                # keep looking until there is nothing we can fold.
+                keep_looking = True
+
+    g.reset_nodes(ops)
 
 def rewrite_constant_fold(g, ops):
     """
@@ -378,6 +407,8 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if target is None:
         target = constants.DEFAULT_TARGET
 
+    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values)
+
     onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, _ = \
         tensorflow_to_onnx(tf_graph, shape_override, const_node_values)
     if not is_subgraph:
@@ -451,6 +482,8 @@ 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 = [rewrite_constant_fold, rewrite_quantize_and_dequantize, rewrite_transpose, rewrite_flatten,