This is an automated email from the ASF dual-hosted git repository.
zha0q1 pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/v1.x by this push:
new c9f111f [v1.x] Add more ONNX export operator support (#19727)
c9f111f is described below
commit c9f111f892bcee760ede967c94eaefe591729374
Author: Joe Evans <[email protected]>
AuthorDate: Wed Jan 6 11:29:51 2021 -0800
[v1.x] Add more ONNX export operator support (#19727)
* Add onnx export support for ones_like operator.
* Clean up dropout, clip and topk export functions.
* Clean up pad export function.
* Add unit test for ones_like onnx export.
* Add onnx export function for arange operator.
* Fix lint.
* Make sure to return all nodes created.
* Extend operator test to work with no inputs, add unit test for arange.
* Extent arange test to also test int32 and int64 dtypes.
* Return scalar nodes in clip conversion function.
* Make sure to return all graph nodes created in export ops.
* Properly obey dtype attribute instead of using input type for arange.
* Use static dtype for parameter to catch errors when dtype != input type.
Co-authored-by: Joe Evans <[email protected]>
---
.../mxnet/contrib/onnx/mx2onnx/_op_translations.py | 224 +++++++++------------
tests/python-pytest/onnx/test_operators.py | 33 ++-
2 files changed, 124 insertions(+), 133 deletions(-)
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index d301975..07537a3 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -504,6 +504,7 @@ def convert_pad(node, **kwargs):
"""Map MXNet's pad operator attributes to onnx's Pad operator
and return the created node.
"""
+ from onnx.helper import make_node
opset_version = kwargs["opset_version"]
name, input_nodes, attrs = get_inputs(node, kwargs)
@@ -515,40 +516,20 @@ def convert_pad(node, **kwargs):
if opset_version >= 11:
# starting with opset 11, pads and constant_value are inputs instead
of attributes
- from onnx.helper import make_tensor, make_tensor_value_info
- initializer = kwargs["initializer"]
- pads_input_name = name + "_pads"
- pads_input_type = onnx.TensorProto.INT64
- pads_input_shape = np.shape(np.array(onnx_pad_width))
- pads_value_node = make_tensor_value_info(pads_input_name,
pads_input_type, pads_input_shape)
- pads_tensor_node = make_tensor(pads_input_name, pads_input_type,
pads_input_shape, onnx_pad_width)
- initializer.append(pads_tensor_node)
- input_nodes.append(pads_input_name)
+ nodes = [
+ create_const_node(name+"_pads", np.array(onnx_pad_width,
dtype='int64'), kwargs)
+ ]
if pad_mode == "constant":
- const_input_name = name + "_constant"
- const_input_type =
onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[pad_value.dtype]
- const_value_node = make_tensor_value_info(const_input_name,
const_input_type, ())
- const_tensor_node = make_tensor(const_input_name,
const_input_type, (), [pad_value])
- initializer.append(const_tensor_node)
- input_nodes.append(const_input_name)
- pad_node = onnx.helper.make_node(
- "Pad",
- input_nodes,
- [name],
- mode=pad_mode,
- name=name
- )
- return [pads_value_node, const_value_node, pad_node]
+ nodes += [
+ create_const_scalar_node(name+"_const", pad_value, kwargs),
+ make_node("Pad", [input_nodes[0], name+"_pads",
name+"_const"], [name], mode=pad_mode, name=name)
+ ]
else:
- pad_node = onnx.helper.make_node(
- "Pad",
- input_nodes,
- [name],
- mode=pad_mode,
- name=name
- )
- return [pads_value_node, pad_node]
+ nodes += [
+ make_node("Pad", [input_nodes[0], name+"_pads"], [name],
mode=pad_mode, name=name)
+ ]
+ return nodes
else:
if pad_mode == "constant":
node = onnx.helper.make_node(
@@ -560,7 +541,6 @@ def convert_pad(node, **kwargs):
pads=onnx_pad_width,
name=name
)
- return [node]
else:
node = onnx.helper.make_node(
'Pad',
@@ -570,7 +550,7 @@ def convert_pad(node, **kwargs):
pads=onnx_pad_width,
name=name
)
- return [node]
+ return [node]
def create_helper_trans_node(node_name, input_node):
@@ -1103,6 +1083,7 @@ def convert_dropout(node, **kwargs):
"""Map MXNet's Dropout operator attributes to onnx's Dropout operator
and return the created node.
"""
+ from onnx.helper import make_node
name, input_nodes, attrs = get_inputs(node, kwargs)
opset_version = kwargs["opset_version"]
@@ -1110,29 +1091,13 @@ def convert_dropout(node, **kwargs):
if opset_version >= 12:
# opset >= 12 requires the ratio to be an input
- initializer = kwargs["initializer"]
- ratio_input_name = name + "_ratio"
- value_node = onnx.helper.make_tensor_value_info(ratio_input_name,
-
onnx.TensorProto.FLOAT, ())
- tensor_node = onnx.helper.make_tensor(ratio_input_name,
onnx.TensorProto.FLOAT,
- (), [probability])
- initializer.append(tensor_node)
- dropout_node = onnx.helper.make_node(
- "Dropout",
- [input_nodes[0], ratio_input_name],
- [name],
- name=name
- )
- return [value_node, dropout_node]
+ nodes = [
+ create_const_scalar_node(name+"_ratio0", np.float32(probability),
kwargs),
+ make_node("Dropout", [input_nodes[0], name+"_ratio0"], [name],
name=name)
+ ]
+ return nodes
else:
- dropout_node = onnx.helper.make_node(
- "Dropout",
- input_nodes,
- [name],
- ratio=probability,
- name=name
- )
- return [dropout_node]
+ return [make_node("Dropout", input_nodes, [name], ratio=probability,
name=name)]
@mx_op.register("Flatten")
@@ -1147,6 +1112,7 @@ def convert_clip(node, **kwargs):
"""Map MXNet's Clip operator attributes to onnx's Clip operator
and return the created node.
"""
+ from onnx.helper import make_node
name, input_nodes, attrs = get_inputs(node, kwargs)
opset_version = kwargs["opset_version"]
@@ -1155,39 +1121,16 @@ def convert_clip(node, **kwargs):
if opset_version >= 11:
# opset >= 11 requires min/max to be inputs
- initializer = kwargs["initializer"]
- min_input_name = name + "_min"
- max_input_name = name + "_max"
- min_value_node = onnx.helper.make_tensor_value_info(min_input_name,
-
onnx.TensorProto.FLOAT, ())
- max_value_node = onnx.helper.make_tensor_value_info(max_input_name,
-
onnx.TensorProto.FLOAT, ())
- min_tensor_node = onnx.helper.make_tensor(min_input_name,
onnx.TensorProto.FLOAT,
- (), [a_min])
- max_tensor_node = onnx.helper.make_tensor(max_input_name,
onnx.TensorProto.FLOAT,
- (), [a_max])
- initializer.append(min_tensor_node)
- initializer.append(max_tensor_node)
- input_nodes.append(min_input_name)
- input_nodes.append(max_input_name)
- clip_node = onnx.helper.make_node(
- "Clip",
- input_nodes,
- [name],
- name=name
- )
- return [min_value_node, max_value_node, clip_node]
-
+ nodes = [
+ create_const_scalar_node(name+"_min", np.float32(a_min), kwargs),
+ create_const_scalar_node(name+"_max", np.float32(a_max), kwargs),
+ make_node("Clip", [input_nodes[0], name+"_min", name+"_max"],
[name], name=name)
+ ]
else:
- clip_node = onnx.helper.make_node(
- "Clip",
- input_nodes,
- [name],
- name=name,
- min=a_min,
- max=a_max
- )
- return [clip_node]
+ nodes = [
+ make_node("Clip", input_nodes, [name], name=name, min=a_min,
max=a_max)
+ ]
+ return nodes
def scalar_op_helper(node, op_name, **kwargs):
@@ -1705,25 +1648,28 @@ def convert_slice_axis(node, **kwargs):
begin = int(attrs.get("begin"))
end = attrs.get("end", None)
- nodes = []
- create_tensor([axis], name+'_axis', kwargs["initializer"])
- create_tensor([begin], name+'_begin', kwargs["initializer"])
+ nodes = [
+ create_tensor([axis], name+'_axis', kwargs["initializer"]),
+ create_tensor([begin], name+'_begin', kwargs["initializer"])
+ ]
if not end or end == 'None':
# ONNX doesn't support None for ends. Since ends=None depicts
# length of dimension, passing dimension in this case.
- create_tensor([axis+1], name+"_axis_plus_1", kwargs["initializer"])
nodes += [
+ create_tensor([axis+1], name+"_axis_plus_1",
kwargs["initializer"]),
make_node('Shape', [input_nodes[0]], [name+"_data_shape"]),
make_node('Slice', [name+'_data_shape', name+'_axis',
name+'_axis_plus_1'],
[name+"_end"])
]
else:
- create_tensor([int(end)], name+'_end', kwargs["initializer"])
+ nodes += [
+ create_tensor([int(end)], name+'_end', kwargs["initializer"])
+ ]
nodes += [
make_node('Slice', [input_nodes[0], name+'_begin', name+'_end',
name+'_axis'],
[name], name=name)
- ]
+ ]
return nodes
@@ -2267,52 +2213,28 @@ def convert_topk(node, **kwargs):
"""Map MXNet's topk operator attributes to onnx's TopK operator
and return the created node.
"""
+ from onnx.helper import make_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')
- dtype = attrs.get('dtype')
- outputs = [name + '_output0']
+ outputs = [name]
if ret_type and ret_type == 'both':
- if dtype and dtype == 'int64':
- outputs.append(name + '_output1')
- else:
- raise NotImplementedError("ONNX expects indices to be of type
int64")
+ outputs.append(name + '_output1')
else:
raise NotImplementedError("ONNX expects both value and indices as
output")
opset_version = kwargs['opset_version']
if opset_version >= 10:
- from onnx.helper import make_tensor, make_tensor_value_info
- initializer = kwargs["initializer"]
- k_input_name = name + "_k"
- k_input_type = onnx.TensorProto.INT64
- k_value_node = make_tensor_value_info(k_input_name, k_input_type, ())
- k_tensor_node = make_tensor(k_input_name, k_input_type, (), k)
- initializer.append(k_tensor_node)
- input_nodes.append(k_input_name)
-
- topk_node = onnx.helper.make_node(
- "TopK",
- input_nodes,
- outputs,
- axis=axis,
- name=name
- )
- return [k_value_node, topk_node]
+ nodes = [
+ create_const_scalar_node(name+"_k", np.int64(k), kwargs),
+ make_node("TopK", [input_nodes[0], name+"_k"], outputs, axis=axis,
name=name)
+ ]
+ return nodes
else:
- topk_node = onnx.helper.make_node(
- "TopK",
- input_nodes,
- outputs,
- axis=axis,
- k=k,
- name=name
- )
-
- return [topk_node]
+ return [make_node("TopK", input_nodes, outputs, axis=axis, k=k,
name=name)]
@mx_op.register("take")
@@ -2525,7 +2447,7 @@ def convert_broadcast_axis(node, **kwargs):
make_node('Shape', [shape_name], [name+'_in_dim']),
make_node('Reshape', [name+'_in_dim', name+'_void'],
[name+'_in_dim_s']),
make_node('Range', [name+'_0_s', name+'_in_dim_s', name+'_1_s'],
[name+'_range']),
- ]
+ ]
for i, axis in enumerate(axis):
if axis not in (0, 1):
@@ -2537,7 +2459,7 @@ def convert_broadcast_axis(node, **kwargs):
make_node('Mul', [name+'_size_'+str(i), name+'_cast_'+str(i)],
[name+'_mul_'+str(i)]),
make_node('Add', [name+'_mul_'+str(i), name+'_1'],
[name+'_add_'+str(i)]),
make_node('Mul', [name+'_add_'+str(i), shape_name],
[name+'_shape_'+str(i+1)])
- ]
+ ]
shape_name = name+'_shape_'+str(i+1)
nodes += [make_node('Expand', [input_nodes[0], shape_name], [name],
name=name)]
@@ -2579,7 +2501,7 @@ def convert_sequencemask(node, **kwargs):
make_node('Range', [name+'_0_s', name+'_in_dim_s', name+'_1_s'],
[name+'_range_0']),
make_node('Less', [name+'_range_0', name+'_2'], [name+'_less_0']),
make_node('Where', [name+'_less_0', name+'_in_shape', name+'_1'],
[name+'_shape_1'])
- ]
+ ]
if(axis == 0):
nodes += [
@@ -2703,6 +2625,22 @@ def convert_zeros_like(node, **kwargs):
return nodes
+@mx_op.register("ones_like")
+def convert_ones_like(node, **kwargs):
+ """Map MXNet's ones_like operator attributes to onnx's ConstantOfShape
operator.
+ """
+ from onnx.helper import make_node, make_tensor
+ name, input_nodes, _ = get_inputs(node, kwargs)
+
+ # create tensor with shape of input
+ tensor_value = make_tensor(name+"_one", kwargs['in_type'], [1], [1])
+ nodes = [
+ make_node("Shape", [input_nodes[0]], [name+"_shape"]),
+ make_node("ConstantOfShape", [name+"_shape"], [name], name=name,
value=tensor_value)
+ ]
+ return nodes
+
+
@mx_op.register("_contrib_arange_like")
def convert_arange_like(node, **kwargs):
"""Map MXNet's arange_like operator attributes to onnx's Range and Reshape
operators.
@@ -2759,3 +2697,31 @@ def convert_arange_like(node, **kwargs):
]
return nodes
+
+@mx_op.register("_arange")
+def convert_arange(node, **kwargs):
+ """Map MXNet's arange operator attributes to onnx's Range operator.
+ """
+ from onnx.helper import make_node
+ name, _, attrs = get_inputs(node, kwargs)
+
+ opset_version = kwargs['opset_version']
+ if opset_version < 11:
+ raise AttributeError("ONNX opset 11 or greater is required to export
this operator")
+
+ start = attrs.get('start', 0.)
+ stop = attrs.get('stop')
+ step = attrs.get('step', 1.)
+ dtype = attrs.get('dtype', 'float32')
+ repeat = int(attrs.get('repeat', 1))
+ if repeat != 1:
+ raise NotImplementedError("arange operator with repeat != 1 not yet
implemented.")
+
+ nodes = [
+ create_const_scalar_node(name+"_start", np.array([start],
dtype=dtype), kwargs),
+ create_const_scalar_node(name+"_stop", np.array([stop], dtype=dtype),
kwargs),
+ create_const_scalar_node(name+"_step", np.array([step], dtype=dtype),
kwargs),
+ make_node("Range", [name+"_start", name+"_stop", name+"_step"], [name])
+ ]
+
+ return nodes
diff --git a/tests/python-pytest/onnx/test_operators.py
b/tests/python-pytest/onnx/test_operators.py
index 057a279..34838a0 100644
--- a/tests/python-pytest/onnx/test_operators.py
+++ b/tests/python-pytest/onnx/test_operators.py
@@ -23,7 +23,7 @@ from mxnet.test_utils import assert_almost_equal
import pytest
import tempfile
-def def_model(op_name, **params):
+def def_model(op_name, dummy_input=False, **params):
class Model(HybridBlock):
def __init__(self, **kwargs):
super(Model, self).__init__(**kwargs)
@@ -33,11 +33,13 @@ def def_model(op_name, **params):
func = F
for name in names:
func = getattr(func, name)
- out = func(*inputs, **params)
- return out
+ if dummy_input:
+ return func(**params), inputs[0]
+ else:
+ return func(*inputs, **params)
return Model
-def op_export_test(model_name, Model, inputs, tmp_path):
+def op_export_test(model_name, Model, inputs, tmp_path, dummy_input=False):
def export_to_onnx(model, model_name, inputs):
model_path = '{}/{}'.format(tmp_path, model_name)
model.export(model_path, epoch=0)
@@ -63,6 +65,8 @@ def op_export_test(model_name, Model, inputs, tmp_path):
pred_nat = 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]
assert_almost_equal(pred_nat, pred_onx)
@@ -94,6 +98,10 @@ def test_onnx_export_zeros_like(tmp_path):
x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype='float32')
op_export_test('zeros_like', M, [x], tmp_path)
+def test_onnx_export_ones_like(tmp_path):
+ M = def_model('ones_like')
+ x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype='float32')
+ op_export_test('ones_like', M, [x], tmp_path)
@pytest.mark.parametrize("dtype", ["float32", "float64"])
@pytest.mark.parametrize("axis", [None,0,1])
@@ -105,6 +113,23 @@ def test_onnx_export_arange_like(tmp_path, dtype, axis,
start, step, test_data):
x = mx.nd.array(test_data, dtype=dtype)
op_export_test('arange_like', M, [x], tmp_path)
+
[email protected]("stop", [2, 50, 5000])
[email protected]("step", [0.25, 0.5, 1, 5])
[email protected]("start", [0., 1.])
[email protected]("dtype", ["float32", "float64", "int32", "int64"])
+def test_onnx_export_arange(tmp_path, dtype, start, stop, step):
+ if "int" in dtype:
+ start = int(start)
+ stop = int(stop)
+ step = int(step)
+ if step == 0:
+ step = 1
+ M = def_model('arange', dummy_input=True, start=start, stop=stop,
step=step, dtype=dtype)
+ x = mx.nd.array([1], dtype='float32')
+ op_export_test('arange', M, [x], tmp_path, dummy_input=True)
+
+
@pytest.mark.parametrize('dtype', ['float32'])
def test_onnx_export_layernorm(tmp_path, dtype):
x = mx.nd.random.uniform(1, 2, (3, 4, 5), dtype=dtype)
