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

[v1.x] ONNX support for Convolution, Pooling 3d case #19855

Merged
Zha0q1 merged 1 commit into from
Feb 6, 2021
Merged

[v1.x] ONNX support for Convolution, Pooling 3d case #19855

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,15 +238,16 @@ def convert_convolution(node, **kwargs):
name, input_nodes, attrs = get_inputs(node, kwargs)

kernel = convert_string_to_list(attrs.get('kernel', '()'))
stride = convert_string_to_list(attrs.get('stride', '(1, 1)'))
dilate = convert_string_to_list(attrs.get('dilate', '(1, 1)'))
pad = convert_string_to_list(attrs.get('pad', '(0, 0)'))
stride = convert_string_to_list(attrs.get('stride', '()'))
dilate = convert_string_to_list(attrs.get('dilate', '()'))
pad = convert_string_to_list(attrs.get('pad', '()'))
num_group = int(attrs.get('num_group', 1))
no_bias = attrs.get('no_bias', 'False')
layout = attrs.get('layout', 'NCHW')

if layout != 'NCHW':
raise NotImplementedError('Pooling currently does not support layout!=\'NCHW\'')
if layout not in ['NCHW', 'NCDHW']:
raise NotImplementedError('Convolution currently does not support layout not in '
'[\'NCHW\', \'NCDHW\']')

if no_bias == 'True':
assert len(input_nodes) == 2, 'Convolution takes 2 input if no_bias==True'
Expand Down Expand Up @@ -707,7 +708,7 @@ def convert_pooling(node, **kwargs):
global_pool = attrs.get('global_pool', 'False')
_ = attrs.get('cudnn_off', 'False')
pooling_convention = attrs.get('pooling_convention', 'valid')
stride = convert_string_to_list(attrs.get('stride', '(1, 1)'))
stride = convert_string_to_list(attrs.get('stride', '()'))
pad = convert_string_to_list(attrs.get('pad', '()'))
p_value = int(attrs.get('p_value', '0'))
count_include_pad = attrs.get('count_include_pad', 'True')
Expand All @@ -721,8 +722,10 @@ def convert_pooling(node, **kwargs):
if pool_type == 'lp' and global_pool == 'False' and pooling_convention != 'valid':
raise NotImplementedError('Pooling currently does not support '
'pooling_convention!=\'valid\' when pool_type==\'lp\' and global_pool==False')
if layout != 'NCHW':
raise NotImplementedError('Pooling currently does not support layout!=\'NCHW\'')

if layout not in ['NCHW', 'NCDHW']:
raise NotImplementedError('Pooling currently does not support layout not in '
'[\'NCHW\', \'NCDHW\']')

kwargs_ = {}
if kernel:
Expand Down
173 changes: 157 additions & 16 deletions tests/python-pytest/onnx/test_operators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def hybrid_forward(self, F, *inputs):
return func(*inputs, **params)
return Model

def op_export_test(model_name, Model, inputs, tmp_path, dummy_input=False, onnx_map=None):
def op_export_test(model_name, Model, inputs, tmp_path, dummy_input=False, onnx_map=None, mx_map=None):
def export_to_onnx(model, model_name, inputs):
model_path = '{}/{}'.format(tmp_path, model_name)
model.export(model_path, epoch=0)
Expand All @@ -62,18 +62,20 @@ def onnx_rt(onnx_file, inputs):
model = Model()
model.initialize(ctx=mx.cpu(0))
model.hybridize()
pred_nat = model(*inputs)
pred_mx = model(*inputs)
onnx_file = export_to_onnx(model, model_name, inputs)
pred_onx = onnx_rt(onnx_file, inputs)
if dummy_input:
pred_nat = pred_nat[0]
if isinstance(pred_nat, list):
for i in range(len(pred_nat)):
pred_mx = pred_mx[0]
if isinstance(pred_mx, list):
for i in range(len(pred_mx)):
pred_onx_i = onnx_map(pred_onx[i]) if onnx_map else pred_onx[i]
assert_almost_equal(pred_nat[i], pred_onx_i, equal_nan=True)
pred_mx_i = mx_map(pred_mx[i]) if mx_map else pred_mx[i]
assert_almost_equal(pred_onx_i, pred_mx_i, equal_nan=True)
else:
pred_onx = onnx_map(pred_onx[0]) if onnx_map else pred_onx[0]
assert_almost_equal(pred_nat, pred_onx, equal_nan=True)
pred_mx = mx_map(pred_mx) if mx_map else pred_mx
assert_almost_equal(pred_onx, pred_mx, equal_nan=True)


