Fix Infinite Loop in constant folding (#2119)

This commit fixes a bug in the function compute_const_folding_using_tf
where an already-folded StridedSlice node could set the 'update'
flag to True, causing an infinite loop.

A regression test was also added. The test checks for the infinite loop
by setting a timeout using the timeout-decorator module from pypi.

Signed-off-by: Ilya Sherstyuk <[email protected]>
Co-authored-by: Ilya Sherstyuk <[email protected]>

Author: ilyasher

ci_build/azure_pipelines/templates/keras2onnx_application_tests.yml

@@ -23,6 +23,7 @@ steps:
       pip install "protobuf<4.21.0"
       pip install h5py==3.7.0
       pip install parameterized
+      pip install timeout-decorator
       pip install coloredlogs flatbuffers
       $(INSTALL_TENSORFLOW)
       $(INSTALL_KERAS)
@@ -91,6 +92,7 @@ steps:
       pip install "protobuf<4.21.0"
       pip install h5py==3.7.0
       pip install parameterized
+      pip install timeout-decorator
       pip install coloredlogs flatbuffers
       %INSTALL_TENSORFLOW%
       %INSTALL_KERAS%

ci_build/azure_pipelines/templates/keras2onnx_unit_test.yml

@@ -23,6 +23,7 @@ steps:
       pip install "protobuf<4.21.0"
       pip install h5py==3.7.0
       pip install parameterized
+      pip install timeout-decorator
       pip install coloredlogs flatbuffers
       pip install $(TENSORFLOW_PATH)
       if [[ ! -z $KERAS ]];
@@ -77,6 +78,7 @@ steps:
       pip install "protobuf<4.21.0"
       pip install h5py==3.7.0
       pip install parameterized
+      pip install timeout-decorator
       pip install coloredlogs flatbuffers
       pip install %TENSORFLOW_PATH%
       IF NOT "%KERAS%"=="" (pip install %KERAS%)

ci_build/azure_pipelines/templates/setup.yml

@@ -3,7 +3,7 @@
 steps:
 - bash: |
     set -ex
-    pip install pytest pytest-cov pytest-runner coverage graphviz requests pyyaml pillow pandas parameterized sympy coloredlogs flatbuffers
+    pip install pytest pytest-cov pytest-runner coverage graphviz requests pyyaml pillow pandas parameterized sympy coloredlogs flatbuffers timeout-decorator
     pip install $(CI_PIP_TF_NAME) $(CI_PIP_ONNX_NAME)
 
     # TF < 2.7 reuires numpy <= 1.19, but onnxruntime >= 1.11 requires numpy >= 1.21

setup.py

@@ -75,7 +75,7 @@ def run(self):
     version=VersionInfo.version,
     description='Tensorflow to ONNX converter',
     setup_requires=['pytest-runner'],
-    tests_require=['graphviz', 'parameterized', 'pytest', 'pytest-cov', 'pyyaml'],
+    tests_require=['graphviz', 'parameterized', 'pytest', 'pytest-cov', 'pyyaml', 'timeout-decorator'],
     cmdclass=cmdclass,
     packages=find_packages(),
     license='Apache License v2.0',

tests/common.py

@@ -11,6 +11,7 @@
 
 from packaging.version import Version
 from parameterized import parameterized
+import timeout_decorator
 import numpy as np
 import tensorflow as tf
 
@@ -50,6 +51,7 @@
     "check_op_count",
     "check_gru_count",
     "check_lstm_count",
+    "timeout",
 ]
 
 
@@ -76,6 +78,10 @@ def __init__(self):
     def is_mac(self):
         return self.platform == "darwin"
 
+    @property
+    def is_windows(self):
+        return self.platform.startswith("win")
+
     @property
     def is_onnxruntime_backend(self):
         return self.backend == "onnxruntime"
@@ -493,3 +499,14 @@ def check_node_domain(node, domain):
     if not domain:
         return not node.domain
     return node.domain == domain
+
+def timeout(seconds):
+    """
+    Decorator for enforcing a time limit on a test.
+
+    NOTE: Please only use for ensuring that a test does not time out.
+    Do not write tests that intentionally time out.
+    """
+    if get_test_config().is_windows:
+        return unittest.skip("timeout testing is unreliable on Windows.")
+    return timeout_decorator.timeout(seconds)

tests/test_backend.py

@@ -1691,6 +1691,23 @@ def graph_validator(g):
             return len(g.get_nodes()) == 2
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, graph_validator=graph_validator)
 
+    @timeout(5)
+    def test_slice_const_fold_halts(self):
+        # Regression test for infinite loop during constant-folding.
+        x_val = np.array([4, 3], dtype=np.int32)
+        x_shape = np.array([-1, 3], dtype=np.int32)
+        def func(x):
+            x_reshaped = tf.reshape(tf.zeros(x), tf.constant(x_shape))
+            s = tf.shape(x_reshaped)
+            const1 = tf.constant([1], dtype=tf.int32)
+            const2 = tf.constant([2], dtype=tf.int32)
+            s_indexed = tf.strided_slice(s, const1, const2, strides=const1, shrink_axis_mask=1)
+            x_indexed = tf.strided_slice(x, const1, const2, strides=const1, shrink_axis_mask=1)
+            mul = tf.multiply(s_indexed, tf.constant(2, dtype=s_indexed.dtype))
+            add = const1 + x_indexed + mul
+            return tf.identity(add, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     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

@@ -210,7 +210,8 @@ def is_huge_shape(x):
             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:
+                                           and output_names[0] not in outputs_to_values \
+                                           and output_names[0] not in unneeded_outputs:
                 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: