diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 61cb353ded4d..035e1f63a08e 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -1772,3 +1772,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] diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py index ef3bda30df38..674b97bedf8e 100644 --- a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py +++ b/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.""" diff --git a/tests/python-pytest/onnx/test_node.py b/tests/python-pytest/onnx/test_node.py index 9b5ff97d1657..df56abb9c6aa 100644 --- a/tests/python-pytest/onnx/test_node.py +++ b/tests/python-pytest/onnx/test_node.py @@ -139,9 +139,18 @@ def test_import_export(self): 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], @@ -216,11 +225,19 @@ def test_imports(self): {'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) ] # test_case = ("test_case_name", "ONNX_op_name", [input_list], np_op, attribute map)