kazum commented on a change in pull request #5052: URL: https://github.com/apache/incubator-tvm/pull/5052#discussion_r437683470
########## File path: python/tvm/contrib/target/onnx.py ########## @@ -0,0 +1,905 @@ +# 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. +# pylint: disable=invalid-name, import-self, len-as-condition, unused-argument, too-many-lines, redefined-builtin +"""Relay to ONNX codegen """ + +import os +import struct +import copy +import numpy +import onnx +import onnx.utils +from onnx import numpy_helper, OperatorSetIdProto, defs +import tvm +from tvm import relay +import tvm._ffi +from tvm.relay.expr_functor import ExprVisitor +from tvm.relay.ty import TupleType, TensorType + +ONNX_OPSET_VERSONS_SUPPORTED = [11] + + +def tvm_array_to_list(arr): + return tuple(x.value for x in arr) + + +def get_onnx_version(): + return onnx.__version__ + + +def infer_type(node): + """A method to infer the type of a relay expression.""" + mod = tvm.IRModule.from_expr(node) + mod = relay.transform.InferType()(mod) + entry = mod["main"] + return entry if isinstance(node, relay.Function) else entry.body + + +def call_node_infer_type(node): + """infer the output types of call node""" + infer_out = infer_type(node) + out_type = infer_out._checked_type_ + types = [] + if isinstance(out_type, TensorType): + types.append(out_type) + elif isinstance(out_type, TupleType): + for tupe_type in out_type.fields: + types.append(tupe_type) Review comment: `types = list(out_type.fields)` looks simpler. ########## File path: tests/python/contrib/test_onnx.py ########## @@ -0,0 +1,467 @@ +# 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. + +"""Relay to ONNX serialization test cases""" +import pytest +pytest.importorskip('onnx') +pytest.importorskip('onnxruntime') + +import numpy as np +import onnxruntime as rt + +import tvm +from tvm import relay +from tvm.contrib.target.onnx import to_onnx + + + +def func_to_onnx(func, name): + mod = tvm.IRModule() + mod['main'] = func + onnx_model = to_onnx(mod, {}, name, path=None) + return onnx_model.SerializeToString() + + +def run_onnx(onnx_model, input_data): + sess = rt.InferenceSession(onnx_model) + input_names = {} + for input, data in zip(sess.get_inputs(), input_data): + input_names[input.name] = data + output_names = [out.name for out in sess.get_outputs()] + res = sess.run(output_names, input_names) + return res + + +def run_relay(func, data_tuple): + target = 'llvm' + ctx = tvm.context('llvm', 0) + intrp = relay.create_executor("graph", ctx=ctx, target=target) + relay_res = intrp.evaluate(func)(*data_tuple) + + result = [] + relay_res = relay_res if isinstance(relay_res, list) else [relay_res] + for res in relay_res: + result.append(res.asnumpy()) + + return result + + +def verify_results(relay_func, indata, test_name, rtol=1e-7, atol=0): + relay_results = run_relay(relay_func, indata) + onnx_results = run_onnx(func_to_onnx(relay_func, test_name), indata) + + for relay_res, onnx_res in zip(relay_results, onnx_results): + np.testing.assert_allclose(relay_res, onnx_res, rtol=rtol, atol=atol) + + +def test_add(): + dtype = 'float32' + t1 = relay.TensorType((5, 10, 5)) + t2 = relay.TensorType((5, 10, 5)) + x = relay.var("x", t1, dtype=dtype) + y = relay.var("y", t2, dtype=dtype) + z = relay.add(x, y) + func = relay.Function([x, y], z) + + x_data = np.random.rand(5, 10, 5).astype(dtype) + y_data = np.random.rand(5, 10, 5).astype(dtype) + + verify_results(func, [x_data, y_data], 'test_add') + + +def test_bias_add(): + for dtype in ['float16', 'float32']: + xshape = (10, 2, 3, 4) + bshape = (2,) + rtol = 1e-2 if dtype == 'float16' else 1e-5 + x = relay.var("x", shape=xshape, dtype=dtype) + bias = relay.var("bias", dtype=dtype) + z = relay.nn.bias_add(x, bias) + func = relay.Function([x, bias], z) + + x_data = np.random.uniform(size=xshape).astype(dtype) + y_data = np.random.uniform(size=bshape).astype(dtype) + + verify_results(func, [x_data, y_data], 'test_bias_add', rtol=rtol) + + +def test_conv2d(): + def verify_conv2d(dtype, scale, dshape, kshape, + padding=(1, 1), + groups=1, + dilation=(1, 1), + **attrs): + x = relay.var("x", shape=dshape, dtype=dtype) + w = relay.var("w", shape=kshape, dtype=dtype) + y = relay.nn.conv2d(x, w, + padding=padding, + dilation=dilation, + groups=groups, + **attrs) + func = relay.Function([x, w], y) + data = np.random.uniform(-scale, scale, size=dshape).astype(dtype) + kernel = np.random.uniform(-scale, scale, size=kshape).astype(dtype) + verify_results(func, [data, kernel], 'test_conv2d', rtol=1e-5, atol=1e-5) + + dshape = (1, 32, 18, 18) + kshape = (32, 1, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=32, groups=32, kernel_size=(3, 3)) + + dshape = (1, 32, 18, 18) + kshape = (32, 4, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=32, groups=8, kernel_size=(3, 3)) + + # also group conv2d + dshape = (1, 32, 18, 18) + kshape = (64, 1, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=64, groups=32, kernel_size=(3, 3)) + + # normal conv2d + dshape = (1, 3, 224, 224) + kshape = (10, 3, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=10, kernel_size=(3, 3)) + + dshape = (1, 3, 224, 224) + kshape = (10, 3, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(2, 2), channels=10, kernel_size=(3, 3)) + + dshape = (1, 3, 18, 18) + kshape = (10, 3, 3, 3) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=10, kernel_size=(3, 3), dilation=(3, 3)) + + dshape = (1, 3, 18, 18) + kshape = (10, 3, 2, 2) + verify_conv2d("float32", 1, dshape, kshape, + padding=(2, 2), channels=10, kernel_size=(2, 2), dilation=(1, 1)) + + dshape = (1, 3, 18, 18) + kshape = (10, 3, 4, 4) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=10, kernel_size=(4, 4)) + + dshape = (1, 3, 18, 18) + kshape = (10, 3, 4, 4) + verify_conv2d("float32", 1, dshape, kshape, + padding=(1, 1), channels=10, kernel_size=(4, 4)) + + +def test_reshape(): + def verify_reshape(shape, newshape): + x = relay.var("x", relay.TensorType(shape, "float32")) + z = relay.reshape(x, newshape=newshape) + + func = relay.Function([x], z) + x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32") + verify_results(func, [x_data], 'test_reshape', rtol=1e-5, atol=1e-5) + + verify_reshape((2, 3, 4), tuple(np.array([4, 2, 3], dtype=np.int64))) + verify_reshape((2, 3, 4), tuple(np.array([2, 0, 0], dtype=np.int64))) + verify_reshape((2, 3, 4), tuple(np.array([0, -1], dtype=np.int64))) + verify_reshape((2, 3, 4), tuple(np.array([-1, 0], dtype=np.int64))) + + +def test_transpose(): + def verify_reshape(shape, newshape): + x = relay.var("x", relay.TensorType(shape, "float32")) + z = relay.transpose(x, newshape) + func = relay.Function([x], z) + x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32") + verify_results(func, [x_data], 'test_transpose', rtol=1e-5, atol=1e-5) + + verify_reshape((1, 2, 3, 4), (0, 2, 3, 1)) + verify_reshape((1, 2, 3, 4), (0, 3, 2, 1)) + + +def test_dense(): + def verify_dense(d_shape, w_shape): + data = relay.var("data", relay.TensorType(d_shape, "float32")) + weight = relay.var("weight", relay.TensorType(w_shape, "float32")) + func = relay.Function([data, weight], relay.nn.dense(data, weight)) + x_data = np.random.uniform(size=d_shape).astype("float32") + w_data = np.random.uniform(size=w_shape).astype("float32") + verify_results(func, [x_data, w_data], 'test_dense', rtol=1e-5, atol=1e-5) + + verify_dense((1, 8), (16, 8)) + verify_dense((1, 4), (3, 4)) + + +def test_max_pool(): + def verify_max_pool(x_shape, pool_size, strides, padding, ceil_mode): + x = relay.var("x", relay.TensorType(x_shape, "float32")) + y = tvm.relay.nn.max_pool2d(x, pool_size=pool_size, strides=strides, padding=padding, + ceil_mode=ceil_mode) + func = relay.Function([x], y) + x_data = np.random.uniform(size=x_shape).astype("float32") + verify_results(func, [x_data], 'test_max_pool', rtol=1e-5, atol=1e-5) + + verify_max_pool((1, 4, 16, 16), pool_size=(2, 2), strides=(2, 2), padding=(0, 0), ceil_mode=False) + + +def test_batch_flatten(): + def verify_test_batch_flatten(d_shape): + data = relay.var("data", relay.TensorType(d_shape, "float32")) + func = relay.Function([data], relay.nn.batch_flatten(data)) + x_data = np.random.uniform(size=d_shape).astype("float32") + verify_results(func, [x_data], 'test_batch_flatten', rtol=1e-5, atol=1e-5) + + verify_test_batch_flatten((1, 2, 3, 4)) + verify_test_batch_flatten((1, 8)) + + +def test_batch_norm(): + def verify_batch_norm(axis=1): + for dtype in ['float16', 'float32']: + data = relay.var("data", relay.TensorType((2, 4, 4, 1), dtype)) + gamma_shape = (data.type_annotation.shape[axis].value,) + beta = relay.var("beta", relay.TensorType(gamma_shape, dtype)) + gamma = relay.var("gamma", relay.TensorType(gamma_shape, dtype)) + moving_mean = relay.var("moving_mean", relay.TensorType(gamma_shape, dtype)) + moving_var = relay.var("moving_var", relay.TensorType(gamma_shape, dtype)) + y = relay.nn.batch_norm(data, gamma, beta, moving_mean, moving_var, axis=axis) + func = relay.Function([data, gamma, beta, moving_mean, moving_var], y[0]) + + x_data = np.random.uniform(size=(2, 4, 4, 1)).astype(dtype) + beta = np.random.uniform(size=gamma_shape).astype(dtype) + gamma = np.random.uniform(size=gamma_shape).astype(dtype) + moving_mean = np.random.uniform(size=gamma_shape).astype(dtype) + moving_var = np.random.uniform(size=gamma_shape).astype(dtype) + verify_results(func, [x_data, gamma, beta, moving_mean, moving_var], 'test_batch_norm', rtol=1e-3, + atol=1e-3) + + verify_batch_norm(axis=1) + verify_batch_norm(axis=3) + + +def test_pad(): + def verify_pad(): + for dtype in ['float16', 'float32']: + dshape = (4, 10, 7, 7) + x = relay.var("x", shape=dshape, dtype=dtype) + y = relay.nn.pad(x, ((1, 1), (2, 2), (3, 3), (4, 4))) + func = relay.Function([x], y) + x_data = np.random.uniform(size=dshape).astype(dtype) + verify_results(func, [x_data], 'test_pad', rtol=1e-5, atol=1e-5) + + verify_pad() + + +def test_sofmax(): + def verify_sofmax(): + for dtype in ['float32']: + shape = (10, 4) + x = relay.var("x", shape=shape, dtype=dtype) + y = relay.nn.softmax(x, axis=1) + func = relay.Function([x], y) + x_data = np.random.uniform(size=shape).astype(dtype) + verify_results(func, [x_data], 'test_softmax', rtol=1e-5, atol=1e-5) + + verify_sofmax() + + +def test_squeeze(): + def verify_squeeze(shape, dtype, axis): + x = relay.var("x", relay.TensorType(shape, dtype)) + z = relay.squeeze(x, axis=axis) + func = relay.Function([x], z) + x_data = np.random.random_sample(shape).astype(dtype) + verify_results(func, [x_data], 'test_squeeze', rtol=1e-5, atol=1e-5) + + verify_squeeze((1, 3, 2, 5), "float32", None) + verify_squeeze((1, 3, 1), "float32", [2, ]) + verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) + + +def test_mean(): + def verify_mean(data_shape, axis, exclude, keepdims): + dtype = "float32" + x = relay.var('x', shape=data_shape, dtype=dtype) + y = relay.mean(x, axis, keepdims, exclude) + func = relay.Function([x], y) + x_data = np.random.uniform(size=data_shape).astype(dtype) + verify_results(func, [x_data], 'test_mean', rtol=1e-5, atol=1e-5) + + verify_mean((1, 2), 0, False, False) + verify_mean((1, 2), 0, True, False) + verify_mean((1, 2), 0, True, True) + verify_mean((1, 2), 1, True, True) + verify_mean((3, 2, 1), 1, False, True) + + +def test_split(): + def verify_split(dshape, indices_or_sections, axis=None): + dtype = "float32" + x = relay.var("x", relay.ty.TensorType(dshape, "float32")) + y = relay.split(x, indices_or_sections, axis=axis) + func = relay.Function([x], y.astuple()) + x_data = np.random.uniform(size=dshape).astype(dtype) + + verify_results(func, [x_data], 'test_split', rtol=1e-5, atol=1e-5) + + verify_split((5, 5, 2, 2), 5, axis=1) + verify_split((5, 5, 2, 2), 5, axis=0) + verify_split((5, 5, 2, 2), [1, 3, 4], axis=0) + verify_split((5, 5, 2, 2), [1, 3, 4], axis=1) + + +def test_concatenate(): + def verify_concatenate(shapes, axis, dtype="float32"): + in_vars = [] + in_data = [] + for i, shape in enumerate(shapes): + in_vars.append(relay.var("x" + str(i), relay.ty.TensorType(shape, dtype))) + in_data.append(np.random.uniform(size=shape).astype(dtype)) + + out_tensor = relay.concatenate(in_vars, axis) + func = relay.Function(in_vars, out_tensor) + verify_results(func, in_data, 'test_concatenate', rtol=1e-5, atol=1e-5) + + verify_concatenate([(2,), (2,), (2,)], -1) + verify_concatenate([(2, 3, 4), (2, 2, 4), (2, 5, 4)], 1) + verify_concatenate([(1, 2, 4), (1, 2, 3), (1, 2, 7), (1, 2, 8), (1, 2, 1)], -1) + verify_concatenate([(5, 6, 7, 3), + (16, 6, 7, 3), + (12, 6, 7, 3), + (8, 6, 7, 3), + (2, 6, 7, 3)], 0) + verify_concatenate([(1, 14400), (1, 2400), (1, 640), (1, 240)], 1) + + +def test_strided_slice(): Review comment: This test fails due to the change in #4312. Could you rebase onto the latest master? ########## File path: python/tvm/contrib/target/onnx.py ########## @@ -0,0 +1,905 @@ +# 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. +# pylint: disable=invalid-name, import-self, len-as-condition, unused-argument, too-many-lines, redefined-builtin +"""Relay to ONNX codegen """ + +import os +import struct +import copy +import numpy +import onnx +import onnx.utils +from onnx import numpy_helper, OperatorSetIdProto, defs +import tvm +from tvm import relay +import tvm._ffi +from tvm.relay.expr_functor import ExprVisitor +from tvm.relay.ty import TupleType, TensorType + +ONNX_OPSET_VERSONS_SUPPORTED = [11] + + +def tvm_array_to_list(arr): + return tuple(x.value for x in arr) + + +def get_onnx_version(): + return onnx.__version__ + + +def infer_type(node): + """A method to infer the type of a relay expression.""" + mod = tvm.IRModule.from_expr(node) + mod = relay.transform.InferType()(mod) + entry = mod["main"] + return entry if isinstance(node, relay.Function) else entry.body + + +def call_node_infer_type(node): + """infer the output types of call node""" + infer_out = infer_type(node) + out_type = infer_out._checked_type_ + types = [] + if isinstance(out_type, TensorType): + types.append(out_type) + elif isinstance(out_type, TupleType): + for tupe_type in out_type.fields: + types.append(tupe_type) + else: + raise RuntimeError("Unsupported output type %s in operator %s" + % (type(out_type), node.op.nae)) + + return types + + +def add_input(data, name, model_container): + dtype = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[data.dtype] + tensor_value_info = onnx.helper.make_tensor_value_info(name, dtype, shape=data.shape) + model_container.add_inputs([tensor_value_info]) + data_tensor = numpy_helper.from_array(data, name) + model_container.add_initializers([data_tensor]) + + +class OpConverter(object): + """ Operator converter Base Class. + """ + + @classmethod + def convert_attributes(cls, attrs): + """convert Relay attributes to ONNX attributes. + The derived classes should implement this method + if attributes are required by the operator + otherwise by default no attributes are passed + """ + return {} + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + onnx_node = onnx.helper.make_node(cls.__name__, + node_entry['input_names'], + node_entry['output_names'], + **attrs) + model_container.add_nodes([onnx_node]) + + +def rename(op_name): + """ This method creates dynamic operator of name op_name with empty attributes + """ + return type(op_name, (OpConverter,), {}) + + +class Reshape(object): + """ Operator converter for Reshape. + """ + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + """Converts Relay operator Reshape to ONNX operator. + Relay operator accepts shape as attribute but ONNX operator + accepts it as a input. + """ + + shape = numpy.asarray([a.value for a in node_entry['relay_node'].attrs.newshape], + dtype=numpy.int64) + input_name = 'shape{}'.format(node_entry['name']) + node = onnx.helper.make_node(cls.__name__, [node_entry['input_names'][0], input_name], + node_entry['output_names']) + model_container.add_nodes([node]) + add_input(shape, input_name, model_container) + + +class Conv(OpConverter): + """ Operator converter for Conv. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'group': attrs.get_int("groups"), + 'pads': attrs.get_int_tuple("padding"), + 'strides': attrs.get_int_tuple("strides"), + 'dilations': attrs.get_int_tuple("dilation"), + 'kernel_shape': attrs.get_int_tuple("kernel_size"), + } + + +class MaxPool(OpConverter): + """ Operator converter for MaxPool. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'pads': attrs.get_int_tuple("padding"), + 'strides': attrs.get_int_tuple("strides"), + 'kernel_shape': attrs.get_int_tuple("pool_size"), + } + + +class Transpose(OpConverter): + """ Operator converter for Transpose. + """ + + @classmethod + def convert_attributes(cls, attrs): + return {'perm': attrs.get_int_tuple("axes")} if attrs["axes"] else {} + + +class MatMul(OpConverter): + """ Operator converter for MatMul. + """ + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + inter_output_name = 'inter{}'.format(node_entry['name']) + transpose_node = onnx.helper.make_node(Transpose.__name__, + [node_entry['input_names'][1]], + [inter_output_name], + perm=(1, 0)) + model_container.add_nodes([transpose_node]) + + inputs = [node_entry['input_names'][0], inter_output_name] + matmul_node = onnx.helper.make_node(cls.__name__, inputs, node_entry['output_names']) + model_container.add_nodes([matmul_node]) + + +class Flatten(OpConverter): + """ Operator converter for Flatten. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'axis': 1, + } + + +class BatchNormalization(OpConverter): + """ Operator converter for BatchNormalization. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'epsilon': float(attrs.get_str('epsilon')), + 'axis': float(attrs.get_int('axis')), + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + """Converts Relay operator batch_norm to ONNX operator. + Relay operator has property axis to handle data in NHWC format. + """ + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + transpose_out_name = node_entry['input_names'][0] + inter_output_names = [node_entry['output_names'][0]] + # axis==3 means channel is specified along the 3rd axis + if attrs['axis'] == 3: + transpose_out_name = 'transpose_{}'.format(node_entry['name']) + node_transposed = onnx.helper.make_node(Transpose.__name__, + [node_entry['input_names'][0]], + [transpose_out_name], + perm=[0, 3, 1, 2]) + model_container.add_nodes([node_transposed]) + inter_output_names = ['batch_norm_{}'.format(node_entry['name'])] + + input_names = [transpose_out_name] + node_entry['input_names'][1:] + batch_norm_node = onnx.helper.make_node(cls.__name__, + input_names, + inter_output_names, + epsilon=attrs['epsilon']) + model_container.add_nodes([batch_norm_node]) + + if attrs['axis'] == 3: + node_transposed = onnx.helper.make_node(Transpose.__name__, + inter_output_names, + [node_entry['output_names'][0]], + perm=[0, 2, 3, 1]) + model_container.add_nodes([node_transposed]) + + +class Dropout(OpConverter): + """ Operator converter for Dropout. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'ratio': float(attrs.get_str('rate')), + } + + +class AveragePool(MaxPool): + """ Operator converter for AveragePool. + """ + + +class Concat(OpConverter): + """ Operator converter for Concat. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'axis': attrs.get_int("axis"), + } + + +class BiasAdd(OpConverter): + """ Operator converter for BiasAdd. + """ + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node_entry can not be a Tuple" + input_node = input_node[0] + data_ndim = len(input_node['types'][0].shape) + axis = node_entry['relay_node'].attrs.get_int("axis") + if axis < 0: + axis = axis + data_ndim + new_axes = data_ndim - axis - 1 + if new_axes: + inter_output_name = 'inter{}'.format(node_entry['name']) + unsqueeze_node = onnx.helper.make_node('Unsqueeze', + [node_entry['input_names'][1]], + [inter_output_name], + axes=tuple(range(1, new_axes + 1))) + model_container.add_nodes([unsqueeze_node]) + else: + inter_output_name = node_entry['input_names'][1] + + inputs = [node_entry['input_names'][0], inter_output_name] + matmul_node = onnx.helper.make_node('Add', inputs, node_entry['output_names']) + model_container.add_nodes([matmul_node]) + + +class ReduceMean(OpConverter): + """ Operator converter for ReduceMean. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'axes': attrs.axis, + 'keepdims': 0 if bool(attrs.get_int("keepdims", 0)) is False else 1 + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node can not be a Tuple" + input_node = input_node[0] + shape = input_node['types'][0].shape + axis = node_entry['relay_node'].attrs.axis + axis = list(range(shape.size())) if not axis else tvm_array_to_list(axis) + exclude = 0 if not bool(node_entry['relay_node'].attrs.exclude) else 1 + keepdims = 0 if not bool(node_entry['relay_node'].attrs.keepdims) else 1 + if exclude: + all_axis = list(range(len(shape))) + axis = set(all_axis) - set(axis) + + node = onnx.helper.make_node(cls.__name__, + node_entry['input_names'], + node_entry['output_names'], + axes=axis, + keepdims=keepdims) + model_container.add_nodes([node]) + + +class Pad(OpConverter): + """ Operator converter for Pad. + """ + + @classmethod + def convert_attributes(cls, attrs): + before = [] + after = [] + for axis_pads in attrs.pad_width: + before.append(axis_pads[0]) + after.append(axis_pads[1]) + pads = before + after + pads = numpy.asarray(pads, dtype=pads[0].dtype) + return { + 'pads': pads, + 'mode': attrs.get_str('pad_mode'), + 'constant_value': attrs.pad_value + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + """Converts Relay operator Pad to ONNX operator. + Relay operator accepts pads as attribute but ONNX operator + accepts it as a input. + """ + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + + name = node_entry['name'] + data = numpy.asarray(attrs['pads'], dtype=attrs['pads'][0].dtype).astype(numpy.int64) + input_name = 'pads_{}'.format(name) + value = numpy.dtype(node_entry['types'][0].dtype).type(attrs['constant_value']) + input_value_name = 'value_{}'.format(name) + add_input(data, input_name, model_container) + add_input(value, input_value_name, model_container) + + input_names = [node_entry['input_names'][0], input_name, input_value_name] + node = onnx.helper.make_node(cls.__name__, input_names, node_entry['output_names']) + model_container.add_nodes([node]) + + +class Softmax(OpConverter): + """ Operator converter for SoftMax. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'axis': attrs.axis, + } + + +class Squeeze(OpConverter): + """ Operator converter for Squeeze. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'axes': attrs.axis, + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node can not be a Tuple" + input_node = input_node[0] + shape = input_node['types'][0].shape + axis = node_entry['relay_node'].attrs.get_int("axis") + if not axis: + axis = [] + for axis_idx, val in enumerate(shape): + if val.value == 1: + axis.append(axis_idx) + else: + axis = node_entry['relay_node'].attrs.get_int_tuple("axis") + + node = onnx.helper.make_node(cls.__name__, + node_entry['input_names'], + node_entry['output_names'], + axes=axis) + model_container.add_nodes([node]) + + +class Slice(OpConverter): + """ Operator converter for Slice. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'starts': attrs.get_int_tuple('begin'), + 'ends': attrs.get_int_tuple('end'), + 'steps': attrs.get_int_tuple('strides') + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node can not be a Tuple" + input_node = input_node[0] + shape = input_node['types'][0].shape + starts = list(attrs['starts']) + ends = list(attrs['ends']) + for i in range(len(starts), len(shape)): + starts.append(0) + for i in range(len(ends), len(shape)): + ends.append(shape[i] + 1) + + name = node_entry['name'] + starts = numpy.asarray(starts).astype(numpy.int64) + starts_name = 'starts_{}'.format(name) + add_input(starts, starts_name, model_container) + + ends = numpy.asarray(ends).astype(numpy.int64) + ends_name = 'ends_{}'.format(name) + add_input(ends, ends_name, model_container) + + input_names = node_entry['input_names'] + [starts_name, ends_name] + + if attrs['steps']: + axes = list(range(len(shape))) + attrs['axes'] = axes + assert len(axes) == len(attrs['steps']), "axes and steps should be of same size" + + steps = numpy.asarray(attrs['steps']).astype(numpy.int64) + steps_name = 'steps_{}'.format(name) + add_input(steps, steps_name, model_container) + + axes = numpy.asarray(attrs['axes']).astype(numpy.int64) + axes_name = 'axes_{}'.format(name) + add_input(axes, axes_name, model_container) + + input_names = input_names + [axes_name, steps_name] + + slice_node = onnx.helper.make_node(cls.__name__, + input_names, + node_entry['output_names']) + model_container.add_nodes([slice_node]) + + +class Split(OpConverter): + """ Operator converter for Split. + """ + + @classmethod + def convert_attributes(cls, attrs): + indices_or_sections = attrs['indices_or_sections'] + + if isinstance(indices_or_sections, (list, tvm.ir.container.Array)): + indices_or_sections = attrs.get_int_tuple('indices_or_sections') + if isinstance(indices_or_sections, tvm.ir.PrimExpr): + indices_or_sections = indices_or_sections.value + + return { + 'indices_or_section': indices_or_sections, + 'axis': attrs.get_int('axis'), + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node can not be a Tuple" + input_node = input_node[0] + shape = input_node['types'][0].concrete_shape + + indices_or_sect = attrs["indices_or_section"] + axis = attrs["axis"] + axis_length = shape[axis] + + if isinstance(indices_or_sect, int): + split = [axis_length // indices_or_sect] * indices_or_sect + else: + split = [] + for i in range(len(indices_or_sect) + 1): + if i == 0: + split.append(indices_or_sect[0]) + elif i == len(indices_or_sect): + split.append(axis_length - indices_or_sect[-1]) + else: + split.append(indices_or_sect[i] - indices_or_sect[i - 1]) + + slice_node = onnx.helper.make_node(cls.__name__, + node_entry['input_names'], + node_entry['output_names'], + split=split, + axis=axis) + model_container.add_nodes([slice_node]) + + +class ConstantOfShapeZeros(OpConverter): + """ Operator converter for ConstantOfShape. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'value': 0 + } + + @classmethod + def convert(cls, node_entry, model_container, node_dict): + attrs = cls.convert_attributes(node_entry['relay_node'].attrs) + input_node = node_dict[node_entry['inputs'][0]] + assert len(input_node) == 1, "input node can not be a Tuple" + input_node = input_node[0] + dtype = input_node['relay_node'].type_annotation.dtype + input_shape_name = 'shape_{}'.format(node_entry['name']) + shape = [val.value for val in input_node['relay_node'].type_annotation.shape] + shape = numpy.asarray(shape).astype(numpy.int64) + add_input(shape, input_shape_name, model_container) + + dtype = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[numpy.dtype(dtype)] + tensor_value = onnx.helper.make_tensor("value", dtype, + [1], [attrs['value']]) + + node = onnx.helper.make_node('ConstantOfShape', + [input_shape_name], + node_entry['output_names'], + value=tensor_value) + model_container.add_nodes([node]) + + +class ConstantOfShapeOnes(ConstantOfShapeZeros): + """ Operator converter for ConstantOfShape. + """ + + @classmethod + def convert_attributes(cls, attrs): + return { + 'value': 1 + } + + +relay_to_onnx_op_mapping = { + 'reshape': Reshape, + 'nn.conv2d': Conv, + 'add': rename('Add'), + 'nn.relu': rename('Relu'), + 'transpose': Transpose, + 'nn.dense': MatMul, + 'nn.max_pool2d': MaxPool, + 'nn.batch_flatten': Flatten, + 'multiply': rename('Mul'), + 'nn.bias_add': BiasAdd, + 'nn.batch_norm': BatchNormalization, + 'nn.global_avg_pool2d': rename('GlobalAveragePool'), + 'concatenate': Concat, + 'nn.dropout': Dropout, + 'nn.avg_pool2d': AveragePool, + 'divide': rename('Div'), + 'mean': ReduceMean, + 'nn.pad': Pad, + 'nn.softmax': Softmax, + 'squeeze': Squeeze, + 'strided_slice': Slice, + 'greater': rename('Greater'), + 'less': rename('Less'), + 'equal': rename('Equal'), + 'zeros_like': ConstantOfShapeZeros, + 'ones_like': ConstantOfShapeOnes, + 'subtract': rename('Sub'), + 'split': Split +} + + +class ModelContainer(object): + """ A container class to hold different attributes of ONNX model graph + """ + + def __init__(self, name, opset_version): + self._name = name + self._opset_version = opset_version + self._inputs = [] + self._outputs = [] + self._nodes = [] + self._initializers = [] + + def add_inputs(self, inputs): + self._inputs.extend(inputs) + + def add_outputs(self, outputs): + self._outputs.extend(outputs) + + def add_nodes(self, nodes): + self._nodes.extend(nodes) + + def add_initializers(self, initializers): + self._initializers.extend(initializers) + + def _get_opsets(self): + opsets = [] + imp = OperatorSetIdProto() + imp.version = self._opset_version + opsets.append(imp) + return opsets + + def make_model(self): + """ Creates the onnx model from the graph """ + onnx_graph = onnx.helper.make_graph( + self._nodes, + self._name, + self._inputs, + self._outputs, + self._initializers + ) + kwargs = {} + kwargs["opset_imports"] = self._get_opsets() + kwargs["producer_name"] = 'TVM Relay' + kwargs["producer_version"] = tvm.__version__ + + return onnx.helper.make_model(onnx_graph, **kwargs) + + +class RelayToONNXConverter(ExprVisitor): + """A helper class to traverse the Relay graph and convert Relay nodes to ONNX model + + Parameters + ---------- + name : str + name of the model + + params : dict + dict of the parameter names and NDarray values + + opset_version : int + target onnx opset version + + """ + + def __init__(self, name, params, opset_version): + super().__init__() + self._name = {} + self._mc = ModelContainer(name, opset_version) + self._params = params + self._node_dict = {} + self._node_count = 0 + self.last_node = None + + @classmethod + def _get_node_entry(cls, relay_node, name, node_index): + return {"relay_node": relay_node, + "inputs": [relay_node], # inputs in the form of relay nodes + "types": [], # output types in case of call nodes else self type + "name": name, # name of the node + "input_names": [name], # input names in case of call nodes else self name + "output_names": [name], # output names in case of call nodes else self name + "op": None, # op name in case of call node else None + "index": node_index Review comment: The key "index" is not used anywhere? ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
