[MXNET-880] ONNX export: Random uniform, Random normal, MaxRoiPool

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

@@ -1738,3 +1738,76 @@ def convert_multinomial(node, **kwargs):
         name=name,
     )
     return [node]
+
+
+@mx_op.register("_random_uniform")
+def convert_random_uniform(node, **kwargs):
+    """Map MXNet's random_uniform operator attributes to onnx's RandomUniform
+    operator and return the created node.
+    """
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    # Converting to float32
+    low = float(attrs.get("low", 0))
+    high = float(attrs.get("high", 1.0))
+    shape = convert_string_to_list(attrs.get('shape', '[]'))
+    dtype = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(attrs.get('dtype', 'float32'))]
+
+    node = onnx.helper.make_node(
+        'RandomUniform',
+        input_nodes,
+        [name],
+        low=low,
+        high=high,
+        dtype=dtype,
+        shape=shape,
+        name=name
+    )
+    return [node]
+
+
+@mx_op.register("_random_normal")
+def convert_random_normal(node, **kwargs):
+    """Map MXNet's random_normal operator attributes to onnx's RandomNormal
+    operator and return the created node.
+    """
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    # Converting to float32
+    mean = float(attrs.get("loc", 0))
+    scale = float(attrs.get("scale", 1.0))
+    shape = convert_string_to_list(attrs.get('shape', '[]'))
+    dtype = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(attrs.get('dtype', 'float32'))]
+
+    node = onnx.helper.make_node(
+        'RandomNormal',
+        input_nodes,
+        [name],
+        mean=mean,
+        scale=scale,
+        dtype=dtype,
+        shape=shape,
+        name=name
+    )
+    return [node]
+
+
+@mx_op.register("ROIPooling")
+def convert_roipooling(node, **kwargs):
+    """Map MXNet's ROIPooling operator attributes to onnx's MaxRoiPool
+    operator and return the created node.
+    """
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    pooled_shape = convert_string_to_list(attrs.get('pooled_size'))
+    scale = float(attrs.get("spatial_scale"))
+
+    node = onnx.helper.make_node(
+        'MaxRoiPool',
+        input_nodes,
+        [name],
+        pooled_shape=pooled_shape,
+        spatial_scale=scale,
+        name=name
+    )
+    return [node]

python/mxnet/contrib/onnx/onnx2mx/_op_translations.py

@@ -29,14 +29,26 @@ def identity(attrs, inputs, proto_obj):
 
 def random_uniform(attrs, inputs, proto_obj):
     """Draw random samples from a uniform distribtuion."""
-    new_attr = translation_utils._remove_attributes(attrs, ['seed'])
-    return 'random_uniform', new_attr, inputs
+    try:
+        from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+    except ImportError:
+        raise ImportError("Onnx and protobuf need to be installed. "
+                          "Instructions to install - /~https://github.com/onnx/onnx")
+    new_attrs = translation_utils._remove_attributes(attrs, ['seed'])
+    new_attrs['dtype'] = TENSOR_TYPE_TO_NP_TYPE[int(new_attrs.get('dtype', 1))]
+    return 'random_uniform', new_attrs, inputs
 
 def random_normal(attrs, inputs, proto_obj):
     """Draw random samples from a Gaussian distribution."""
+    try:
+        from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+    except ImportError:
+        raise ImportError("Onnx and protobuf need to be installed. "
+                          "Instructions to install - /~https://github.com/onnx/onnx")
     new_attr = translation_utils._remove_attributes(attrs, ['seed'])
-    new_attr = translation_utils._fix_attribute_names(new_attr, {'mean' : 'loc'})
-    return 'random_uniform', new_attr, inputs
+    new_attr = translation_utils._fix_attribute_names(new_attr, {'mean': 'loc'})
+    new_attr['dtype'] = TENSOR_TYPE_TO_NP_TYPE[int(new_attr.get('dtype', 1))]
+    return 'random_normal', new_attr, inputs
 
 def sample_multinomial(attrs, inputs, proto_obj):
     """Draw random samples from a multinomial distribution."""

tests/python-pytest/onnx/test_node.py

@@ -139,9 +139,18 @@ def get_onnx_graph(testname, input_names, inputs, output_name, output_shape, att
             test_name, mxnet_op, onnx_name, inputs, attrs, mxnet_specific, fix_attrs, check_value, check_shape = test
             with self.subTest(test_name):
                 names, input_tensors, inputsym = get_input_tensors(inputs)
-                test_op = mxnet_op(*inputsym, **attrs)
-                mxnet_output = forward_pass(test_op, None, None, names, inputs)
-                outputshape = np.shape(mxnet_output)
+                if inputs:
+                    test_op = mxnet_op(*inputsym, **attrs)
+                    mxnet_output = forward_pass(test_op, None, None, names, inputs)
+                    outputshape = np.shape(mxnet_output)
+                else:
+                    test_op = mxnet_op(**attrs)
+                    shape = attrs.get('shape', (1,))
+                    x = mx.nd.zeros(shape, dtype='float32')
+                    xgrad = mx.nd.zeros(shape, dtype='float32')
+                    exe = test_op.bind(ctx=mx.cpu(), args={'x': x}, args_grad={'x': xgrad})
+                    mxnet_output = exe.forward(is_train=False)[0].asnumpy()
+                    outputshape = np.shape(mxnet_output)
 
                 if mxnet_specific:
                     onnxmodelfile = onnx_mxnet.export_model(test_op, {}, [np.shape(ip) for ip in inputs],
@@ -204,11 +213,19 @@ def get_onnx_graph(testname, input_names, inputs, output_name, output_shape, att
      {'kernel': (4, 5), 'pad': (0, 0), 'stride': (1, 1), 'p_value': 2, 'pool_type': 'lp', 'global_pool': True}, False,
      {'modify': {'p_value': 'p'},
       'remove': ['pool_type', 'kernel', 'pad', 'stride', 'global_pool']}, True, False),
+    ("test_roipool", mx.sym.ROIPooling, "MaxRoiPool",
+     [[[get_rnd(shape=(8, 6), low=1, high=100, dtype=np.int32)]], [[0, 0, 0, 4, 4]]],
+     {'pooled_size': (2, 2), 'spatial_scale': 0.7}, False,
+     {'modify': {'pooled_size': 'pooled_shape'}}, True, False),
 
     # since results would be random, checking for shape alone
     ("test_multinomial", mx.sym.sample_multinomial, "Multinomial",
      [np.array([0, 0.1, 0.2, 0.3, 0.4]).astype("float32")],
-     {'shape': (10,)}, False, {'modify': {'shape': 'sample_size'}}, False, True)
+     {'shape': (10,)}, False, {'modify': {'shape': 'sample_size'}}, False, True),
+    ("test_random_normal", mx.sym.random_normal, "RandomNormal", [],
+     {'shape': (2, 2), 'loc': 0, 'scale': 1}, False, {'modify': {'loc': 'mean'}}, False, True),
+    ("test_random_uniform", mx.sym.random_uniform, "RandomUniform", [],
+     {'shape': (2, 2), 'low': 0.5, 'high': 1.0}, False, {}, False, True)
 ]
 
 if __name__ == '__main__':