diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 7e84cea5ead5..988092910778 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -1545,3 +1545,10 @@ def convert_sum(node, **kwargs): name=name ) return [node] + +@mx_op.register("size_array") +def convert_size(node, **kwargs): + """Map MXNet's size_array operator attributes to onnx's Size operator + and return the created node. + """ + return create_basic_op_node('Size', node, kwargs) diff --git a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py index f61910f838ea..2ceabaec1dcd 100644 --- a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py +++ b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py @@ -21,7 +21,7 @@ from ._op_translations import identity, random_uniform, random_normal from ._op_translations import add, subtract, multiply, divide, absolute, negative, add_n from ._op_translations import tanh, arccos, arcsin, arctan, _cos, _sin, _tan -from ._op_translations import softplus, shape, gather, lp_pooling +from ._op_translations import softplus, shape, gather, lp_pooling, size from ._op_translations import ceil, floor, hardsigmoid, global_lppooling from ._op_translations import concat from ._op_translations import leaky_relu, _elu, _prelu, _selu, softmax, fully_connected @@ -139,6 +139,7 @@ 'Softplus' : softplus, 'Tan' : _tan, 'Shape' : shape, + 'Size' : size, 'Gather' : gather, 'HardSigmoid' : hardsigmoid, 'LpPool' : lp_pooling, diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py index 368b98d662b1..702832529314 100644 --- a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py +++ b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py @@ -642,6 +642,10 @@ def shape(attrs, inputs, proto_obj): """Returns shape of input array.""" return 'shape_array', attrs, inputs +def size(attrs, inputs, proto_obj): + """Returns array containing size of data.""" + return "size_array", attrs, inputs + def reduce_l2(attrs, inputs, proto_obj): """Reduce input tensor by l2 normalization.""" new_attrs = translation_utils._fix_attribute_names(attrs, {'axes':'axis'}) diff --git a/tests/python-pytest/onnx/export/onnx_backend_test.py b/tests/python-pytest/onnx/export/onnx_backend_test.py index 678435d92357..872fe95b5eb5 100644 --- a/tests/python-pytest/onnx/export/onnx_backend_test.py +++ b/tests/python-pytest/onnx/export/onnx_backend_test.py @@ -94,7 +94,8 @@ 'test_operator_permute2', 'test_clip' 'test_cast', - 'test_depthtospace' + 'test_depthtospace', + 'test_size' ] BASIC_MODEL_TESTS = [ diff --git a/tests/python-pytest/onnx/import/test_cases.py b/tests/python-pytest/onnx/import/test_cases.py index 7f34247c94e2..c407f91c7a64 100644 --- a/tests/python-pytest/onnx/import/test_cases.py +++ b/tests/python-pytest/onnx/import/test_cases.py @@ -86,7 +86,8 @@ 'test_operator_maxpool', 'test_operator_params', 'test_operator_permute2', - 'test_depthtospace' + 'test_depthtospace', + 'test_size' ] BASIC_MODEL_TESTS = [