opset 9 scan support

Author: pengwa

tests/common.py

@@ -15,8 +15,9 @@
 
 __all__ = ["TestConfig", "get_test_config", "unittest_main",
            "check_tf_min_version", "skip_tf_versions",
-           "check_opset_min_version", "check_target", "skip_onnxruntime_backend", "skip_caffe2_backend",
-           "check_onnxruntime_incompatibility", "validate_const_node", "group_nodes_by_type"]
+           "check_opset_min_version", "check_target", "skip_caffe2_backend", "skip_onnxruntime_backend",
+           "skip_specific_opset_version", "check_onnxruntime_incompatibility", "validate_const_node",
+           "group_nodes_by_type"]
 
 
 # pylint: disable=missing-docstring
@@ -155,6 +156,13 @@ def check_opset_min_version(min_required_version, message=""):
     return unittest.skipIf(config.opset < min_required_version, reason)
 
 
+def skip_specific_opset_version(opset_v, message=""):
+    """ Skip if opset = opset_v """
+    config = get_test_config()
+    reason = _append_message("conversion requires opset != {}".format(opset_v), message)
+    return unittest.skipIf(config.opset == opset_v, reason)
+
+
 def check_target(required_target, message=""):
     """ Skip if required_target is NOT specified """
     config = get_test_config()

tests/test_backend.py

@@ -1398,8 +1398,8 @@ def test_erf(self):
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val}, rtol=0.01)
 
