Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
[v1.x]Onnx support for upsampling (#19795)
Browse files Browse the repository at this point in the history
* bare bone implementation

* Update _op_translations.py

* Update _op_translations.py

* Update _op_translations.py
  • Loading branch information
Zha0q1 committed Feb 1, 2021
1 parent 5abbb28 commit 76c8c7d
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 0 deletions.
29 changes: 29 additions & 0 deletions python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
Original file line number Diff line number Diff line change
Expand Up @@ -3400,6 +3400,35 @@ def convert_gather_nd(node, **kwargs):
return nodes


@mx_op.register('UpSampling')
def convert_upsampling(node, **kwargs):
"""Map MXNet's UpSampling operator to onnx.
"""
from onnx.helper import make_node
name, input_nodes, attrs = get_inputs(node, kwargs)

scale = int(attrs.get('scale', '1'))
sample_type = attrs.get('sample_type')
num_args = int(attrs.get('num_args', '1'))

if num_args > 1:
raise NotImplementedError('Upsampling conversion does not currently support num_args > 1')

if sample_type != 'nearest':
raise NotImplementedError('Upsampling conversion does not currently support \
sample_type != nearest')

nodes = [
create_tensor([], name+'_roi', kwargs['initializer'], dtype='float32'),
create_tensor([1, 1, scale, scale], name+'_scales', kwargs['initializer'],
dtype='float32'),
make_node('Resize', [input_nodes[0], name+'_roi', name+'_scales'], [name], mode='nearest',
coordinate_transformation_mode='half_pixel')
]

return nodes


@mx_op.register('SwapAxis')
def convert_swapaxis(node, **kwargs):
"""Map MXNet's SwapAxis operator
Expand Down
9 changes: 9 additions & 0 deletions tests/python-pytest/onnx/test_operators.py
Original file line number Diff line number Diff line change
Expand Up @@ -595,6 +595,15 @@ def test_onnx_export_gather_nd(tmp_path, dtype):
op_export_test('gather_nd2', M2, [x2, y2], tmp_path)


@pytest.mark.parametrize('dtype', ['float16', 'float32'])
@pytest.mark.parametrize('shape', [(3, 4, 5, 6), (1, 1, 1, 1)])
@pytest.mark.parametrize('scale', [1, 2, 3])
def test_onnx_export_upsampling(tmp_path, dtype, shape, scale):
A = mx.random.uniform(0, 1, shape).astype(dtype)
M = def_model('UpSampling', scale=scale, sample_type='nearest', num_args=1)
op_export_test('UpSampling', M, [A], tmp_path)


@pytest.mark.parametrize('dtype', ['int32', 'int64', 'float16', 'float32', 'float64'])
@pytest.mark.parametrize('params', [((4, 5, 6), (0, 2)), ((4, 5, 6), (0, 1)),
((1, 2, 3, 4, 1), (0, 4)),
Expand Down

0 comments on commit 76c8c7d

Please sign in to comment.