add reverse support for dynamic axi

Author: RandySheriffH

tests/test_backend.py

@@ -20,6 +20,7 @@
 from common import *  # pylint: disable=wildcard-import,unused-wildcard-import
 from tf2onnx import constants, utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+from tensorflow.python.ops import init_ops
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument
 
@@ -3010,6 +3011,24 @@ def test_hashtable_lookup(self):
         self._run_test_case([lookup_results.name], {_INPUT: query})
         os.remove(filnm)
 
+    @check_opset_min_version(11, "GRU")
+    def test_cudnngru(self):
+        seq_length = 3
+        batch_size = 5
+        input_size = 2
+        num_layers = 2
+        num_units = 2
+        num_dirs = 2
+        initializer = init_ops.constant_initializer(0.5)
+        x = np.random.randint(0, 100, [seq_length, batch_size, input_size]).astype(np.float32)
+        h = np.random.randint(0, 100, [num_layers * num_dirs, batch_size, num_units]).astype(np.float32).reshape(
+            [num_layers * num_dirs, batch_size, num_units])
+        cudnngru = tf.contrib.cudnn_rnn.CudnnGRU(num_layers, num_units, 'linear_input', 'bidirectional',
+                                                 kernel_initializer=initializer, bias_initializer=initializer)
+        cudnngru.build([seq_length, batch_size, input_size])
+        outputs = cudnngru.call(x, tuple([h]))
+        self.run_test_case({}, [], [outputs[0].name], rtol=1e-05, atol=1e-04)
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/rnn.py

@@ -175,12 +175,6 @@ def version_7(cls, ctx, node, **kwargs):
 class CudnnRNN:
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
-        #print ("CudnnRNN captured")
-        #print (node.attr["direction"].s)
-        #ii = input(os.getpid())
-        #print ("input\n", node.input)
-        #print ("\noutput\n", node.output)
-        #print ("\nattr\n", node.attr)
         X = node.input[0]
         X_shape = ctx.get_shape(X)
         H = node.input[1]
@@ -190,6 +184,14 @@ def version_11(cls, ctx, node, **kwargs):
             node.attr["rnn_mode"].s == b"gru",
             "rnn mode other than gru are not supported yet"
         )
+        utils.make_sure(
+            node.attr["dropout"].f == 0,
+            "dropout not supported yet"
+        )
+        utils.make_sure(
+            node.attr["input_mode"].s == b"linear_input",
+            "input mode must be linear input"
+        )
         num_dirs = 1 if node.attr["direction"].s == b"unidirectional" else 2
         num_layers = int(H_shape[0]/num_dirs)
         num_units = hidden_size = H_shape[2]
@@ -217,26 +219,24 @@ def NM(nm):
         W_flattened = ctx.make_node('Slice', [P, zero_const.output[0], w_end_const.output[0]])
         R_flattened = ctx.make_node('Slice', [P, w_end_const.output[0], r_end_const.output[0]])
         B_flattened = ctx.make_node('Slice', [P, r_end_const.output[0], b_end_const.output[0]])
-        #W = utils.make_name('W')
-        #R = utils.make_name('R')
-        #B = utils.make_name('B')
-        W = ctx.make_node('Reshape', [W_flattened.output[0], w_shape_const.output[0]])
-        R = ctx.make_node('Reshape', [R_flattened.output[0], r_shape_const.output[0]])
-        B = ctx.make_node('Reshape', [B_flattened.output[0], b_shape_const.output[0]])
-        ctx.make_node('Split', [W.output[0]], outputs = WS)
-        ctx.make_node('Split', [R.output[0]], outputs = RS)
-        ctx.make_node('Split', [B.output[0]], outputs = BS)
+        W = utils.make_name('W')
+        R = utils.make_name('R')
+        B = utils.make_name('B')
+        ctx.make_node('Reshape', [W_flattened.output[0], w_shape_const.output[0]], outputs=[W])
+        ctx.make_node('Reshape', [R_flattened.output[0], r_shape_const.output[0]], outputs=[R])
+        ctx.make_node('Reshape', [B_flattened.output[0], b_shape_const.output[0]], outputs=[B])
+        ctx.make_node('Split', [W], outputs = WS)
+        ctx.make_node('Split', [R], outputs = RS)
+        ctx.make_node('Split', [B], outputs = BS)
         ctx.make_node('Split', [H], outputs = HS)
         XNF = XNB = X
-        gru_nodes = []
-        squeeze_nodes = []
         for i in range(num_layers):