def test_onnx_export_abs(tmp_path):
Expand Down Expand Up @@ -780,7 +782,7 @@ def test_onnx_export_pooling_avg(tmp_path, dtype, shape, count_include_pad, pool
if pooling_convention == 'full' and count_include_pad == True:
return
# onnxruntime requires that pad is smaller than kernel
if pad and pad[0] >= kernel[0] and pad[1] >= kernel[1]:
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
Expand All @@ -791,12 +793,46 @@ def test_onnx_export_pooling_avg(tmp_path, dtype, shape, count_include_pad, pool
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', count_include_pad=count_include_pad, pool_type='avg',
pooling_convention=pooling_convention, **kwargs)
pooling_convention=pooling_convention, layout='NCHW', **kwargs)
# Note here we use np.nan_to_num to map the onnx output because onnxruntime AveragePool will
# output NaN in some edge cases where mxnet outputs 0
op_export_test('pooling_avg', M, [x], tmp_path, onnx_map=np.nan_to_num)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 16, 16, 16), (1, 1, 10, 18, 18)])
@pytest.mark.parametrize('count_include_pad', [True, False])
@pytest.mark.parametrize('pooling_convention', ['full', 'valid'])
@pytest.mark.parametrize('kernel', [(1, 1, 1), (3, 3, 3), (1, 7, 7)])
@pytest.mark.parametrize('stride', [None, (1, 1, 1), (1, 2, 3)])
@pytest.mark.parametrize('pad', [None, (0, 1, 1), (1, 2, 3)])
def test_onnx_export_pooling_avg_3d(tmp_path, dtype, shape, count_include_pad, pooling_convention,
kernel, stride, pad):
# mxnet and onnxruntime has different implementation of count_include_pad on the left column
# and bottom row
if pooling_convention == 'full' and count_include_pad == True:
return
# onnxruntime requires that pad is smaller than kernel
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1] or pad[2] >= kernel[2]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
if kernel:
kwargs['kernel'] = kernel
if stride:
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', count_include_pad=count_include_pad, pool_type='avg',
pooling_convention=pooling_convention, layout='NCDHW', **kwargs)
# Note here we use np.nan_to_num to map the onnx output because onnxruntime AveragePool will
# output NaN in some edge cases where mxnet outputs 0
def mx_nan_to_num(a):
return np.nan_to_num(a.asnumpy())
op_export_test('pooling_avg_3d', M, [x], tmp_path, onnx_map=np.nan_to_num, mx_map=mx_nan_to_num)



@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 64, 64), (2, 1, 60, 60)])
@pytest.mark.parametrize('pooling_convention', ['full', 'valid'])
Expand All @@ -805,7 +841,7 @@ def test_onnx_export_pooling_avg(tmp_path, dtype, shape, count_include_pad, pool
@pytest.mark.parametrize('pad', [None, (1, 1), (3, 4), (4, 5)])
def test_onnx_export_pooling_max(tmp_path, dtype, shape, pooling_convention, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and pad[0] >= kernel[0] and pad[1] >= kernel[1]:
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
Expand All @@ -815,10 +851,34 @@ def test_onnx_export_pooling_max(tmp_path, dtype, shape, pooling_convention, ker
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', pool_type='max', pooling_convention=pooling_convention, **kwargs)
M = def_model('Pooling', pool_type='max', pooling_convention=pooling_convention,
layout='NCHW', **kwargs)
op_export_test('pooling_max', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 16, 16, 16), (1, 1, 10, 18, 18)])
@pytest.mark.parametrize('pooling_convention', ['full', 'valid'])
@pytest.mark.parametrize('kernel', [(1, 1, 1), (3, 3, 3), (1, 7, 7)])
@pytest.mark.parametrize('stride', [None, (1, 1, 1), (1, 2, 3)])
@pytest.mark.parametrize('pad', [None, (0, 1, 1), (1, 2, 3)])
def test_onnx_export_pooling_max_3d(tmp_path, dtype, shape, pooling_convention, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1] or pad[2] >= kernel[2]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
if kernel:
kwargs['kernel'] = kernel
if stride:
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', pool_type='max', pooling_convention=pooling_convention,
layout='NCDHW', **kwargs)
op_export_test('pooling_max_3d', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 64, 64), (2, 1, 60, 60)])
@pytest.mark.parametrize('p_value', [1, 2])
Expand All @@ -827,7 +887,7 @@ def test_onnx_export_pooling_max(tmp_path, dtype, shape, pooling_convention, ker
@pytest.mark.parametrize('pad', [None, (1, 1), (3, 4), (4, 5)])
def test_onnx_export_pooling_lp(tmp_path, dtype, shape, p_value, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and pad[0] >= kernel[0] and pad[1] >= kernel[1]:
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
Expand All @@ -838,10 +898,33 @@ def test_onnx_export_pooling_lp(tmp_path, dtype, shape, p_value, kernel, stride,
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', pool_type='lp', pooling_convention='valid',
p_value=p_value, **kwargs)
p_value=p_value, layout='NCHW', **kwargs)
op_export_test('pooling_lp', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 16, 16, 16), (1, 1, 10, 18, 18)])
@pytest.mark.parametrize('p_value', [1, 2])
@pytest.mark.parametrize('kernel', [(1, 1, 1), (3, 3, 3), (1, 7, 7)])
@pytest.mark.parametrize('stride', [None, (1, 1, 1), (1, 2, 3)])
@pytest.mark.parametrize('pad', [None, (0, 1, 1), (1, 2, 3)])
def test_onnx_export_pooling_lp_3d(tmp_path, dtype, shape, p_value, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1] or pad[2] >= kernel[2]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
if kernel:
kwargs['kernel'] = kernel
if stride:
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
M = def_model('Pooling', pool_type='lp', pooling_convention='valid',
p_value=p_value, layout='NCDHW', **kwargs)
op_export_test('pooling_lp_3d', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 64, 64), (2, 1, 60, 60)])
@pytest.mark.parametrize('pool_type', ['avg', 'max', 'lp'])
Expand All @@ -851,7 +934,7 @@ def test_onnx_export_pooling_lp(tmp_path, dtype, shape, p_value, kernel, stride,
@pytest.mark.parametrize('pad', [None, (3, 4)])
def test_onnx_export_pooling_global(tmp_path, dtype, shape, pool_type, p_value, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and pad[0] >= kernel[0] and pad[1] >= kernel[1]:
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
Expand All @@ -863,10 +946,35 @@ def test_onnx_export_pooling_global(tmp_path, dtype, shape, pool_type, p_value,
kwargs['pad'] = pad
# kernel, stride, and pad should have no effect on the results
M = def_model('Pooling', global_pool=True, pool_type=pool_type, pooling_convention='valid',
p_value=p_value, **kwargs)
p_value=p_value, layout='NCHW', **kwargs)
op_export_test('pooling_global', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 3, 16, 16, 16), (1, 1, 10, 18, 18)])
@pytest.mark.parametrize('pool_type', ['avg', 'max', 'lp'])
@pytest.mark.parametrize('p_value', [1, 2])
@pytest.mark.parametrize('kernel', [(1, 1, 1), (3, 3, 3)])
@pytest.mark.parametrize('stride', [None, (1, 1, 1)])
@pytest.mark.parametrize('pad', [None, (0, 1, 1)])
def test_onnx_export_pooling_global_3d(tmp_path, dtype, shape, pool_type, p_value, kernel, stride, pad):
# onnxruntime requires that pad is smaller than kernel
if pad and (pad[0] >= kernel[0] or pad[1] >= kernel[1] or pad[2] >= kernel[2]):
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
kwargs = {}
if kernel:
kwargs['kernel'] = kernel
if stride:
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
# kernel, stride, and pad should have no effect on the results
M = def_model('Pooling', global_pool=True, pool_type=pool_type, pooling_convention='valid',
p_value=p_value, layout='NCDHW', **kwargs)
op_export_test('pooling_global_3d', M, [x], tmp_path)


