Merge pull request #442 from lucienwang1009/reverse_sequence

reverse sequence

Author: nbcsm

tests/common.py

@@ -13,8 +13,8 @@
 from parameterized import parameterized
 from tf2onnx import constants, utils
 
-__all__ = ["TestConfig", "get_test_config", "unittest_main",
-           "check_tf_min_version", "skip_tf_versions",
+__all__ = ["TestConfig", "get_test_config", "unittest_main", "check_onnxruntime_backend",
+           "check_tf_min_version", "skip_tf_versions", "check_onnxruntime_min_version",
            "check_opset_min_version", "check_target", "skip_caffe2_backend", "skip_onnxruntime_backend",
            "skip_opset", "check_onnxruntime_incompatibility", "validate_const_node",
            "group_nodes_by_type", "test_ms_domain", "check_node_domain"]
@@ -177,6 +177,21 @@ def skip_onnxruntime_backend(message=""):
     return unittest.skipIf(config.is_onnxruntime_backend, reason)
 
 
+def check_onnxruntime_backend(message=""):
+    """ Skip if backend is NOT onnxruntime """
+    config = get_test_config()
+    reason = _append_message("only supported by onnxruntime", message)
+    return unittest.skipIf(not config.is_onnxruntime_backend, reason)
+
+
+def check_onnxruntime_min_version(min_required_version, message=""):
+    """ Skip if onnxruntime version < min_required_version """
+    config = get_test_config()
+    reason = _append_message("conversion requires onnxruntime >= {}".format(min_required_version), message)
+    return unittest.skipIf(config.is_onnxruntime_backend and
+                           config.backend_version < LooseVersion(min_required_version), reason)
+
+
 def skip_caffe2_backend(message=""):
     """ Skip if backend is caffe2 """
     config = get_test_config()
@@ -253,6 +268,7 @@ def check_gru_count(graph, expected_count):
 def test_ms_domain(versions=None):
     """ Parameterize test case to apply ms opset(s) as extra_opset. """
 
+    @check_onnxruntime_backend()
     def _custom_name_func(testcase_func, param_num, param):
         del param_num
         arg = param.args[0]

tests/run_pretrained_models.py

@@ -155,10 +155,10 @@ def run_tensorflow(self, sess, inputs):
             self.tf_runtime = time.time() - start
         return result
 
-    def to_onnx(self, tf_graph, opset=None, shape_override=None, input_names=None):
+    def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None):
         """Convert graph to tensorflow."""
         return process_tf_graph(tf_graph, continue_on_error=False, verbose=True, opset=opset,
-                                target=Test.target, shape_override=shape_override,
+                                extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
                                 input_names=input_names, output_names=self.output_names)
 
     def run_caffe2(self, name, model_proto, inputs):
@@ -207,7 +207,8 @@ def create_onnx_file(name, model_proto, inputs, outdir):
         utils.save_protobuf(model_path, model_proto)
         print("\tcreated", model_path)
 
-    def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=None, perf=None, fold_const=None):
+    def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=None, extra_opset=None,
+                 perf=None, fold_const=None):
         """Run complete test against backend."""
         print(name)
         self.perf = perf
@@ -267,8 +268,8 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
             model_proto = None
             try:
                 # convert model to onnx
-                onnx_graph = self.to_onnx(sess.graph, opset=opset, shape_override=shape_override,
-                                          input_names=inputs.keys())
+                onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
+                                          shape_override=shape_override, input_names=inputs.keys())
                 model_proto = onnx_graph.make_model("converted from tf2onnx")
                 new_model_proto = optimizer.optimize_graph(onnx_graph, debug=debug).make_model("optimized")
                 if new_model_proto:
@@ -328,6 +329,8 @@ def get_args():
                         choices=["caffe2", "onnxmsrtnext", "onnxruntime"], help="backend to use")
     parser.add_argument("--verbose", help="verbose output", action="store_true")
     parser.add_argument("--opset", type=int, default=None, help="opset to use")
+    parser.add_argument("--extra_opset", default=None,
+                        help="extra opset with format like domain:version, e.g. com.microsoft:1")
     parser.add_argument("--debug", help="debug vlog", action="store_true")
     parser.add_argument("--list", help="list tests", action="store_true")
     parser.add_argument("--onnx-file", help="create onnx file in directory")
@@ -338,6 +341,11 @@ def get_args():
     args = parser.parse_args()
 
     args.target = args.target.split(",")
+    if args.extra_opset:
+        tokens = args.extra_opset.split(':')
+        if len(tokens) != 2:
+            raise ValueError("invalid extra_opset argument")
+        args.extra_opset = [utils.make_opsetid(tokens[0], int(tokens[1]))]
     return args
 
 
@@ -385,7 +393,8 @@ def main():
         count += 1
         try:
             ret = t.run_test(test, backend=args.backend, debug=args.debug, onnx_file=args.onnx_file,
-                             opset=args.opset, perf=args.perf, fold_const=args.fold_const)
+                             opset=args.opset, extra_opset=args.extra_opset, perf=args.perf,
+                             fold_const=args.fold_const)
         except Exception as ex:
             ret = None
             print(ex)

tests/test_backend.py

@@ -1539,17 +1539,19 @@ def test_erf(self):
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val}, rtol=0.01)
 
