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

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

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



##########
File path: python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
##########
@@ -4064,3 +4064,94 @@ 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)
+    input_type = kwargs['in_type']
+
+    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_size = int(attrs.get('state_size'))
+    state_outputs = str(attrs.get('state_outputs', 'False'))
+
+    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([2], name+'_2', kwargs['initializer'])
+    create_tensor([3], name+'_3', 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+'_shape']),

Review comment:
       Thanks for the suggestion

##########
File path: python/mxnet/contrib/onnx/mx2onnx/export_onnx.py
##########
@@ -141,7 +141,12 @@ def get_outputs(sym, params, in_shape, in_label, in_type):
 
         out_names = list()
         for name in sym.list_outputs():
-            if name.endswith('_output'):
+            # handel special cases for RNN operator

Review comment:
       Done




----------------------------------------------------------------
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