@pytest.mark.parametrize('dtype', ['float16', 'float32'])
def test_onnx_export_log2(tmp_path, dtype):
x = mx.random.normal(0, 10, (2, 3, 4, 5)).astype(dtype)
Expand Down Expand Up @@ -924,7 +1032,40 @@ def test_onnx_export_convolution(tmp_path, dtype, shape, num_filter, num_group,
if dilate:
kwargs['dilate'] = dilate
M = def_model('Convolution', num_filter=num_filter, num_group=num_group, no_bias=no_bias,
**kwargs)
layout='NCHW', **kwargs)
inputs = [x, w] if no_bias else [x, w, b]
op_export_test('convolution', M, inputs, tmp_path)


@pytest.mark.parametrize('dtype', ['float32'])
@pytest.mark.parametrize('shape', [(1, 4, 16, 16, 16), (1, 3, 10, 18, 18)])
@pytest.mark.parametrize('num_filter', [2, 4, 32])
@pytest.mark.parametrize('num_group', [1, 2])
@pytest.mark.parametrize('no_bias', [True, False])
@pytest.mark.parametrize('kernel', [(3, 3, 3), (1, 1, 1), (1, 7, 7)])
@pytest.mark.parametrize('stride', [None, (1, 1, 1), (1, 2, 3)])
@pytest.mark.parametrize('pad', [None, (0, 1, 1), (1, 2, 3)])
@pytest.mark.parametrize('dilate', [None, [2, 2, 2]])
def test_onnx_export_convolution_3D(tmp_path, dtype, shape, num_filter, num_group, no_bias,
kernel, stride, pad, dilate):
if shape[1] % num_group:
return
x = mx.random.uniform(0, 1, shape, dtype=dtype)
w_shape = (num_filter,) + (shape[1] // num_group,) + kernel
w = mx.random.uniform(0, 1, w_shape, dtype=dtype)
b_shape = (num_filter)
b = mx.random.uniform(0, 1, b_shape, dtype=dtype)
kwargs = {}
if kernel:
kwargs['kernel'] = kernel
if stride:
kwargs['stride'] = stride
if pad:
kwargs['pad'] = pad
if dilate:
kwargs['dilate'] = dilate
M = def_model('Convolution', num_filter=num_filter, num_group=num_group, no_bias=no_bias,
layout='NCDHW', **kwargs)
inputs = [x, w] if no_bias else [x, w, b]
op_export_test('convolution', M, inputs, tmp_path)

Expand Down