-    @check_opset_min_version(8, "Scan")
-    @skip_opset(9, "ReverseSequence")
-    def test_reverse_sequence_batch_major(self):
+    def _test_reverse_sequence_batch_major(self, extra_opset=None):
+        process_args = {}
+        if extra_opset is not None:
+            process_args["extra_opset"] = [extra_opset]
+
         x_val = np.array([[[1, 2, 3], [4, 5, 6], [0, 0, 0]],
                           [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                           [[1, 2, 3], [0, 0, 0], [0, 0, 0]]],
                          dtype=np.float32)
         x = tf.placeholder(tf.float32, [None, 3, 3], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=1, batch_axis=0, seq_lengths=[2, 3, 1])
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
         tf.reset_default_graph()
 
         x_val = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3],
@@ -1560,19 +1562,21 @@ def test_reverse_sequence_batch_major(self):
         x = tf.placeholder(tf.float32, [None, 3], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=1, batch_axis=0, seq_lengths=[3] * 9)
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
         tf.reset_default_graph()
 
         x_val_shape = [5, 5, 7, 8, 9]
         x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
         x = tf.placeholder(tf.float32, [None, 5, 7, 8, 9], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=1, batch_axis=0, seq_lengths=[5, 5, 5, 5, 5])
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
+
+    def _test_reverse_sequence_time_major(self, extra_opset=None):
+        process_args = {}
+        if extra_opset is not None:
+            process_args["extra_opset"] = [extra_opset]
 
-    @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]],
                           [[0, 0, 0], [7, 8, 9], [0, 0, 0]]
@@ -1581,7 +1585,7 @@ def test_reverse_sequence_time_major(self):
         x = tf.placeholder(tf.float32, [3, None, 3], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=0, batch_axis=1, seq_lengths=[2, 3, 1])
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
         tf.reset_default_graph()
 
         x_val = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3],
@@ -1592,15 +1596,36 @@ def test_reverse_sequence_time_major(self):
         x = tf.placeholder(tf.float32, [9, None], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=0, batch_axis=1, seq_lengths=[9, 9, 9])
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
         tf.reset_default_graph()
 
         x_val_shape = [5, 5, 7, 8, 9]
         x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
         x = tf.placeholder(tf.float32, [5, None, 7, 8, 9], name=_TFINPUT)
         x_ = tf.reverse_sequence(x, seq_axis=0, batch_axis=1, seq_lengths=[5, 5, 5, 5, 5])
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        self._run_test_case([_OUTPUT], {_INPUT: x_val}, process_args=process_args)
+
+    @check_opset_min_version(8, "Scan")
+    @skip_opset(9, "ReverseSequence")
+    def test_reverse_sequence_batch_major(self):
+        self._test_reverse_sequence_batch_major()
+
+    @check_opset_min_version(8, "Scan")
+    @skip_opset(9, "ReverseSequence")
+    def test_reverse_sequence_time_major(self):
+        self._test_reverse_sequence_time_major()
+
+    # only support onnxruntime with version larger than 0.4.0
+    @test_ms_domain()
+    @check_onnxruntime_min_version("0.4.0")
+    def test_ms_reverse_sequence_batch_major(self, extra_opset):
+        self._test_reverse_sequence_batch_major(extra_opset)
+
+    @test_ms_domain()
+    @check_onnxruntime_min_version("0.4.0")
+    def test_ms_reverse_sequence_time_major(self, extra_opset):
+        self._test_reverse_sequence_time_major(extra_opset)
 
     @check_opset_min_version(8, "where")
     def test_where(self):

tf2onnx/custom_opsets/ms.py

@@ -36,3 +36,41 @@ def version_1(cls, ctx, node, **kwargs):
         utils.make_sure(dtype is not None, "Tidx of %s is None", node.name)
         ctx.remove_node(node.name)
         make_range(ctx, node.input[0], node.input[1], node.input[2], node.output[0], node.name, shape, dtype)
+
+
+@tf_op("ReverseSequence", domain=constants.MICROSOFT_DOMAIN)
+class ReverseSequence:
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        """ReverseSequence"""
+        # T output = ReverseSequence(T input, int32|int64 seq_lengths, @int seq_dim, @int batch_dim)
+        # T output = ReverseSequence(T input, int32 seqence_lens, @int time_axis, @int batch_axis)
+        seq_dim = node.get_attr("seq_dim")
+        utils.make_sure(seq_dim is not None, "sequence dim must be given in {}".format(node.name))
+        seq_dim = seq_dim.i
+        batch_dim = node.get_attr("batch_dim")
+        if batch_dim is not None:
+            batch_dim = batch_dim.i
+        else:
+            batch_dim = 0
+
+        output_dtypes = node.output_dtypes
+        output_shapes = node.output_shapes
+        ctx.remove_node(node.name)
+        node = ctx.make_node(
+            "ReverseSequence",
+            node.input,
+            outputs=node.output,
+            shapes=output_shapes,
+            dtypes=output_dtypes,
+            domain=constants.MICROSOFT_DOMAIN,
+            attr={"time_axis": seq_dim, "batch_axis": batch_dim}
+        )
+
+        seq_len_dtype = ctx.get_dtype(node.input[1])
+        utils.make_sure(seq_len_dtype is not None, "dtype of {} is None".format(node.input[1]))
+        target_dtype = TensorProto.INT32
+        if seq_len_dtype != target_dtype:
+            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=target_dtype)
+            ctx.copy_shape(cast_node.input[0], cast_node.output[0])
+            ctx.set_dtype(cast_node.output[0], target_dtype)