From 7832bb1cbae043e25a6c287558f4ba96f7e40c25 Mon Sep 17 00:00:00 2001 From: Joe Evans Date: Wed, 24 Feb 2021 19:13:10 +0000 Subject: [PATCH 1/4] Add onnx export support and unit tests for one_hot. --- .../contrib/onnx/mx2onnx/_op_translations.py | 21 +++++++++++++++++++ tests/python-pytest/onnx/test_operators.py | 11 +++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index d75adff8916f..95ae5891bfc2 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -3999,3 +3999,24 @@ def convert_argsort(node, **kwargs): ] return nodes + + +@mx_op.register('one_hot') +def convert_one_hot(node, **kwargs): + """Map MXNet's one_hot operator attributes to onnx's OneHot operator + """ + from onnx.helper import make_node + name, input_nodes, attrs = get_inputs(node, kwargs) + + depth = int(attrs.get('depth')) + on_value = float(attrs.get('on_value', 1.)) + off_value = float(attrs.get('off_value', 0.)) + dtype = attrs.get('dtype', 'float32') + + create_tensor([off_value, on_value], name+'_values', kwargs['initializer'], dtype=dtype) + create_tensor([depth], name+'_depth', kwargs['initializer']) + nodes = [ + make_node('OneHot', [input_nodes[0], name+'_depth', name+'_values'], [name], name=name) + ] + + return nodes diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py index 9c4326159eb3..29ff357d428f 100644 --- a/tests/python-pytest/onnx/test_operators.py +++ b/tests/python-pytest/onnx/test_operators.py @@ -1167,4 +1167,13 @@ def test_onnx_export_take_raise(tmp_path, dtype, axis): x = mx.nd.random.normal(0, 10, (3, 4, 5)).astype(dtype) y = mx.random.randint(0, 3, (6, 7)).astype(dtype) M = def_model('take', axis=axis, mode='raise') - op_export_test('take', M, [x, y], tmp_path) \ No newline at end of file + op_export_test('take', M, [x, y], tmp_path) + + +# onnxruntime currently only supports float32 and int64 +@pytest.mark.parametrize("dtype", ["float32", "int64"]) +@pytest.mark.parametrize("depth", [1, 3, 5, 10]) +def test_onnx_export_one_hot(tmp_path, dtype, depth): + M = def_model('one_hot', depth=depth, dtype=dtype) + x = mx.random.randint(0, 10, (depth * depth)).astype('int64') + op_export_test('one_hot', M, [x], tmp_path) From b7a2fd73d4eb86b3c6aad3d2aae75608c51bad1f Mon Sep 17 00:00:00 2001 From: Joe Evans Date: Wed, 24 Feb 2021 19:41:07 +0000 Subject: [PATCH 2/4] Add onnx export function for random_uniform_like. --- .../contrib/onnx/mx2onnx/_op_translations.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 95ae5891bfc2..a12e360d8a40 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -4020,3 +4020,22 @@ def convert_one_hot(node, **kwargs): ] return nodes + +@mx_op.register('_random_uniform_like') +def convert_one_hot(node, **kwargs): + """Map MXNet's random_uniform_like operator attributes to onnx's RandomUniformLike operator + """ + from onnx.helper import make_node + name, input_nodes, attrs = get_inputs(node, kwargs) + + low = float(attrs.get('low', 0.)) + high = float(attrs.get('high', 1.)) + dtype = attrs.get('dtype', 'float32') + + nodes = [ + make_node('RandomUniformLike', [input_nodes[0]], [name], name=name, + dtype=onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(dtype)], + low=low, high=high) + ] + + return nodes From 5089fe9d9a5026ef89c504d30917c8a0c36940a1 Mon Sep 17 00:00:00 2001 From: Joe Evans Date: Wed, 24 Feb 2021 20:59:09 +0000 Subject: [PATCH 3/4] Fix lint. --- python/mxnet/contrib/onnx/mx2onnx/_op_translations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index a12e360d8a40..7b3fa08898f0 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -4022,7 +4022,7 @@ def convert_one_hot(node, **kwargs): return nodes @mx_op.register('_random_uniform_like') -def convert_one_hot(node, **kwargs): +def convert_random_uniform_like(node, **kwargs): """Map MXNet's random_uniform_like operator attributes to onnx's RandomUniformLike operator """ from onnx.helper import make_node @@ -4033,7 +4033,7 @@ def convert_one_hot(node, **kwargs): dtype = attrs.get('dtype', 'float32') nodes = [ - make_node('RandomUniformLike', [input_nodes[0]], [name], name=name, + make_node('RandomUniformLike', [input_nodes[0]], [name], name=name, dtype=onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(dtype)], low=low, high=high) ] From 1fec6a6e921ca1f6def206af291aaffd7c5bd87d Mon Sep 17 00:00:00 2001 From: Joe Evans Date: Tue, 2 Mar 2021 00:00:33 +0000 Subject: [PATCH 4/4] Update tests and use correct dtype for on and off values. --- python/mxnet/contrib/onnx/mx2onnx/_op_translations.py | 2 +- tests/python-pytest/onnx/test_operators.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 28621a734947..2521cf5adb37 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -4047,7 +4047,7 @@ def convert_one_hot(node, **kwargs): off_value = float(attrs.get('off_value', 0.)) dtype = attrs.get('dtype', 'float32') - create_tensor([off_value, on_value], name+'_values', kwargs['initializer'], dtype=dtype) + create_tensor([off_value, on_value], name+'_values', kwargs['initializer'], dtype=np.dtype(dtype)) create_tensor([depth], name+'_depth', kwargs['initializer']) nodes = [ make_node('OneHot', [input_nodes[0], name+'_depth', name+'_values'], [name], name=name) diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py index aca4a6806735..f73672b59cc7 100644 --- a/tests/python-pytest/onnx/test_operators.py +++ b/tests/python-pytest/onnx/test_operators.py @@ -1183,12 +1183,13 @@ def test_onnx_export_take_raise(tmp_path, dtype, axis): op_export_test('take', M, [x, y], tmp_path) -# onnxruntime currently only supports float32 and int64 -@pytest.mark.parametrize("dtype", ["float32", "int64"]) +# onnxruntime currently does not support int32 +@pytest.mark.parametrize("dtype", ["float16", "float32", "int64"]) @pytest.mark.parametrize("depth", [1, 3, 5, 10]) -def test_onnx_export_one_hot(tmp_path, dtype, depth): +@pytest.mark.parametrize("shape", [(1,1), (1,5), (5,5), (3,4,5)]) +def test_onnx_export_one_hot(tmp_path, dtype, depth, shape): M = def_model('one_hot', depth=depth, dtype=dtype) - x = mx.random.randint(0, 10, (depth * depth)).astype('int64') + x = mx.random.randint(0, 10, shape).astype('int64') op_export_test('one_hot', M, [x], tmp_path)