-            suffix = '_' + str(i*2)
-            gru_nodes.append(ctx.make_node('GRU', [XNF, NM('W' + suffix), NM('R' + suffix), NM('B' + suffix), '', NM('H'+ suffix)],
-                             outputs = [NM('Y' + suffix), NM('YH' + suffix)],
-                             attr={'direction':'forward', 'hidden_size':num_units}))
+            suffix = '_' + str(i*num_dirs)
+            ctx.make_node('GRU', [XNF, NM('W' + suffix), NM('R' + suffix), NM('B' + suffix), '', NM('H'+ suffix)],
+                          outputs = [NM('Y' + suffix), NM('YH' + suffix)],
+                          attr={'direction':'forward', 'hidden_size':num_units})
             XNF = NM(X + suffix)
-            squeeze_nodes.append(ctx.make_node('Squeeze', [NM('Y' + suffix)], outputs = [XNF], attr={'axes': [1]}))
+            ctx.make_node('Squeeze', [NM('Y' + suffix)], outputs = [XNF], attr={'axes': [1]})
             if num_dirs == 2:
                 suffix = '_' + str(i*2+1)
                 ctx.make_node('GRU', [XNB, NM('W' + suffix), NM('R' + suffix), NM('B' + suffix), '', NM('H'+ suffix)],
@@ -249,5 +249,4 @@ def NM(nm):
             ctx.make_node('Concat', [XNF, XNB], outputs = [node.output[0]], attr={'axis': -1})
         else:
             identity_0 = ctx.make_node('Identity', [XNF], outputs = [node.output[0]])
-        concat_0 = ctx.make_node('Concat', YHS, outputs = [node.output[1]], attr={'axis': 0})
-        #print ("Done")
+        concat_0 = ctx.make_node('Concat', YHS, outputs = [node.output[1]], attr={'axis': 0})

tf2onnx/onnx_opset/tensor.py

@@ -1647,23 +1647,35 @@ def version_10(cls, ctx, node, **kwargs):
                     inputs = [new_node.output[0]]
 
                 # Add a Constant node (seq_len) for ReverseSequence.
-
-                # Index 1 for the shape should not return 0
-                # since the input must have rank >= 2.
-                rs_batch_size = ctx.get_shape(inputs[-1])[1]
-
-                # Make sure rs_batch_size and input_shape[axis] are not -1 each
-                utils.make_sure(input_shape[axis] is not -1 \
-                                , "shape of axis {} is unknown".format(axis))
-                utils.make_sure(rs_batch_size is not -1 \
-                                , "ReverseSequence batch size for axis {} is unknown".format(axis))
-
-                seq_list = [input_shape[axis]] * rs_batch_size
-                seq_array = np.asarray(seq_list, dtype=np.int64)  # dtype should be int64
-
-                const_seq_name = utils.make_name(const_name_root)
-                new_node = ctx.make_const(name=const_seq_name, np_val=seq_array)
-                inputs.append(new_node.output[0])
+                if ctx.opset >= 11:
+                    batch_shape = ctx.make_node("Shape", [inputs[-1]])
+                    const_one = ctx.make_const(utils.make_name(node.name + "_const_one"), np.array([1], dtype=np.int64))
+                    const_two = ctx.make_const(utils.make_name(node.name + "_const_two"), np.array([2], dtype=np.int64))
+                    batch_size = ctx.make_node("Slice",
+                                               [batch_shape.output[0], const_one.output[0], const_two.output[0]])
+                    input_shape = ctx.make_node("Shape", [node.input[0]])
+                    const_axis = ctx.make_const(utils.make_name(node.name + "_const_axis"),
+                                                np.array([axis], dtype=np.int64))
+                    const_axis_next = ctx.make_const(utils.make_name(node.name + "_const_axis_next"),
+                                                     np.array([axis + 1], dtype=np.int64))
+                    input_axis = ctx.make_node("Slice",
+                                               [input_shape.output[0], const_axis.output[0], const_axis_next.output[0]])
+                    seq_array = ctx.make_node("Expand", [input_axis.output[0], batch_size.output[0]])
+                    inputs.append(seq_array.output[0])
+                else:
+                    # Index 1 for the shape should not return 0
+                    # since the input must have rank >= 2.
+                    rs_batch_size = ctx.get_shape(inputs[-1])[1]
+                    # Make sure rs_batch_size and input_shape[axis] are not -1 each
+                    utils.make_sure(input_shape[axis] is not -1 \
+                                    , "shape of axis {} is unknown".format(axis))
+                    utils.make_sure(rs_batch_size is not -1 \
+                                    , "ReverseSequence batch size for axis {} is unknown".format(axis))
+                    seq_list = [input_shape[axis]] * rs_batch_size
+                    seq_array = np.asarray(seq_list, dtype=np.int64)  # dtype should be int64
+                    const_seq_name = utils.make_name(const_name_root)
+                    new_node = ctx.make_const(name=const_seq_name, np_val=seq_array)
+                    inputs.append(new_node.output[0])
 
                 # Add a ReverseSequence node.