diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 8fdc5944fb91..c9a21becf0b5 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -2055,7 +2055,6 @@ def convert_topk(node, **kwargs): and return the created node. """ name, input_nodes, attrs = get_inputs(node, kwargs) - axis = int(attrs.get('axis', '-1')) k = int(attrs.get('k', '1')) ret_type = attrs.get('ret_typ') @@ -2080,3 +2079,26 @@ def convert_topk(node, **kwargs): ) return [topk_node] + + +@mx_op.register("UpSampling") +def convert_upsample(node, **kwargs): + """Map MXNet's UpSampling operator attributes to onnx's Upsample operator + sample_type = attrs.get('sample_type', 'nearest') + sample_type = 'linear' if sample_type == 'bilinear' else sample_type + scale = convert_string_to_list(attrs.get('scale')) + scaleh = float(scale[0]) + scalew = float(scale[0]) + if len(scale) > 1: + scalew = float(scale[1]) + scale = [1.0, 1.0, scaleh, scalew] + + node = onnx.helper.make_node( + 'Upsample', + input_nodes, + [name], + scales=scale, + mode=sample_type, + name=name + ) + return [node] \ No newline at end of file diff --git a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py index bd48d26ef6ba..7f6e2ab5b036 100644 --- a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py +++ b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py @@ -23,7 +23,7 @@ from ._op_translations import tanh, arccos, arcsin, arctan, _cos, _sin, _tan from ._op_translations import softplus, shape, gather, lp_pooling, size from ._op_translations import ceil, floor, hardsigmoid, global_lppooling -from ._op_translations import concat, hardmax, topk +from ._op_translations import concat, hardmax, topk, upsampling from ._op_translations import leaky_relu, _elu, _prelu, _selu, softmax, fully_connected from ._op_translations import global_avgpooling, global_maxpooling, linalg_gemm from ._op_translations import sigmoid, pad, relu, matrix_multiplication, batch_norm @@ -148,5 +148,6 @@ 'SpaceToDepth' : spacetodepth, 'Hardmax' : hardmax, 'LpNormalization' : lpnormalization, - 'TopK' : topk -} + 'TopK' : topk, + 'Upsample' : upsampling, +} \ No newline at end of file diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py index 627181d6ae21..a1d8b8ee54b0 100644 --- a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py +++ b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py @@ -787,3 +787,24 @@ def topk(attrs, inputs, proto_obj): {'ret_typ': 'both', 'dtype': 'int64'}) return 'topk', new_attrs, inputs + + +def upsampling(attrs, inputs, proto_obj): + """Rearranges blocks of spatial data into depth.""" + new_attrs = translation_utils._fix_attribute_names(attrs, {'scales': 'scale', + 'mode': 'sample_type'}) + sample_type = new_attrs.get('sample_type', 'nearest') + if sample_type != 'nearest': + raise NotImplementedError("Operator {} in ONNX supports 'linear' mode " + "for linear, bilinear, trilinear etc. There is no " + "way to distinguish these so far. Therefore, supporting " + "import of only nearest neighbor upsampling for now. " + "https://github.com/onnx/onnx/issues/1774. " + "Use contrib.BilinearResize2D for bilinear mode." + .format('UpSample')) + + scale = tuple(new_attrs.get('scale'))[2:] + scale = tuple([int(s) for s in scale]) + mx_op = symbol.UpSampling(inputs[0], scale=scale, sample_type=sample_type) + + return mx_op, new_attrs, inputs \ No newline at end of file diff --git a/tests/python-pytest/onnx/test_cases.py b/tests/python-pytest/onnx/test_cases.py index eb6b71896e28..59585986d966 100644 --- a/tests/python-pytest/onnx/test_cases.py +++ b/tests/python-pytest/onnx/test_cases.py @@ -79,7 +79,8 @@ 'test_softplus', 'test_reduce_', 'test_split_equal', - 'test_top_k' + 'test_top_k', + 'test_upsample_n' ], 'import': ['test_gather', 'test_softsign',