[GitHub] [incubator-mxnet] Zha0q1 commented on a change in pull request #19958: [v1.x] ONNX export support for RNN

[GitBox <gi...@apache.org>]

Zha0q1 commented on a change in pull request #19958:
URL: https://github.com/apache/incubator-mxnet/pull/19958#discussion_r584045592



##########
File path: python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
##########
@@ -4064,3 +4064,90 @@ def convert_sequence_reverse(node, **kwargs):
         ]
 
     return nodes
+
+
+@mx_op.register("RNN")
+def convert_RNN(node, **kwargs):
+    """Map MXNet's RNN operator attributes to onnx's operators
+    and return the created node.
+    """
+    from onnx.helper import make_node
+    from onnx import TensorProto
+
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    mode = str(attrs.get('mode'))
+    if mode != 'lstm':
+        raise NotImplementedError('Currently RNN onnx export only supports lstm mode')
+
+    bidirectional = str(attrs.get('bidirectional', 'False'))
+    if bidirectional != 'False':
+        raise NotImplementedError('Currently RNN onnx export only supports bidirectional is False')
+
+    num_layers = int(attrs.get('num_layers', '1'))
+    if num_layers != 1:
+        raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1')
+
+    p = float(attrs.get('p', '0'))
+    if p != 0:
+        raise NotImplementedError('Currently RNN onnx export only supports p equals to 0')
+
+    use_sequence_length = str(attrs.get('use_sequence_length', 'False'))
+    if use_sequence_length != 'False':
+        raise NotImplementedError('Currently RNN onnx export only supports use_sequence_length equals to False')
+
+    projection_size = str(attrs.get('projection_size', 'None'))
+    if projection_size != 'None':
+        raise NotImplementedError('Currently RNN onnx export only supports projection_size equals to None')
+
+    state_outputs = str(attrs.get('state_outputs', 'False'))
+    if state_outputs != 'True':
+        raise NotImplementedError('Currently RNN onnx export only supports state_outputs equals to True')
+
+    state_size = int(attrs.get('state_size'))
+    data = input_nodes[0]
+    param = input_nodes[1]
+    initial_h = input_nodes[2]
+    initial_c = input_nodes[3]
+
+    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'])
+    create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
+    create_tensor([1, 4*state_size, state_size], name+'_R_shape', kwargs['initializer'])
+    create_tensor([1, 8*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']),
+        # get W
+        make_node('Mul', [name+'_4*state_size', name+'_input_size'], [name+'_mul0']),
+        make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
+        make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2', name+'_W3']),
+        make_node('Concat', [name+'_W0', name+'_W3', name+'_W1', name+'_W2'], [name+'_W_'], axis=0),

Review comment:
       We might consider using GatherND if favor of Split + Concat. This way performance might be better




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org