Merge pull request #1145 from onnx/tom/foldSliceFromConstShape

Added constant folding for Slice nodes following a shape node of know…

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1213,6 +1213,22 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val1})
 
+    def test_slice_from_shape_const_fold(self):
+        x_val = np.array([4, 3], dtype=np.int64)
+        x_shape = np.array([-1, 3], dtype=np.int64)
+        def func(x):
+            z = tf.zeros(x)
+            x = tf.reshape(z, tf.constant(x_shape))
+            s = tf.shape(x)
+            t1 = tf.constant([1], dtype=tf.int32)
+            t2 = tf.constant([2], dtype=tf.int32)
+            y = tf.strided_slice(s, t1, t2, shrink_axis_mask=1)
+            return tf.identity(y, name=_TFOUTPUT)
+        def graph_validator(g):
+            # After constant folding just an input and const output node remain
+            return len(g.get_nodes()) == 2
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, graph_validator=graph_validator)
+
     def test_slice(self):
         x_val = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.float32)
         def func(x):

tf2onnx/tf_utils.py

@@ -20,7 +20,7 @@
 
 from onnx import helper, onnx_pb, numpy_helper
 
-from tf2onnx.utils import make_sure, is_tf_const_op, port_name
+from tf2onnx.utils import make_sure, is_tf_const_op, port_name, map_onnx_to_numpy_type
 from . import logging
 
 logger = logging.getLogger(__name__)
@@ -140,15 +140,44 @@ def compress_graph_def(graph_def):
                 tensor.tensor_content = b''
     return const_node_values
 
-def compute_const_folding_using_tf(g, const_node_values):
+def get_index_from_strided_slice_of_shape(node, outputs_to_values):
+    """Returns the  index of the dimension that the strided slice is reading from the shape node or None"""
+    attr_vals = {
+        'shrink_axis_mask': 1,
+        'ellipsis_mask': 0,
+        'begin_mask': 0,
+        'new_axis_mask': 0,
+        'end_mask': 0
+    }
+    for a in node.node_def.attr:
+        if a in attr_vals:
+            i = get_tf_node_attr(node, a)
+            if i != attr_vals[a]:
+                return None
+    i1 = outputs_to_values.get(node.inputs[1].name)
+    i2 = outputs_to_values.get(node.inputs[2].name)
+    i3 = outputs_to_values.get(node.inputs[3].name)
+    if i1 is None or i2 is None or i3 is None:
+        return None
+    if i1.shape != (1,) or i2.shape != (1,) or i3.shape != (1,):
+        return None
+    i1, i2, i3 = i1[0], i2[0], i3[0]
+    if i1 + 1 != i2 or i3 != 1:
+        return None
+    return i1
+
+def compute_const_folding_using_tf(g, const_node_values, graph_outputs):
     """Find nodes with constant inputs and compute their values using TF"""
     if const_node_values is None:
         const_node_values = {}
+    graph_outputs = set(graph_outputs)
     from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
 
     ops = g.get_operations()
     outputs_to_values = {}
     outputs_to_dtypes = {}
+    outputs_to_shapes = {}
+    shape_node_outputs = {}
 
     for node in ops:
         # Load values of constants. Use const_node_values if possible
@@ -158,6 +187,14 @@ def compute_const_folding_using_tf(g, const_node_values):
                 tensor.tensor_content = const_node_values[node.name]
             outputs_to_values[node.outputs[0].name] = get_tf_tensor_data(tensor)
             outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
+        for out in node.outputs:
+            outputs_to_shapes[out.name] = get_tf_tensor_shape(out)
+
+    for node in ops:
+        if node.type == "Shape":
+            shape = outputs_to_shapes.get(node.inputs[0].name)
+            if shape is not None:
+                shape_node_outputs[node.outputs[0].name] = shape
 
     unneeded_outputs = set()
     progress = True
@@ -167,12 +204,21 @@ 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]
+            if node.type == 'StridedSlice' and input_names[0] in shape_node_outputs \
+                                           and output_names[0] not in outputs_to_values:
+                shape = shape_node_outputs[input_names[0]]
+                i = get_index_from_strided_slice_of_shape(node, outputs_to_values)
+                if i is not None and 0 <= i < len(shape) and shape[i] is not None:
+                    np_dtype = map_onnx_to_numpy_type(map_tf_dtype(node.outputs[0].dtype))
+                    outputs_to_values[output_names[0]] = np.array(shape[i], dtype=np_dtype)
+                    outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
+                    progress = True
             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
+            can_fold = can_fold and output_names[0] not in unneeded_outputs and output_names[0] not in graph_outputs
             if can_fold:
                 # Make a mini graph containing just the node to fold
                 g2 = tf.Graph()

tf2onnx/tfonnx.py

@@ -409,7 +409,7 @@ 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)
+    outputs_to_values, outputs_to_dtypes = compute_const_folding_using_tf(tf_graph, const_node_values, output_names)
 
     onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, _ = \
         tensorflow_to_onnx(tf_graph, shape_override, const_node_values)