-    # @unittest.skipIf(OPSET < 8, "supported with opset 8 or better")
-    @unittest.skip("FIXME: the newest onnxruntime wheel hasn't been published to PYPI, so scan op is not supported")
+    @check_opset_min_version(8, "Scan")
+    @skip_specific_opset_version(9, "ReverseSequence not supported")
     def test_reverse_sequence_batch_major(self):
         x_val = np.array([[[1, 2, 3], [4, 5, 6], [0, 0, 0]],
                           [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -1429,8 +1429,8 @@ def test_reverse_sequence_batch_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    # @unittest.skipIf(OPSET < 8, "supported with opset 8 or better")
-    @unittest.skip("FIXME: the newest onnxruntime wheel hasn't been published to PYPI, so scan op is not supported")
+    @check_opset_min_version(8, "Scan")
+    @skip_specific_opset_version(9, "ReverseSequence not supported")
     def test_reverse_sequence_time_major(self):
         x_val = np.array([[[1, 2, 3], [1, 2, 3], [1, 2, 3]],
                           [[4, 5, 6], [4, 5, 6], [0, 0, 0]],
@@ -1461,8 +1461,7 @@ def test_reverse_sequence_time_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    # @unittest.skipIf(OPSET < 8, "supported with opset 8 or better")
-    @unittest.skip("FIXME: the newest onnxruntime wheel hasn't been published to PYPI, so Select op is not supported")
+    @check_opset_min_version(8, "where")
     def test_where(self):
         x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.int32)
         true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],

tests/test_custom_rnncell.py

@@ -13,7 +13,7 @@
 from tensorflow.contrib import rnn
 from tensorflow.python.ops import init_ops
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import check_tf_min_version, check_opset_min_version, unittest_main
+from common import check_tf_min_version, check_opset_min_version, unittest_main, skip_specific_opset_version
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -284,6 +284,7 @@ def test_multi_rnn_lstm(self, state_is_tuple=True):
 
     @check_opset_min_version(8, "Scan")
     @check_tf_min_version("1.8")
+    @skip_specific_opset_version(9, "ReverseSequence cannot be efficient mapped in OPSET 9")
     def test_bidrectional_attention_wrapper_lstm_encoder(self):
         size = 30
         time_step = 3

tf2onnx/rewriter/custom_rnn_rewriter.py

@@ -105,13 +105,19 @@ def rewrite(self, context):
 
             state_inputs_initial_values = []
             for state_input in scan_props.state_inputs_initial_values:
-                nodes = self._adapt_scan_sequence_input_or_output("input", state_input, False)
-                state_inputs_initial_values.append(nodes[-1].output[0])
+                if self.g.opset == 8:
+                    nodes = self._adapt_scan_sequence_input_or_output("input", state_input, False)
+                    state_inputs_initial_values.append(nodes[-1].output[0])
+                else:
+                    state_inputs_initial_values.append(state_input)
 
             scan_inputs_initial_values = []
             for scan_input in scan_props.scan_inputs_initial_values:
-                nodes = self._adapt_scan_sequence_input_or_output("input", scan_input, False)
-                scan_inputs_initial_values.append(nodes[-1].output[0])
+                if self.g.opset == 8:
+                    nodes = self._adapt_scan_sequence_input_or_output("input", scan_input, False)
+                    scan_inputs_initial_values.append(nodes[-1].output[0])
+                else:
+                    scan_inputs_initial_values.append(scan_input)
 
             cell_g_info = context.cell_graph
             scan_body_g = LoopRewriterBase.construct_graph_from_nodes(self.g, cell_g_info.nodes, cell_g_info.outputs)
@@ -155,17 +161,24 @@ def _create_scan_node(self, context, scan_props, init_values):
                 n = self.g.get_node_by_output(tensor_value_info.id)
                 self.g.remove_node(n.name)
             else:
-                loop_outputs_shapes.append(None)
+                loop_outputs_shapes.append([-1])
                 loop_outputs_dtypes.append(None)
 
-        # here we did not give the sequence_length, because
-        # current batch size is 1, not original batch size
-        # original seq_length will be used by the loop body of Scan op.
-        scan_node = self.g.make_node("Scan", [""] + init_values, op_name_scope="custom_rnn_scan",
-                                     attr={"num_scan_inputs": len(scan_props.scan_inputs)},
-                                     output_count=len(scan_props.state_outputs + scan_props.scan_outputs),
-                                     shapes=loop_outputs_shapes, dtypes=loop_outputs_dtypes,
-                                     skip_conversion=False)
+        if self.g.opset == 8:
+            # here we did not give the sequence_length, because
+            # current batch size is 1, not original batch size
+            # original seq_length will be used by the loop body of Scan op.
+            scan_node = self.g.make_node("Scan", [""] + init_values, op_name_scope="custom_rnn_scan",
+                                         attr={"num_scan_inputs": len(scan_props.scan_inputs)},
+                                         output_count=len(scan_props.state_outputs + scan_props.scan_outputs),
+                                         shapes=loop_outputs_shapes, dtypes=loop_outputs_dtypes,
+                                         skip_conversion=False)
+        else:
+            scan_node = self.g.make_node("Scan", init_values, op_name_scope="custom_rnn_scan",
+                                         attr={"num_scan_inputs": len(scan_props.scan_inputs)},
+                                         output_count=len(scan_props.state_outputs + scan_props.scan_outputs),
+                                         shapes=loop_outputs_shapes, dtypes=loop_outputs_dtypes,
+                                         skip_conversion=False)
 
         return scan_node
 
@@ -175,17 +188,22 @@ def _connect_scan_with_output(self, context, scan_node):
         index = 0
         for out_tensor_value_info in context.loop_properties.state_outputs_exits:
             if out_tensor_value_info.id:
-                nodes = self._adapt_scan_sequence_input_or_output("state_output_reshape",
-                                                                  scan_node.output[index], True)
-                self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
-
+                if self.g.opset == 8:
+                    nodes = self._adapt_scan_sequence_input_or_output("state_output_reshape",
+                                                                        scan_node.output[index], True)
+                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
+                else:
+                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
             index += 1
 
         for out_tensor_value_info in context.loop_properties.scan_outputs_exits:
             if out_tensor_value_info.id:
-                nodes = self._adapt_scan_sequence_input_or_output("scan_output_reshape",
-                                                                  scan_node.output[index], True)
-                self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
+                if self.g.opset == 8:
+                    nodes = self._adapt_scan_sequence_input_or_output("scan_output_reshape",
+                                                                        scan_node.output[index], True)
+                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, nodes[-1].output[0])
+                else:
+                    self.g.replace_all_inputs(self.g.get_nodes(), out_tensor_value_info.id, scan_node.output[index])
             index += 1
 
 

tf2onnx/tfonnx.py

@@ -1454,6 +1454,15 @@ def reverse_op8(ctx, node, name, args):
     node.input[0] = node.input[1]
     node.input[1] = tmp
 
+def reverse_op9(ctx, node, name, args):
+    # T output = ReverseSequence(T input, int32|int64 seq_lengths, @int seq_dim, @int batch_dim)
+    # we cannot easily construct reverse_sequence equivalence in opset 9, so we will not support it
+    # here. Actually using loops to do that is kind of meaningless since there will be performance
+    # issue there for sure.
+
+    raise RuntimeError("ReverseSequence is not supported to convert in OPSET9,"
+                       " if possible please try use OPSET 8 instead.")
+
 
 def shape_op(ctx, node, name, args):
     # out_type output = Shape(T input, @int32|int64 out_type), out_type by default int32
@@ -1817,6 +1826,7 @@ def where_op(ctx, node, name, args):
     "Less": (logical_compare_op, []),
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
+    "ReverseSequence": (reverse_op9, []),
     "Sign": (sign_op9, []),
     "Sinh": (direct_op, []),
     "Where": (where_op, []),