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 b219ac2 [v1.x] Add more ONNX export support to operators (#19625)
b219ac2 is described below
commit b219ac2e72e8a2a3e2fa776683bf9a7cd1e5405a
Author: Joe Evans <[email protected]>
AuthorDate: Fri Dec 11 10:22:19 2020 -0800
[v1.x] Add more ONNX export support to operators (#19625)
---
ci/docker/runtime_functions.sh | 1 +
.../mxnet/contrib/onnx/mx2onnx/_op_translations.py | 186 +++++++++++++++++++++
tests/python-pytest/onnx/test_operators.py | 134 +++++++++++++++
3 files changed, 321 insertions(+)
diff --git a/ci/docker/runtime_functions.sh b/ci/docker/runtime_functions.sh
index 68ecf30..c387236 100755
--- a/ci/docker/runtime_functions.sh
+++ b/ci/docker/runtime_functions.sh
@@ -1281,6 +1281,7 @@ integrationtest_ubuntu_cpu_onnx() {
pytest tests/python-pytest/onnx/mxnet_export_test.py
pytest tests/python-pytest/onnx/test_models.py
pytest tests/python-pytest/onnx/test_node.py
+ pytest tests/python-pytest/onnx/test_operators.py
pytest tests/python-pytest/onnx/test_onnxruntime.py
}
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index f03fabb..efb58c0 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -154,6 +154,29 @@ def create_basic_op_node(op_name, node, kwargs):
)
return [node]
+def create_const_scalar_node(input_name, value, kwargs):
+ """Helper function to create a tensor value node and a
+ initializer tensor node with constant value."""
+ from onnx.helper import make_tensor, make_tensor_value_info
+ initializer = kwargs["initializer"]
+ input_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[value.dtype]
+ value_node = make_tensor_value_info(input_name, input_type, ())
+ tensor_node = make_tensor(input_name, input_type, (), (value,))
+ initializer.append(tensor_node)
+ return value_node
+
+def create_const_node(input_name, value, kwargs):
+ """Helper function to create a tensor value node and a
+ initializer tensor node with constant value."""
+ from onnx.helper import make_tensor, make_tensor_value_info
+ initializer = kwargs["initializer"]
+ input_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[value.dtype]
+ input_shape = value.shape
+ value_node = make_tensor_value_info(input_name, input_type, input_shape)
+ tensor_node = make_tensor(input_name, input_type, input_shape, value)
+ initializer.append(tensor_node)
+ return value_node
+
@mx_op.register("null")
def convert_weights_and_inputs(node, **kwargs):
"""Helper function to convert weights and inputs.
@@ -802,6 +825,7 @@ def convert_leakyrelu(node, **kwargs):
"""Map MXNet's LeakyReLU operator attributes to onnx's Elu/LeakyRelu/PRelu
operators
based on the input node's attributes and return the created node.
"""
+ from onnx.helper import make_node
name, input_nodes, attrs = get_inputs(node, kwargs)
act_type = attrs.get("act_type", "leaky")
@@ -816,6 +840,19 @@ def convert_leakyrelu(node, **kwargs):
inputs=input_nodes,
outputs=[name],
name=name)
+ elif act_type in ('gelu'):
+ sqrt2 = np.float32(1.4142135623730951)
+ nodes = [
+ create_const_scalar_node(name+"_sqrt2", sqrt2, kwargs),
+ make_node("Div", [input_nodes[0], name+"_sqrt2"],
[name+"_div0_out"]),
+ make_node("Erf", [name+"_div0_out"], [name+"_erf0_out"]),
+ create_const_scalar_node(name+"_one", np.float32(1.0), kwargs),
+ create_const_scalar_node(name+"_half", np.float32(0.5), kwargs),
+ make_node("Add", [name+"_erf0_out", name+"_one"],
[name+"_add0_out"]),
+ make_node("Mul", [input_nodes[0], name+"_add0_out"],
[name+"_mul0_out"]),
+ make_node("Mul", [name+"_mul0_out", name+"_half"], [name])
+ ]
+ return nodes
else:
node = onnx.helper.make_node(
act_name[act_type],
@@ -2214,3 +2251,152 @@ def convert_take(node, **kwargs):
name=name,
)
return [node]
+
+
+@mx_op.register("LayerNorm")
+def convert_layer_norm(node, **kwargs):
+ """Map MXNet's LayerNorm operator attributes to onnx operators.
+ """
+ from onnx.helper import make_node
+ name, input_nodes, attrs = get_inputs(node, kwargs)
+
+ in_shape = kwargs['in_shape']
+ axes = [-i for i in range(len(in_shape[0]), 0, -1)]
+ eps = attrs.get('eps')
+ nodes = [
+ make_node("ReduceMean", [input_nodes[0]], [name+"_rm0_out"],
axes=axes),
+ make_node("Sub", [input_nodes[0], name+"_rm0_out"],
[name+"_sub0_out"]),
+ create_const_scalar_node(name+"_two", np.float32(2.), kwargs),
+ make_node("Pow", [name+"_sub0_out", name+"_two"], [name+"_pow0_out"]),
+ make_node("ReduceMean", [name+"_pow0_out"], [name+"_rm1_out"],
axes=axes),
+ create_const_scalar_node(name+"_eps", np.float32(eps), kwargs),
+ make_node("Add", [name+"_rm1_out", name+"_eps"], [name+"_add0_out"]),
+ make_node("Sqrt", [name+"_add0_out"], [name+"_sqrt0_out"]),
+ make_node("Div", [name+"_sub0_out", name+"_sqrt0_out"],
[name+"_div0_out"]),
+ make_node("Mul", [name+"_div0_out", input_nodes[1]],
[name+"_mul0_out"]),
+ make_node("Add", [name+"_mul0_out", input_nodes[2]], [name], name)
+ ]
+
+ return nodes
+
+
+@mx_op.register("Embedding")
+def convert_embedding(node, **kwargs):
+ """Map MXNet's Embedding operator attributes to onnx's
+ Gather operator."""
+ name, input_nodes, attrs = get_inputs(node, kwargs)
+ axis = int(attrs.get('axis', 0))
+ node = onnx.helper.make_node(
+ "Gather",
+ input_nodes,
+ [name],
+ axis=axis,
+ name=name
+ )
+ return [node]
+
+
+@mx_op.register("stack")
+def convert_stack(node, **kwargs):
+ """Map MXNet's stack operator to onnx operators.
+ """
+ name, input_nodes, attrs = get_inputs(node, kwargs)
+ axis = int(attrs.get('axis', 0))
+ idx = 0
+ nodes = []
+ for input_node in input_nodes:
+ nodes.append(onnx.helper.make_node(
+ "Unsqueeze",
+ inputs=[input_node],
+ outputs=[name+"_unsqueeze"+str(idx)],
+ axes=[axis]
+ ))
+ idx += 1
+
+ nodes.append(onnx.helper.make_node(
+ "Concat",
+ inputs=[name+"_unsqueeze"+str(i) for i in range(len(nodes))],
+ outputs=[name],
+ name=name,
+ axis=axis
+ ))
+ return nodes
+
+
+@mx_op.register("slice")
+def convert_slice(node, **kwargs):
+ """Map MXNet's slice operator to onnx Slice operator."""
+ name, input_nodes, attrs = get_inputs(node, kwargs)
+ starts = convert_string_to_list(attrs.get("begin"))
+ ends = convert_string_to_list(attrs.get("end"))
+ steps = attrs.get("step", [])
+ nodes = [
+ create_const_node(name+"_begin", np.array(starts), kwargs),
+ create_const_node(name+"_end", np.array(ends), kwargs)
+ ]
+ inputs = [input_nodes[0], name+"_begin", name+"_end"]
+ if len(steps) > 0:
+ nodes.append(create_const_node(name+"_steps", np.array(steps,
dtype='int64'), kwargs))
+ inputs.append(name+"_steps")
+ nodes.append(onnx.helper.make_node("Slice", inputs, [name], name=name))
+ return nodes
+
+
+@mx_op.register("zeros_like")
+def convert_zeros_like(node, **kwargs):
+ """Map MXNet's zeros_like operator attributes to onnx's ConstantOfShape
operator.
+ """
+ from onnx.helper import make_node, make_tensor
+ name, _, _ = get_inputs(node, kwargs)
+
+ # create tensor with shape of input
+ create_const_node(name+"_shape", np.array(kwargs['in_shape'][0],
dtype='int64'), kwargs)
+ tensor_value = make_tensor(name+"_zero", kwargs['in_type'], [1], [0])
+ nodes = [
+ make_node("ConstantOfShape", [name+"_shape"], [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.
+ """
+ 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")
+
+ input_type = kwargs['in_type']
+ dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
+ in_shape = kwargs['in_shape']
+ axis = attrs.get('axis')
+
+ if axis is None:
+ # output will be same shape as input
+ output_shape = in_shape[0]
+ else:
+ # determine shape of axis
+ output_shape = [in_shape[0][int(axis)]]
+
+ start = np.array([attrs.get('start', 0.)], dtype=dtype)
+ step = np.array([attrs.get('step', 1.)], dtype=dtype)
+ repeat = np.array([attrs.get('repeat', 1)], dtype=dtype)
+ if repeat != 1:
+ raise NotImplementedError("arange_like operator with repeat != 1 not
yet implemented.")
+
+ tot_elements = np.prod(output_shape)
+ limit = np.array([start + (tot_elements * step)], dtype=dtype)
+
+ # create constant inputs
+ nodes = [
+ create_const_scalar_node(name+"_start", start, kwargs),
+ create_const_scalar_node(name+"_limit", limit, kwargs),
+ create_const_scalar_node(name+"_step", step, kwargs),
+ create_const_node(name+"_shape", np.array(output_shape,
dtype='int64'), kwargs),
+ make_node("Range", [name+"_start", name+"_limit", name+"_step"],
[name+"_range0_out"]),
+ make_node("Reshape", [name+"_range0_out", name+"_shape"], [name])
+ ]
+ return nodes
diff --git a/tests/python-pytest/onnx/test_operators.py
b/tests/python-pytest/onnx/test_operators.py
new file mode 100644
index 0000000..aa66858
--- /dev/null
+++ b/tests/python-pytest/onnx/test_operators.py
@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import mxnet as mx
+from mxnet.gluon import HybridBlock, nn
+import numpy as np
+import onnxruntime as rt
+from mxnet.test_utils import assert_almost_equal
+import pytest
+import tempfile
+
+def op_export_test(op_name, Model, inputs, tmp_path):
+ def export_to_onnx(model, op_name, inputs):
+ model_path = '{}/{}'.format(tmp_path, op_name)
+ model.export(model_path, epoch=0)
+ sym_file = '{}-symbol.json'.format(model_path)
+ params_file = '{}-0000.params'.format(model_path)
+ dtype = inputs[0].dtype
+ onnx_file = '{}/{}.onnx'.format(tmp_path, op_name)
+ mx.contrib.onnx.export_model(sym_file, params_file, [i.shape for i in
inputs],
+ dtype, onnx_file)
+ return onnx_file
+ def onnx_rt(onnx_file, inputs):
+ sess = rt.InferenceSession(onnx_file)
+ input_dict = dict((sess.get_inputs()[i].name, inputs[i].asnumpy()) for
i in range(len(inputs)))
+ pred = sess.run(None, input_dict)[0]
+ return pred
+
+ # create a new model
+ model = Model()
+ model.initialize(ctx=mx.cpu(0))
+ model.hybridize()
+ pred_nat = model(*inputs)
+ onnx_file = export_to_onnx(model, op_name, inputs)
+ pred_onx = onnx_rt(onnx_file, inputs)
+ assert_almost_equal(pred_nat, pred_onx)
+
+
+def test_onnx_export_abs():
+ with tempfile.TemporaryDirectory() as tmp_path:
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x):
+ out = F.abs(x)
+ return out
+ x = mx.nd.array([[-2, -1], [0, 99]], dtype='float32')
+ op_export_test('abs', Model, [x], tmp_path)
+
+def test_onnx_export_slice():
+ with tempfile.TemporaryDirectory() as tmp_path:
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x):
+ out = F.slice(x, begin=(0,1), end=(2,4))
+ return out
+ x = mx.nd.array([[1,2,3,4],[5,6,7,8],[9,10,11,12]], dtype='float32')
+ op_export_test('slice', Model, [x], tmp_path)
+
+def test_onnx_export_stack():
+ with tempfile.TemporaryDirectory() as tmp_path:
+ dtype = 'float32'
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x, y):
+ out = F.stack(x, y)
+ return out
+ x = mx.nd.array([1, 2], dtype=dtype)
+ y = mx.nd.array([3, 4], dtype=dtype)
+ op_export_test('stack', Model, [x, y], tmp_path)
+
+def test_onnx_export_zeros_like():
+ with tempfile.TemporaryDirectory() as tmp_path:
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x):
+ out = F.zeros_like(x)
+ return out
+ x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype='float32')
+ op_export_test('zeros_like', Model, [x], tmp_path)
+
[email protected]("dtype", ["float32", "double"])
+def test_onnx_export_arange_like(dtype):
+ with tempfile.TemporaryDirectory() as tmp_path:
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x):
+ out = F.contrib.arange_like(x)
+ return out
+ x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype=dtype)
+ op_export_test('arange_like', Model, [x], tmp_path)
+
+def test_onnx_export_layernorm():
+ with tempfile.TemporaryDirectory() as tmp_path:
+ dtype = 'float32'
+ class Model(HybridBlock):
+ def __init__(self, **kwargs):
+ super(Model, self).__init__(**kwargs)
+ def hybrid_forward(self, F, x, gamma, beta):
+ out = F.LayerNorm(x, gamma, beta, axis=1)
+ return out
+ x = mx.nd.array([[1,3],[2,4]], dtype=dtype)
+ gamma = mx.random.uniform(0, 1, x[0].shape).astype(dtype)
+ beta = mx.random.uniform(0, 1, x[0].shape).astype(dtype)
+ op_export_test('LayerNorm', Model, [x, gamma, beta], tmp_path)
+
+
+if __name__ == '__main__':
+ test_onnx_export_abs()
+ test_onnx_export_slice()
+ test_onnx_export_stack()
+ test_onnx_export_zeros_like()
+ test_onnx_export_arange_like('float32')
+ test_onnx_export_arange_like('double')
+ test_onnx_export_layernorm()
+
