[v1.x] ONNX export support for RNN and sum_axis (#20226)

* export support RNN

* add sum_axis

* fix sanity

* fix sanity

* fix sanity

* change regiester sum_axis

Co-authored-by: Wei Chu <weichu@amazon.com>

Author: waytrue17

python/mxnet/onnx/mx2onnx/_op_translations/_op_translations_opset12.py

@@ -2137,7 +2137,9 @@ def convert_square(node, **kwargs):
     )
     return [tensor_node, node]
 
+# sum_axis is equivalent to sum in MXNet
 @mx_op.register("sum")
+@mx_op.register("sum_axis")
 def convert_sum(node, **kwargs):
     """Map MXNet's sum operator attributes to onnx's ReduceSum operator
     and return the created node.
@@ -4476,12 +4478,12 @@ def convert_RNN(node, **kwargs):
     initial_h = input_nodes[2]
 
     nodes = []
+    create_tensor([0], name+'_0', kwargs['initializer'])
 
     mode = str(attrs.get('mode'))
     if mode == 'lstm':
         initial_c = input_nodes[3]
         if num_layers == 2:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
             create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
             create_tensor([1, 4*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
@@ -4553,7 +4555,6 @@ def convert_RNN(node, **kwargs):
                 make_node('Concat', [name+'_lstm0_c', name+'_lstm1_c'], [name+'2'], axis=0),
             ]
         elif num_layers == 1:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([1], name+'_1', kwargs['initializer'])
             create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
             create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
@@ -4598,7 +4599,6 @@ def convert_RNN(node, **kwargs):
 
     elif mode == 'gru':
         if num_layers == 2:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
             create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
             create_tensor([1, 3*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
@@ -4669,7 +4669,7 @@ def convert_RNN(node, **kwargs):
             ]
 
         elif num_layers == 1:
-            create_tensor([0], name+'_0', kwargs['initializer'])
+
             create_tensor([1], name+'_1', kwargs['initializer'])
             create_tensor([3*state_size], name+'_3*state_size', kwargs['initializer'])
             create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
@@ -4712,6 +4712,100 @@ def convert_RNN(node, **kwargs):
         else:
             raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
 
+    elif mode in ['rnn_tanh', 'rnn_relu']:
+        activations = ['Tanh']
+        if mode == 'rnn_relu':
+            activations = ['Relu']
+        if num_layers == 2:
+
+            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
+            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
+            create_tensor([1, state_size, state_size], name+'_WR_shape', kwargs['initializer'])
+            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
+            create_tensor([4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])
+
+            nodes += [
+                make_node('Shape', [data], [name+'_data_shape']),
+                make_node('Split', [name+'_data_shape'], [name+'_seq_length', name+'_batch_size', name+'_input_size']),
+
+                # Layer 0
+                # get W
+                make_node('Slice', [param, name+'_0', name+'_state_size^2'], [name+'_W0_1d']),
+                make_node('Reshape', [name+'_W0_1d', name+'_WR_shape'], [name+'_W0']),
+                # get R
+                make_node('Add', [name+'_state_size^2', name+'_state_size^2'], [name+'_R0_offset']),
+                make_node('Slice', [param, name+'_state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
+                make_node('Reshape', [name+'_R0_1d', name+'_WR_shape'], [name+'_R0']),
+                # get B
+                make_node('Add', [name+'_WR_offset', name+'_2*state_size'], [name+'_B0_offset']),
+                make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
+                make_node('Reshape', [name+'_B0_1d', name+'_B_shape'], [name+'_B0']),
+                # get initial states
+                make_node('Split', [initial_h], [name+'_initial_h0', name+'_initial_h1'], axis=0),
+                # get seq_len
+                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
+                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
+                # Layer 0 RNN
+                make_node('RNN', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len', name+'_initial_h0'],
+                          [name+'_rnn0_out_', name+'_rnn0_h'], hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'_rnn0_out_'], [name+'_rnn0_out'], axes=[1]),
+
+                # Layer 1
+                # get W
+                make_node('Add', [name+'_R0_offset', name+'_state_size^2'], [name+'_W1_offset']),
+                make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
+                make_node('Reshape', [name+'_W1_1d', name+'_WR_shape'], [name+'_W1']),
+                # get R
+                make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
+                make_node('Reshape', [name+'_R1_1d', name+'_WR_shape'], [name+'_R1']),
+                # get B
+                make_node('Add', [name+'_B0_offset', name+'_2*state_size'], [name+'_B1_offset']),
+                make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
+                make_node('Reshape', [name+'_B1_1d', name+'_B_shape'], [name+'_B1']),
+                # Layer 1 RNN
+                make_node('RNN', [name+'_rnn0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
+                                  name+'_initial_h1'], [name+'_rnn1_out_', name+'_rnn1_h'],
+                          hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'_rnn1_out_'], [name], axes=[1]),
+                make_node('Concat', [name+'_rnn0_h', name+'_rnn1_h'], [name+'1'], axis=0)
+            ]
+
+        elif num_layers == 1:
+
+            create_tensor([1], name+'_1', kwargs['initializer'])
+            create_tensor([state_size], name+'_state_size', kwargs['initializer'])
+            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
+            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
+            create_tensor([1, state_size, state_size], name+'_R_shape', kwargs['initializer'])
+            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
+
+            nodes += [
+                make_node('Shape', [data], [name+'_data_shape']),
+                make_node('Split', [name+'_data_shape'],
+                          [name+'_seq_length', name+'_batch_size', name+'_input_size'], name='split0'),
+                # get W
+                make_node('Mul', [name+'_state_size', name+'_input_size'], [name+'_mul0']),
+                make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
+                make_node('Concat', [name+'_1', name+'_state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
+                make_node('Reshape', [name+'_W_1d', name+'_W_shape'], [name+'_W']),
+                # get R
+                make_node('Add', [name+'_mul0', name+'_state_size^2'], [name+'_add0']),
+                make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
+                make_node('Reshape', [name+'_R_1d', name+'_R_shape'], [name+'_R']),
+                # get B
+                make_node('Add', [name+'_add0', name+'_2*state_size'], [name+'_add1']),
+                make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
+                make_node('Reshape', [name+'_B_1d', name+'_B_shape'], [name+'_B']),
+                # get seq_len
+                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
+                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
+                # compute RNN
+                make_node('RNN', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', initial_h],
+                          [name+'0_', name+'1'], hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'0_'], [name], axes=[1]),
+            ]
+        else:
+            raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
     else:
         raise NotImplementedError(f"Currently RNN onnx export does not support {mode} mode")
     return nodes

python/mxnet/onnx/mx2onnx/_op_translations/_op_translations_opset13.py

@@ -1047,11 +1047,12 @@ def convert_RNN(node, **kwargs):
     nodes = []
 
     mode = str(attrs.get('mode'))
+    create_tensor([0], name+'_0', kwargs['initializer'])
     create_tensor([1], name+'_1', kwargs['initializer'])
+
     if mode == 'lstm':
         initial_c = input_nodes[3]
         if num_layers == 2:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
             create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
             create_tensor([1, 4*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
@@ -1123,7 +1124,6 @@ def convert_RNN(node, **kwargs):
                 make_node('Concat', [name+'_lstm0_c', name+'_lstm1_c'], [name+'2'], axis=0),
             ]
         elif num_layers == 1:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
             create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
             create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
@@ -1167,7 +1167,6 @@ def convert_RNN(node, **kwargs):
 
     elif mode == 'gru':
         if num_layers == 2:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
             create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
             create_tensor([1, 3*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
@@ -1238,7 +1237,6 @@ def convert_RNN(node, **kwargs):
             ]
 
         elif num_layers == 1:
-            create_tensor([0], name+'_0', kwargs['initializer'])
             create_tensor([3*state_size], name+'_3*state_size', kwargs['initializer'])
             create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
             create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
@@ -1272,14 +1270,106 @@ def convert_RNN(node, **kwargs):
                 # get seq_len
                 make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                 make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
-                # compute LSTM
+                # compute GRU
                 make_node('GRU', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', initial_h],
                           [name+'0_', name+'1'], hidden_size=state_size, linear_before_reset=1),
                 make_node('Squeeze', [name+'0_', name+'_1'], [name]),
             ]
         else:
             raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
 
+    elif mode in ['rnn_tanh', 'rnn_relu']:
+        activations = ['Tanh']
+        if mode == 'rnn_relu':
+            activations = ['Relu']
+        if num_layers == 2:
+            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
+            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
+            create_tensor([1, state_size, state_size], name+'_WR_shape', kwargs['initializer'])
+            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
+            create_tensor([4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])
+
+            nodes += [
+                make_node('Shape', [data], [name+'_data_shape']),
+                make_node('Split', [name+'_data_shape'], [name+'_seq_length', name+'_batch_size', name+'_input_size']),
+
+                # Layer 0
+                # get W
+                make_node('Slice', [param, name+'_0', name+'_state_size^2'], [name+'_W0_1d']),
+                make_node('Reshape', [name+'_W0_1d', name+'_WR_shape'], [name+'_W0']),
+                # get R
+                make_node('Add', [name+'_state_size^2', name+'_state_size^2'], [name+'_R0_offset']),
+                make_node('Slice', [param, name+'_state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
+                make_node('Reshape', [name+'_R0_1d', name+'_WR_shape'], [name+'_R0']),
+                # get B
+                make_node('Add', [name+'_WR_offset', name+'_2*state_size'], [name+'_B0_offset']),
+                make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
+                make_node('Reshape', [name+'_B0_1d', name+'_B_shape'], [name+'_B0']),
+                # get initial states
+                make_node('Split', [initial_h], [name+'_initial_h0', name+'_initial_h1'], axis=0),
+                # get seq_len
+                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
+                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
+                # Layer 0 RNN
+                make_node('RNN', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
+                                  name+'_initial_h0'], [name+'_rnn0_out_', name+'_rnn0_h'],
+                          hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'_rnn0_out_', name+'_1'], [name+'_rnn0_out']),
+
+                # Layer 1
+                # get W
+                make_node('Add', [name+'_R0_offset', name+'_state_size^2'], [name+'_W1_offset']),
+                make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
+                make_node('Reshape', [name+'_W1_1d', name+'_WR_shape'], [name+'_W1']),
+                # get R
+                make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
+                make_node('Reshape', [name+'_R1_1d', name+'_WR_shape'], [name+'_R1']),
+                # get B
+                make_node('Add', [name+'_B0_offset', name+'_2*state_size'], [name+'_B1_offset']),
+                make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
+                make_node('Reshape', [name+'_B1_1d', name+'_B_shape'], [name+'_B1']),
+                # Layer 1 RNN
+                make_node('RNN', [name+'_rnn0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
+                                  name+'_initial_h1'], [name+'_rnn1_out_', name+'_rnn1_h'],
+                          hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'_rnn1_out_', name+'_1'], [name]),
+                make_node('Concat', [name+'_rnn0_h', name+'_rnn1_h'], [name+'1'], axis=0)
+            ]
+
+        elif num_layers == 1:
+            create_tensor([state_size], name+'_state_size', kwargs['initializer'])
+            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
+            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
+            create_tensor([1, state_size, state_size], name+'_R_shape', kwargs['initializer'])
+            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
+
+            nodes += [
+                make_node('Shape', [data], [name+'_data_shape']),
+                make_node('Split', [name+'_data_shape'], [name+'_seq_length', name+'_batch_size',
+                                                          name+'_input_size'], name='split0'),
+                # get W
+                make_node('Mul', [name+'_state_size', name+'_input_size'], [name+'_mul0']),
+                make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
+                make_node('Concat', [name+'_1', name+'_state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
+                make_node('Reshape', [name+'_W_1d', name+'_W_shape'], [name+'_W']),
+                # get R
+                make_node('Add', [name+'_mul0', name+'_state_size^2'], [name+'_add0']),
+                make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
+                make_node('Reshape', [name+'_R_1d', name+'_R_shape'], [name+'_R']),
+                # get B
+                make_node('Add', [name+'_add0', name+'_2*state_size'], [name+'_add1']),
+                make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
+                make_node('Reshape', [name+'_B_1d', name+'_B_shape'], [name+'_B']),
+                # get seq_len
+                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
+                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
+                # compute RNN
+                make_node('RNN', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', initial_h],
+                          [name+'0_', name+'1'], hidden_size=state_size, activations=activations),
+                make_node('Squeeze', [name+'0_', name+'_1'], [name]),
+            ]
+        else:
+            raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
     else:
         raise NotImplementedError(f"Currently RNN onnx export does not support {mode} mode")
     return nodes