Merge pull request #384 from pengwa/opset_9_scan

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_opset", "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_opset(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

@@ -1402,8 +1402,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_opset(9, "ReverseSequence")
     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]],
@@ -1433,8 +1433,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_opset(9, "ReverseSequence")
     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]],
@@ -1465,8 +1465,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_opset
 
 
 # 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_opset(9, "ReverseSequence")
     def test_bidrectional_attention_wrapper_lstm_encoder(self):
         size = 30
         time_step = 3

tf2onnx/graph.py

@@ -942,6 +942,9 @@ def replace_all_inputs(ops, old_input, new_input):
             return
 
         for node in ops:
+            if old_input in node.input and new_input in node.output:
+                raise RuntimeError("creating a circle in the graph is not allowed: " + node.name)
+
             for i, input_name in enumerate(node.input):
                 if input_name == old_input:
                     node.input[i] = new_input

tf2onnx/optimizer/transpose_optimizer.py

@@ -45,8 +45,6 @@ def __init__(self, graph, output_names, debug=False):
         self._handler_map = {}
         self._force_stop = {}
 
-        # make sure all proto of nodes or attribtues are update to date
-        self._g.update_proto()
         self._initialize_handlers()
         self.pre_optimize_action()
 
@@ -74,9 +72,6 @@ def pre_optimize_action(self):
                         if name == output_name:
                             child.input[i] = const_name
             self._g.make_const(const_name, new_data)
-
-            # need call this to make input update synced to protobuf val
-            self._g.update_proto()
             self._g.remove_node(reshape_op.name)
             self._g.topological_sort(self._g.get_nodes())
 
@@ -111,7 +106,6 @@ def post_optimize_action(self):
                                           name=op_name)
                     else:
                         self._remove_useless_tranpose(op)
-        self._g.update_proto()
         self._g.topological_sort(self._g.get_nodes())
 
     def merge_duplicated_transposes(self):

tf2onnx/rewriter/bigru_rewriter.py

@@ -116,7 +116,6 @@ def process_bigru(g, bi_grus):
             raise ValueError(
                 "Reverse is still used by GRU as input, cannot remove")
 
-    g.update_proto()
     return g.get_nodes()
 
 

tf2onnx/rewriter/bilstm_rewriter.py

@@ -101,8 +101,6 @@ def process_bilstm(g, bi_lstms):
         else:
             raise ValueError("Reverse is still used by LSTM as input, cannot remove")
 
-
-    g.update_proto()
     return g.get_nodes()
 
 

tf2onnx/rewriter/custom_rnn_rewriter.py

@@ -85,6 +85,10 @@ def _parse_rnn_loop(self, context):
         return True
 
     def need_rewrite(self, context):
+        if self.g.opset < 8:
+            log.debug("skip the custom_rnn_rewriter due to lower opset version %s", self.g.opset)
+            return False
+
         context.rnn_scope = self._get_rnn_scope_name(context.while_context_scope)
 
         if not self._parse_rnn_loop(context):
@@ -105,13 +109,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:  # since opset 9
+                    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:  # since opset 9
+                    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 +165,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 +192,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:  # since opset 9
+                    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:  # since opset 9
+                    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 NotImplementedError("ReverseSequence is not supported to convert in OPSET 9,"
+                              " if possible please try using 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
@@ -1824,6 +1833,7 @@ def where_op(ctx, node, name, args):
     "IsNan": (direct_op, ["IsNaN"]),
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
+    "ReverseSequence": (reverse_op9, []),
     "Sign": (sign_op9, []),
     "Sinh": (direct_op, []),
     "Where": (where_op, []),
@@ -2451,8 +2461,6 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     # onnx requires topological sorting
     topological_sort(g, continue_on_error)
 
-    g.update_proto()
-
     if verbose:
         print("tensorflow ops: {}".format(op_cnt))
         print("tensorflow attr: {}".format(attr_cnt))