| # 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 |
| # pylint: disable=import-outside-toplevel |
| """ONNX: Open Neural Network Exchange frontend for Relay.""" |
| import copy |
| import math |
| import warnings |
| from typing import Optional |
| |
| import numpy as np |
| import tvm |
| from tvm import relay |
| from tvm.ir import IRModule |
| from tvm.topi.utils import get_const_tuple |
| |
| from ... import nd as _nd |
| from .. import analysis |
| from .. import expr as _expr |
| from .. import function as _function |
| from .. import loops as _loops |
| from .. import op as _op |
| from .. import qnn as _qnn |
| from .. import random as _random |
| from .. import ty as _ty |
| from .. import vision as _vision |
| from .common import ( |
| AttrCvt, |
| Renamer, |
| autopad, |
| ensure_scalar_shape, |
| fold_constant, |
| get_name, |
| get_relay_op, |
| gru_cell, |
| infer_channels, |
| infer_shape, |
| infer_type, |
| infer_value, |
| lstm_cell, |
| new_var, |
| rnn_cell, |
| shape_of, |
| try_resolve_var_to_const, |
| unbind, |
| ) |
| |
| __all__ = ["from_onnx"] |
| |
| # The default configurations of Relay ONNX frontend. |
| ONNX_DEFAULT_CONFIGS = { |
| # By default, TVM converts qualified onnx `matmul` to `transpose(weight) + nn.batch_matmul_NT`. |
| # Change this flag to False to directly convert to `nn.batch_matmul`. |
| # Note that `nn.batch_matmul` with format other than NT is in experimental, it may have some |
| # performance issues. |
| "use_nt_batch_matmul": True, |
| } |
| |
| |
| class onnx_input(list): |
| """A helper extension to list that returns None for out of bound indices.""" |
| |
| def __getitem__(self, item): |
| if isinstance(item, slice): |
| if item.stop is None: |
| stop = len(self) |
| else: |
| stop = item.stop |
| indices = list(range(stop)[item]) |
| return [self[i] for i in indices] |
| if isinstance(item, int): |
| return list(self)[item] if item < len(self) else None |
| raise TypeError("list indices must be integers or slices, not %s" % type(item).__name__) |
| |
| |
| def get_numpy(tensor_proto): |
| """Grab data in TensorProto and convert to numpy array.""" |
| try: |
| from onnx.numpy_helper import to_array |
| except ImportError as e: |
| raise ImportError("Unable to import onnx which is required {}".format(e)) |
| return to_array(tensor_proto) |
| |
| |
| def get_type(elem_type): |
| """Converts onnx integer datatype to numpy datatype""" |
| # If a string was passed instead of a tensor type, it does not need |
| # conversion and can be returned. |
| if isinstance(elem_type, str): |
| return elem_type |
| |
| try: |
| from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE |
| except ImportError as e: |
| raise ImportError("Unable to import onnx which is required {}".format(e)) |
| |
| try: |
| from onnx import TensorProto |
| except ImportError as e: |
| raise ImportError("Unable to import TensorProto from onnx {}".format(e)) |
| |
| # Onnx mapping converts bfloat16 to float16 because |
| # numpy does not have a bfloat16 data type. However, |
| # tvm has one, so we force the return type to be bfloat16 |
| if elem_type == int(TensorProto.BFLOAT16): |
| return "bfloat16" |
| return str(TENSOR_TYPE_TO_NP_TYPE[elem_type]) |
| |
| |
| def get_info(info_proto): |
| """Extract the shape from a ValueInfoProto.""" |
| shape = [] |
| shape_name = [] |
| for dim in info_proto.type.tensor_type.shape.dim: |
| name = dim.dim_param |
| value = dim.dim_value |
| if value is None or value == 0: |
| value = _ty.Any() |
| shape_name.append(name) |
| else: |
| shape_name.append(value) |
| shape.append(value) |
| |
| name = info_proto.name |
| if info_proto.type.tensor_type.elem_type: |
| dtype = get_type(info_proto.type.tensor_type.elem_type) |
| else: |
| dtype = None |
| return name, shape, dtype, shape_name |
| |
| |
| def dimension_picker(prefix, suffix=""): |
| """Check that dimensions are supported.""" |
| |
| def _impl(attr): |
| kernel = attr["kernel_shape"] |
| if len(kernel) == 1: |
| return prefix + "1d" + suffix |
| if len(kernel) == 2: |
| return prefix + "2d" + suffix |
| if len(kernel) == 3: |
| return prefix + "3d" + suffix |
| msg = "Only 1D, 2D, and 3D kernels are supported for operator {}." |
| op_name = prefix + "1d/2d/3d" |
| raise tvm.error.OpAttributeInvalid(msg.format(op_name)) |
| |
| return _impl |
| |
| |
| def revert_caffe2_pad(pads): |
| """Caffe2 requires two times the normal padding.""" |
| if len(pads) == 4: |
| pads = pads[:2] |
| elif len(pads) == 2: |
| pass |
| else: |
| raise tvm.error.OpAttributeInvalid("Number of pads must be either 2 or 4.") |
| return pads |
| |
| |
| def get_pad_pair(input1d, kernel1d, stride1d, mode): |
| """infer pad size""" |
| if input1d % stride1d == 0: |
| pad = max(kernel1d - stride1d, 0) |
| else: |
| pad = max(kernel1d - (input1d % stride1d), 0) |
| pad_before = pad // 2 |
| pad_after = pad - pad_before |
| if "LOWER" in mode: |
| return [pad_after, pad_before] |
| return [pad_before, pad_after] |
| |
| |
| def onnx_default_layout(dims, op_name): |
| if dims == 1: |
| return "NCW" |
| if dims == 2: |
| return "NCHW" |
| if dims == 3: |
| return "NCDHW" |
| |
| msg = "Only 1D, 2D and 3D layouts are currently supported for operator {}." |
| raise tvm.error.OpAttributeInvalid(msg.format(op_name)) |
| |
| |
| def onnx_storage_order2layout(storage_order, dims, op_name): |
| """converter of onnx storage order parameter to tvm storage order format""" |
| if storage_order not in (0, 1): |
| raise tvm.error.OpAttributeInvalid("Mode of storage_order must be either 0 or 1") |
| |
| if dims == 1: |
| return "NCW" if storage_order == 0 else "NWC" |
| if dims == 2: |
| return "NCHW" if storage_order == 0 else "NHWC" |
| if dims == 3: |
| return "NCDHW" if storage_order == 0 else "NDHWC" |
| |
| msg = "Only 1D, 2D and 3D layouts are currently supported for operator {}." |
| raise tvm.error.OpAttributeInvalid(msg.format(op_name)) |
| |
| |
| def dimension_constraint(): |
| def _dim_check(attrs): |
| if len(attrs["kernel_shape"]) in [1, 2, 3]: |
| return True |
| return False |
| |
| return _dim_check, "Only 1d, 2d and 3d kernel supported." |
| |
| |
| def get_scalar(x, params, dtype="float32"): |
| """Helper to get a scalar value for Quantized operators.""" |
| if isinstance(x, _expr.Var) and x.name_hint in params: |
| return _op.const(params[x.name_hint].numpy(), dtype) |
| rank = len(infer_shape(x)) |
| assert rank <= 1, "scale and zero_point input must be scalars" |
| if rank == 1: |
| x = _op.squeeze(x, [0]) |
| return _op.cast(x, dtype) |
| |
| |
| def get_scalar_or_1d_tensor(x, params, dtype="float32"): |
| """Helper to get a scalar value or 1D tensor for Quantized operators.""" |
| if isinstance(x, _expr.Var) and x.name_hint in params: |
| return _op.const(params[x.name_hint].numpy(), dtype) |
| rank = len(infer_shape(x)) |
| assert rank <= 1, "scale and zero_point input must be scalars or 1D tensors" |
| return _op.cast(x, dtype) |
| |
| |
| def matmul_out_dtype(inputs, out_dtype): |
| """Common function to handle MatMul and MatMulInteger16""" |
| a_shape = shape_of(inputs[0]) |
| a_rank = infer_shape(a_shape)[0] |
| b_shape = shape_of(inputs[1]) |
| b_rank = infer_shape(b_shape)[0] |
| if a_rank > 2 or b_rank > 2: |
| |
| def flatten_to_nd(x, x_shape, nd=3): |
| ndims = infer_shape(x_shape)[0] |
| if ndims == nd: |
| return x |
| newshape = _op.concatenate( |
| [ |
| _expr.const([-1], dtype=infer_type(x_shape).checked_type.dtype), |
| _op.strided_slice(x_shape, [ndims - nd + 1], [ndims]), |
| ], |
| 0, |
| ) |
| out = _op.reshape(x, fold_constant(newshape)) |
| return out |
| |
| # Determine the output batch dimension. |
| new_a_shape = a_shape |
| new_b_shape = b_shape |
| if a_rank > b_rank: |
| rank_diff = a_rank - b_rank |
| new_b_shape = _op.concatenate( |
| [ |
| _expr.const([1] * rank_diff, dtype=infer_type(b_shape).checked_type.dtype), |
| b_shape, |
| ], |
| 0, |
| ) |
| elif a_rank < b_rank: |
| rank_diff = b_rank - a_rank |
| new_a_shape = _op.concatenate( |
| [ |
| _expr.const([1] * rank_diff, dtype=infer_type(a_shape).checked_type.dtype), |
| a_shape, |
| ], |
| 0, |
| ) |
| else: |
| pass |
| |
| out_batch = _op.concatenate( |
| [ |
| _op.maximum( |
| _op.strided_slice(new_b_shape, [i], [i + 1]), |
| _op.strided_slice(new_a_shape, [i], [i + 1]), |
| ) |
| for i in range(max(a_rank, b_rank) - 2) |
| ], |
| 0, |
| ) |
| |
| b_type = infer_type(inputs[1]) |
| # Convert to dense if the second matrix is 2d and non-dynamic |
| if b_rank == 2 and not _ty.is_dynamic(b_type.checked_type): |
| a = flatten_to_nd(inputs[0], a_shape, 2) |
| b = _op.transpose(inputs[1]) |
| output = _op.nn.dense(a, b, out_dtype=out_dtype) |
| else: |
| a = inputs[0] |
| b = inputs[1] |
| # broadcast a and b |
| a_broadcasted_shape = fold_constant( |
| _op.concatenate( |
| [ |
| out_batch, |
| _op.strided_slice(a_shape, [a_rank - 2], [a_rank]), |
| ], |
| 0, |
| ) |
| ) |
| b_broadcasted_shape = fold_constant( |
| _op.concatenate( |
| [ |
| out_batch, |
| _op.strided_slice(b_shape, [b_rank - 2], [b_rank]), |
| ], |
| 0, |
| ) |
| ) |
| if not tvm.ir.structural_equal(a_shape, a_broadcasted_shape): |
| a = _op.transform.broadcast_to(a, a_broadcasted_shape) |
| if not tvm.ir.structural_equal(b_shape, b_broadcasted_shape): |
| b = _op.transform.broadcast_to(b, b_broadcasted_shape) |
| # Convert a and b into 3 dimensional tensors. |
| a = flatten_to_nd(a, shape_of(a), 3) |
| b = flatten_to_nd(b, shape_of(b), 3) |
| if ONNX_DEFAULT_CONFIGS["use_nt_batch_matmul"]: |
| # Transpose matrix dimensions of b. |
| bt = _op.transpose(b, [0, 2, 1]) |
| # Perform a NT batch matmul. |
| output = _op.nn.batch_matmul(a, bt, out_dtype=out_dtype) |
| else: |
| # Perform a NN batch matmul. |
| output = _op.nn.batch_matmul(a, b, out_dtype=out_dtype, transpose_b=False) |
| # Reshape output to original dimensions. |
| final_shape = _op.concatenate( |
| [ |
| out_batch, |
| _op.strided_slice( |
| a_shape, [infer_shape(a_shape)[0] - 2], [infer_shape(a_shape)[0] - 1] |
| ), |
| _op.strided_slice( |
| b_shape, [infer_shape(b_shape)[0] - 1], [infer_shape(b_shape)[0]] |
| ), |
| ], |
| 0, |
| ) |
| return _op.reshape(output, fold_constant(final_shape)) |
| |
| if a_rank == 1: |
| return _op.squeeze(_op.nn.matmul(_op.expand_dims(inputs[0], axis=0), inputs[1]), axis=[0]) |
| |
| # Otherwise a simple dense op will get the job done. |
| input_1_t = _op.transpose(inputs[1], axes=(1, 0)) |
| return _op.nn.dense(inputs[0], input_1_t, out_dtype=out_dtype) |
| |
| |
| def layer_norm(x, eps, gamma, beta): |
| """Common function to handle layer norm""" |
| eps_dtype = infer_type(x).checked_type.dtype |
| |
| u, s = _op.mean_variance(x, axis=-1, keepdims=True) |
| output = _op.divide( |
| _op.subtract(x, u), |
| _op.sqrt(_op.add(s, _op.const(eps, dtype=eps_dtype))), |
| ) |
| output = _op.multiply(output, gamma) |
| if beta is not None: |
| output = _op.add(output, beta) |
| |
| return output |
| |
| |
| class OnnxOpConverter(object): |
| """A helper class for holding onnx op converters.""" |
| |
| @classmethod |
| def get_converter(cls, opset): |
| """Get converter matches given opset. |
| |
| Parameters |
| ---------- |
| opset: int |
| opset from model. |
| |
| Returns |
| ------- |
| converter, which should be `_impl_vx`. Number x is the biggest |
| number smaller than or equal to opset belongs to all support versions. |
| """ |
| versions = [int(d.replace("_impl_v", "")) for d in dir(cls) if "_impl_v" in d] |
| versions = sorted(versions + [opset]) |
| version = versions[max([i for i, v in enumerate(versions) if v == opset]) - 1] |
| if hasattr(cls, "_impl_v{}".format(version)): |
| return getattr(cls, "_impl_v{}".format(version)) |
| raise NotImplementedError( |
| "opset version {} of {} not implemented".format(version, cls.__name__) |
| ) |
| |
| |
| class Unary(OnnxOpConverter): |
| """A helper class for unary op converters.""" |
| |
| name = "" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 1, "Unary math op {} takes 1 input, {} given".format( |
| cls.name, len(inputs) |
| ) |
| op_name = cls.name |
| return get_relay_op(op_name)(*inputs) |
| |
| |
| class Elemwise(OnnxOpConverter): |
| """A helper class for elemwise op converters.""" |
| |
| name = "" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 2, "Math op {} take 2 inputs, {} given".format(cls.name, len(inputs)) |
| op_name = cls.name |
| conv_ops = ["conv2d", "conv2d_transpose"] |
| if attr.get("broadcast", 0) and any(x in str(inputs[0]) for x in conv_ops): |
| # TODO(zhreshold): remove hard coded infershape |
| axis = int(attr.get("axis", 0)) |
| inputs[1] = _op.expand_dims(inputs[1], axis=axis, num_newaxis=2) |
| return get_relay_op(op_name)(*inputs) |
| |
| |
| class Pool(OnnxOpConverter): |
| """A helper class for pool op converters.""" |
| |
| name = "" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| input_shape = infer_shape(data) |
| ndim = len(input_shape) |
| |
| attr_cvt, data = cls._run_calculation(inputs, attr, params) |
| out = attr_cvt([data], attr, params) |
| |
| if ndim - len(attr["kernel_shape"]) == 1: |
| out = _op.squeeze(out, axis=[0]) |
| return out |
| |
| @classmethod |
| def _run_calculation(cls, inputs, attr, params): |
| """Helper method to return the processed input data and AttrCvt object""" |
| |
| data = inputs[0] |
| input_shape = infer_shape(data) |
| input_dtype = infer_type(data).checked_type.dtype |
| ndim = len(input_shape) |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| if cls.name == "avg_pool": |
| pad_tuple = [] |
| for axis in range(len(input_shape) - 2): |
| axis_shape = input_shape[2 + axis] |
| stride = attr.get("strides", [1] * ndim)[axis] |
| kernel = attr["kernel_shape"][axis] |
| pad = get_pad_pair(axis_shape, kernel, stride, attr["auto_pad"]) |
| pad_tuple.append(pad) |
| pad_tuple = tuple([val for pair in zip(*pad_tuple) for val in pair]) |
| attr["pads"] = pad_tuple |
| else: |
| # Warning: Pool does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| if "int" in input_dtype: |
| pad_val = np.iinfo(np.dtype(input_dtype)).min |
| else: |
| pad_val = np.finfo(np.dtype(input_dtype)).min |
| data = autopad( |
| data, |
| attr.get("strides", [1] * (ndim - 2)), |
| attr["kernel_shape"], |
| [1] * ndim, |
| pad_value=pad_val, |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator {} is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"], cls.name)) |
| attr.pop("auto_pad") |
| |
| if "storage_order" in attr: |
| attr["layout"] = onnx_storage_order2layout( |
| attr["storage_order"], dims=(len(input_shape) - 2), op_name=cls.name |
| ) |
| else: |
| if ndim - len(attr["kernel_shape"]) == 1: |
| data = _op.expand_dims(data, axis=0) |
| input_shape = [1] + list(input_shape) |
| |
| attr["layout"] = onnx_default_layout(dims=(len(input_shape) - 2), op_name=cls.name) |
| |
| return ( |
| AttrCvt( |
| op_name=dimension_picker(cls.name), |
| transforms={ |
| "kernel_shape": "pool_size", |
| "pads": ("padding", 0), |
| "dilations": ("dilation", 1), |
| }, |
| ignores=["storage_order"], |
| custom_check=dimension_constraint(), |
| ), |
| data, |
| ) |
| |
| |
| class Absolute(Unary): |
| """Operator converter for Absolute.""" |
| |
| name = "abs" |
| |
| |
| class Add(Elemwise): |
| """Operator converter for Add.""" |
| |
| name = "add" |
| |
| |
| class AveragePool(Pool): |
| """Operator converter for AveragePool.""" |
| |
| name = "avg_pool" |
| |
| |
| class QLinearAveragePool(Pool): |
| """Operator converter for QLinearAveragePool from Microsoft onnxruntime contrib opset.""" |
| |
| name = "avg_pool" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| x_scale = get_scalar(inputs[1], params) |
| x_zero_point = get_scalar(inputs[2], params, dtype="int32") |
| y_scale = fold_constant(get_scalar(inputs[3], params)) |
| y_zero_point = get_scalar(inputs[4], params, dtype="int32") |
| |
| attr_cvt, data = cls._run_calculation(inputs, attr, params) |
| |
| input_dtype = infer_type(data).checked_type.dtype |
| # Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 |
| # and then requantize afer (according to documentation below) |
| # https://github.com/microsoft/onnxruntime/blob/master/docs/ContribOperators.md#com.microsoft.QLinearAveragePool |
| float_node = _qnn.op.dequantize(data, x_scale, x_zero_point) |
| out = attr_cvt([float_node], attr, params) |
| return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=input_dtype) |
| |
| |
| class BatchNorm(OnnxOpConverter): |
| """Operator converter for BatchNorm.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # TODO(zhreshold): 'spatial' is not properly handled here. |
| # TODO(vvchernov): 'training_mode' (onnx tag) is not correctly handled, ignore for now |
| out = AttrCvt( |
| op_name="batch_norm", |
| ignores=["spatial", "is_test", "consumed_inputs", "momentum", "training_mode"], |
| )(inputs, attr, params) |
| # We only support test mode, so we return data, moving_mean, moving_var, |
| # and then moving_mean and moving_var again as placeholders for |
| # the expected "saved_mean", "saved_var". |
| return _expr.TupleWrapper(_expr.Tuple((*out, out[1], out[2])), 5) |
| |
| |
| class InstanceNorm(OnnxOpConverter): |
| """Operator converter for BatchNorm.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return AttrCvt(op_name="instance_norm")(inputs, attr, params) |
| |
| |
| class Conv(OnnxOpConverter): |
| """Operator converter for Conv.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # Use shape of input to determine convolution type. |
| data = inputs[0] |
| kernel = inputs[1] |
| input_shape = infer_shape(data) |
| ndim = len(input_shape) |
| |
| kernel_type = infer_type(inputs[1]) |
| kernel_shapes = [get_const_tuple(kernel_type.checked_type.shape)] |
| |
| if "kernel_shape" not in attr: |
| attr["kernel_shape"] = kernel_shapes[0][2:] |
| |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: Convolution does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| data = autopad( |
| data, |
| attr.get("strides", [1] * (ndim - 2)), |
| attr["kernel_shape"], |
| attr.get("dilations", [1] * (ndim - 2)), |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = [0 for i in range(ndim - 2)] |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"])) |
| attr.pop("auto_pad") |
| |
| attr["channels"] = kernel_shapes[0][0] |
| out = AttrCvt( |
| op_name=dimension_picker("conv"), |
| transforms={ |
| "kernel_shape": "kernel_size", |
| "dilations": ("dilation", 1), |
| "pads": ("padding", 0), |
| "group": ("groups", 1), |
| }, |
| custom_check=dimension_constraint(), |
| )([data, kernel], attr, params) |
| |
| use_bias = len(inputs) == 3 |
| if use_bias: |
| out = _op.nn.bias_add(out, inputs[2]) |
| return out |
| |
| |
| class ConvTranspose(OnnxOpConverter): |
| """Operator converter for ConvTranspose.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # get number of channels |
| out_type = infer_type(inputs[1]) |
| kernel_shape = [get_const_tuple(out_type.checked_type.shape)] |
| out_channels = kernel_shape[0][1] * attr.get("group", 1) |
| attr["channels"] = out_channels |
| groups = attr.get("group", 1) |
| |
| if "kernel_shape" not in attr: |
| attr["kernel_shape"] = kernel_shape[0][2:] |
| |
| attr["groups"] = groups |
| # infer pads for auto_pad |
| data = inputs[0] |
| input_shape = infer_shape(data) |
| ndim = len(input_shape) |
| if "auto_pad" in attr or "output_shape" in attr: |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if "output_shape" in attr or attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: Convolution does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| kernel_shape = attr["kernel_shape"] |
| kndim = len(kernel_shape) |
| dilations = attr.get("dilations", [1] * kndim) |
| output_padding = attr.get("output_padding", [0] * kndim) |
| strides = attr["strides"] |
| total_pad = [0] * kndim |
| # https://github.com/onnx/onnx/blob/main/docs/Operators.md#ConvTranspose |
| if "output_shape" in attr: |
| for i in range(kndim): |
| total_pad[i] = ( |
| strides[i] * (input_shape[ndim - kndim + i] - 1) |
| + output_padding[i] |
| + ((kernel_shape[i] - 1) * dilations[i] + 1) |
| - attr["output_shape"][i] |
| ) |
| left = [p // 2 for p in total_pad] |
| right = [total_pad[i] - left[i] for i in range(kndim)] |
| if "output_shape" in attr and "auto_pad" not in attr: |
| pad = right + left |
| elif "LOWER" in attr["auto_pad"]: |
| pad = left + right |
| else: |
| pad = right + left |
| attr["pads"] = pad |
| else: |
| data = autopad( |
| data, |
| attr.get("strides", [1] * (ndim - 2)), |
| attr["kernel_shape"], |
| attr.get("dilations", [1] * (ndim - 2)), |
| deconv=True, |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"])) |
| if "auto_pad" in attr: |
| attr.pop("auto_pad") |
| |
| out = AttrCvt( |
| op_name=dimension_picker("conv", "_transpose"), |
| transforms={ |
| "kernel_shape": "kernel_size", |
| "dilations": ("dilation", 1), |
| "pads": ("padding", 0), |
| "group": ("groups", 1), |
| }, |
| disables=["output_shape"], |
| custom_check=dimension_constraint(), |
| )([data, inputs[1]], attr, params) |
| use_bias = len(inputs) == 3 |
| if use_bias: |
| out = _op.nn.bias_add(out, inputs[2]) |
| return out |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| # get number of channels |
| out_type = infer_type(inputs[1]) |
| kernel_shape = [get_const_tuple(out_type.checked_type.shape)] |
| out_channels = kernel_shape[0][1] * attr.get("group", 1) |
| attr["channels"] = out_channels |
| groups = attr.get("group", 1) |
| |
| if "kernel_shape" not in attr: |
| attr["kernel_shape"] = kernel_shape[0][2:] |
| |
| attr["groups"] = groups |
| # infer pads for auto_pad |
| data = inputs[0] |
| input_shape = infer_shape(data) |
| ndim = len(input_shape) |
| if "auto_pad" in attr or "output_shape" in attr: |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if "output_shape" in attr or attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: Convolution does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| kernel_shape = attr["kernel_shape"] |
| kndim = len(kernel_shape) |
| dilations = attr.get("dilations", [1] * kndim) |
| output_padding = attr.get("output_padding", [0] * kndim) |
| strides = attr["strides"] |
| total_pad = [0] * kndim |
| # https://github.com/onnx/onnx/blob/main/docs/Operators.md#ConvTranspose |
| if "output_shape" in attr: |
| for i in range(kndim): |
| total_pad[i] = ( |
| strides[i] * (input_shape[ndim - kndim + i] - 1) |
| + output_padding[i] |
| + ((kernel_shape[i] - 1) * dilations[i] + 1) |
| - attr["output_shape"][i] |
| ) |
| else: |
| for i in range(kndim): |
| total_pad[i] = ( |
| output_padding[i] |
| + ((kernel_shape[i] - 1) * dilations[i] + 1) |
| - strides[i] |
| ) |
| left = [p // 2 for p in total_pad] |
| right = [total_pad[i] - left[i] for i in range(kndim)] |
| if "output_shape" in attr and "auto_pad" not in attr: |
| pad = right + left |
| elif "LOWER" in attr["auto_pad"]: |
| pad = left + right |
| else: |
| pad = right + left |
| attr["pads"] = pad |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"])) |
| if "auto_pad" in attr: |
| attr.pop("auto_pad") |
| |
| out = AttrCvt( |
| op_name=dimension_picker("conv", "_transpose"), |
| transforms={ |
| "kernel_shape": "kernel_size", |
| "dilations": ("dilation", 1), |
| "pads": ("padding", 0), |
| "group": ("groups", 1), |
| }, |
| disables=["output_shape"], |
| custom_check=dimension_constraint(), |
| )([data, inputs[1]], attr, params) |
| use_bias = len(inputs) == 3 |
| if use_bias: |
| out = _op.nn.bias_add(out, inputs[2]) |
| return out |
| |
| |
| class GlobalAveragePool(OnnxOpConverter): |
| """Operator converter for GlobalAveragePool""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| rank = len(infer_shape(inputs[0])) |
| if rank == 3: |
| return _op.nn.global_avg_pool1d(inputs[0]) |
| if rank == 4: |
| return _op.nn.global_avg_pool2d(inputs[0]) |
| if rank == 5: |
| return _op.nn.global_avg_pool3d(inputs[0]) |
| raise NotImplementedError( |
| "Global average pooling is only implemented for 1D, 2D, and 3D kernels, got %dD." |
| % (rank - 2), |
| ) |
| |
| |
| class QLinearGlobalAveragePool(OnnxOpConverter): |
| "Operator converter for QLinearGlobalAveragePool from Microsoft onnxruntime contrib opset." |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| rank = len(infer_shape(inputs[0])) |
| |
| x_scale = get_scalar(inputs[1], params) |
| x_zero_point = get_scalar(inputs[2], params, dtype="int32") |
| y_scale = fold_constant(get_scalar(inputs[3], params)) |
| y_zero_point = get_scalar(inputs[4], params, dtype="int32") |
| |
| input_dtype = infer_type(inputs[0]).checked_type.dtype |
| |
| # Onnxruntime documentation does not mention that this global avg_pool should follow the |
| # sequence dequantize -> float op -> quantize, but that is how QLinearAveragePool is done. |
| # |
| # This op also follows the same pattern since qnn op is not available right now. |
| # TODO: Generate QNN op to perform quantized operation instead of dequant -> op -> quant |
| x = _qnn.op.dequantize(inputs[0], x_scale, x_zero_point) |
| if rank == 3: |
| out = _op.nn.global_avg_pool1d(x) |
| elif rank == 4: |
| out = _op.nn.global_avg_pool2d(x) |
| elif rank == 5: |
| out = _op.nn.global_avg_pool3d(x) |
| else: |
| raise NotImplementedError( |
| "Global average pooling is only implemented for 1D, 2D, and 3D kernels, got %dD." |
| % (rank - 2), |
| ) |
| return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=input_dtype) |
| |
| |
| class GlobalMaxPool(OnnxOpConverter): |
| """Operator converter for GlobalMaxPool""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| rank = len(infer_shape(inputs[0])) |
| if rank == 3: |
| return _op.nn.global_max_pool1d(inputs[0]) |
| if rank == 4: |
| return _op.nn.global_max_pool2d(inputs[0]) |
| if rank == 5: |
| return _op.nn.global_max_pool3d(inputs[0]) |
| raise NotImplementedError( |
| "Global max pooling is only implemented for 1D, 2D, and 3D kernels, got %dD." |
| % (rank - 2), |
| ) |
| |
| |
| class Div(Elemwise): |
| """Operator converter for Divide.""" |
| |
| name = "divide" |
| |
| |
| class Elu(OnnxOpConverter): |
| """Operator converter for Elu.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = float(attr.get("alpha", 1.0)) |
| return _expr.const(-alpha) * _op.nn.relu( |
| _expr.const(1.0) - _op.exp(inputs[0]) |
| ) + _op.nn.relu(inputs[0]) |
| |
| |
| class Gelu(OnnxOpConverter): |
| """Operator converter for Gelu from Microsoft onnxruntime contrib opset. |
| |
| gelu(x) = 0.5x(1 + erf(x/sqrt(2))) |
| """ |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| x = inputs[0] |
| |
| # Declare consts |
| const_dtype = infer_type(x).checked_type.dtype |
| half = _expr.const(0.5, dtype=const_dtype) |
| one = _expr.const(1.0, dtype=const_dtype) |
| sqrt2 = _expr.const(math.sqrt(2), dtype=const_dtype) |
| |
| # Compute gelu |
| term1 = _op.multiply(half, x) |
| erf = _op.erf(_op.divide(x, sqrt2)) |
| term2 = _op.add(one, erf) |
| return _op.multiply(term1, term2) |
| |
| |
| class BiasGelu(OnnxOpConverter): |
| """Operator converter for BiasGelu from Microsoft onnxruntime contrib opset. |
| |
| bias_gelu(x, b) = 0.5(x, b)(1 + erf((x + b)/sqrt(2))) |
| """ |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| x = inputs[0] |
| b = inputs[1] |
| |
| b_shape = infer_shape(b) |
| assert len(b_shape) == 1, "BiasGelu bias term must be a 1D tensor" |
| |
| inp = _op.add(x, b) |
| return Gelu._impl_v1([inp], attr, params) |
| |
| |
| class EmbedLayerNormalization(OnnxOpConverter): |
| """Operator converter for EmbedLayerNormalization from Microsoft onnxruntime contrib opset. |
| |
| This layer embeds the input tokens, sums them, and applies layer normalization. |
| """ |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| input_ids = inputs[0] |
| segment_ids = inputs[1] |
| word_emb = inputs[2] |
| pos_emb = inputs[3] |
| segment_emb = inputs[4] |
| gamma = inputs[5] |
| beta = inputs[6] |
| |
| mask = inputs[7] |
| pos_ids = inputs[8] |
| |
| eps = attr.get("epsilon", 1e-12) |
| |
| (batch_size, seq_len) = infer_shape(input_ids) |
| |
| if segment_ids: |
| assert segment_emb |
| |
| if pos_ids is None: |
| pos_ids = _op.const([list(range(seq_len))] * batch_size, dtype="int32") |
| |
| word_vec = _op.take(word_emb, input_ids, axis=0) |
| segment_vec = _op.take(segment_emb, segment_ids, axis=0) |
| pos_vec = _op.take(pos_emb, pos_ids, axis=0) |
| |
| vec_sum = _op.add(word_vec, pos_vec) |
| if segment_ids: |
| vec_sum = _op.add(vec_sum, segment_vec) |
| |
| ln = layer_norm(vec_sum, eps, gamma, beta) |
| |
| mask_index = _op.const(np.zeros((batch_size,), dtype="int32")) |
| if mask: |
| # calculate number of words per sentence |
| mask_index = _op.sum(mask, axis=1) |
| |
| # TODO(@anwang2009): onnxruntime v1.10.0 requires a third output of vec_sum |
| return _expr.TupleWrapper(_expr.Tuple([ln, mask_index]), 2) |
| |
| |
| class SkipLayerNormalization(OnnxOpConverter): |
| """Operator converter for SkipLayerNormalization from Microsoft onnxruntime contrib opset. |
| |
| This layer sums the two input tensors (along with optional bias), and applies layer |
| normalization. |
| """ |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| skip = inputs[1] |
| gamma = inputs[2] |
| beta = inputs[3] |
| bias = inputs[4] |
| |
| assert ( |
| beta is not None and bias is not None |
| ), "SkipLayerNormalization import currently only supports required beta and bias" |
| |
| eps = attr.get("epsilon", 1e-12) |
| |
| x = _op.add(data, skip) |
| if bias is not None: |
| x = _op.add(x, bias) |
| |
| output = layer_norm(x, eps, gamma, beta) |
| |
| # onnxruntime doesn't compute the other outputs, despite the documentation |
| placeholder = _op.const(0, dtype="float32") |
| |
| return _expr.TupleWrapper(_expr.Tuple([output, placeholder, placeholder]), 3) |
| |
| |
| class Attention(OnnxOpConverter): |
| """Operator converter for Attention from Microsoft onnxruntime contrib opset. |
| |
| This is the self-attention mechanism used in transformer models. |
| """ |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| num_heads = attr["num_heads"] |
| assert ( |
| "qkv_hidden_sizes" not in attr |
| ), "different hidden sizes for Q, K, V are not currently supported" |
| assert "unidirectional" not in attr, "unidirectional attention not current supported" |
| |
| # (batch, seq, in_hidden) |
| input_emb = inputs[0] |
| |
| # (in_hidden, 3 * out_hidden), where out_hidden = num_heads * head_size |
| weight = inputs[1] |
| |
| # (3 * out_hidden,) |
| bias = inputs[2] |
| |
| # 1. ( batch, 1, max_seq, max_seq) |
| # 2. ( batch, past_seq + seq,) |
| # 3. ( batch, seq, past_seq + seq,) |
| # 4. ( batch,) |
| # 5. (2 * batch,) |
| # For now, we only support case 2. |
| mask_index = inputs[3] |
| |
| # (2, batch, num_heads, past_seq, head_size) |
| past = inputs[4] |
| |
| # (batch, num_heads, seq, seq) |
| extra_add = inputs[5] |
| |
| (batch_size, seq_len, _) = infer_shape(input_emb) |
| (out_hidden_x3,) = infer_shape(bias) |
| assert out_hidden_x3 % 3 == 0, "bias shape should be divisible by 3" |
| out_hidden = out_hidden_x3 // 3 |
| assert ( |
| out_hidden % num_heads == 0 |
| ), "output hidden size should be divisible by number of attention heads" |
| head_size = out_hidden // num_heads |
| |
| assert ( |
| mask_index is not None |
| ), "Attention import currently only supports required mask_index" |
| mask_index_shape = infer_shape(mask_index) |
| assert ( |
| len(mask_index_shape) == 2 |
| and mask_index_shape[0] == batch_size |
| and mask_index_shape[1] == seq_len |
| ), "currently only support (batch_size, sequence_length) mask index" |
| |
| assert past is None, "past K, V state is not currently supported" |
| assert extra_add is None, "extra add to QxK not currently supported" |
| |
| # split weight and biases and do the matmuls |
| w_Q, w_K, w_V = _op.split(weight, 3, axis=1) |
| b_Q, b_K, b_V = _op.split(bias, 3, axis=0) |
| # need to merge batch dimensions since TVM matmul is 2D |
| input_emb = _op.reverse_reshape(input_emb, (-1, 0)) |
| Q = _op.add(_op.nn.matmul(input_emb, w_Q), b_Q) |
| K = _op.add(_op.nn.matmul(input_emb, w_K), b_K) |
| V = _op.add(_op.nn.matmul(input_emb, w_V), b_V) |
| |
| # massage tensors in preparation for batched matmul |
| def massage(tensor): |
| tensor = _op.reshape(tensor, (batch_size, seq_len, num_heads, head_size)) |
| |
| # (batch_size, num_heads, seq_len, head_size) |
| tensor = _op.transpose(tensor, axes=[0, 2, 1, 3]) |
| |
| # (batch_size * num_heads, seq_len, head_size) |
| return _op.reverse_reshape(tensor, (-1, 0, 0)) |
| |
| Q = massage(Q) |
| K = massage(K) |
| V = massage(V) |
| |
| K_present = _op.reshape(K, (batch_size, num_heads, seq_len, head_size)) |
| V_present = _op.reshape(V, (batch_size, num_heads, seq_len, head_size)) |
| present = _op.stack([K_present, V_present], axis=0) |
| |
| att_scores = _op.nn.batch_matmul(Q, K, transpose_a=False, transpose_b=True) |
| score_dtype = infer_type(att_scores).checked_type.dtype |
| att_scores = _op.divide( |
| att_scores, |
| _op.const(np.sqrt(head_size), dtype=infer_type(att_scores).checked_type.dtype), |
| ) |
| att_scores = _op.reshape(att_scores, (batch_size, num_heads, seq_len, seq_len)) |
| |
| # build the attention mask |
| att_mask = _op.cast(mask_index, score_dtype) |
| att_mask = _op.expand_dims(att_mask, 1, num_newaxis=2) |
| att_mask = _op.subtract(_op.const(1, dtype=score_dtype), att_mask) |
| att_mask = _op.multiply(att_mask, _op.const(-10000, dtype=score_dtype)) |
| |
| # apply the mask |
| att_scores = _op.add(att_scores, att_mask) |
| att_scores = _op.reshape(att_scores, (batch_size * num_heads, seq_len, seq_len)) |
| |
| att_probs = _op.nn.softmax(att_scores, axis=-1) |
| |
| output = _op.nn.batch_matmul(att_probs, V, transpose_a=False, transpose_b=False) |
| output = _op.reverse_reshape(output, (-1, num_heads, 0, 0)) |
| output = _op.transpose(output, axes=[0, 2, 1, 3]) |
| output = _op.reshape(output, (0, 0, out_hidden)) |
| |
| return _expr.TupleWrapper(_expr.Tuple([output, present]), 2) |
| |
| |
| class Gemm(OnnxOpConverter): |
| """Operator converter for Gemm.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 3 or len(inputs) == 2, "Gemm op take 2 or 3 inputs, {} given".format( |
| len(inputs) |
| ) |
| input0_state = infer_type(inputs[0]) |
| dtype = input0_state.checked_type.dtype |
| # Y = alpha * A * B + beta * C |
| alpha = float(attr.get("alpha", 1.0)) |
| beta = float(attr.get("beta", 1.0)) |
| transA = int(attr.get("transA", 0)) |
| transB = int(attr.get("transB", 0)) |
| # get number of channels |
| channels = infer_channels(inputs[1], not transB) |
| if transA: |
| inputs[0] = _op.transpose(inputs[0], axes=(1, 0)) |
| if not transB: |
| inputs[1] = _op.transpose(inputs[1], axes=(1, 0)) |
| if len(input0_state.checked_type.shape) != 2: |
| inputs[0] = _op.nn.batch_flatten(inputs[0]) |
| if alpha != 1.0: |
| inputs[0] *= _expr.const(alpha, dtype=dtype) |
| out = _op.nn.dense(inputs[0], inputs[1], units=channels) |
| if len(inputs) == 3: |
| out = out + _expr.const(beta, dtype=dtype) * inputs[2] |
| return out |
| |
| |
| class MatMul(OnnxOpConverter): |
| """Operator converter for MatMul.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 2, "MatMul op take 2 inputs, {} given".format(len(inputs)) |
| # Need to check input shape as batch matmul must be supported. |
| return matmul_out_dtype(inputs, out_dtype=infer_type(inputs[0]).checked_type.dtype) |
| |
| |
| class MatMulInteger16(OnnxOpConverter): |
| """Operator converter for MatMulInteger16 from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| assert len(inputs) == 2, "MatMulInteger16 op take 2 inputs, {} given".format(len(inputs)) |
| a_dtype = infer_type(inputs[0]).checked_type.dtype |
| b_dtype = infer_type(inputs[1]).checked_type.dtype |
| # Check input data types |
| assert a_dtype in ("int16", "uint16"), "MatMulInteger16: invalid dtype for first input" |
| assert b_dtype in ("int16", "uint16"), "MatMulInteger16: invalid dtype for second input" |
| out_dtype = "int32" |
| if a_dtype == "uint16" and b_dtype == "uint16": |
| out_dtype = "uint32" |
| return matmul_out_dtype(inputs, out_dtype) |
| |
| |
| class Mod(OnnxOpConverter): |
| """Operator converter for Mod.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 2, "Mod op take 2 inputs, {} given".format(len(inputs)) |
| |
| # Note: attr['fmod'] determines whether the operator should behave like np.fmod or np.mod. |
| # attr['fmod'] == 0 will behave as np.mod and attr['fmod'] == 1 will force fmod treatment. |
| # The relay equivalent of np.fmod is relay.mod and np.mod is relay.floor_mod |
| if attr.get("fmod", 0) == 0: |
| op_name = "floor_mod" |
| else: |
| op_name = "mod" |
| |
| return AttrCvt(op_name)(inputs, {}, params) |
| |
| |
| class MaxPool(Pool): |
| """Operator converter for MaxPool""" |
| |
| name = "max_pool" |
| |
| |
| class MaxUnpool(OnnxOpConverter): |
| """Operator converter for MaxUnpool""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| # Unpack inputs and attributes |
| data = inputs[0] |
| data_type = infer_type(data).checked_type.dtype |
| indices = inputs[1] |
| output_shape = inputs[2] |
| kernel_shape = attr.get("kernel_shape") |
| pads = attr.get("pads", None) |
| strides = attr.get("strides", [1] * len(kernel_shape)) |
| |
| # Compute the proper output shape before padding. |
| multiplier = _op.concatenate( |
| [_expr.const([1, 1], dtype="int64"), _expr.const(list(strides), dtype="int64")], axis=0 |
| ) |
| total_output_shape = multiplier * shape_of(data, dtype="int64") |
| # Add extra dimensions from kernel size and stride mismatch |
| total_output_shape += _op.concatenate( |
| [_expr.const([0, 0], "int64"), _expr.const(list(kernel_shape), "int64")], axis=0 |
| ) - _op.concatenate( |
| [_expr.const([0, 0], "int64"), _expr.const(list(strides), "int64")], axis=0 |
| ) |
| |
| # Compute padding amount if output shape is specified. |
| if output_shape is not None: |
| total_output_shape = output_shape |
| |
| elif pads is not None: |
| # Get pads in the proper format for relay. |
| pads = _op.concatenate( |
| [_expr.const([0, 0, 0, 0], "int64"), _expr.const(list(pads), "int64")], axis=0 |
| ) |
| pads = _op.reshape(pads, [-1, 2]) |
| # Compute the total padding per axis. |
| total_pad = _op.sum(pads, axis=-1) |
| # Reversing maxpool means that padding actually makes our output smaller. |
| total_output_shape = total_output_shape - total_pad |
| |
| # Create a tensor of zeros then scatter our data through it. |
| zeros_tensor = _op.zeros(total_output_shape, data_type) |
| # We need to flatten all our tensors before scattering. |
| flat_tensor = _op.scatter( |
| _op.reshape(zeros_tensor, [-1]), |
| _op.reshape(indices, [-1]), |
| _op.reshape(data, [-1]), |
| axis=0, |
| ) |
| # Now reshape back to prepadded shape. |
| output_tensor = _op.reshape(flat_tensor, total_output_shape) |
| |
| return output_tensor |
| |
| |
| class LpPool(OnnxOpConverter): |
| """A helper class for lppool op converters.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| data = inputs[0] |
| input_shape = infer_shape(data) |
| ndim = len(input_shape) |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: LpPool does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| data = autopad( |
| data, |
| attr["strides"], |
| attr["kernel_shape"], |
| [1] * ndim, |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator {} is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"], "LpPool")) |
| attr.pop("auto_pad") |
| |
| if "storage_order" in attr: |
| attr["layout"] = onnx_storage_order2layout( |
| attr["storage_order"], dims=(len(input_shape) - 2), op_name="LpPool" |
| ) |
| else: |
| attr["layout"] = onnx_default_layout(dims=(len(input_shape) - 2), op_name="LpPool") |
| |
| p_value = attr.get("p", 2) |
| p = _expr.const(p_value, dtype) |
| reci_p = _expr.const(1.0 / p_value, dtype) |
| data = _op.power(data, p) |
| |
| out = AttrCvt( |
| op_name=dimension_picker("avg_pool"), |
| transforms={"kernel_shape": "pool_size", "pads": ("padding", 0)}, |
| extras={"count_include_pad": True}, |
| ignores=["p"], |
| custom_check=dimension_constraint(), |
| )([data], attr, params) |
| kernels = attr["kernel_shape"] |
| out = _op.abs(out) * _expr.const(np.prod(kernels).astype(dtype)) |
| return _op.power(out, reci_p) |
| |
| |
| class GlobalLpPool(OnnxOpConverter): |
| """Operator converter for GlobalLpPool.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # TODO: GlobalLpPool does not yet support dynamic shapes |
| in_shape = infer_shape(inputs[0]) |
| attr["kernel_shape"] = in_shape[2:] |
| |
| return LpPool._impl_v1(inputs, attr, params) |
| |
| |
| class Mul(Elemwise): |
| """Operator converter for Multiply.""" |
| |
| name = "multiply" |
| |
| |
| class Pad(OnnxOpConverter): |
| """Operator converter for Pad.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| pad_width = [] |
| pads = attr.pop("paddings") |
| dims = int(len(pads) / 2) |
| for i in range(dims): |
| pad_width.append((pads[i], pads[i + dims])) |
| attr["pad_width"] = pad_width |
| pad_mode = attr.get("mode", b"constant").decode("utf-8") |
| if pad_mode in ["constant", "edge", "reflect"]: |
| attr["pad_mode"] = pad_mode |
| attr.pop("mode", None) |
| else: |
| raise tvm.error.OpAttributeInvalid( |
| "Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.' |
| ) |
| |
| return AttrCvt( |
| _op.nn.pad, |
| transforms={ |
| "value": "pad_value", |
| }, |
| )(inputs, attr, params) |
| |
| @classmethod |
| def _impl_v2(cls, inputs, attr, params): |
| pad_width = [] |
| pads = attr.pop("pads") |
| dims = int(len(pads) / 2) |
| for i in range(dims): |
| pad_width.append((pads[i], pads[i + dims])) |
| attr["pad_width"] = pad_width |
| pad_mode = attr.get("mode", b"constant").decode("utf-8") |
| if pad_mode in ["constant", "edge", "reflect"]: |
| attr["pad_mode"] = pad_mode |
| attr.pop("mode", None) |
| else: |
| raise tvm.error.OpAttributeInvalid( |
| "Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.' |
| ) |
| |
| return AttrCvt( |
| "pad", |
| transforms={ |
| "value": "pad_value", |
| }, |
| )(inputs, attr, params) |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| pads = inputs[1] |
| if len(inputs) == 3: |
| value = fold_constant(_op.take(inputs[2], _op.const(0))) |
| else: |
| value = 0.0 |
| |
| pad_width_expr = fold_constant(_op.transpose(_op.reshape(pads, (2, -1)))) |
| pad_mode = attr.get("mode", b"constant").decode("utf-8") |
| if not pad_mode in ["constant", "edge", "reflect"]: |
| raise tvm.error.OpAttributeInvalid( |
| "Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.' |
| ) |
| |
| return _op.nn.pad(inputs[0], pad_width_expr, value, pad_mode=pad_mode) |
| |
| |
| class ParametricSoftPlus(OnnxOpConverter): |
| """Operator converter for ParametricSoftPlus.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = _expr.const(float(attr.get("alpha", 1.0))) |
| beta = _expr.const(float(attr.get("beta", 1.0))) |
| return _op.log(_op.exp(beta * inputs[0]) + _expr.const(1.0)) * alpha |
| |
| |
| class Pow(OnnxOpConverter): |
| """Operator converter for Pow.""" |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| x = inputs[0] |
| y = inputs[1] |
| |
| x_type = infer_type(x).checked_type.dtype |
| output_type = x_type |
| y_type = infer_type(y).checked_type.dtype |
| |
| if not x_type.startswith("float"): |
| x_type = "float32" |
| x = _op.cast(x, x_type) |
| |
| if x_type != y_type: |
| y = _op.cast(y, x_type) |
| |
| # TODO: come up with good default integer pow() func for common backends |
| result = _op.power(x, y) |
| if x_type != output_type: |
| return _op.cast(result, output_type) |
| return result |
| |
| |
| class Prelu(OnnxOpConverter): |
| """Operator converter for Prelu.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 2, "Prelu need 2 inputs, {} given".format(len(inputs)) |
| input_shape = shape_of(inputs[0]) |
| alpha = _op.broadcast_to_like(inputs[1], inputs[0]) |
| alpha = _op.reshape(alpha, [-1]) |
| output = _op.nn.prelu(_op.reshape(inputs[0], [-1]), alpha, axis=0) |
| return _op.reshape(output, input_shape) |
| |
| |
| class Reciprocal(OnnxOpConverter): |
| """Operator converter for Reciprocal.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| return _expr.const(1.0, dtype=dtype) / inputs[0] |
| |
| |
| class Flatten(OnnxOpConverter): |
| """Operator converter for Flatten.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axis", 1) |
| ishape = shape_of(inputs[0]) |
| ndim = infer_shape(ishape)[0] |
| if axis < 0: |
| axis = axis + ndim |
| |
| if axis == 1: |
| out = _op.nn.batch_flatten(inputs[0]) |
| else: |
| pre_shape = _op.prod(_op.strided_slice(ishape, [0], [axis], [1]), keepdims=True) |
| post_shape = _op.prod(_op.strided_slice(ishape, [axis], [ndim], [1]), keepdims=True) |
| newshape = fold_constant(_op.concatenate([pre_shape, post_shape], axis=0)) |
| out = _op.reshape(inputs[0], newshape) |
| return out |
| |
| |
| class Reshape(OnnxOpConverter): |
| """Operator converter for Reshape.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return _op.reshape(inputs[0], attr["shape"]) |
| |
| @classmethod |
| def _impl_v5(cls, inputs, attr, params): |
| allowzero = attr.get("allowzero", False) |
| if get_name(inputs[1]) in params: |
| shape = tuple(params[inputs[1].name_hint].numpy().astype("int32")) |
| out = _op.reshape(inputs[0], shape, allowzero=allowzero) |
| else: |
| out = _op.reshape(*inputs, allowzero=allowzero) |
| return out |
| |
| |
| class DepthToSpace(OnnxOpConverter): |
| """Operator converter for DepthToSpace.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| block_size = int(attr["blocksize"]) |
| mode = attr.get("mode", b"DCR").decode("utf-8") |
| return _op.nn.depth_to_space(inputs[0], block_size, mode=mode) |
| |
| |
| class SpaceToDepth(OnnxOpConverter): |
| """Operator converter for SpaceToDepth.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| |
| block_size = int(attr["blocksize"]) |
| return _op.nn.space_to_depth(inputs[0], block_size) |
| |
| |
| class Concat(OnnxOpConverter): |
| """Operator converter for Concat.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, args, params): |
| return AttrCvt(op_name="concatenate")((inputs,), args) |
| |
| |
| class Scale(OnnxOpConverter): |
| """Operator converter for Scale.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| scale = float(attr.get("scale", 1.0)) |
| return inputs[0] * _expr.const(scale) |
| |
| |
| class Selu(OnnxOpConverter): |
| """Operator converter for Selu.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = float(attr.get("alpha", 1.67326319217681884765625)) |
| gamma = float(attr.get("gamma", 1.05070102214813232421875)) |
| return _expr.const(gamma) * ( |
| _expr.const(-alpha) * _op.nn.relu(_expr.const(1.0) - _op.exp(inputs[0])) |
| + _op.nn.relu(inputs[0]) |
| ) |
| |
| |
| class ScaledTanh(OnnxOpConverter): |
| """Operator converter for ScaledTanh.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = float(attr.get("alpha", 1.0)) |
| beta = float(attr.get("beta", 1.0)) |
| return _op.tanh(_expr.const(beta) * inputs[0]) * _expr.const(alpha) |
| |
| |
| class Shrink(OnnxOpConverter): |
| """Operator converter for Shrink.""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| x = inputs[0] |
| dtype = infer_type(x).checked_type.dtype |
| lambd = _op.const(attr.get("lambd", 0.5), dtype=dtype) |
| bias = _op.const(attr.get("bias", 0.0), dtype=dtype) |
| |
| zeros = _op.zeros_like(x) |
| return _op.where(x < -lambd, x + bias, zeros) + _op.where(x > lambd, x - bias, zeros) |
| |
| |
| class Softsign(OnnxOpConverter): |
| """Operator converter for Softsign.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return inputs[0] / (_expr.const(1.0) + Absolute.get_converter(1)(inputs, attr, params)) |
| |
| |
| class Sub(Elemwise): |
| """Operator converter for Subtract.""" |
| |
| name = "subtract" |
| |
| |
| class Sum(OnnxOpConverter): |
| """Operator converter for Sum.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # Onnx Sum Operator |
| for in_index in range(len(inputs) - 1): |
| inputs[in_index + 1] = _op.add(inputs[in_index], inputs[in_index + 1]) |
| |
| return inputs[len(inputs) - 1] |
| |
| |
| class Affine(OnnxOpConverter): |
| """Operator converter for Affine transformation.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = _expr.const(attr.get("alpha", 1.0)) |
| beta = _expr.const(attr.get("beta", 0.0)) |
| return (alpha * inputs[0]) + beta |
| |
| |
| class ThresholdedRelu(OnnxOpConverter): |
| """Operator converter for ThresholdedRelu.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = float(attr.get("alpha", 1.0)) |
| alpha_tensor = _op.full_like(inputs[0], fill_value=_expr.const(alpha)) |
| mask = _op.greater(inputs[0], alpha_tensor).astype("float32") |
| return inputs[0] * mask |
| |
| |
| def _broadcast_constraint(): |
| def _broadcast_check(attrs): |
| if attrs.get("axis", None): |
| return False |
| return True |
| |
| return _broadcast_check, "Specifying broadcast axis not allowed." |
| |
| |
| def _fully_connected(opset): |
| def _impl(inputs, attr, params): |
| # get number of channels |
| channels = infer_channels(inputs[1], params) |
| attr["units"] = channels |
| return AttrCvt("dense", ignores=["axis", "axis_w"])(inputs, attr) |
| |
| return _impl |
| |
| |
| class Upsample(OnnxOpConverter): |
| """Operator converter for Upsample (nearest mode).""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| scales = attr.get("scales") |
| |
| input_shape = infer_shape(inputs[0]) |
| dims = len(input_shape) |
| |
| if not scales: |
| # Here we are going to higher OPSET version. |
| assert len(inputs) == 2, "Upsample op takes 2 inputs, {} given".format(len(inputs)) |
| |
| if get_name(inputs[1]) in params: |
| scales = params[inputs[1].name_hint].numpy() |
| else: |
| scales = inputs[1] |
| if isinstance(scales, _expr.Constant): |
| scales = list(scales.data.numpy()) |
| if not isinstance(scales, _expr.Expr): |
| assert scales[0] == 1.0 and scales[1] == 1.0 |
| |
| mode = attr.get("mode") |
| if mode == b"nearest": |
| method = "nearest_neighbor" |
| elif mode == b"linear": |
| method = "trilinear" if dims == 5 else "bilinear" |
| else: |
| raise tvm.error.OpAttributeInvalid( |
| 'Value {} in attribute "mode" of operator Upsample is not valid.'.format(mode) |
| ) |
| |
| # in 3d case, we use the purely static op |
| if dims == 5: |
| if isinstance(scales, _expr.Expr): |
| scale_h = _op.take(scales, _op.const(3)) |
| scale_w = _op.take(scales, _op.const(4)) |
| scale_d = _op.take(scales, _op.const(1)) |
| else: |
| assert len(scales) == 5 |
| scale_h = scales[-2] |
| scale_w = scales[-1] |
| scale_d = scales[-3] |
| |
| layout = "NCDHW" |
| out = _op.nn.upsampling3d( |
| inputs[0], |
| scale_d, |
| scale_h, |
| scale_w, |
| layout=layout, |
| method=method, |
| coordinate_transformation_mode="asymmetric", |
| ) |
| # in 2d case, use dynamic op |
| else: |
| if isinstance(scales, _expr.Expr): |
| scale_h = _op.take(scales, _op.const(3)) |
| scale_w = _op.take(scales, _op.const(4)) |
| else: |
| assert len(scales) == 4 |
| scale_h = scales[-2] |
| scale_w = scales[-1] |
| layout = "NCHW" |
| |
| out = _op.nn.upsampling( |
| inputs[0], |
| scale_h, |
| scale_w, |
| layout=layout, |
| method=method, |
| align_corners=False, |
| ) |
| return out |
| |
| |
| class Shape(OnnxOpConverter): |
| """Operator converter for Shape.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return shape_of(inputs[0], "int64") |
| |
| @classmethod |
| def _impl_v15(cls, inputs, attr, params): |
| start = attr.get("start") |
| end = attr.get("end") |
| return shape_of(inputs[0], dtype="int64", start=start, end=end) |
| |
| |
| class CumSum(OnnxOpConverter): |
| """Operator converter for CumSum.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| dim = inputs[1] |
| |
| if dim is not None: |
| dim = int(infer_value(dim, params).numpy()) |
| |
| exclusive = attr.get("exclusive", 0) |
| reverse = attr.get("reverse", 0) |
| |
| if reverse != 0: |
| out = _op.reverse(data, axis=dim) |
| out = _op.cumsum(out, axis=dim, exclusive=exclusive) |
| return _op.reverse(out, axis=dim) |
| |
| return _op.cumsum(data, axis=dim, exclusive=exclusive) |
| |
| |
| class Cast(OnnxOpConverter): |
| """Operator converter for Cast.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return AttrCvt(op_name="cast", transforms={"to": "dtype"})(inputs, attr) |
| |
| @classmethod |
| def _impl_v5(cls, inputs, attr, params): |
| try: |
| from onnx import TensorProto |
| except ImportError as e: |
| raise ImportError("Unable to import TensorProto from onnx {}".format(e)) |
| |
| # If onnx mapping is used, bfloat16 gets converted to float16 |
| # which is not the desired behavior |
| if attr["to"] == int(TensorProto.BFLOAT16): |
| attr["to"] = "bfloat16" |
| else: |
| try: |
| from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE |
| |
| attr["to"] = str(TENSOR_TYPE_TO_NP_TYPE[attr["to"]]) |
| except ImportError as e: |
| raise ImportError("Unable to import onnx.mapping which is required {}".format(e)) |
| return AttrCvt(op_name="cast", transforms={"to": "dtype"})(inputs, attr) |
| |
| |
| class Unsqueeze(OnnxOpConverter): |
| """Operator converter for Unsqueeze.""" |
| |
| @classmethod |
| def run_calculation(cls, tensor, axes): |
| axes = sorted(axes) |
| for axis in axes: |
| tensor = _op.expand_dims(tensor, axis=axis, num_newaxis=1) |
| return tensor |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return cls.run_calculation(inputs[0], attr["axes"]) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| if isinstance(inputs[1], _expr.Constant): |
| constant_axes = list(inputs[1].data.numpy()) |
| constant_axes = list(map(int, constant_axes)) |
| return cls.run_calculation(inputs[0], constant_axes) |
| |
| rank_input = len(infer_type(inputs[0]).checked_type.shape) |
| num_new_axis = int(infer_type(inputs[1]).checked_type.shape[0]) |
| axes = relay.sort(inputs[1]) |
| axes = relay.split(axes, num_new_axis).astuple() |
| result = inputs[0] |
| |
| # TODO (AndrewZhaoLuo): investigate performance issues with consecutive |
| # dynamic expand_dims on non-llvm targets. |
| for i in range(num_new_axis): |
| axis = relay.TupleGetItem(axes, i) |
| # Unpack scalar |
| axis = relay.reshape(axis, []) |
| axis = relay.where( |
| axis >= relay.const(0, "int64"), axis, axis + relay.const(rank_input, "int64") |
| ) |
| result = _op.expand_dims(result, axis) |
| return result |
| |
| |
| class Squeeze(OnnxOpConverter): |
| """Operator converter for Squeeze.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axes", None) |
| return _op.squeeze(inputs[0], axis) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| ishape = infer_shape(inputs[0]) |
| axis = inputs[1] |
| |
| if axis is None: |
| # If axes is not provided, all the single dimensions will be removed from the shape. |
| if not ishape: # scalar |
| return inputs[0] |
| |
| axis = [i for i in range(len(ishape)) if ishape[i] == 1] |
| axis = _op.const(axis) |
| |
| dtype = infer_type(axis).checked_type.dtype |
| |
| if isinstance(axis, _expr.Constant): |
| constant_axes = list(axis.data.numpy()) |
| constant_axes = list(map(int, constant_axes)) |
| return _op.squeeze(inputs[0], constant_axes) |
| |
| rank = _op.shape_of(_op.shape_of(inputs[0], dtype), dtype) |
| axis = _op.where(axis < _op.const(0, dtype), axis + rank, axis) |
| return _op.squeeze(inputs[0], fold_constant(axis)) |
| |
| |
| class Split(OnnxOpConverter): |
| """Operator converter for Split.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| splits = attr.get("split", None) |
| if splits is not None and len(splits) > 1: |
| indices = [] |
| index = 0 |
| for i in splits[:-1]: |
| index += i |
| indices.append(index) |
| # When splits isnt specified divide evenly over axis. |
| else: |
| indices = attr["tvm_custom"]["num_outputs"] |
| output = _op.split(inputs[0], indices, attr.get("axis", 0)) |
| # If the output of split is a single value, unpack if from the TupleWrapper |
| if len(output) == 1: |
| output = output[0] |
| return output |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| splits = inputs[1] |
| splits_rank = None |
| if splits is not None: |
| splits_rank = len(infer_shape(splits)) |
| if splits is not None and splits_rank > 0: |
| if isinstance(splits, _expr.Constant): |
| splits = splits.data.asnumpy() |
| indices = [] |
| index = 0 |
| for i in splits[:-1]: |
| index += i |
| indices.append(index) |
| else: |
| raise ValueError("Dynamic Split not yet supported") |
| # When splits isnt specified divide evenly over axis. |
| else: |
| indices = attr["tvm_custom"]["num_outputs"] |
| output = _op.split(inputs[0], indices, attr.get("axis", 0)) |
| # If the output of split is a single value, unpack if from the TupleWrapper |
| if len(output) == 1: |
| output = output[0] |
| return output |
| |
| |
| class Slice(OnnxOpConverter): |
| """Operator converter for Slice.""" |
| |
| @classmethod |
| def _common(cls, starts, ends, axes): |
| N = max(axes) + 1 |
| new_axes = list(range(N)) |
| new_starts = [0] * N |
| new_ends = [np.iinfo(np.int32).max] * N |
| for i, axis in enumerate(axes): |
| new_starts[axis] = starts[i] |
| new_ends[axis] = ends[i] |
| return new_starts, new_ends, new_axes |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if isinstance(attr["starts"], int): |
| attr["starts"] = (attr["starts"],) |
| attr["ends"] = (attr["ends"],) |
| |
| try: |
| # Update the starts and ends according to axes if required. |
| if isinstance(attr["axes"], int): |
| attr["axes"] = (attr["axes"],) |
| new_starts, new_ends, new_axes = cls._common(attr["starts"], attr["ends"], attr["axes"]) |
| attr["axes"] = new_axes |
| attr["starts"] = new_starts |
| attr["ends"] = new_ends |
| except KeyError: |
| pass |
| begin = list(attr["starts"]) |
| end = list(attr["ends"]) |
| |
| return _op.strided_slice(inputs[0], begin=begin, end=end) |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| starts = inputs[1] |
| ends = inputs[2] |
| axes = inputs[3] |
| steps = inputs[4] |
| |
| ishape = infer_shape(inputs[0]) |
| data_rank = len(ishape) |
| |
| if axes is not None: |
| # Normalize for negative axes |
| axes_dtype = infer_type(axes).checked_type.dtype |
| axes = fold_constant( |
| _op.where( |
| axes < _op.const(0, axes_dtype), axes + _op.const(data_rank, axes_dtype), axes |
| ) |
| ) |
| |
| def has_static_axes(): |
| return ( |
| isinstance(axes, _expr.Constant) |
| and isinstance(starts, _expr.Constant) |
| and isinstance(ends, _expr.Constant) |
| and (steps is None or isinstance(steps, _expr.Constant)) |
| ) |
| |
| if axes is not None and has_static_axes(): |
| axes_np = axes.data.numpy().astype("int64") |
| begin_np = starts.data.numpy().astype("int64") |
| end_np = ends.data.numpy().astype("int64") |
| if steps is None: |
| strides_np = np.ones_like(begin_np).astype("int64") |
| else: |
| strides_np = steps.data.numpy().astype("int64") |
| if all([isinstance(ishape[i], int) for i in axes_np]): |
| return _op.strided_slice( |
| inputs[0], list(begin_np), list(end_np), list(strides_np), axes=list(axes_np) |
| ) |
| |
| # Update the starts and ends according to axes if required. |
| if axes is not None: |
| data_shape = shape_of(inputs[0], dtype=infer_type(ends).checked_type.dtype) |
| starts = _op.scatter( |
| _op.const([0] * data_rank, dtype=infer_type(starts).checked_type.dtype), |
| axes, |
| starts, |
| axis=0, |
| ) |
| ends = _op.scatter(data_shape, axes, ends, axis=0) |
| if steps is not None: |
| steps = _op.scatter( |
| _op.const([1] * data_rank, dtype=infer_type(steps).checked_type.dtype), |
| axes, |
| steps, |
| axis=0, |
| ) |
| |
| if steps is None: |
| steps = _op.const([1] * data_rank, dtype=infer_type(starts).checked_type.dtype) |
| |
| return _op.strided_slice( |
| inputs[0], fold_constant(starts), fold_constant(ends), fold_constant(steps) |
| ) |
| |
| |
| def normalize_gather_indices(data, indices, axis): |
| """Make sure gather indicies aren't negative""" |
| ind_dtype = infer_type(indices).checked_type.dtype |
| # Normalize the indices to a positive range |
| s = _op.take(_op.shape_of(data, dtype=ind_dtype), _op.const(axis, dtype="int64")) |
| cond = fold_constant(indices < _op.const(0, ind_dtype)) |
| if isinstance(cond, _expr.Constant): |
| val = cond.data.numpy() |
| if val.size == 1: |
| cond = val.item() |
| if cond: |
| indices = indices + s |
| return indices |
| indices = _op.where(cond, indices + s, indices) |
| return indices |
| |
| |
| class Gather(OnnxOpConverter): |
| """Operator converter for Gather.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axis", 0) |
| data = inputs[0] |
| indices = inputs[1] |
| indices = normalize_gather_indices(data, indices, axis) |
| return _op.take(data, indices, axis) |
| |
| |
| class GatherElements(OnnxOpConverter): |
| """Operator converter for GatherElements.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| indices = inputs[1] |
| axis = attr.get("axis", 0) |
| indices = normalize_gather_indices(data, indices, axis) |
| return _op.gather(data, axis, indices) |
| |
| |
| class GatherND(OnnxOpConverter): |
| """Operator converter for GatherND.""" |
| |
| @classmethod |
| def _impl_common(cls, data, indices, batch_dims=0): |
| indices_dims = len(infer_shape(indices)) |
| indices_shape = infer_shape(indices) |
| indices = _op.transpose(indices, axes=[-1] + list(range(indices_dims - 1))) |
| index_rank = indices_shape[-1] |
| return _op.gather_nd( |
| data, |
| indices, |
| batch_dims=batch_dims, |
| index_rank=index_rank, |
| ) |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return cls._impl_common(inputs[0], inputs[1]) |
| |
| @classmethod |
| def _impl_v12(cls, inputs, attr, params): |
| batch_dims = attr.get("batch_dims", 0) |
| return cls._impl_common(inputs[0], inputs[1], batch_dims) |
| |
| |
| class Compress(OnnxOpConverter): |
| """Operator converter for compress""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| input_tensor, condition_tensor = inputs |
| |
| axis = attr.get("axis", None) |
| |
| # Change one hot tensor to indices e.g. [0, 1, 1, 0, 1] -> [1, 2, 4] |
| condition_tensor = _op.reshape(_op.argwhere(condition_tensor), (-1,)) |
| |
| if axis is not None: |
| return _op.take(input_tensor, condition_tensor, axis=axis) |
| |
| # if axis is None, flatten input tensor before selection |
| input_tensor = _op.reshape(input_tensor, (-1,)) |
| return _op.take(input_tensor, condition_tensor, axis=0) |
| |
| |
| class Scatter(OnnxOpConverter): |
| """Operator converter for Scatter.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axis", 0) |
| return _op.scatter(inputs[0], inputs[1], inputs[2], axis) |
| |
| |
| class ScatterND(OnnxOpConverter): |
| """Operator converter for ScatterND.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| indices_dim = len(infer_shape(inputs[1])) |
| axes = list(range(indices_dim)) |
| return _op.scatter_nd( |
| inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2], "update" |
| ) |
| |
| |
| class EyeLike(OnnxOpConverter): |
| """Operator converter for EyeLike.""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| dtype = attr.get("dtype", None) |
| if dtype is None: |
| in_checked_type = infer_type(inputs[0]).checked_type |
| in_dtype = in_checked_type.dtype |
| dtype = in_dtype |
| else: |
| dtype = get_type(dtype) |
| |
| in_shape = _op.shape_of(inputs[0]) |
| zeros = _op.zeros(in_shape, dtype) |
| |
| dim = _op.take(in_shape, _op.const(0)) |
| |
| indices = _op.arange(_op.const(0), dim, dtype="int32") |
| ones = _op.full(_op.const(1), _op.reshape(dim, (1,)), dtype=dtype) |
| k = _op.const(attr.get("k", 0), dtype="int32") |
| return _op.scatter_nd(zeros, _op.stack([indices, indices + k], axis=0), ones, "update") |
| |
| |
| class LRN(OnnxOpConverter): |
| """Operator converter for Local Response Normalization.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| """LRN support only NCHW format |
| https://github.com/onnx/onnx/blob/main/docs/Operators.md#LRN |
| """ |
| axis = 1 |
| alpha = attr.get("alpha", 0.0001) |
| beta = attr.get("beta", 0.75) |
| bias = attr.get("bias", 1.0) |
| nsize = attr.get("size") |
| attr = {"size": nsize, "axis": axis, "alpha": alpha, "beta": beta, "bias": bias} |
| return AttrCvt("lrn")(inputs, attr) |
| |
| |
| class Maximum(OnnxOpConverter): |
| """Operator converter for Maximum.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) == 1: |
| return inputs[0] |
| _max = inputs[0] |
| for i in range(1, len(inputs)): |
| _max = AttrCvt("maximum")([_max, inputs[i]], {}) |
| return _max |
| |
| |
| class Minimum(OnnxOpConverter): |
| """Operator converter for Minimum.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) == 1: |
| return inputs[0] |
| _min = inputs[0] |
| for i in range(1, len(inputs)): |
| _min = AttrCvt("minimum")([_min, inputs[i]], {}) |
| return _min |
| |
| |
| class Mean(OnnxOpConverter): |
| """Operator converter for Mean.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) == 1: |
| return inputs[0] |
| # avoid overflow |
| concat = _op.concatenate([_op.expand_dims(x, axis=0) for x in inputs], axis=0) |
| return _op.mean(concat, axis=0, keepdims=False) |
| |
| |
| class MeanVarianceNormalization(OnnxOpConverter): |
| """Operator converter for MeanVarianceNormalization.""" |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| axis = attr.get("axes", (0, 2, 3)) |
| data_mean = _op.mean(inputs[0], axis=axis, keepdims=True) |
| data_mean_squared = _op.power(data_mean, _expr.const(2, "float32")) |
| data_squared = _op.power(inputs[0], _expr.const(2, "float32")) |
| data_squared_mean = _op.mean(data_squared, axis=axis, keepdims=True) |
| return (inputs[0] - data_mean) / _op.sqrt(data_squared_mean - data_mean_squared) |
| |
| |
| class HardSigmoid(OnnxOpConverter): |
| """Operator converter for HardSigmoid.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = attr.get("alpha", 0.2) |
| beta = attr.get("beta", 0.5) |
| transformX = (inputs[0] * _expr.const(alpha)) + _expr.const(beta) |
| attr = {"a_min": 0, "a_max": 1} |
| return AttrCvt("clip")([transformX], attr) |
| |
| |
| class HardSwish(OnnxOpConverter): |
| """Operator converter for HardSwish.""" |
| |
| @classmethod |
| def _impl_v14(cls, inputs, attr, params): |
| alpha = attr.get("alpha", 1 / 6) |
| beta = attr.get("beta", 0.5) |
| transformX = inputs[0] * _expr.const(alpha) + _expr.const(beta) |
| attr = {"a_min": 0, "a_max": 1} |
| return inputs[0] * AttrCvt("clip")([transformX], attr) |
| |
| |
| class Reduce(OnnxOpConverter): |
| """Operator converter for reduce ops.""" |
| |
| name = "" |
| |
| @classmethod |
| def run_calculation(cls, inputs, axis, keepdims): |
| attr = {"axis": axis, "keepdims": keepdims} |
| return AttrCvt(cls.name)(inputs, attr) |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if "axes" in attr: |
| axis = attr.get("axes", 0) |
| else: |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| |
| return cls.run_calculation(inputs, axis, attr.get("keepdims", True)) |
| |
| @classmethod |
| def _impl_v12(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if len(inputs) == 2: |
| if isinstance(inputs[1], _expr.Constant): |
| # Get axis and unpack scalar |
| constant_axis = int(inputs[1].data.numpy()[0]) |
| return cls.run_calculation([inputs[0]], constant_axis, attr.get("keepdims", True)) |
| |
| raise ValueError("Dynamic Reduce is not supported yet!") |
| |
| return cls._impl_v1(inputs, attr, params) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| noop_with_empty_axes = attr.get("noop_with_empty_axes", 0) |
| num_axis = int(infer_type(inputs[1]).checked_type.shape[0]) if inputs[1] is not None else 0 |
| |
| if noop_with_empty_axes and num_axis == 0: |
| return inputs[0] |
| |
| if len(inputs) == 2: |
| if isinstance(inputs[1], _expr.Constant): |
| # Get axis and unpack scalar |
| constant_axis = int(inputs[1].data.numpy()[0]) |
| return cls.run_calculation([inputs[0]], constant_axis, attr.get("keepdims", True)) |
| |
| if num_axis > 0: |
| raise ValueError("Dynamic Reduce is not supported yet!") |
| |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| return cls.run_calculation([inputs[0]], axis, attr.get("keepdims", True)) |
| |
| return cls._impl_v1(inputs, attr, params) |
| |
| |
| class ReduceMax(Reduce): |
| """Operator converter for ReduceMax.""" |
| |
| name = "max" |
| |
| |
| class ReduceMin(Reduce): |
| """Operator converter for ReduceMin.""" |
| |
| name = "min" |
| |
| |
| class ReduceSum(Reduce): |
| """Operator converter for ReduceSum.""" |
| |
| name = "sum" |
| |
| |
| class ReduceMean(Reduce): |
| """Operator converter for ReduceMean.""" |
| |
| name = "mean" |
| |
| |
| class ReduceProd(Reduce): |
| """Operator converter for ReduceProd.""" |
| |
| name = "prod" |
| |
| |
| class ReduceLogSumExp(Reduce): |
| """Operator converter for ReduceLogSumExp.""" |
| |
| name = "logsumexp" |
| |
| |
| class ReduceSumSquare(OnnxOpConverter): |
| """Operator converter for ReduceSumSquare.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if "axes" in attr: |
| axis = attr.get("axes", 0) |
| else: |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| attr = {"axis": axis, "keepdims": attr.get("keepdims", True)} |
| inputs[0] = inputs[0] * inputs[0] |
| |
| return AttrCvt("sum")(inputs, attr) |
| |
| |
| class ReduceL1(OnnxOpConverter): |
| """Operator converter for ReduceL1.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if "axes" in attr: |
| axis = attr.get("axes", 0) |
| else: |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| attr = {"axis": axis, "keepdims": attr.get("keepdims", True)} |
| inputs[0] = _op.abs(inputs[0]) |
| |
| return AttrCvt("sum")(inputs, attr) |
| |
| |
| class ReduceL2(OnnxOpConverter): |
| """Operator converter for ReduceL2.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if "axes" in attr: |
| axis = attr.get("axes", 0) |
| else: |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| attr = {"axis": axis, "keepdims": attr.get("keepdims", True)} |
| inputs[0] = inputs[0] * inputs[0] |
| out = AttrCvt("sum")(inputs, attr) |
| |
| return _op.sqrt(out) |
| |
| |
| class ReduceLogSum(OnnxOpConverter): |
| """Operator converter for ReduceLogSum.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if not infer_shape(inputs[0]): # promote scalar to 1-D tensor |
| inputs[0] = _op.expand_dims(inputs[0], axis=0) |
| |
| if "axes" in attr: |
| axis = attr.get("axes", 0) |
| else: |
| axis_len = len(infer_shape(inputs[0])) |
| axis = list(range(axis_len)) |
| attr = {"axis": axis, "keepdims": attr.get("keepdims", True)} |
| out = AttrCvt("sum")(inputs, attr) |
| |
| return _op.log(out) |
| |
| |
| class ArgMax(OnnxOpConverter): |
| """Operator converter for ArgMax.""" |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| axis = attr.get("axis", 0) |
| keepdims = attr.get("keepdims", True) |
| select_last_index = attr.get("select_last_index", False) |
| attr = {"axis": axis, "keepdims": keepdims, "select_last_index": select_last_index} |
| return _op.cast(AttrCvt("argmax")(inputs, attr), "int64") |
| |
| |
| class ArgMin(OnnxOpConverter): |
| """Operator converter for ArgMin.""" |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| axis = attr.get("axis", 0) |
| keepdims = attr.get("keepdims", True) |
| select_last_index = attr.get("select_last_index", False) |
| attr = {"axis": axis, "keepdims": keepdims, "select_last_index": select_last_index} |
| return _op.cast(AttrCvt("argmin")(inputs, attr), "int64") |
| |
| |
| class Softmax(OnnxOpConverter): |
| """Operator converter for Softmax.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axis", 1) |
| in_shape = infer_shape(inputs[0]) |
| ndim = len(in_shape) |
| if axis < 0: |
| axis += ndim |
| if axis == 0: |
| reshape_shape = [-1] |
| elif axis == ndim - 1: |
| return _op.nn.softmax(inputs[0], axis=axis) |
| else: |
| axis_val = [in_shape[i] for i in range(axis)] |
| reshape_shape = [np.prod(axis_val)] + [-1] |
| data_reshape = _op.reshape(inputs[0], newshape=reshape_shape) |
| out = _op.nn.softmax(data_reshape, axis=-1) |
| out = _op.reshape(out, newshape=in_shape) |
| return out |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, _): |
| axis = attr.get("axis", -1) |
| ndim = len(infer_shape(inputs[0])) |
| if axis < 0: |
| axis += ndim |
| return _op.nn.softmax(inputs[0], axis=axis) |
| |
| |
| class LogSoftmax(OnnxOpConverter): |
| """Operator converter for Softmax.""" |
| |
| @classmethod |
| def run_calculation(cls, inputs, attr, params, opset): |
| """Run the calculation for Log Softmax calculation.""" |
| res = Softmax.get_converter(opset)(inputs, attr, params) |
| return _op.log(res) |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return cls.run_calculation(inputs, attr, params, opset=1) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| return cls.run_calculation(inputs, attr, params, opset=13) |
| |
| |
| class Hardmax(OnnxOpConverter): |
| """Operator converter for Hardmax.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| axis = attr.get("axis", 1) |
| ndim = len(infer_shape(inputs[0])) |
| if axis < 0: |
| axis += ndim |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| |
| if axis == 0: |
| pre = _op.const([1], "int64") |
| else: |
| pre = _op.prod( |
| _op.strided_slice(shape_of(inputs[0]), [0], [axis], [1]), axis=0, keepdims=True |
| ) |
| post = _op.prod( |
| _op.strided_slice(shape_of(inputs[0]), [axis], [2147483647], [1]), axis=0, keepdims=True |
| ) |
| newshape = _op.concatenate([pre, post], axis=0) |
| x = _op.reshape(inputs[0], fold_constant(newshape)) |
| argmax = _op.argmax(x, axis=1) |
| onehot = _op.one_hot( |
| argmax, |
| _op.const(1.0, dtype), |
| _op.const(0.0, dtype), |
| fold_constant(_op.take(shape_of(x), _op.const([1], "int64"))), |
| 1, |
| dtype, |
| ) |
| return _op.reshape(onehot, shape_of(inputs[0])) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params) -> relay.Expr: |
| inferred_type = infer_type(inputs[0]) |
| dtype = inferred_type.checked_type.dtype |
| ndim = len(inferred_type.checked_type.shape) |
| axis = attr.get("axis", -1) % ndim |
| |
| argmax = _op.argmax(inputs[0], axis=axis) |
| return _op.one_hot( |
| argmax, |
| _op.const(1.0, dtype), |
| _op.const(0.0, dtype), |
| fold_constant(_op.take(shape_of(inputs[0]), _op.const([axis], "int64"))), |
| axis, |
| dtype, |
| ) |
| |
| |
| class OneHot(OnnxOpConverter): |
| """Operator converter for OneHot.""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| # Extract relay one_hot inputs. |
| indices, depth, values = inputs |
| ndim = len(infer_shape(indices)) |
| # Split onnx on off values into two separate expressions. |
| off_value, on_value = _op.take(values, _op.const(0)), _op.take(values, _op.const(1)) |
| # Extract the datatype of the output from on_value. |
| dtype = infer_type(on_value).checked_type.dtype |
| ind_dtype = infer_type(indices).checked_type.dtype |
| # Normalize the indices to a positive range |
| indices = _op.where( |
| indices < _op.const(0, ind_dtype), indices + _op.cast(depth, ind_dtype), indices |
| ) |
| # set default value when axis is not set in the model |
| if "axis" not in attr: |
| attr["axis"] = -1 |
| axis = attr["axis"] |
| if axis < 0: |
| axis += ndim + 1 |
| |
| return _op.one_hot(indices, on_value, off_value, depth, axis, dtype=dtype) |
| |
| |
| class ConstantOfShape(OnnxOpConverter): |
| """Operator converter for ConstantOfShape.""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| if "value" in attr: |
| np_value = get_numpy(attr.pop("value"))[0] |
| value = _expr.const(np_value) |
| dtype = np_value.dtype.name |
| else: |
| value = _expr.const(0) |
| dtype = "float32" |
| output = _op.full(value, inputs[0], dtype=dtype) |
| return output |
| |
| |
| class Constant(OnnxOpConverter): |
| """Operator converter for ConstantOfShape.""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| if "value" not in attr: |
| raise tvm.errors.OpAttributeRequired("no value in Constant") |
| value = attr.pop("value") |
| # Constants may rarely have string types. These are likely exported |
| # from other frameworks and not actually used in TVM. We'll just use |
| # a zero valued constant for compatibility. |
| if isinstance(value, bytes): |
| np_value = np.asarray([0]).astype("int64") |
| else: |
| np_value = get_numpy(value) |
| dtype = np_value.dtype.name |
| value = _expr.const(np_value, dtype) |
| return value |
| |
| |
| class Sign(OnnxOpConverter): |
| """Operator converter for Sign.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return _op.sign(inputs[0]) |
| |
| |
| class Equal(Elemwise): |
| """Operator converter for Equal.""" |
| |
| name = "equal" |
| |
| |
| class Not(Elemwise): |
| """Operator converter for Not.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return _op.logical_not(inputs[0]) |
| |
| |
| class And(Elemwise): |
| """Operator converter for And.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return _op.logical_and(inputs[0], inputs[1]) |
| |
| |
| class Tile(Elemwise): |
| """Operator converter for Tile""" |
| |
| @classmethod |
| def _impl_v6(cls, inputs, attr, params): |
| return _op.tile(inputs[0], inputs[1]) |
| |
| |
| class Erf(OnnxOpConverter): |
| """Operator converter for Erf""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| return _op.erf(inputs[0]) |
| |
| |
| class Where(OnnxOpConverter): |
| """Operator converter for Where""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| return _op.where(*inputs) |
| |
| |
| class Or(Elemwise): |
| """Operator converter for Or.""" |
| |
| @classmethod |
| def _impl_v7(cls, inputs, attr, params): |
| return _op.logical_or(inputs[0], inputs[1]) |
| |
| |
| class Expand(OnnxOpConverter): |
| """Operator converter for Expand.""" |
| |
| @classmethod |
| def _impl_v8(cls, inputs, attr, params): |
| dtype = infer_type(inputs[1]).checked_type.dtype |
| in_shape = shape_of(inputs[0], dtype=dtype) |
| shape = inputs[1] |
| |
| # Currently 'op.broadcast_to' expect the rank of the given 'shape' |
| # (the 2nd input) is always higher than that of the given 'input' (the 1st input) |
| # However, ONNX Expand supports multi-directional broadcasting, which allows |
| # above pattern and also some extent of 'shape' can be smaller than the corresponding |
| # extent of 'input'. In this case, the extent of 'shape' must be 1. |
| # https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md |
| # In above cases, we cannot directorly apply 'op.broadcast_to' instead of 'expand' |
| # so, here we solved this problem by expanding the given 'shape' itself. |
| def expand_shape(in_shape, shape): |
| """A function expands the shape when the rank is lower than that of the given |
| intput. Also it replaces the extent of the shape with the corresponding extent |
| of the intput when it is 1. |
| """ |
| in_dims = infer_shape(in_shape)[0] |
| new_dims = infer_shape(shape)[0] |
| |
| if in_dims < new_dims: |
| in_shape = _op.concatenate( |
| [ |
| _expr.const( |
| [ |
| 1, |
| ] |
| * (new_dims - in_dims), |
| dtype=dtype, |
| ), |
| in_shape, |
| ], |
| axis=0, |
| ) |
| elif new_dims < in_dims: |
| shape = _op.concatenate( |
| [ |
| _expr.const( |
| [ |
| 1, |
| ] |
| * (in_dims - new_dims), |
| dtype=dtype, |
| ), |
| shape, |
| ], |
| axis=0, |
| ) |
| new_shape = _op.maximum(in_shape, shape) |
| return new_shape |
| |
| shape = fold_constant(expand_shape(in_shape, shape)) |
| return _op.broadcast_to(inputs[0], shape=shape) |
| |
| |
| class RNN(OnnxOpConverter): |
| """Operator converter for RNNs such as RNN, LSTM and GRU.""" |
| |
| @classmethod |
| def _activation_helper(cls, activation, alpha, beta): |
| convert_map = _get_convert_map(1) |
| attrs = {} |
| if alpha is not None: |
| attrs["alpha"] = alpha |
| if beta is not None: |
| attrs["beta"] = beta |
| return lambda x: convert_map[activation.decode("utf-8")]([x], attrs, {}) |
| |
| @classmethod |
| def _activation_needs_alpha(cls, activation): |
| needs_alpha = [ |
| "Affine", |
| "LeakyRelu", |
| "ThresholdedRelu", |
| "ScaledTanh", |
| "HardSigmoid", |
| "Elu", |
| ] |
| return activation.decode("utf-8") in needs_alpha |
| |
| @classmethod |
| def _activation_needs_beta(cls, activation): |
| needs_beta = [ |
| "Affine", |
| "ScaledTanh", |
| "HardSigmoid", |
| ] |
| return activation.decode("utf-8") in needs_beta |
| |
| @classmethod |
| def bidir_rnn_cell( |
| cls, |
| input_seqs, |
| weight_dicts, |
| acts, |
| ): |
| """ |
| Bidirectional RNN cell |
| """ |
| seq_len = len(input_seqs) |
| forward_outputs, fw_H_t = rnn_cell( |
| input_seqs, |
| **weight_dicts[0], |
| act=acts[0], |
| ) |
| |
| reverse_outputs, rev_H_t = rnn_cell( |
| input_seqs, |
| **weight_dicts[1], |
| act=acts[1], |
| backwards=True, |
| ) |
| |
| final_outputs = [] |
| for i in range(seq_len): |
| final_outputs.append( |
| _op.stack([forward_outputs[i], reverse_outputs[seq_len - 1 - i]], axis=0) |
| ) |
| |
| return ( |
| _op.stack(final_outputs, axis=0), |
| _op.stack([fw_H_t, rev_H_t], axis=0), |
| ) |
| |
| @classmethod |
| def _default_activations(cls, num_directions): |
| return [_op.tanh] * num_directions |
| |
| @classmethod |
| def _get_activations(cls, attr, multiplier, num_directions, rnn_type): |
| """ |
| Activation functions |
| """ |
| if "activations" in attr: |
| activations = attr["activations"] |
| if len(activations) != multiplier * num_directions: |
| raise NotImplementedError( |
| "{} assumes {} * num_directions activation functions are provided".format( |
| rnn_type, multiplier |
| ) |
| ) |
| alpha_loc = 0 |
| alphas = attr.get("activation_alpha", []) |
| if isinstance(alphas, float): |
| alphas = [alphas] |
| beta_loc = 0 |
| betas = attr.get("activation_beta", []) |
| if isinstance(betas, float): |
| betas = [betas] |
| acts = [] |
| for i in range(multiplier * num_directions): |
| alpha = None |
| beta = None |
| activation = activations[i] |
| if cls._activation_needs_alpha(activation) and len(alphas) > alpha_loc: |
| alpha = alphas[alpha_loc] |
| alpha_loc += 1 |
| if cls._activation_needs_beta(activation) and len(betas) > beta_loc: |
| beta = betas[beta_loc] |
| beta_loc += 1 |
| acts.append(cls._activation_helper(activation, alpha, beta)) |
| else: |
| acts = cls._default_activations(num_directions) |
| return acts |
| |
| @classmethod |
| def _inputs_helper(cls, inputs, layout): |
| """ |
| Process inputs |
| """ |
| # Unpack inputs, note that if optional and not provided then value will be None. |
| X = inputs[0] |
| Wp = inputs[1] |
| Rp = inputs[2] |
| Bp = inputs[3] |
| # Sequence length currently unused as it can be inferred from shapes. |
| # sequence_lens = inputs['sequence_lens'] |
| Hp_0 = inputs[5] |
| |
| num_directions = infer_shape(Wp)[0] |
| |
| if num_directions not in [1, 2]: |
| raise ValueError("num_directions must be either 1 or 2!") |
| |
| if layout == 1: |
| X = _op.transpose(X, axes=(1, 0)) |
| |
| # Initialize state if not provided. |
| if Hp_0 is None: |
| W_dtype = infer_type(Wp).checked_type.dtype |
| X_shape = infer_shape(X) |
| hidden_size = infer_shape(Rp)[-1] |
| batch_size = X_shape[1] |
| Hp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) |
| elif layout == 1: |
| Hp_0 = _op.transpose(Hp_0, axes=(1, 0)) |
| |
| # TODO (vvchernov): It can be replaced by _op.split if issue #8412 is resolved |
| X_steps = unbind(X, axis=0) |
| |
| H_ts = _op.split(Hp_0, num_directions) |
| Ws = _op.split(Wp, num_directions) |
| Rs = _op.split(Rp, num_directions) |
| |
| Bs = None |
| if Bp is not None: |
| Bs = _op.split(Bp, num_directions) |
| return X_steps, H_ts, Ws, Rs, Bs, num_directions |
| |
| @classmethod |
| def _impl_common(cls, inputs, attr, layout): |
| X_steps, H_ts, Ws, Rs, Bs, num_directions = cls._inputs_helper(inputs, layout) |
| acts = cls._get_activations(attr, 1, num_directions, "RNN") |
| |
| weights_dicts = [] |
| for i in range(num_directions): |
| weights_dict = {} |
| |
| weights_dict["hidden_state"] = _op.squeeze(H_ts[i], axis=[0]) |
| |
| weights_dict["w_inp"] = _op.squeeze(Ws[i], axis=[0]) |
| weights_dict["w_hid"] = _op.squeeze(Rs[i], axis=[0]) |
| if Bs is not None: |
| Bi, Bh = _op.split(Bs[i], 2, -1) |
| weights_dict["b_inp"] = _op.squeeze(Bi, axis=[0]) |
| weights_dict["b_hid"] = _op.squeeze(Bh, axis=[0]) |
| weights_dicts.append(weights_dict) |
| |
| if num_directions == 2: |
| output, H = RNN.bidir_rnn_cell( |
| input_seqs=X_steps, |
| weight_dicts=weights_dicts, |
| acts=acts, |
| ) |
| else: |
| # outputs shape = [seqs_num, (batch_size, hidden_size)] |
| outputs, H = rnn_cell( |
| input_seqs=X_steps, |
| **weights_dicts[0], |
| act=acts[0], |
| ) |
| |
| # output shape = (seqs_num, num_directions, batch_size, hidden_size) |
| output = _op.expand_dims(_op.stack(outputs, axis=0), axis=1) |
| H = _op.expand_dims(H, axis=0) |
| |
| if layout == 1: |
| output = _op.transpose(output, axes=(1, 0)) |
| H = _op.transpose(H, axes=(1, 0)) |
| return _expr.TupleWrapper(_expr.Tuple((output, H)), 2) |
| |
| @classmethod |
| def _impl_v7(cls, inputs, attr, params): |
| return cls._impl_common(inputs, attr, 0) |
| |
| @classmethod |
| def _impl_v14(cls, inputs, attr, params): |
| layout = attr.get("layout", 0) |
| return cls._impl_common(inputs, attr, layout) |
| |
| |
| class LSTM(RNN): |
| """Operator converter for LSTM""" |
| |
| @classmethod |
| def bidir_lstm_cell( |
| cls, |
| input_seqs, |
| weight_dicts, |
| acts, |
| ): |
| """ |
| Bidirectional LSTM cell |
| """ |
| seq_len = len(input_seqs) |
| forward_outputs, fw_H_t, fw_C_t = lstm_cell( |
| input_seqs, |
| **weight_dicts[0], |
| f_act=acts[0], |
| g_act=acts[1], |
| h_act=acts[2], |
| ) |
| |
| reverse_outputs, rev_H_t, rev_C_t = lstm_cell( |
| input_seqs, |
| **weight_dicts[1], |
| f_act=acts[3], |
| g_act=acts[4], |
| h_act=acts[5], |
| backwards=True, |
| ) |
| |
| final_outputs = [] |
| for i in range(seq_len): |
| final_outputs.append( |
| _op.stack([forward_outputs[i], reverse_outputs[seq_len - 1 - i]], axis=0) |
| ) |
| |
| return ( |
| _op.stack(final_outputs, axis=0), |
| _op.stack([fw_H_t, rev_H_t], axis=0), |
| _op.stack([fw_C_t, rev_C_t], axis=0), |
| ) |
| |
| @classmethod |
| def _default_activations(cls, num_directions): |
| return [_op.sigmoid, _op.tanh, _op.tanh] * num_directions |
| |
| @classmethod |
| def _impl_common(cls, inputs, attr, layout): |
| X_steps, H_ts, Ws, Rs, Bs, num_directions = cls._inputs_helper(inputs, layout) |
| acts = cls._get_activations(attr, 3, num_directions, "LSTM") |
| |
| # cell state |
| Cp_0 = inputs[6] |
| if Cp_0 is None: |
| C_ts = _expr.TupleWrapper( |
| _expr.Tuple([_op.zeros_like(H_ts[i]) for i in range(num_directions)]), |
| num_directions, |
| ) |
| else: |
| if layout == 1: |
| Cp_0 = _op.transpose(Cp_0, axes=(1, 0)) |
| C_ts = _op.split(Cp_0, num_directions) |
| |
| # peepholes |
| Pp = inputs[7] |
| if Pp is not None: |
| p_i, p_o, p_f = _op.split(Pp, 3, axis=1) |
| |
| p_is = _op.split(p_i, num_directions) |
| p_fs = _op.split(p_f, num_directions) |
| p_os = _op.split(p_o, num_directions) |
| |
| weights_dicts = [] |
| for i in range(num_directions): |
| weights_dict = {} |
| |
| weights_dict["hidden_state"] = _op.squeeze(H_ts[i], axis=[0]) |
| weights_dict["cell_state"] = _op.squeeze(C_ts[i], axis=[0]) |
| |
| # Weights permutation: onnx format i-o-f-c, lstm cell format i-f-c-o |
| mati, mato, matf, matc = _op.split(_op.squeeze(Ws[i], axis=[0]), 4) |
| weights_dict["w_inp"] = _op.concatenate([mati, matf, matc, mato], axis=0) |
| mati, mato, matf, matc = _op.split(_op.squeeze(Rs[i], axis=[0]), 4) |
| weights_dict["w_hid"] = _op.concatenate([mati, matf, matc, mato], axis=0) |
| if Bs is not None: |
| Bi, Bh = _op.split(Bs[i], 2, -1) |
| mati, mato, matf, matc = _op.split(_op.squeeze(Bi, axis=[0]), 4) |
| weights_dict["b_inp"] = _op.concatenate([mati, matf, matc, mato], axis=0) |
| mati, mato, matf, matc = _op.split(_op.squeeze(Bh, axis=[0]), 4) |
| weights_dict["b_hid"] = _op.concatenate([mati, matf, matc, mato], axis=0) |
| if Pp is not None: |
| weights_dict["p_i"] = _op.squeeze(p_is[i], axis=[0]) |
| weights_dict["p_f"] = _op.squeeze(p_fs[i], axis=[0]) |
| weights_dict["p_o"] = _op.squeeze(p_os[i], axis=[0]) |
| weights_dicts.append(weights_dict) |
| |
| if num_directions == 2: |
| output, H, C = LSTM.bidir_lstm_cell( |
| input_seqs=X_steps, |
| weight_dicts=weights_dicts, |
| acts=acts, |
| ) |
| else: |
| # outputs shape = [seqs_num, (batch_size, hidden_size)] |
| outputs, H, C = lstm_cell( |
| input_seqs=X_steps, |
| **weights_dicts[0], |
| f_act=acts[0], |
| g_act=acts[1], |
| h_act=acts[2], |
| ) |
| |
| # output shape = (seqs_num, num_directions, batch_size, hidden_size) |
| output = _op.expand_dims(_op.stack(outputs, axis=0), axis=1) |
| H = _op.expand_dims(H, axis=0) |
| C = _op.expand_dims(C, axis=0) |
| |
| if layout == 1: |
| output = _op.transpose(output, axes=(1, 0)) |
| H = _op.transpose(H, axes=(1, 0)) |
| C = _op.transpose(C, axes=(1, 0)) |
| return _expr.TupleWrapper(_expr.Tuple((output, H, C)), 3) |
| |
| |
| class GRU(RNN): |
| """Operator convert for GRU""" |
| |
| @classmethod |
| def bidir_gru_cell( |
| cls, |
| input_seqs, |
| weight_dicts, |
| acts, |
| ): |
| """ |
| Bidirectional GRU cell |
| """ |
| seq_len = len(input_seqs) |
| forward_outputs, fw_H_t = gru_cell( |
| input_seqs, |
| **weight_dicts[0], |
| rz_act=acts[0], |
| n_act=acts[1], |
| ) |
| |
| reverse_outputs, rev_H_t = gru_cell( |
| input_seqs, |
| **weight_dicts[1], |
| rz_act=acts[2], |
| n_act=acts[3], |
| backwards=True, |
| ) |
| |
| final_outputs = [] |
| for i in range(seq_len): |
| final_outputs.append( |
| _op.stack([forward_outputs[i], reverse_outputs[seq_len - 1 - i]], axis=0) |
| ) |
| |
| return ( |
| _op.stack(final_outputs, axis=0), |
| _op.stack([fw_H_t, rev_H_t], axis=0), |
| ) |
| |
| @classmethod |
| def _default_activations(cls, num_directions): |
| return [_op.sigmoid, _op.tanh] * num_directions |
| |
| @classmethod |
| def _impl_common(cls, inputs, attr, layout): |
| X_steps, H_ts, Ws, Rs, Bs, num_directions = cls._inputs_helper(inputs, layout) |
| acts = cls._get_activations(attr, 2, num_directions, "GRU") |
| linear_before_reset = attr.get("linear_before_reset", 0) |
| |
| weights_dicts = [] |
| for i in range(num_directions): |
| weights_dict = {} |
| |
| weights_dict["hidden_state"] = _op.squeeze(H_ts[i], axis=[0]) |
| weights_dict["linear_before_reset"] = linear_before_reset |
| |
| # Weights permutation: onnx format i-o-f-c, lstm cell format i-f-c-o |
| matz, matr, matn = _op.split(_op.squeeze(Ws[i], axis=[0]), 3) |
| weights_dict["w_inp"] = _op.concatenate([matr, matz, matn], axis=0) |
| matz, matr, matn = _op.split(_op.squeeze(Rs[i], axis=[0]), 3) |
| weights_dict["w_hid"] = _op.concatenate([matr, matz, matn], axis=0) |
| if Bs is not None: |
| Bi, Bh = _op.split(Bs[i], 2, -1) |
| matz, matr, matn = _op.split(_op.squeeze(Bi, axis=[0]), 3) |
| weights_dict["b_inp"] = _op.concatenate([matr, matz, matn], axis=0) |
| matz, matr, matn = _op.split(_op.squeeze(Bh, axis=[0]), 3) |
| weights_dict["b_hid"] = _op.concatenate([matr, matz, matn], axis=0) |
| weights_dicts.append(weights_dict) |
| |
| if num_directions == 2: |
| output, H = GRU.bidir_gru_cell( |
| input_seqs=X_steps, |
| weight_dicts=weights_dicts, |
| acts=acts, |
| ) |
| else: |
| # outputs shape = [seqs_num, (batch_size, hidden_size)] |
| outputs, H = gru_cell( |
| input_seqs=X_steps, |
| **weights_dicts[0], |
| rz_act=acts[0], |
| n_act=acts[1], |
| ) |
| |
| # output shape = (seqs_num, num_directions, batch_size, hidden_size) |
| output = _op.expand_dims(_op.stack(outputs, axis=0), axis=1) |
| H = _op.expand_dims(H, axis=0) |
| |
| if layout == 1: |
| output = _op.transpose(output, axes=(1, 0)) |
| H = _op.transpose(H, axes=(1, 0)) |
| return _expr.TupleWrapper(_expr.Tuple((output, H)), 2) |
| |
| |
| class Resize(OnnxOpConverter): |
| """Operator converter for Resize""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| mode = attr.get("mode").decode("ascii") |
| if mode == "nearest": |
| method = "nearest_neighbor" |
| elif mode == "linear": |
| method = "linear" |
| elif mode == "cubic": |
| method = "cubic" |
| else: |
| raise tvm.error.OpAttributeInvalid( |
| 'Value {} in attribute "mode" of operator Resize is not valid.'.format(mode) |
| ) |
| |
| scale = inputs[1] |
| size = _op.cast(shape_of(inputs[0]), infer_type(scale).checked_type.dtype) * scale |
| ndims = len(infer_shape(inputs[0])) |
| out = None |
| if ndims == 3: |
| out_size = fold_constant(_op.strided_slice(size, [2], [3])) |
| out = _op.image.resize1d(inputs[0], out_size, None, "NCW", method, "asymmetric") |
| elif ndims == 4: |
| out_size = fold_constant(_op.strided_slice(size, [2], [4])) |
| out = _op.image.resize2d(inputs[0], out_size, None, "NCHW", method, "asymmetric") |
| elif ndims == 5: |
| out_size = fold_constant(_op.strided_slice(size, [2], [5])) |
| out = _op.image.resize3d(inputs[0], out_size, None, "NCDHW", method, "asymmetric") |
| else: |
| raise NotImplementedError("Resize only supports 3, 4, or 5 dims") |
| return out |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| scale = inputs[2] |
| scale_shape = infer_shape(scale) |
| if len(inputs) == 4: |
| assert ( |
| len(scale_shape) == 0 or scale_shape[0] == 0 |
| ), "One of scale or size should be passed, not both." |
| size = inputs[3] |
| else: |
| assert len(scale_shape) != 0, "One of scale or size should be passed." |
| size = _op.cast(shape_of(inputs[0]), infer_type(scale).checked_type.dtype) * scale |
| return cls.v11_13_common(inputs, size, attr, params) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| scale = inputs[2] |
| size = inputs[3] |
| |
| # Some versions of onnx exporters produce an opset 13 model with the opset 11 |
| # resize op, handle that edge case |
| if scale is not None and size is not None: |
| return cls._impl_v11(inputs, attr, params) |
| |
| if size is not None: |
| assert scale is None, "One of scale or size should be passed, not both." |
| else: |
| scale_type = infer_type(scale) |
| scale_shape = scale_type.checked_type.shape |
| scale_dtype = scale_type.checked_type.dtype |
| assert len(scale_shape) != 0, "One of scale or size should be passed." |
| size = _op.cast(shape_of(inputs[0]), scale_dtype) * scale |
| |
| return cls.v11_13_common(inputs, size, attr, params) |
| |
| @classmethod |
| def v11_13_common(cls, inputs, size, attr, params): |
| """ |
| Resize v11 and Resize v13 are identical except in how |
| they handle the passing of scale and size. This utility |
| provides the implementation for both |
| """ |
| roi = inputs[1] |
| if roi is not None and infer_shape(roi)[0] == 0: |
| roi = None |
| ndims = len(infer_shape(inputs[0])) |
| mode = attr.get("mode").decode("ascii") |
| if mode == "nearest": |
| method = "nearest_neighbor" |
| elif mode == "linear": |
| method = "linear" |
| elif mode == "cubic": |
| method = "cubic" |
| else: |
| raise tvm.error.OpAttributeInvalid( |
| 'Value {} in attribute "mode" of operator Resize is not valid.'.format(mode) |
| ) |
| |
| coord_trans = attr.get("coordinate_transformation_mode", b"half_pixel").decode("ascii") |
| nearest_mode = attr.get("nearest_mode", b"round_prefer_floor").decode("ascii") |
| alpha = attr.get("cubic_coeff_a", -0.75) |
| exclude = attr.get("exclude_outside", 0) |
| extrapolation_value = attr.get("extrapolation_value", 0.0) |
| |
| if roi is not None: |
| roi = fold_constant( |
| _op.concatenate( |
| [ |
| _op.strided_slice(roi, [2], [ndims]), |
| _op.strided_slice(roi, [ndims + 2], [2 * ndims]), |
| ], |
| axis=0, |
| ) |
| ) |
| |
| out_size = fold_constant(_op.strided_slice(size, [2], [ndims])) |
| |
| out = None |
| if ndims == 3: |
| out = _op.image.resize1d( |
| inputs[0], |
| out_size, |
| roi, |
| "NCW", |
| method, |
| coord_trans, |
| nearest_mode, |
| alpha, |
| exclude, |
| extrapolation_value, |
| ) |
| elif ndims == 4: |
| out = _op.image.resize2d( |
| inputs[0], |
| out_size, |
| roi, |
| "NCHW", |
| method, |
| coord_trans, |
| nearest_mode, |
| alpha, |
| exclude, |
| extrapolation_value, |
| ) |
| elif ndims == 5: |
| out = _op.image.resize3d( |
| inputs[0], |
| out_size, |
| roi, |
| "NCDHW", |
| method, |
| coord_trans, |
| nearest_mode, |
| alpha, |
| exclude, |
| extrapolation_value, |
| ) |
| else: |
| raise NotImplementedError("Resize only supports 3, 4, or 5 dims") |
| |
| return out |
| |
| |
| class NonZero(OnnxOpConverter): |
| """Operator converter for NonZero""" |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| if len(inputs) > 1: |
| raise ValueError("Expect 1 input only") |
| |
| output = AttrCvt(op_name="argwhere")(inputs, attr, params) |
| # ONNX NonZero always outputs int64 |
| output = _op.cast(output, "int64") |
| return _op.transpose(output, axes=(1, 0)) |
| |
| |
| class ReverseSequence(OnnxOpConverter): |
| """Operator converter for ReverseSequence""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| |
| return _op.reverse_sequence(inputs[0], inputs[1], attr["time_axis"], attr["batch_axis"]) |
| |
| |
| class TopK(OnnxOpConverter): |
| """Operator converter for TopK""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) != 2: |
| raise ValueError("Expect 2 input only") |
| axis = attr.get("axis", -1) |
| largest = attr.get("largest", 1) |
| |
| if largest == 0: |
| # TODO(mbrookhart): optimize this by adding a smallest attribute to topi if this |
| # ever becomes a bottleneck |
| ndim = len(infer_shape(inputs[0])) |
| if axis < 0: |
| axis += ndim |
| sort = _op.sort(inputs[0], axis=axis) |
| argsort = _op.argsort(inputs[0], axis=axis, dtype="int64") |
| begin = [0] * ndim |
| stride = [1] * ndim |
| end = _op.concatenate( |
| [ |
| _op.const([np.iinfo(np.int64).max] * axis, dtype="int64"), |
| inputs[1], |
| _op.const([np.iinfo(np.int64).max] * (ndim - axis - 1), dtype="int64"), |
| ], |
| axis=0, |
| ) |
| return _expr.TupleWrapper( |
| _expr.Tuple( |
| [ |
| _op.strided_slice(sort, begin, end, stride), |
| _op.strided_slice(argsort, begin, end, stride), |
| ] |
| ), |
| 2, |
| ) |
| |
| return _op.topk(inputs[0], inputs[1], axis=axis, dtype="int64") |
| |
| |
| class Range(OnnxOpConverter): |
| """Operator converter for Range""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) != 3: |
| raise ValueError("Expect 3 input only") |
| |
| return _op.arange( |
| inputs[0], inputs[1], inputs[2], dtype=infer_type(inputs[0]).checked_type.dtype |
| ) |
| |
| |
| class IsInf(OnnxOpConverter): |
| """Operator converter for IsInf""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| detect_negative = attr.get("detect_negative", 1) |
| detect_positive = attr.get("detect_positive", 1) |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| isinf = _op.isinf(inputs[0]) |
| if not detect_negative: |
| isinf = isinf * (inputs[0] > _op.const(0, dtype)) |
| if not detect_positive: |
| isinf = isinf * (inputs[0] < _op.const(0, dtype)) |
| return isinf |
| |
| |
| class Celu(OnnxOpConverter): |
| """Operator convereter for celu""" |
| |
| @classmethod |
| def _impl_v12(cls, inputs, attr, params): |
| x = inputs[0] |
| dtype = infer_type(x).checked_type.dtype |
| alpha = _op.const(attr.get("alpha", 1.0), dtype) |
| zero = _op.const(0, dtype) |
| one = _op.const(1, dtype) |
| out = _op.maximum(zero, x) + _op.minimum(zero, alpha * (_op.exp(x / alpha) - one)) |
| return out |
| |
| |
| class MaxRoiPool(OnnxOpConverter): |
| """Operator converter for MaxRoiPool.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| assert len(inputs) == 2, "MMaxRoiPool op take 2 inputs, {} given".format(len(inputs)) |
| |
| data = inputs[0] |
| rois = inputs[1] |
| pooled_shape = attr.get("pooled_shape") |
| spatial_scale = attr.get("spatial_scale", 1.0) |
| |
| return _vision.roi_pool(data, rois, pooled_shape, spatial_scale) |
| |
| |
| class RoiAlign(OnnxOpConverter): |
| """Operator converter for RoiAlign.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if len(inputs) != 3: |
| raise ValueError("Expect 3 inputs only") |
| x = inputs[0] |
| rois = inputs[1] |
| batch_indices = inputs[2] |
| mode = attr.get("mode", b"avg") |
| if mode not in (b"avg", b"max"): |
| raise NotImplementedError("RoiAlign in Relay only uses avg and max modes") |
| output_height = attr.get("output_height", 1) |
| output_width = attr.get("output_width", 1) |
| |
| sampling_ratio = attr.get("sampling_ratio", 0) |
| spatial_scale = attr.get("spatial_scale", 1.0) |
| |
| batch_indices = _op.expand_dims(batch_indices, axis=1, num_newaxis=1) |
| batch_indices = _op.cast(batch_indices, infer_type(rois).checked_type.dtype) |
| rois = _op.concatenate([batch_indices, rois], 1) |
| |
| return _vision.roi_align( |
| x, rois, [output_height, output_width], spatial_scale, sampling_ratio, mode=mode |
| ) |
| |
| |
| class Clip(OnnxOpConverter): |
| """Operator converter for Clip.""" |
| |
| @staticmethod |
| def convert_attributes(inputs, attr, params): |
| convert = AttrCvt("clip", transforms={"min": "a_min", "max": "a_max"}) |
| return convert(inputs, attr, params) |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| if "min" not in attr: |
| attr["min"] = -np.inf |
| if "max" not in attr: |
| attr["max"] = np.inf |
| return Clip.convert_attributes(inputs, attr, params) |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| if len(inputs) == 3 and isinstance(inputs[2], _expr.Constant): |
| attr["max"] = inputs[2].data.numpy().item() |
| inputs = inputs[0:2] |
| if len(inputs) >= 2 and isinstance(inputs[1], _expr.Constant): |
| attr["min"] = inputs[1].data.numpy().item() |
| inputs = inputs[0:1] |
| if "min" in attr and "max" in attr: |
| return Clip.convert_attributes(inputs, attr, params) |
| |
| assert len(inputs) <= 3, "Clip-11 takes up to 3 inputs, input, min, max" |
| result = inputs[0] |
| for i, op in enumerate([_op.tensor.maximum, _op.tensor.minimum]): |
| if i < len(inputs) - 1: |
| if inputs[i + 1] is not None: |
| result = op(result, inputs[i + 1]) |
| return result |
| |
| |
| class Softplus(OnnxOpConverter): |
| """Operator converter for Softplus.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| data_dtype = infer_type(data).checked_type.dtype |
| data = _op.exp(data) + _expr.const(1, dtype=data_dtype) |
| return _op.log(data) |
| |
| |
| class Loop(OnnxOpConverter): |
| """Operator converter for Loop""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| max_loop_count = inputs[0] |
| cond = inputs[1] |
| loop_deps = inputs[2:] |
| num_deps = len(loop_deps) |
| # Create a copy of the body function to prevent the original |
| # from being modified. |
| body = copy.copy(attr["body"]) |
| iter_dtype = infer_type(max_loop_count).checked_type.dtype |
| |
| # Determine what condition mode we're in. |
| assert cond is not None or max_loop_count is not None |
| is_for_loop = max_loop_count is not None and cond is None |
| is_condition_for_loop = cond is not None and max_loop_count is not None |
| |
| # Loop inputs will be packed as |
| # [iter_count, max_count, condition, loop_deps, scan_outputs] |
| def cond_fn(*loop_inputs): |
| i = loop_inputs[0] |
| max_count = loop_inputs[1] |
| w = loop_inputs[2] |
| |
| if cond is not None: |
| out_while = _op.equal(w, _expr.const(True, "bool")) |
| if max_loop_count is not None: |
| out_loop = _op.less(i, max_count) |
| |
| if is_condition_for_loop: |
| return _op.logical_and(out_while, out_loop) |
| if is_for_loop: |
| return out_loop |
| return out_while |
| |
| # Get the current graph proto and create a clone for the subgraph |
| graph_scope = GraphProto.current |
| subgraph_scope = GraphProto( |
| graph_scope._shape, graph_scope._dtype, graph_scope._freeze_params |
| ) |
| # Load nodes from outer graph into inner graph. |
| subgraph_scope._nodes = graph_scope._nodes.copy() |
| |
| # Create a list of variables for each value updated in the loop. |
| def get_var(name, val, scan=False): |
| checked_type = infer_type(val) |
| if hasattr(checked_type, "type_annotation"): |
| checked_type = checked_type.type_annotation |
| if hasattr(checked_type, "checked_type"): |
| checked_type = checked_type.checked_type |
| shape = get_const_tuple(checked_type.shape) |
| actual_shape = [] |
| for dim in shape: |
| if isinstance(dim, int) and dim == 0: |
| actual_shape.append(_ty.Any()) |
| else: |
| actual_shape.append(dim) |
| if scan: |
| return _expr.var(name, shape=[_ty.Any()] + actual_shape, dtype=checked_type.dtype) |
| |
| return _expr.var(name, shape=actual_shape, dtype=checked_type.dtype) |
| |
| loop_vars = [ |
| _expr.var(body.input[0].name, shape=(), dtype=iter_dtype), # iteration count |
| _expr.var("max_count", shape=(), dtype=iter_dtype), # iteration count |
| get_var(body.input[1].name, cond), # exit condition |
| ] |
| loop_vars += [get_var(body.input[i + 2].name, v) for i, v in enumerate(loop_deps)] |
| loop_var_names = [v.name_hint for v in loop_vars] |
| |
| num_scan_outputs = len(body.output) - (1 + num_deps) |
| |
| # Construct variables and initial empty tensors for any scan outputs. |
| # To do this, we'll figure out the output shapes of the body subgraph by importing |
| # it and doing type inference. |
| scan_output_vars = [] |
| scan_output_init = [] |
| if num_scan_outputs > 0: |
| with subgraph_scope: |
| loop_outputs = subgraph_scope.from_onnx( |
| body, graph_scope.opset, get_output_expr=True |
| ) |
| loop_outputs = _expr.TupleWrapper(loop_outputs, len(body.output)) |
| |
| for i in range(num_scan_outputs): |
| name, _, _, _ = get_info(body.output[i + 1 + num_deps]) |
| output_node = infer_type(loop_outputs[i + 1 + num_deps]) |
| shape = get_const_tuple(output_node.checked_type.shape) |
| dtype = output_node.checked_type.dtype |
| scan_output_vars.append( |
| _expr.var(name, shape=([_ty.Any()] * (len(shape) + 1)), dtype=dtype) |
| ) |
| scan_output_init.append( |
| _op.reshape(_expr.const(np.array([]).astype(dtype)), [0] + [1] * len(shape)) |
| ) |
| |
| # Now we can remove loop iter variables from our inner loop's inputs. |
| # This is kind of a hack since we have graph inputs that we don't |
| # want to treat as actual inputs. |
| while len(body.input) != 0: |
| body.input.pop(0) |
| |
| # Define the loop body, in this function we need to unpack loop inputs, |
| # convert the loop subgraph, and pack outputs for the next iteration. |
| def body_fn(*loop_inputs): |
| # Unpack inputs |
| loop_count = loop_inputs[0] |
| max_count = loop_inputs[1] |
| cond = loop_inputs[2] |
| current_vars = list(loop_inputs[3 : (3 + num_deps)]) |
| scan_outputs = loop_inputs[(3 + num_deps) :] |
| |
| # Prepare body inputs by adding them to node dictionary. |
| new_inputs = [loop_count, max_count, cond] + current_vars |
| for i, inp in enumerate(new_inputs): |
| subgraph_scope._nodes[loop_var_names[i]] = inp |
| |
| # Get the output of the current loop using the updated inputs. |
| with subgraph_scope: |
| loop_outputs = subgraph_scope.from_onnx( |
| body, graph_scope.opset, get_output_expr=True |
| ) |
| # Unpack the body outputs and prepare variables for next iteration. |
| new_cond = loop_outputs[0] |
| new_loop_vars = [loop_outputs[i] for i in range(1, 1 + num_deps)] |
| new_scan_outputs = [loop_outputs[i] for i in range(1 + num_deps, len(loop_outputs))] |
| |
| # Add new scan outputs to tracking |
| combined_scan_outputs = [] |
| for i, scan in enumerate(scan_outputs): |
| rank = len(infer_shape(scan)) - 1 |
| new_scan = new_scan_outputs[i] |
| expand_scan = _op.expand_dims(new_scan, axis=0) |
| # For non scalar outputs we need to broadcast the initial value. |
| if rank > 0: |
| new_scan_shape = shape_of(new_scan, dtype=iter_dtype) |
| scan_broadcast = _op.concatenate( |
| [_op.reshape(loop_count, [1]), new_scan_shape], axis=0 |
| ) |
| scan = _op.broadcast_to(scan, scan_broadcast) |
| combined_scan = _op.concatenate([scan, expand_scan], axis=0) |
| combined_scan_outputs.append(combined_scan) |
| |
| # Increment counter. |
| if max_loop_count is not None: |
| incr = _expr.const(1, dtype=iter_dtype) |
| loop_count = loop_count + incr |
| |
| # Pack loop outputs for next iteration |
| # [iter_count, cond, loop_deps, loop_scans] |
| return [loop_count, max_count, new_cond] + new_loop_vars + combined_scan_outputs |
| |
| # Create the loop function. |
| loop = fold_constant(_loops.while_loop(cond_fn, loop_vars + scan_output_vars, body_fn)) |
| |
| # Now need to run initial values through the graph. |
| init_count = _expr.const(0, dtype=iter_dtype) |
| loop_vals = loop(init_count, max_loop_count, cond, *loop_deps, *scan_output_init) |
| |
| # Extract final iteration outputs. |
| if num_deps + num_scan_outputs == 1: |
| outputs = _expr.TupleGetItem(loop_vals, 3) |
| else: |
| outputs = _expr.TupleWrapper( |
| _expr.Tuple( |
| [ |
| _expr.TupleGetItem(loop_vals, i + 3) |
| for i in range(num_deps + num_scan_outputs) |
| ] |
| ), |
| num_deps + num_scan_outputs, |
| ) |
| |
| # Update outer graph with constants found in the subgraph. |
| free_vars = analysis.free_vars(loop) |
| graph_scope._params.update(subgraph_scope._params) |
| graph_scope._nodes.update(subgraph_scope._nodes) |
| for var in free_vars: |
| graph_scope._nodes.update({var.name_hint: var}) |
| return outputs |
| |
| |
| class If(OnnxOpConverter): |
| """Operator converter for If""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| cond = inputs[0] |
| # Convert array to bool if needed. |
| if len(infer_shape(cond)) > 0: |
| cond = _op.take(cond, _expr.const(0, dtype="int64")) |
| then_branch = attr.get("then_branch", None) |
| else_branch = attr.get("else_branch", None) |
| assert then_branch is not None and else_branch is not None |
| |
| # Create graph converters for both branches. |
| graph_scope = GraphProto.current |
| then_graph = GraphProto(graph_scope._shape, graph_scope._dtype, graph_scope._freeze_params) |
| then_graph._nodes = graph_scope._nodes.copy() |
| else_graph = GraphProto(graph_scope._shape, graph_scope._dtype, graph_scope._freeze_params) |
| else_graph._nodes = graph_scope._nodes.copy() |
| |
| # Convert each branch to a relay expression. |
| with then_graph: |
| then_expr = then_graph.from_onnx(then_branch, graph_scope.opset, get_output_expr=True) |
| with else_graph: |
| else_expr = else_graph.from_onnx(else_branch, graph_scope.opset, get_output_expr=True) |
| |
| # Add constants from both branches to parent graph. |
| graph_scope._params.update(then_graph._params) |
| graph_scope._nodes.update(then_graph._nodes) |
| then_free_vars = analysis.free_vars(then_expr) |
| for var in then_free_vars: |
| graph_scope._nodes.update({var.name_hint: var}) |
| graph_scope._params.update(else_graph._params) |
| graph_scope._nodes.update(else_graph._nodes) |
| else_free_vars = analysis.free_vars(else_expr) |
| for var in else_free_vars: |
| graph_scope._nodes.update({var.name_hint: var}) |
| |
| # Now we can construct the relay if statement and return. |
| ret = _expr.If(cond, then_expr, else_expr) |
| if len(then_branch.output) > 1: |
| ret = _expr.TupleWrapper(ret, len(then_branch.output)) |
| return ret |
| |
| |
| class Scan(OnnxOpConverter): |
| """Operator converter for Scan""" |
| |
| @classmethod |
| def _impl_v8(cls, inputs, attr, params): |
| new_inputs = inputs[1:] |
| batch_num = infer_shape(inputs[1])[0] |
| out = [] |
| for i in range(batch_num): |
| v9_inputs = [ |
| _op.take(new_inputs[j], _expr.const(i), axis=0) for j in range(len(new_inputs)) |
| ] |
| results = cls._impl_v9(v9_inputs, attr, params) |
| results = [_op.expand_dims(results[j], axis=0) for j in range(len(results))] |
| if i == 0: |
| out = results |
| else: |
| out = [_op.concatenate([out[j], results[j]], axis=0) for j in range(len(results))] |
| |
| out = _expr.TupleWrapper(_expr.Tuple(out), len(out)) |
| return out |
| |
| @classmethod |
| def _impl_v9(cls, inputs, attr, params): |
| body = attr.get("body") |
| num_scan_inputs = attr.get("num_scan_inputs") |
| num_all_inputs = len(inputs) |
| num_state_inputs = len(body.input) - num_scan_inputs |
| num_state_outputs = num_state_inputs |
| num_all_outputs = len(body.output) |
| num_scan_outputs = num_all_outputs - num_state_outputs |
| scan_input_axes = attr.get("scan_input_axes", [0] * num_scan_inputs) |
| scan_input_directions = attr.get("scan_input_directions", [0] * num_scan_inputs) |
| scan_output_axes = list(attr.get("scan_output_axes", [0] * num_scan_outputs)) |
| scan_output_directions = attr.get("scan_output_directions", [0] * num_scan_outputs) |
| # loop count are the same for all scan inputs, so get loop count by first input scan |
| # strided_slice not support dynamic axes, so assume input shape are static |
| max_loop_count = infer_shape(inputs[num_state_inputs])[scan_input_axes[0]] |
| |
| # Create a copy of the body function to prevent the original |
| # from being modified. |
| body = copy.copy(attr["body"]) |
| |
| # Loop inputs will be packed as |
| # [iter_count, loop_deps, scan_outputs] |
| def cond_fn(*loop_inputs): |
| i = loop_inputs[0] |
| return _op.less(i, relay.const(max_loop_count, "int32")) |
| |
| # Get the current graph proto and create a clone for the subgraph |
| graph_scope = GraphProto.current |
| subgraph_scope = GraphProto( |
| graph_scope._shape, graph_scope._dtype, graph_scope._freeze_params |
| ) |
| # Load nodes from outer graph into inner graph. |
| subgraph_scope._nodes = graph_scope._nodes.copy() |
| |
| # Create a list of variables for each value updated in the loop. |
| def get_var(name, val, scan=False): |
| checked_type = infer_type(val) |
| if hasattr(checked_type, "type_annotation"): |
| checked_type = checked_type.type_annotation |
| if hasattr(checked_type, "checked_type"): |
| checked_type = checked_type.checked_type |
| shape = get_const_tuple(checked_type.shape) |
| actual_shape = [] |
| for dim in shape: |
| if isinstance(dim, int) and dim == 0: |
| actual_shape.append(_ty.Any()) |
| else: |
| actual_shape.append(dim) |
| if scan: |
| return _expr.var(name, shape=[_ty.Any()] + actual_shape, dtype=checked_type.dtype) |
| |
| return _expr.var(name, shape=actual_shape, dtype=checked_type.dtype) |
| |
| # Construct variables and initial empty tensors for any scan outputs. |
| # To do this, we'll figure out the output shapes of the body subgraph by importing |
| # it and doing type inference. |
| scan_output_vars = [] |
| scan_output_init = [] |
| if num_scan_outputs > 0: |
| with subgraph_scope: |
| loop_outputs = subgraph_scope.from_onnx( |
| body, graph_scope.opset, get_output_expr=True |
| ) |
| loop_outputs = _expr.TupleWrapper(loop_outputs, len(body.output)) |
| |
| for i in range(num_scan_outputs): |
| name, _, _, _ = get_info(body.output[i + num_state_outputs]) |
| output_node = infer_type(loop_outputs[i + num_state_outputs]) |
| shape = list(get_const_tuple(output_node.checked_type.shape)) |
| if scan_output_axes[i] < 0: |
| scan_output_axes[i] = len(shape) + scan_output_axes[i] + 1 |
| shape.insert(scan_output_axes[i], max_loop_count) |
| dtype = output_node.checked_type.dtype |
| scan_output_vars.append(_expr.var(name, shape=shape, dtype=dtype)) |
| scan_output_init.append(_op.zeros(shape, dtype)) |
| |
| # loop vars = [iter_count, scan_state, scan_out] |
| loop_vars = [ |
| _expr.var("iter", shape=(), dtype="int32"), # iteration count |
| ] |
| loop_vars += [ |
| get_var(body.input[i].name, v) for i, v in enumerate(inputs) if i < num_state_inputs |
| ] |
| loop_vars += scan_output_vars |
| body_input_var_names = ["iter"] + [body.input[i].name for i in range(len(body.input))] |
| |
| # # Now we can remove loop iter variables from our inner loop's inputs. |
| # # This is kind of a hack since we have graph inputs that we don't |
| # # want to treat as actual inputs. |
| while len(body.input) != 0: |
| body.input.pop(0) |
| |
| # Define the loop body, in this function we need to unpack loop inputs, |
| # convert the loop subgraph, and pack outputs for the next iteration. |
| def body_fn(*loop_inputs): |
| # Unpack inputs |
| loop_count = loop_inputs[0] |
| state_vars = list(loop_inputs[1 : 1 + num_state_inputs]) |
| scan_vars = list(loop_inputs[1 + num_state_inputs :]) |
| # body take scan graph scan inputs as original input |
| input_scan_exprs = [] |
| for i in range(num_state_inputs, num_all_inputs): |
| if scan_input_directions[i - num_state_inputs] != 0: |
| input_scan_exprs.append( |
| relay.take( |
| inputs[i], |
| relay.const(max_loop_count - 1, "int32") - loop_count, |
| axis=scan_input_axes[i - num_state_inputs], |
| ) |
| ) |
| else: |
| input_scan_exprs.append( |
| relay.take( |
| inputs[i], |
| loop_count, |
| axis=scan_input_axes[i - num_state_inputs], |
| ) |
| ) |
| |
| # Prepare body inputs by adding them to node dictionary. |
| body_inputs = [loop_count] + state_vars + input_scan_exprs |
| for i, inp in enumerate(body_inputs): |
| subgraph_scope._nodes[body_input_var_names[i]] = inp |
| |
| # Get the output of the current loop using the updated inputs. |
| with subgraph_scope: |
| loop_outputs = subgraph_scope.from_onnx( |
| body, graph_scope.opset, get_output_expr=True |
| ) |
| # Unpack the body outputs and prepare variables for next iteration. |
| new_state_vars = [loop_outputs[i] for i in range(num_state_outputs)] |
| new_scan_vars = [loop_outputs[i] for i in range(num_state_outputs, num_all_outputs)] |
| |
| # Add new scan outputs to tracking |
| combined_scan_outputs = [] |
| for i in range(num_scan_outputs): |
| if scan_output_directions[i] == 0: |
| # append new scan output |
| combined_scan = _op.concatenate( |
| [scan_vars[i], _op.expand_dims(new_scan_vars[i], axis=scan_output_axes[i])], |
| axis=scan_output_axes[i], |
| ) |
| # pop head scan output |
| combined_scan = _op.strided_slice( |
| combined_scan, |
| begin=[1], |
| end=[max_loop_count + 1], |
| strides=[1], |
| axes=[scan_output_axes[i]], |
| ) |
| else: |
| # prepend new scan output |
| combined_scan = _op.concatenate( |
| [_op.expand_dims(new_scan_vars[i], axis=scan_output_axes[i]), scan_vars[i]], |
| axis=scan_output_axes[i], |
| ) |
| # pop tail scan output |
| combined_scan = _op.strided_slice( |
| combined_scan, |
| begin=[0], |
| end=[max_loop_count], |
| strides=[1], |
| axes=[scan_output_axes[i]], |
| ) |
| combined_scan_outputs.append(combined_scan) |
| |
| incr = _expr.const(1, dtype="int32") |
| loop_count = loop_count + incr |
| |
| # Pack loop outputs for next iteration |
| # [iter_count, state_var, scan_var] |
| return [loop_count] + new_state_vars + combined_scan_outputs |
| |
| # Create the loop function. |
| loop = fold_constant(_loops.while_loop(cond_fn, loop_vars, body_fn)) |
| |
| # Now need to run initial values through the graph. |
| init_count = _expr.const(0, dtype="int32") |
| |
| input_states = [inputs[i] for i in range(num_state_inputs)] |
| loop_vals = loop(init_count, *input_states, *scan_output_init) |
| |
| outputs = _expr.TupleWrapper( |
| _expr.Tuple([_expr.TupleGetItem(loop_vals, i + 1) for i in range(num_all_outputs)]), |
| num_all_outputs, |
| ) |
| |
| # Update outer graph with constants found in the subgraph. |
| free_vars = analysis.free_vars(loop) |
| graph_scope._params.update(subgraph_scope._params) |
| graph_scope._nodes.update(subgraph_scope._nodes) |
| for var in free_vars: |
| graph_scope._nodes.update({var.name_hint: var}) |
| return outputs |
| |
| |
| class LinearRegressor(OnnxOpConverter): |
| """Operator converter for LinearRegressor.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| data = inputs[0] |
| coefficients = attr.get("coefficients", 0) |
| data_shape = infer_shape(data) |
| targets = attr.get("targets", 1) |
| coefficients = _expr.const(list(coefficients), dtype="float32") |
| coefficients_shape = infer_shape(coefficients) |
| |
| coefficients = _op.reshape(coefficients, (targets, coefficients_shape[0] // targets)) |
| if coefficients_shape[0] // targets < data_shape[-1]: |
| data = _op.split(data, [coefficients_shape[0] // targets], -1)[0] |
| |
| mm_out = _op.nn.dense(data, coefficients) |
| |
| if "intercepts" in attr: |
| intercepts = attr.get("intercepts", 0) |
| intercepts = _expr.const(list(intercepts), dtype="float32") |
| |
| if targets == 1: |
| return _op.nn.bias_add(mm_out, intercepts, axis=-1) |
| return get_relay_op("add")(mm_out, intercepts) |
| |
| return mm_out |
| |
| |
| class NonMaxSuppression(OnnxOpConverter): |
| """Operator converter for NonMaxSuppression.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| # Get parameter values |
| boxes = inputs[0] |
| scores = inputs[1] |
| max_output_boxes_per_class = inputs[2] |
| iou_threshold = inputs[3] |
| score_threshold = inputs[4] |
| |
| boxes_dtype = infer_type(boxes).checked_type.dtype |
| |
| if attr.get("center_point_box", 0) != 0: |
| xc, yc, w, h = _op.split(boxes, 4, axis=2) |
| half_w = w / _expr.const(2.0, boxes_dtype) |
| half_h = h / _expr.const(2.0, boxes_dtype) |
| x1 = xc - half_w |
| x2 = xc + half_w |
| y1 = yc - half_h |
| y2 = yc + half_h |
| boxes = _op.concatenate([y1, x1, y2, x2], axis=2) |
| |
| if iou_threshold is None: |
| iou_threshold = _expr.const(0.0, dtype="float32") |
| if score_threshold is None: |
| score_threshold = _expr.const(0.0, dtype="float32") |
| |
| def conditionally_squeeze_scalar(x): |
| rank = len(infer_shape(x)) |
| assert rank <= 1, "nms thresholds must be scalars" |
| if rank == 1: |
| return _op.squeeze(x, [0]) |
| return x |
| |
| max_output_boxes_per_class = conditionally_squeeze_scalar(max_output_boxes_per_class) |
| iou_threshold = conditionally_squeeze_scalar(iou_threshold) |
| score_threshold = conditionally_squeeze_scalar(score_threshold) |
| |
| nms_out = _op.vision.all_class_non_max_suppression( |
| boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold |
| ) |
| |
| return _op.strided_slice(nms_out[0], _op.const([0], dtype="int64"), nms_out[1]) |
| |
| |
| class ATen(OnnxOpConverter): |
| """Operator converter for Pytorch ATen ops.""" |
| |
| @classmethod |
| def _op_dispatch(cls, operator, inputs, attr, params): |
| op_map = { |
| "size": cls._size, |
| "arange": cls._arange, |
| "index_put": cls._index_put, |
| "reshape": cls._reshape, |
| "embedding_bag": cls._embedding_bag, |
| } |
| assert operator in op_map, "Operator %s is not supported." % operator |
| return op_map[operator](inputs, attr, params) |
| |
| @classmethod |
| def _size(cls, inputs, attr, params): |
| return _op.take( |
| _op.shape_of(inputs[0], dtype="int64"), |
| _expr.const(-1, dtype="int64"), |
| axis=0, |
| mode="wrap", |
| ) |
| |
| @classmethod |
| def _arange(cls, inputs, attr, params): |
| return _op.arange(inputs[0], inputs[1], inputs[2], dtype="int64") |
| |
| @classmethod |
| def _check_index(cls, indices, values): |
| def unfolding_indices(indices, values): |
| n = len(indices) |
| flatten_indices = [] |
| slices_size = [] |
| for index in indices: |
| flatten_indices.append(_op.reshape(index, _op.const([-1]))) |
| slices_size.append(infer_shape(flatten_indices[-1])[0]) |
| repeat_size = [1] |
| tile_size = [1] |
| for i in range(1, n): |
| repeat_size.append(slices_size[-i] * repeat_size[-1]) |
| tile_size.append(slices_size[i - 1] * tile_size[-1]) |
| repeat_size.reverse() |
| unflod_slices = [] |
| for i in range(n): |
| unflod_slices.append( |
| fold_constant( |
| _op.repeat(_op.tile(flatten_indices[i], (tile_size[i],)), repeat_size[i], 0) |
| ) |
| ) |
| return unflod_slices, _op.reshape(values, _op.const([-1])) |
| |
| values_shape = infer_shape(values) |
| if len(values_shape) != 1: |
| return unfolding_indices(indices, values) |
| return indices, values |
| |
| @classmethod |
| def _index_put(cls, inputs, attr, params): |
| in_tensor = inputs[0] |
| indices, values = cls._check_index(inputs[1 : len(inputs) - 2], inputs[len(inputs) - 2]) |
| accumulate = inputs[len(inputs) - 1].data.asnumpy() != 0 |
| if not accumulate: |
| mode = "update" |
| else: |
| mode = "add" |
| index_tensor = _op.stack(indices, axis=0) |
| return _op.transform.scatter_nd(in_tensor, index_tensor, values, mode) |
| |
| @classmethod |
| def _reshape(cls, inputs, attr, params): |
| return _op.reshape(inputs[0], inputs[1]) |
| |
| @classmethod |
| def _embedding_bag(cls, inputs, attr, params): |
| mode_map = {0: _op.sum, 1: _op.mean, 2: _op.max} |
| |
| mode = attr.get("mode", 1) |
| reduction_fn = mode_map[mode] |
| weights, indices, offsets = inputs[0], inputs[1], inputs[2] |
| offsets_shape = _op.shape_of(offsets, dtype="int64") |
| indices_shape = _op.stack( |
| [ |
| _op.take(offsets_shape, _expr.const(0, dtype="int64")), |
| _expr.const(-1, dtype="int64"), |
| ], |
| axis=0, |
| ) |
| indices = _op.reshape(indices, indices_shape) |
| embedding = _op.take(weights, indices.astype("int64"), axis=0) |
| rembedding = reduction_fn(embedding, axis=1) |
| # EmbeddingBag has 4 outputs for some reason despite only one ever being used. |
| # Fill the rest with 0s. |
| unused_output = _expr.const(0, dtype="float32") |
| return _expr.TupleWrapper( |
| _expr.Tuple((rembedding, unused_output, unused_output, unused_output)), 4 |
| ) |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| operator = attr.get("operator", None).decode("utf-8") |
| assert operator, "ATen Operator not found" |
| return cls._op_dispatch(operator, inputs, attr, params) |
| |
| |
| class QuantizeLinear(OnnxOpConverter): |
| """Operator converter for QuantizeLinear.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| data, scale, zp = inputs |
| out_dtype = infer_type(zp).checked_type.dtype |
| return _qnn.op.quantize(data, scale, _op.cast(zp, "int32"), 0, out_dtype) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| data, scale, zp = inputs |
| out_dtype = infer_type(zp).checked_type.dtype |
| axis = attr.get("axis", 1) |
| if len(infer_shape(data)) < 2: |
| axis = 0 |
| return _qnn.op.quantize(data, scale, _op.cast(zp, "int32"), axis, out_dtype) |
| |
| |
| class DequantizeLinear(OnnxOpConverter): |
| """Operator converter for QuantizeLinear.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| data, scale, zp = inputs |
| return _qnn.op.dequantize(data, scale, _op.cast(zp, "int32"), 0) |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| data, scale, zp = inputs |
| axis = attr.get("axis", 1) |
| if len(infer_shape(data)) <= 1: |
| axis = 0 |
| return _qnn.op.dequantize(data, scale, _op.cast(zp, "int32"), axis) |
| |
| |
| class DynamicQuantizeLinear(OnnxOpConverter): |
| """Operator converter for QuantizeLinear.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| """This op is deprecated an only supports uint8""" |
| data = inputs[0] |
| data_dtype = infer_type(data).checked_type.dtype |
| zero = _op.const(0, dtype=data_dtype) |
| maximum = _op.maximum(zero, _op.max(data)) |
| minimum = _op.minimum(zero, _op.min(data)) |
| scale = (maximum - minimum) / _op.const(255, dtype=data_dtype) |
| zp = zero - _op.min(data) / scale |
| zp = _op.cast(_op.round(_op.clip(zp, 0, 255)), "uint8") |
| return _expr.TupleWrapper( |
| _expr.Tuple( |
| [_qnn.op.quantize(data, scale, _op.cast(zp, "int32"), 0, "uint8"), scale, zp] |
| ), |
| size=3, |
| ) |
| |
| |
| class QLinearConv(OnnxOpConverter): |
| """Operator converter for QLinearConv.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| data = inputs[0] |
| x_scale = get_scalar(inputs[1], params) |
| x_zero_point = get_scalar(inputs[2], params, "int32") |
| weight = inputs[3] |
| w_scale = get_scalar_or_1d_tensor(inputs[4], params) |
| w_zero_point = get_scalar_or_1d_tensor(inputs[5], params, "int32") |
| y_scale = fold_constant(get_scalar(inputs[6], params)) |
| y_zero_point = get_scalar(inputs[7], params, "int32") |
| |
| # Check shapes for per channel quantization |
| w_scale_shape = infer_shape(w_scale) |
| w_zero_point_shape = infer_shape(w_zero_point) |
| if len(w_scale_shape) == 1 or len(w_zero_point_shape) == 1: |
| m = infer_shape(weight)[0] |
| if m != w_scale_shape[0] or m != w_zero_point_shape[0]: |
| raise tvm.error.OpAttributeInvalid( |
| "The number of elements should be equal to the number of output channels" |
| ) |
| |
| input_shape = infer_shape(data) |
| |
| ndim = len(input_shape) |
| kernel_type = infer_type(weight) |
| kernel_shapes = [get_const_tuple(kernel_type.checked_type.shape)] |
| if "kernel_shape" not in attr: |
| attr["kernel_shape"] = kernel_shapes[0][2:] |
| |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: Convolution does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| zp = fold_constant(x_zero_point) |
| assert isinstance(zp, relay.Constant), "Zero point expected to be a constant" |
| data = autopad( |
| data, |
| attr.get("strides", [1] * (ndim - 2)), |
| attr["kernel_shape"], |
| attr.get("dilations", [1] * (ndim - 2)), |
| pad_value=zp.data, |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"])) |
| attr.pop("auto_pad") |
| |
| out_channels = kernel_shapes[0][0] |
| dilation = attr.get("dilations", [1] * (ndim - 2)) |
| strides = attr.get("strides", [1] * (ndim - 2)) |
| padding = attr["pads"] if "pads" in attr else 0 |
| groups = attr["group"] if "group" in attr else 1 |
| |
| if ndim != 4: |
| raise tvm.error.OpAttributeInvalid( |
| "Only 2D kernels are supported for operator QLinearConv." |
| ) |
| |
| out = _qnn.op.conv2d( |
| data, |
| weight, |
| x_zero_point, |
| w_zero_point, |
| x_scale, |
| w_scale, |
| kernel_size=attr["kernel_shape"], |
| channels=out_channels, |
| strides=strides, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| use_bias = len(inputs) == 9 |
| if use_bias: |
| out = _op.nn.bias_add(out, inputs[8]) |
| |
| out_dtype = infer_type(inputs[7]).checked_type.dtype |
| requantize_scale = _op.multiply(x_scale, w_scale) |
| |
| # requantize requires y_scale to be constant, |
| # if y_scale is not constant, doing dequantize -> quantize |
| if isinstance(y_scale, _expr.Constant): |
| out = _qnn.op.requantize( |
| out, |
| requantize_scale, |
| _op.const(0, dtype="int32"), |
| y_scale, |
| y_zero_point, |
| out_dtype=out_dtype, |
| axis=1, |
| ) |
| else: |
| out = _qnn.op.dequantize(out, requantize_scale, _op.const(0, dtype="int32"), axis=1) |
| out = _qnn.op.quantize(out, y_scale, y_zero_point, axis=1, out_dtype=out_dtype) |
| return out |
| |
| |
| class QLinearAdd(OnnxOpConverter): |
| """Operator converter for QLinearAdd from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| a = inputs[0] |
| a_scale = get_scalar(inputs[1], params) |
| a_zero_point = get_scalar(inputs[2], params, "int32") |
| b = inputs[3] |
| b_scale = get_scalar(inputs[4], params) |
| b_zero_point = get_scalar(inputs[5], params, "int32") |
| c_scale = get_scalar(inputs[6], params) |
| c_zero_point = get_scalar(inputs[7], params, "int32") |
| |
| dtype = infer_type(a).checked_type.dtype |
| |
| ## Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 |
| ## and then requantize afer |
| ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qladd.cpp |
| a = _qnn.op.dequantize( |
| inputs[0], a_scale, a_zero_point |
| ) # , c_scale, c_zero_point, out_dtype = dtype) |
| b = _qnn.op.dequantize( |
| inputs[3], b_scale, b_zero_point |
| ) # , c_scale, c_zero_point, out_dtype = dtype) |
| out = _op.add(a, b) |
| return _qnn.op.quantize(out, c_scale, c_zero_point, out_dtype=dtype) |
| |
| |
| class QLinearMatMul(OnnxOpConverter): |
| """ |
| Operator converter for QLinearMatMul from Microsoft onnxruntime contrib opset. |
| |
| Limitations: |
| - Only supports 2D input tensors. |
| - Not guaranteed to meet the integer-overflow behavior stipulated in the |
| ONNX documentation for this operator. |
| |
| The QLinearMatMul converter is re-used for MatMulInteger and is adapted for |
| the latter with the optional `expected_out_dtypes` argument. |
| """ |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params, expected_out_dtypes=None): |
| if expected_out_dtypes is None: |
| # The default QLinearMatMul converter is expected to have one of |
| # these output dtypes. |
| expected_out_dtypes = ["int8", "uint8"] |
| |
| # Some of the ops used below take scalar-like inputs, and may require either |
| # of the following: |
| # |
| # - the input is Const node (not merely an expression that *could* be reduced |
| # to a single Const at graph-compilation time) |
| # |
| # - the input has a specific dtype |
| # |
| # This function attempts to present 'x' in a form that meets both of those |
| # requirements. |
| def try_resolve_to_const(x, dtype_override=None): |
| x2 = try_resolve_var_to_const(x, params) |
| num_elem = np.prod(infer_shape(x)) |
| if num_elem == 1: |
| x2 = ensure_scalar_shape(x2) |
| x_dtype = infer_type(x).checked_type.dtype |
| if (dtype_override is not None) and (dtype_override != x_dtype): |
| x2 = _op.cast(x2, dtype_override) |
| x3 = fold_constant(x2) |
| return x3 |
| |
| # Unpack the inputs and obtain some type info... |
| a, a_scale, a_zp, b, b_scale, b_zp, y_scale, y_zp = inputs |
| |
| a_type = infer_type(a).checked_type # 'T1' in ONNX doc for this op |
| a_scale_type = infer_type(a_scale).checked_type |
| a_zp_type = infer_type(a_zp).checked_type |
| |
| b_type = infer_type(b).checked_type # 'T2' in ONNX doc for this op |
| b_scale_type = infer_type(b_scale).checked_type |
| b_zp_type = infer_type(b_zp).checked_type |
| |
| y_scale_type = infer_type(y_scale).checked_type |
| y_zp_type = infer_type(y_zp).checked_type # 'T3' in ONNX doc for this op |
| |
| a_shape = infer_shape(a) |
| b_shape = infer_shape(b) |
| |
| # Verify type assumptions, based on the ONNX doc for this op... |
| assert a_type.dtype in ["int8", "uint8"] |
| assert a_scale_type.dtype == "float32" |
| assert a_zp_type.dtype == a_type.dtype |
| |
| assert b_type.dtype in ["int8", "uint8"] |
| assert b_scale_type.dtype == "float32" |
| assert b_zp_type.dtype == b_type.dtype |
| |
| assert y_scale_type.dtype == "float32" |
| assert y_zp_type.dtype in expected_out_dtypes |
| |
| # TODO: relax this limitation in a future version of this importer. |
| a_rank = len(a_shape) |
| b_rank = len(b_shape) |
| assert (a_rank == 2) and (b_rank == 2), ( |
| "QLinearMatMul importer currently requires both 'a' and 'b' tensors to be 2D, but" |
| " rank(a)={}, rank(b)={}".format(a_rank, b_rank) |
| ) |
| |
| # _qnn.op.dense requires the zero-point values to have dtype int32. |
| a_scale_scalar = try_resolve_to_const(a_scale) |
| a_zp_scalar = try_resolve_to_const(a_zp, "int32") |
| |
| b_scale_scalar = try_resolve_to_const(b_scale) |
| b_zp_scalar = try_resolve_to_const(b_zp, "int32") |
| |
| y_scale_scalar = try_resolve_to_const(y_scale) |
| y_zp_scalar = try_resolve_to_const(y_zp, "int32") |
| |
| # TODO: Confirm that we're using 'num_hidden_units' correctly / as intended with |
| # the '_qnn.op.dense' instance below. |
| num_hidden_units = infer_shape(b)[-1] |
| |
| # - Specify the matmul result dtype as int32, so that hopefully the matmul will use |
| # a 32-bit accumulator as seems to be required by the ONNX op's documentation. |
| # |
| # TL;DR: |
| # The ONNX documentation for this op is clear about acceptable overflow |
| # behavior during the matmul operation: |
| # - The scalar multiplication ops MAY NOT overflow. |
| # - The scalar addition ops, which sum the results of the scalar multiplication, |
| # MAY overflow, but if they do so, it must behave as one would expect during |
| # 32-bit integer-addition overflow. |
| # As of this writing, Relay's qnn.op.dense operator doesn't expose a way for us to |
| # express these constraints. |
| # |
| # TODO: Extend TVM / Relay / TIR / etc. to allow this kind of constraint to be |
| # expressed in a Relay graph. And then update this importer and various TVM |
| # backends accordingly. |
| matmul_result_dtype = "int32" |
| |
| matmul_result = _qnn.op.dense( |
| a, |
| _op.transpose(b), |
| a_zp_scalar, |
| b_zp_scalar, |
| a_scale_scalar, |
| b_scale_scalar, |
| num_hidden_units, |
| matmul_result_dtype, |
| ) |
| |
| # This information might only be found in the C++ code-comments for the |
| # dense.matmul op, but the quantized tensor returned by _qnn.op.dense |
| # has scale==(a_scale_scalar * b_scale_scalar), and zero_point==0. |
| # |
| # 'matmul_result_zp_scalar' has type 'int32' to satisfy input requirements |
| # of the [de/re]quantize ops below. |
| matmul_result_scale_scalar = fold_constant(_op.multiply(a_scale_scalar, b_scale_scalar)) |
| matmul_result_zp_scalar = _op.const(0, dtype="int32") |
| |
| if "int32" in expected_out_dtypes: |
| # This is the adaptation of the QLinearMatMul converter for MatMulInteger, |
| # in the MatMulInteger case we skip the unnecessary requantization step. |
| return matmul_result |
| |
| # requantize requires y_scale to be constant, |
| # if y_scale is not constant, doing dequantize -> quantize |
| if isinstance(y_scale_scalar, _expr.Constant): |
| y = _qnn.op.requantize( |
| matmul_result, |
| matmul_result_scale_scalar, |
| matmul_result_zp_scalar, |
| y_scale_scalar, |
| y_zp_scalar, |
| axis=-1, |
| rounding="TONEAREST", |
| out_dtype=y_zp_type.dtype, |
| ) |
| else: |
| matmul_result_deq = _qnn.op.dequantize( |
| matmul_result, matmul_result_scale_scalar, matmul_result_zp_scalar, axis=0 |
| ) |
| |
| y = _qnn.op.quantize( |
| matmul_result_deq, y_scale_scalar, y_zp_scalar, axis=0, out_dtype=y_zp_type.dtype |
| ) |
| |
| return y |
| |
| |
| class MatMulInteger(OnnxOpConverter): |
| """Operator converter for MatMulInteger.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| a = inputs[0] |
| b = inputs[1] |
| |
| a_dtype = infer_type(a).checked_type.dtype |
| b_dtype = infer_type(b).checked_type.dtype |
| |
| assert a_dtype in ("int8", "uint8"), "MatMulInteger: invalid dtype for first input" |
| assert b_dtype in ("int8", "uint8"), "MatMulInteger: invalid dtype for second input" |
| |
| assert a_dtype == b_dtype, "MatMulInteger: input dtypes must match" |
| |
| a_scale = _op.const(1.0, dtype="float32") |
| b_scale = _op.const(1.0, dtype="float32") |
| out_scale = _op.const(1.0, dtype="float32") |
| |
| a_zero_point = _op.const(0.0, dtype=a_dtype) |
| b_zero_point = _op.const(0.0, dtype=b_dtype) |
| out_zero_point = _op.const(0.0, dtype="int32") |
| |
| if len(inputs) == 4: |
| a_zero_point = inputs[2] |
| b_zero_point = inputs[3] |
| |
| a_zp_dtype = infer_type(a_zero_point).checked_type.dtype |
| b_zp_dtype = infer_type(b_zero_point).checked_type.dtype |
| assert ( |
| a_zp_dtype == a_dtype and b_zp_dtype == b_dtype |
| ), "MatMulInteger: input dtype doesn't match zero point dtype" |
| elif len(inputs) != 2: |
| raise AssertionError( |
| "MatMulInteger op takes 2 or 4 inputs, {} given".format(len(inputs)) |
| ) |
| |
| inputs = [ |
| a, |
| a_scale, |
| a_zero_point, |
| b, |
| b_scale, |
| b_zero_point, |
| out_scale, |
| out_zero_point, |
| ] |
| |
| return QLinearMatMul.get_converter(10)(inputs, attr, params, expected_out_dtypes=["int32"]) |
| |
| |
| class QLinearMul(OnnxOpConverter): |
| """Operator converter for QLinearMul from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| a = inputs[0] |
| a_scale = get_scalar(inputs[1], params) |
| a_zero_point = get_scalar(inputs[2], params, "int32") |
| b = inputs[3] |
| b_scale = get_scalar(inputs[4], params) |
| b_zero_point = get_scalar(inputs[5], params, "int32") |
| y_scale = fold_constant(get_scalar(inputs[6], params)) |
| y_zero_point = get_scalar(inputs[7], params, "int32") |
| |
| dtype = infer_type(a).checked_type.dtype |
| |
| ## Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 |
| ## and then requantize afer |
| ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qlmul.cpp |
| a = _qnn.op.dequantize(inputs[0], a_scale, a_zero_point) |
| b = _qnn.op.dequantize(inputs[3], b_scale, b_zero_point) |
| out = _op.multiply(a, b) |
| return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype) |
| |
| |
| class QLinearLeakyRelu(OnnxOpConverter): |
| """Operator converter for QLinearLeakyRelu from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| |
| a_scale = get_scalar(inputs[1], params) |
| a_zero_point = get_scalar(inputs[2], params, "int32") |
| y_scale = fold_constant(get_scalar(inputs[3], params)) |
| y_zero_point = get_scalar(inputs[4], params, "int32") |
| alpha = float(attr.get("alpha", 1.0)) |
| |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| |
| # Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 |
| # and then requantize afer (according to documentation below) |
| # https://github.com/microsoft/onnxruntime/blob/master/docs/ContribOperators.md#com.microsoft.QLinearLeakyRelu |
| a = _qnn.op.dequantize(inputs[0], a_scale, a_zero_point) |
| out = _op.nn.leaky_relu(a, alpha) |
| return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype) |
| |
| |
| class QLinearSigmoid(OnnxOpConverter): |
| """Operator converter for QLinearSigmoid from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| x = inputs[0] |
| x_scale = get_scalar(inputs[1], params) |
| x_zero_point = get_scalar(inputs[2], params, "int32") |
| y_scale = fold_constant(get_scalar(inputs[3], params)) |
| y_zero_point = get_scalar(inputs[4], params, "int32") |
| |
| dtype = infer_type(x).checked_type.dtype |
| |
| ## Apparently, onnxruntime doesn't do this op in integer, they dequantize to fp32 |
| ## and then requantize after: |
| ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/ |
| ## providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp#L245 |
| x = _qnn.op.dequantize(x, x_scale, x_zero_point) |
| out = _op.sigmoid(x) |
| return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype) |
| |
| |
| class QLinearConcat(OnnxOpConverter): |
| """Operator converter for QLinearConcat from Microsoft onnxruntime contrib opset.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| # which axis to concat on |
| axis = attr["axis"] |
| |
| y_scale = fold_constant(get_scalar(inputs[0], params)) |
| y_zero_point = get_scalar(inputs[1], params, "int32") |
| |
| # input tensors, scales, zero_points |
| assert ( |
| len(inputs) % 3 == 2 |
| ), "Additional input count must be a multiple of 3 -- tensor/scale/zero_point tuples" |
| tensors = [] |
| scales = [] |
| zero_points = [] |
| for i in range(2, len(inputs), 3): |
| tensors.append(inputs[i]) |
| scales.append(get_scalar(inputs[i + 1], params)) |
| zero_points.append(get_scalar(inputs[i + 2], params, "int32")) |
| |
| return _qnn.op.concatenate(tensors, scales, zero_points, y_scale, y_zero_point, axis) |
| |
| |
| class ConvInteger(OnnxOpConverter): |
| """Operator converter for ConvInteger.""" |
| |
| @classmethod |
| def _impl_v10(cls, inputs, attr, params): |
| data = inputs[0] |
| weight = inputs[1] |
| data_zp = inputs[2] |
| weight_zp = inputs[3] |
| if data_zp is None: |
| data_zp = _expr.const(0, "int32") |
| if weight_zp is None: |
| weight_zp = _expr.const(0, "int32") |
| |
| input_type = infer_type(data) |
| input_shape = get_const_tuple(input_type.checked_type.shape) |
| |
| ndim = len(input_shape) |
| kernel_type = infer_type(weight) |
| kernel_shape = get_const_tuple(kernel_type.checked_type.shape) |
| if "kernel_shape" not in attr: |
| attr["kernel_shape"] = kernel_shape[2:] |
| |
| if "auto_pad" in attr: |
| attr["auto_pad"] = attr["auto_pad"].decode("utf-8") |
| if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"): |
| # Warning: Convolution does not yet support dynamic shapes, |
| # one will need to run dynamic_to_static on this model after import |
| data = autopad( |
| data, |
| attr.get("strides", [1] * (ndim - 2)), |
| attr["kernel_shape"], |
| attr.get("dilations", [1] * (ndim - 2)), |
| pad_value=data_zp, |
| mode=attr["auto_pad"], |
| ) |
| elif attr["auto_pad"] == "VALID": |
| attr["pads"] = tuple([0 for i in range(ndim - 2)]) |
| elif attr["auto_pad"] == "NOTSET": |
| pass |
| else: |
| msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' |
| raise tvm.error.OpAttributeInvalid(msg.format(attr["auto_pad"])) |
| attr.pop("auto_pad") |
| |
| out_channels = kernel_shape[0] |
| dilation = attr.get("dilations", [1] * (ndim - 2)) |
| strides = attr.get("strides", [1] * (ndim - 2)) |
| padding = attr["pads"] if "pads" in attr else 0 |
| groups = attr["group"] if "group" in attr else 1 |
| |
| if ndim != 4: |
| raise tvm.error.OpAttributeInvalid( |
| "Only 2D kernels are supported for operator ConvInteger." |
| ) |
| |
| return _qnn.op.conv2d( |
| data, |
| weight, |
| _op.cast(data_zp, "int32"), |
| _op.cast(weight_zp, "int32"), |
| _expr.const(1.0, "float32"), |
| _expr.const(1.0, "float32"), |
| kernel_size=attr["kernel_shape"], |
| channels=out_channels, |
| strides=strides, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| |
| |
| class BitShift(OnnxOpConverter): |
| """Operator converter for NonZero""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| if len(inputs) != 2: |
| raise ValueError("Bitshift expects 2 inputs") |
| |
| direction = attr.get("direction", "LEFT").decode("ascii") |
| if direction == "LEFT": |
| out = _op.left_shift(*inputs) |
| elif direction == "RIGHT": |
| out = _op.right_shift(*inputs) |
| else: |
| raise ValueError("Unsupported Shift Direction: " + direction) |
| return out |
| |
| |
| class Unique(OnnxOpConverter): |
| """Operator converter for unique""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| if len(inputs) != 1: |
| raise ValueError("Unique expects 1 input") |
| |
| data = inputs[0] |
| axis = attr.get("axis", None) |
| if axis is None: # If axis is None, flatten the input before calling unique |
| data = _op.reshape(data, _op.const([-1])) |
| else: |
| data_shape = infer_shape(data) |
| if len(data_shape) != 1: |
| raise ValueError("TVM only supports 1D Unique operator.") |
| is_sorted = attr.get("sorted", 1) # sorted is 0 or 1, 1 by default |
| |
| # ONNX documentation lists return_counts as optional but there is no input to specify |
| # whether it is returned. Therefore we'll just always return it. |
| unique = _op.unique(data, is_sorted=(is_sorted == 1), return_counts=True) |
| num_unique = unique[3] |
| |
| trim_unique_lambda = lambda input: _op.strided_slice(input, _op.const([0]), num_unique) |
| |
| unique_vals = trim_unique_lambda(unique[0]) |
| indices = _op.cast(trim_unique_lambda(unique[1]), "int64") # ONNX always returns int64 |
| inverse_indices = _op.cast(unique[2], "int64") # ONNX always returns int64 |
| counts = _op.cast(trim_unique_lambda(unique[4]), "int64") # ONNX always returns int64 |
| # ONNX unique returns unique, indices, inverse_indices, (optional) counts |
| return _expr.TupleWrapper(_expr.Tuple([unique_vals, indices, inverse_indices, counts]), 4) |
| |
| |
| class Einsum(OnnxOpConverter): |
| """Operator converter for Einsum""" |
| |
| @classmethod |
| def _impl_v12(cls, inputs, attr, params): |
| equation = attr["equation"].decode("utf-8") |
| return _op.einsum(inputs, equation) |
| |
| |
| class Trilu(OnnxOpConverter): |
| """Operator converter for Trilu""" |
| |
| @classmethod |
| def _impl_v14(cls, inputs, attr, params): |
| upper = attr.get("upper", True) |
| if len(inputs) == 2: |
| data, k = inputs |
| else: |
| data = inputs[0] |
| k = 0 |
| return _op.trilu(data, k, upper) |
| |
| |
| class RandomNormal(OnnxOpConverter): |
| """Operator converter for random_normal""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = get_type(attr.get("dtype", 1)) |
| mean = attr.get("mean", 0.0) |
| scale = attr.get("scale", 1.0) |
| seed = attr.get("seed", None) |
| shape = attr["shape"] |
| |
| assert dtype in [ |
| "float32", |
| "float64", |
| ], "Only float random value generation is currently supported." |
| |
| if seed is None: |
| seed = np.random.randint(1e6) |
| else: |
| seed = int(seed) |
| key = _random.threefry_key(seed) |
| output = _op.random.normal(key, shape, dtype=dtype, mean=mean, scale=scale) |
| _, vals = _expr.TupleWrapper(output, 2) |
| return vals |
| |
| |
| class RandomNormalLike(OnnxOpConverter): |
| """Operator converter for random_normal_like""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = attr.get("dtype", None) |
| scale = attr.get("scale", 1.0) |
| mean = attr.get("mean", 0.0) |
| seed = attr.get("seed", None) |
| shape = infer_shape(inputs[0]) |
| if dtype is None: |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| else: |
| dtype = get_type(dtype) |
| |
| assert dtype in [ |
| "float32", |
| "float64", |
| ], "Only float random value generation is currently supported." |
| |
| if seed is None: |
| seed = np.random.randint(1e6) |
| else: |
| seed = int(seed) |
| key = _random.threefry_key(seed) |
| output = _op.random.normal(key, shape, dtype=dtype, mean=mean, scale=scale) |
| _, vals = _expr.TupleWrapper(output, 2) |
| return vals |
| |
| |
| class RandomUniform(OnnxOpConverter): |
| """Operator converter for random_uniform""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = get_type(attr.get("dtype", 1)) |
| high = attr.get("high", 1.0) |
| low = attr.get("low", 0.0) |
| seed = attr.get("seed", None) |
| shape = attr["shape"] |
| |
| assert dtype in [ |
| "float32", |
| "float64", |
| ], "Only float random value generation is currently supported." |
| |
| if seed is None: |
| seed = np.random.randint(1e6) |
| else: |
| seed = int(seed) |
| key = _random.threefry_key(seed) |
| output = _op.random.uniform(key, shape, dtype=dtype, low=low, high=high) |
| _, vals = _expr.TupleWrapper(output, 2) |
| return vals |
| |
| |
| class RandomUniformLike(OnnxOpConverter): |
| """Operator converter for random_uniform_like""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| dtype = attr.get("dtype", None) |
| high = attr.get("high", 1.0) |
| low = attr.get("low", 0.0) |
| seed = attr.get("seed", None) |
| shape = infer_shape(inputs[0]) |
| if dtype is None: |
| dtype = infer_type(inputs[0]).checked_type.dtype |
| else: |
| dtype = get_type(dtype) |
| |
| assert dtype in [ |
| "float32", |
| "float64", |
| ], "Only float random value generation is currently supported." |
| |
| if seed is None: |
| seed = np.random.randint(1e6) |
| else: |
| seed = int(seed) |
| key = _random.threefry_key(seed) |
| output = _op.random.uniform(key, shape, dtype=dtype, low=low, high=high) |
| _, vals = _expr.TupleWrapper(output, 2) |
| return vals |
| |
| |
| class Multinomial(OnnxOpConverter): |
| """Operator converter for multinomial""" |
| |
| @classmethod |
| def _impl_v7(cls, inputs, attr, params): |
| dtype = attr.get("dtype", "int64") |
| sample_size = attr.get("sample_size", 1) |
| seed = attr.get("seed", None) |
| if seed is None: |
| seed = np.random.randint(1e6) |
| key = _op.random.threefry_key(seed) |
| output = _op.random.multinomial(key, inputs[0], sample_size) |
| _, indices = _expr.TupleWrapper(output, 2) |
| return _op.cast(indices, get_type(dtype)) |
| |
| |
| class NegativeLogLikelihoodLoss(OnnxOpConverter): |
| """Operator converter for NegativeLogLikehoodLoss""" |
| |
| VALID_REDUCTIONS = {"mean", "sum", "none"} |
| |
| @classmethod |
| def run_calculation( |
| cls: "NegativeLogLikelihoodLoss", |
| input_tensor: relay.Expr, |
| target_tensor: relay.Expr, |
| weight_tensor: Optional[relay.Expr], |
| ignore_index: int, |
| ): |
| """Run calculation for NegativeLogLikelihood, returning output tensor and |
| weight tensor used for mean-style reductions. |
| """ |
| # Convert negative indices --> positive indices for gather ops, note we have to |
| # use the original target tensor to interact with ignore_index to have proper behavior. |
| normalized_target_tensor = normalize_gather_indices(input_tensor, target_tensor, 1) |
| |
| if weight_tensor is None: |
| channels = infer_shape(input_tensor)[1] |
| weight_tensor = relay.ones( |
| [channels], |
| dtype=infer_type(input_tensor).checked_type.dtype, |
| ) |
| |
| loss = -relay.gather( |
| input_tensor, |
| axis=1, |
| indices=relay.expand_dims(normalized_target_tensor, 1), |
| ) |
| loss = relay.squeeze(loss, axis=[1]) |
| |
| expanded_normalized_target_tensor = relay.expand_dims(normalized_target_tensor, 0) |
| expanded_normalized_target_tensor = relay.nn.batch_flatten( |
| expanded_normalized_target_tensor |
| ) |
| flattened_weights = relay.gather_nd(weight_tensor, expanded_normalized_target_tensor) |
| select_weights = relay.reshape_like(flattened_weights, loss) |
| loss *= select_weights |
| |
| if ignore_index is not None: |
| # "Ignore" values whose target is the ignore_index |
| mask_tensor = relay.equal( |
| target_tensor, relay.const(ignore_index, dtype=target_tensor.type_annotation.dtype) |
| ) |
| mask_tensor = relay.const(1, dtype="int8") - relay.cast(mask_tensor, "int8") |
| loss = relay.where( |
| mask_tensor, loss, relay.const(0, infer_type(loss).checked_type.dtype) |
| ) |
| |
| # This is not explained super clearly in the onnx spec, but masked values don't |
| # contribute toward the final value in reduction |
| select_weights *= relay.cast_like(mask_tensor, select_weights) |
| |
| weight_total = relay.sum(select_weights) |
| return loss, weight_total |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| ignore_index = attr.get("ignore_index", None) |
| reduction = attr.get("reduction", b"mean").decode("utf-8") |
| |
| if reduction not in cls.VALID_REDUCTIONS: |
| raise ValueError( |
| f"Unknown reduction type {reduction}, choices are {cls.VALID_REDUCTIONS}" |
| ) |
| |
| input_tensor, target_tensor = inputs[0], inputs[1] |
| if len(inputs) == 3: |
| weight_tensor = inputs[2] |
| else: |
| weight_tensor = None |
| |
| loss, weight_total = cls.run_calculation( |
| input_tensor, |
| target_tensor, |
| weight_tensor=weight_tensor, |
| ignore_index=ignore_index, |
| ) |
| if reduction == "mean": |
| return relay.sum(loss) / weight_total |
| if reduction == "sum": |
| return relay.sum(loss) |
| # Case reduction == 'none' |
| return loss |
| |
| |
| class SoftmaxCrossEntropyLoss(OnnxOpConverter): |
| """Operator converter for SCE_loss""" |
| |
| @classmethod |
| def _impl_v13(cls, inputs, attr, params): |
| ignore_index = attr.get("ignore_index", None) |
| reduction = attr.get("reduction", b"mean").decode("utf-8") |
| input_tensor, target_tensor = inputs[0], inputs[1] |
| if len(inputs) == 3: |
| weight_tensor = inputs[2] |
| else: |
| weight_tensor = None |
| |
| get_log_prob = attr["tvm_custom"]["num_outputs"] == 2 |
| log_softmax_attr = {"axis": 1} |
| log_softmax_tensor = LogSoftmax.get_converter(13)([input_tensor], log_softmax_attr, None) |
| |
| loss, weight_total = NegativeLogLikelihoodLoss.run_calculation( |
| log_softmax_tensor, |
| target_tensor, |
| weight_tensor, |
| ignore_index=ignore_index, |
| ) |
| |
| if reduction == "mean": |
| loss = relay.sum(loss) / weight_total |
| elif reduction == "sum": |
| loss = relay.sum(loss) |
| |
| if get_log_prob: |
| return relay.TupleWrapper(relay.Tuple((loss, log_softmax_tensor)), 2) |
| return loss |
| |
| |
| class Adagrad(OnnxOpConverter): |
| """Operator converter for adagrad op.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| decay_factor = attr.get("decay_factor", 0.0) |
| epsilon = attr.get("epsilon", 0.0) |
| norm_coefficient = attr.get("norm_coefficient", 0.0) |
| |
| R = inputs[0] |
| T = inputs[1] |
| |
| # convert attributes to constants, proper types |
| dtype_inputs = infer_type(inputs[3]).checked_type.dtype |
| decay_factor = relay.const(decay_factor, dtype=dtype_inputs) |
| epsilon = relay.const(epsilon, dtype=dtype_inputs) |
| norm_coefficient = relay.const(norm_coefficient, dtype=dtype_inputs) |
| T = relay.cast_like(T, inputs[3]) |
| |
| assert ( |
| len(inputs) - 2 |
| ) % 3 == 0, f"Expect triplets for remaining inputs, found {len(inputs) - 2}" |
| |
| # Remaining inputs are: |
| # [x_1, x_2 ..., x_1_gradient, x_2_gradient, ... x_1_sq_g, x_2_sq_g...] |
| num_input_tensors = (len(inputs) - 2) // 3 |
| output_tensors = [] |
| output_accumulated_squared_gradients = [] |
| for i in range(num_input_tensors): |
| x = inputs[i + 2] |
| gradient = inputs[i + 2 + num_input_tensors] |
| accumulated_squared_gradient = inputs[i + 2 + 2 * num_input_tensors] |
| |
| r = R / (relay.const(1.0, dtype=dtype_inputs) + T * decay_factor) |
| g_regularized = norm_coefficient * x + gradient |
| new_accumulated_squared_gradient = ( |
| accumulated_squared_gradient + g_regularized * g_regularized |
| ) |
| h_adaptive = relay.sqrt(new_accumulated_squared_gradient) + epsilon |
| |
| x_new = x - r * g_regularized / h_adaptive |
| |
| output_tensors.append(x_new) |
| output_accumulated_squared_gradients.append(new_accumulated_squared_gradient) |
| |
| # append lists together, momentums come after result tensors |
| result = output_tensors + output_accumulated_squared_gradients |
| return _expr.TupleWrapper(_expr.Tuple(result), len(result)) |
| |
| |
| class Adam(OnnxOpConverter): |
| """Operator converter for Adam op.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = attr.get("alpha", 0.9) |
| beta = attr.get("beta", 0.999) |
| |
| # Note in the docs epsilon default is 0.0 but in the tests it is set to 1e-2: |
| # https://git.io/Ju5C4 |
| epsilon = attr.get("epsilon", 1e-2) |
| norm_coefficient = attr.get("norm_coefficient", 0.0) |
| norm_coefficient_post = attr.get("norm_coefficient_post", 0.0) |
| |
| R = inputs[0] |
| T = inputs[1] |
| |
| assert ( |
| len(inputs) - 2 |
| ) % 4 == 0, f"Expect 4-lets for remaining inputs, found {len(inputs) - 2}" |
| |
| # convert attributes to constants, proper types |
| dtype_inputs = infer_type(inputs[3]).checked_type.dtype |
| inverse_alpha = relay.const(1 - alpha, dtype=dtype_inputs) |
| alpha = relay.const(alpha, dtype=dtype_inputs) |
| inverse_beta = relay.const(1 - beta, dtype=dtype_inputs) |
| beta = relay.const(beta, dtype=dtype_inputs) |
| epsilon = relay.const(epsilon, dtype=dtype_inputs) |
| norm_coefficient = relay.const(norm_coefficient, dtype=dtype_inputs) |
| norm_coefficient_post = relay.const(norm_coefficient_post, dtype=dtype_inputs) |
| one = relay.const(1, dtype=dtype_inputs) |
| T = relay.cast_like(T, inputs[3]) |
| |
| # Remaining inputs are: |
| # [x_1, x_2 ..., x_1_grad, x_2_grad, ... x_1_g_accum, x_2_g_accum..., x_1_g_sq_accum, ...] |
| num_input_tensors = (len(inputs) - 2) // 4 |
| output_tensors = [] |
| output_accumulated_gradients = [] |
| output_accumulated_squared_gradients = [] |
| for i in range(num_input_tensors): |
| x = inputs[i + 2] |
| g = inputs[i + 2 + num_input_tensors] |
| v = inputs[i + 2 + 2 * num_input_tensors] |
| h = inputs[i + 2 + 3 * num_input_tensors] |
| |
| g_regularized = norm_coefficient * x + g |
| v_new = alpha * v + inverse_alpha * g_regularized |
| h_new = beta * h + inverse_beta * g_regularized * g_regularized |
| h_sqrt = relay.sqrt(h_new) + epsilon |
| |
| true_branch = R * relay.sqrt(one - relay.power(beta, T)) / (one - relay.power(alpha, T)) |
| R_adjusted = relay.If(T > relay.const(0, dtype=dtype_inputs), true_branch, R) |
| |
| x_new = x - R_adjusted * (v_new / h_sqrt) |
| x_result = (one - norm_coefficient_post) * x_new |
| |
| output_tensors.append(x_result) |
| output_accumulated_gradients.append(v_new) |
| output_accumulated_squared_gradients.append(h_new) |
| |
| # append lists together to get final result |
| result = ( |
| output_tensors + output_accumulated_gradients + output_accumulated_squared_gradients |
| ) |
| return _expr.TupleWrapper(_expr.Tuple(result), len(result)) |
| |
| |
| class Momentum(OnnxOpConverter): |
| """Operator converter for Momentum op.""" |
| |
| @classmethod |
| def _impl_v1(cls, inputs, attr, params): |
| alpha = attr["alpha"] |
| beta = attr["beta"] |
| mode = attr["mode"].decode("utf-8") |
| norm_coefficient = attr["norm_coefficient"] |
| |
| assert mode in ["nesterov", "standard"], f"Unknown momentum mode {mode}" |
| R = inputs[0] |
| T = inputs[1] |
| |
| assert ( |
| len(inputs) - 2 |
| ) % 3 == 0, f"Expect triplets for remaining inputs, found {len(inputs) - 2}" |
| # Remaining inputs are: |
| # [x_1, x_2 ..., x_1_gradient, x_2_gradient, ... x_1_momentum, x_2_momentum...] |
| num_input_tensors = (len(inputs) - 2) // 3 |
| |
| # convert attributes to constants |
| dtype_inputs = infer_type(inputs[3]).checked_type.dtype |
| alpha = relay.const(alpha, dtype=dtype_inputs) |
| beta = relay.const(beta, dtype=dtype_inputs) |
| norm_coefficient = relay.const(norm_coefficient, dtype=dtype_inputs) |
| default_beta = relay.const(1.0, dtype=dtype_inputs) |
| |
| # Calculate updated values for every input |
| output_tensors = [] |
| output_momentums = [] |
| for i in range(num_input_tensors): |
| x = inputs[i + 2] |
| gradient = inputs[i + 2 + num_input_tensors] |
| momentum = inputs[i + 2 + 2 * num_input_tensors] |
| g_regularized = norm_coefficient * x + gradient |
| beta_adjusted = relay.If(T > relay.const(0, dtype="int64"), beta, default_beta) |
| new_momentum = alpha * momentum + beta_adjusted * g_regularized |
| |
| if mode == "standard": |
| x_output = x - R * new_momentum |
| else: |
| # mode == 'nesterov' |
| x_output = x - R * (g_regularized + alpha * new_momentum) |
| |
| output_tensors.append(x_output) |
| output_momentums.append(new_momentum) |
| |
| # append lists together, momentums come after result tensors |
| result = output_tensors + output_momentums |
| return _expr.TupleWrapper(_expr.Tuple(result), len(result)) |
| |
| |
| class Round(OnnxOpConverter): |
| """Operator converter for round op.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| # Onnx round uses Banker's rounding which rounds .5 to the nearest even integer |
| |
| x = inputs[0] |
| dtype = infer_type(x).checked_type.dtype |
| half = _expr.const(0.5, dtype=dtype) |
| one = _expr.const(1, dtype=dtype) |
| two = _expr.const(2, dtype=dtype) |
| |
| rounded = _op.ceil(x - half) |
| bankers_mask = one - (_op.ceil(x + half) - _op.floor(x + half)) |
| non_even = _op.abs(_op.mod(rounded, two)) |
| return rounded + (bankers_mask * non_even) |
| |
| |
| class SequenceConstruct(OnnxOpConverter): |
| """Operator converter for sequence construction op.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| # Construct a tuple from input tensors. |
| return _expr.Tuple(inputs) |
| |
| |
| class SequenceInsert(OnnxOpConverter): |
| """Operator converter for sequence insert op.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| # Insert a new tensor into a tuple of tensors. |
| input_sequence = inputs[0] |
| new_tensor = inputs[1] |
| |
| if len(inputs) == 3: |
| position = inputs[2] |
| # Non constant position is not supported. |
| if isinstance(position, _expr.Constant): |
| position = position.data.numpy() |
| elif position.name_hint in params: |
| position = params[position.name_hint].numpy() |
| else: |
| raise NotImplementedError("Position must be a constant.") |
| else: |
| position = -1 |
| |
| if position < 0: |
| position = len(input_sequence) + position + 1 |
| # Convert sequence to a list, insert new tensor, and repackage as Tuple. |
| tensor_list = [input_sequence[i] for i in range(len(input_sequence))] |
| # Insert new tensor. |
| tensor_list.insert(position, new_tensor) |
| # Create new tuple and return. |
| return _expr.Tuple(tensor_list) |
| |
| |
| class ConcatFromSequence(OnnxOpConverter): |
| """Operator converter for sequence concatenation op.""" |
| |
| @classmethod |
| def _impl_v11(cls, inputs, attr, params): |
| axis = attr.get("axis", 0) |
| new_axis = attr.get("new_axis", 0) |
| |
| # If a new axis should be created, just stack input tensors. |
| if new_axis == 1: |
| return _op.stack(inputs[0], axis=axis) |
| |
| return _op.concatenate(inputs[0], axis=axis) |
| |
| |
| # compatible operators that do NOT require any conversion. |
| _identity_list = [] |
| |
| # _convert_map defines maps of name to converter functor(callable) |
| # for 1 to 1 mapping, use Renamer if nothing but name is different |
| # use AttrCvt if attributes need to be converted |
| # for 1 to N mapping(composed), use custom callable functions |
| # for N to 1 mapping, currently not supported(?) |
| def _get_convert_map(opset): |
| return { |
| # defs/experimental |
| "Identity": Renamer("copy"), |
| "Affine": Affine.get_converter(opset), |
| "BitShift": BitShift.get_converter(opset), |
| "ThresholdedRelu": ThresholdedRelu.get_converter(opset), |
| "ScaledTanh": ScaledTanh.get_converter(opset), |
| "ParametricSoftplus": ParametricSoftPlus.get_converter(opset), |
| "Constant": Constant.get_converter(opset), |
| "ConstantOfShape": ConstantOfShape.get_converter(opset), |
| # 'GivenTensorFill' |
| "FC": AttrCvt("dense", ignores=["axis", "axis_w"]), |
| "Scale": Scale.get_converter(opset), |
| # 'GRUUnit' |
| # 'ATen' |
| # 'ImageScaler' |
| "MeanVarianceNormalization": MeanVarianceNormalization.get_converter(opset), |
| # 'Crop' |
| # 'Embedding' |
| "Upsample": Upsample.get_converter(opset), |
| "SpatialBN": BatchNorm.get_converter(opset), |
| # defs/generator |
| # 'Constant' # Implemented |
| # 'RandomUniform' |
| # 'RandomNormal' |
| # 'RandomUniformLike' |
| # 'RandomNormalLike' |
| # defs/logical |
| # defs/math |
| "Add": Add.get_converter(opset), |
| "Sub": Sub.get_converter(opset), |
| "Mul": Mul.get_converter(opset), |
| "Div": Div.get_converter(opset), |
| "Neg": Renamer("negative"), |
| "Abs": Absolute.get_converter(opset), |
| "Reciprocal": Reciprocal.get_converter(opset), |
| "Floor": Renamer("floor"), |
| "Ceil": Renamer("ceil"), |
| "Round": Round.get_converter(opset), |
| "IsInf": IsInf.get_converter(opset), |
| "IsNaN": Renamer("isnan"), |
| "Sqrt": Renamer("sqrt"), |
| "Relu": Renamer("relu"), |
| "Celu": Celu.get_converter(opset), |
| "LeakyRelu": Renamer("leaky_relu"), |
| "Selu": Selu.get_converter(opset), |
| "Elu": Elu.get_converter(opset), |
| "Gelu": Gelu.get_converter(opset), |
| "BiasGelu": BiasGelu.get_converter(opset), |
| # TODO: We need a better way to handle different domains, in case |
| # of name collisions. EmbedLayerNormalization, SkipLayerNormalization, and Attention |
| # are in the `com.microsoft` domain. |
| "EmbedLayerNormalization": EmbedLayerNormalization.get_converter(opset), |
| "SkipLayerNormalization": SkipLayerNormalization.get_converter(opset), |
| "Attention": Attention.get_converter(opset), |
| "Exp": Renamer("exp"), |
| "Greater": Renamer("greater"), |
| "GreaterOrEqual": Renamer("greater_equal"), |
| "Less": Renamer("less"), |
| "LessOrEqual": Renamer("less_equal"), |
| "Log": Renamer("log"), |
| "Acos": Renamer("acos"), |
| "Acosh": Renamer("acosh"), |
| "Asin": Renamer("asin"), |
| "Asinh": Renamer("asinh"), |
| "Atan": Renamer("atan"), |
| "Atanh": Renamer("atanh"), |
| "Cos": Renamer("cos"), |
| "Cosh": Renamer("cosh"), |
| "Sin": Renamer("sin"), |
| "Sinh": Renamer("sinh"), |
| "Tan": Renamer("tan"), |
| "Tanh": Renamer("tanh"), |
| "Pow": Pow.get_converter(opset), |
| "PRelu": Prelu.get_converter(opset), |
| "Sigmoid": Renamer("sigmoid"), |
| "HardSigmoid": HardSigmoid.get_converter(opset), |
| "HardSwish": HardSwish.get_converter(opset), |
| "Max": Maximum.get_converter(opset), |
| "Min": Minimum.get_converter(opset), |
| "Sum": Sum.get_converter(opset), |
| "Mean": Mean.get_converter(opset), |
| "Clip": Clip.get_converter(opset), |
| "Softplus": Softplus.get_converter(opset), |
| # softmax default axis is different in onnx |
| "Softmax": Softmax.get_converter(opset), |
| "LogSoftmax": LogSoftmax.get_converter(opset), |
| "OneHot": OneHot.get_converter(opset), |
| "Hardmax": Hardmax.get_converter(opset), |
| "Shrink": Shrink.get_converter(opset), |
| "Softsign": Softsign.get_converter(opset), |
| "Gemm": Gemm.get_converter(opset), |
| "MatMul": MatMul.get_converter(opset), |
| "MatMulInteger": MatMulInteger.get_converter(opset), |
| "MatMulInteger16": MatMulInteger16.get_converter(opset), |
| "Mod": Mod.get_converter(opset), |
| "Xor": Renamer("logical_xor"), |
| # defs/nn |
| "AveragePool": AveragePool.get_converter(opset), |
| "LpPool": LpPool.get_converter(opset), |
| "GlobalLpPool": GlobalLpPool.get_converter(opset), |
| "MaxPool": MaxPool.get_converter(opset), |
| "MaxUnpool": MaxUnpool.get_converter(opset), |
| "Conv": Conv.get_converter(opset), |
| "ConvTranspose": ConvTranspose.get_converter(opset), |
| "GlobalAveragePool": GlobalAveragePool.get_converter(opset), |
| "GlobalMaxPool": GlobalMaxPool.get_converter(opset), |
| "BatchNormalization": BatchNorm.get_converter(opset), |
| "InstanceNormalization": InstanceNorm.get_converter(opset), |
| # 'LpNormalization' |
| "Dropout": AttrCvt("dropout", {"ratio": "rate"}, ignores=["is_test"]), |
| "Flatten": Flatten.get_converter(opset), |
| "LRN": LRN.get_converter(opset), |
| # Recurrent Layers |
| "RNN": RNN.get_converter(opset), |
| "LSTM": LSTM.get_converter(opset), |
| "GRU": GRU.get_converter(opset), |
| # defs/vision |
| "MaxRoiPool": MaxRoiPool.get_converter(opset), |
| "RoiAlign": RoiAlign.get_converter(opset), |
| "NonMaxSuppression": NonMaxSuppression.get_converter(opset), |
| # defs/reduction |
| "ReduceMax": ReduceMax.get_converter(opset), |
| "ReduceMin": ReduceMin.get_converter(opset), |
| "ReduceSum": ReduceSum.get_converter(opset), |
| "ReduceMean": ReduceMean.get_converter(opset), |
| "ReduceProd": ReduceProd.get_converter(opset), |
| "ReduceLogSumExp": ReduceLogSumExp.get_converter(opset), |
| "ReduceLogSum": ReduceLogSum.get_converter(opset), |
| "ReduceSumSquare": ReduceSumSquare.get_converter(opset), |
| "ReduceL1": ReduceL1.get_converter(opset), |
| "ReduceL2": ReduceL2.get_converter(opset), |
| # defs/sorting |
| "ArgMax": ArgMax.get_converter(opset), |
| "ArgMin": ArgMin.get_converter(opset), |
| "TopK": TopK.get_converter(opset), |
| # defs/tensor |
| "Cast": Cast.get_converter(opset), |
| "Reshape": Reshape.get_converter(opset), |
| "Expand": Expand.get_converter(opset), |
| "Concat": Concat.get_converter(opset), |
| "Split": Split.get_converter(opset), |
| "Slice": Slice.get_converter(opset), |
| "Transpose": AttrCvt("transpose", {"perm": "axes"}), |
| "DepthToSpace": DepthToSpace.get_converter(opset), |
| "SpaceToDepth": SpaceToDepth.get_converter(opset), |
| "Gather": Gather.get_converter(opset), |
| "GatherElements": GatherElements.get_converter(opset), |
| "GatherND": GatherND.get_converter(opset), |
| "Compress": Compress.get_converter(opset), |
| "Size": AttrCvt("ndarray_size", extras={"dtype": "int64"}), |
| "Scatter": Scatter.get_converter(opset), |
| "ScatterElements": Scatter.get_converter(opset), |
| "ScatterND": ScatterND.get_converter(opset), |
| "EyeLike": EyeLike.get_converter(opset), |
| "Squeeze": Squeeze.get_converter(opset), |
| "Unsqueeze": Unsqueeze.get_converter(opset), |
| "Pad": Pad.get_converter(opset), |
| "Shape": Shape.get_converter(opset), |
| "Sign": Sign.get_converter(opset), |
| "Equal": Equal.get_converter(opset), |
| "Not": Not.get_converter(opset), |
| "And": And.get_converter(opset), |
| "Tile": Tile.get_converter(opset), |
| "Erf": Erf.get_converter(opset), |
| "Where": Where.get_converter(opset), |
| "Or": Or.get_converter(opset), |
| "Resize": Resize.get_converter(opset), |
| "NonZero": NonZero.get_converter(opset), |
| "Range": Range.get_converter(opset), |
| "CumSum": CumSum.get_converter(opset), |
| "Unique": Unique.get_converter(opset), |
| "Einsum": Einsum.get_converter(opset), |
| "Trilu": Trilu.get_converter(opset), |
| # defs/control_flow |
| "Loop": Loop.get_converter(opset), |
| "If": If.get_converter(opset), |
| # Torch ATen Dispatcher. |
| "ATen": ATen.get_converter(opset), |
| # Quantization |
| "QuantizeLinear": QuantizeLinear.get_converter(opset), |
| "DequantizeLinear": DequantizeLinear.get_converter(opset), |
| "DynamicQuantizeLinear": DynamicQuantizeLinear.get_converter(opset), |
| "ReverseSequence": ReverseSequence.get_converter(opset), |
| "QLinearConv": QLinearConv.get_converter(opset), |
| "QLinearConcat": QLinearConcat.get_converter(opset), |
| "QLinearAdd": QLinearAdd.get_converter(opset), |
| "QLinearMatMul": QLinearMatMul.get_converter(opset), |
| "QLinearMul": QLinearMul.get_converter(opset), |
| "QLinearSigmoid": QLinearSigmoid.get_converter(opset), |
| "ConvInteger": ConvInteger.get_converter(opset), |
| "QLinearAveragePool": QLinearAveragePool.get_converter(opset), |
| "QLinearGlobalAveragePool": QLinearGlobalAveragePool.get_converter(opset), |
| "QLinearLeakyRelu": QLinearLeakyRelu.get_converter(opset), |
| # Random number generation. |
| "RandomNormal": RandomNormal.get_converter(opset), |
| "RandomNormalLike": RandomNormalLike.get_converter(opset), |
| "RandomUniform": RandomUniform.get_converter(opset), |
| "RandomUniformLike": RandomUniformLike.get_converter(opset), |
| "Multinomial": Multinomial.get_converter(opset), |
| # Loss functions / training |
| "NegativeLogLikelihoodLoss": NegativeLogLikelihoodLoss.get_converter(opset), |
| "SoftmaxCrossEntropyLoss": SoftmaxCrossEntropyLoss.get_converter(opset), |
| "Adagrad": Adagrad.get_converter(opset), |
| "Adam": Adam.get_converter(opset), |
| "Momentum": Momentum.get_converter(opset), |
| "Scan": Scan.get_converter(opset), |
| # ML |
| "LinearRegressor": LinearRegressor.get_converter(opset), |
| # Sequence operators |
| "SequenceConstruct": SequenceConstruct.get_converter(opset), |
| "SequenceInsert": SequenceInsert.get_converter(opset), |
| "ConcatFromSequence": ConcatFromSequence.get_converter(opset), |
| } |
| |
| |
| class GraphProto: |
| """A helper class for handling Relay expression copying from pb2.GraphProto. |
| Definition: https://github.com/onnx/onnx/blob/main/onnx/onnx.proto |
| |
| Parameters |
| ---------- |
| shape : dict of str to tuple, optional |
| The input shape to the graph |
| |
| dtype : str or dict of str to str |
| The input types to the graph |
| |
| freeze_params: bool |
| If this parameter is true, the importer will take any provided |
| onnx input values (weights, shapes, etc) and embed them into the relay model |
| as Constants instead of variables. This allows more aggressive optimizations |
| at compile time and helps in making models static if certain inputs represent |
| attributes relay would traditionally consider compile-time constants. |
| |
| """ |
| |
| current = None |
| |
| def __init__(self, shape, dtype, freeze_params=False): |
| self._nodes = {} |
| self._params = {} |
| self._inputs = {} |
| self._renames = {} |
| self._num_input = 0 |
| self._num_param = 0 |
| self._shape = shape.copy() if shape else {} |
| self._input_names = [] |
| self._dtype = dtype |
| self.opset = None |
| self._freeze_params = freeze_params |
| |
| def __enter__(self): |
| self._old_manager = GraphProto.current |
| GraphProto.current = self |
| return self |
| |
| def __exit__(self, ptype, value, trace): |
| GraphProto.current = self._old_manager |
| |
| def freeze(self, func, params): |
| bind_map = {} |
| for name in params.keys(): |
| if name in self._nodes.keys(): |
| bind_map[self._nodes[name]] = _expr.const(params[name]) |
| body = _expr.bind(func.body, bind_map) |
| fn = _function.Function(analysis.free_vars(body), body) |
| return fn, {} |
| |
| def from_onnx(self, graph, opset, get_output_expr=False): |
| """Construct Relay expression from ONNX graph. |
| |
| Onnx graph is a python protobuf object. |
| The companion parameters will be handled automatically. |
| However, the input names from onnx graph is vague, mixing inputs and |
| network weights/bias such as "1", "2"... |
| For convenience, we rename the `real` input names to "input_0", |
| "input_1"... And renaming parameters to "param_0", "param_1"... |
| |
| Parameters |
| ---------- |
| graph : onnx protobuf object |
| The loaded onnx graph |
| |
| opset : opset version |
| |
| get_output_expr: bool |
| If set to true, this conversion will return each output expression rather |
| than a packaged module. This can be useful when converting subgraphs to |
| relay. |
| |
| Returns |
| ------- |
| mod : tvm.IRModule |
| The returned relay module |
| |
| params : dict |
| A dict of name: tvm.nd.array pairs, used as pretrained weights |
| """ |
| self.opset = opset |
| self._parse_graph_initializers(graph) |
| self._parse_graph_input(graph) |
| self._check_user_inputs_in_outermost_graph_scope() |
| self._check_for_unsupported_ops(graph) |
| self._construct_nodes(graph) |
| |
| # now return the outputs |
| outputs = [self._nodes[self._parse_value_proto(i)] for i in graph.output] |
| outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs) |
| # If requested, directly return the converted expressions. |
| if get_output_expr: |
| return outputs |
| ## Maintain the order of inputs and parameters from the ONNX graph, but only include |
| ## those parameters that are needed to execute the relay graph |
| free_vars = analysis.free_vars(outputs) |
| nodes = {v: k for k, v in self._nodes.items()} |
| free_vars = [nodes[var] for var in free_vars] |
| for i_name in self._params: |
| if i_name in free_vars and i_name not in self._inputs: |
| self._inputs[i_name] = self._nodes[i_name] |
| # Create a function from our output expression and all input variables. |
| func = _function.Function([v for k, v in self._inputs.items()], outputs) |
| return IRModule.from_expr(func), self._params |
| |
| def _parse_graph_initializers(self, graph): |
| """Parse network inputs to relay, aka parameters.""" |
| for init_tensor in graph.initializer: |
| if not init_tensor.name.strip(): |
| raise ValueError("Tensor's name is required.") |
| array = self._parse_array(init_tensor) |
| if self._freeze_params: |
| self._nodes[init_tensor.name] = _expr.const(array) |
| else: |
| self._params[init_tensor.name] = array |
| self._nodes[init_tensor.name] = new_var( |
| init_tensor.name, |
| shape=self._params[init_tensor.name].shape, |
| dtype=self._params[init_tensor.name].dtype, |
| ) |
| |
| def _parse_graph_input(self, graph): |
| for i in graph.input: |
| # from onnx v0.2, GraphProto.input has type ValueInfoProto, |
| # and the name is 'i.name' |
| i_name, i_shape, d_type, i_shape_name = get_info(i) |
| if i_name in self._params: |
| # i is a param instead of input |
| self._num_param += 1 |
| self._nodes[i_name] = new_var( |
| i_name, shape=self._params[i_name].shape, dtype=self._params[i_name].dtype |
| ) |
| elif i_name in self._nodes: |
| continue |
| else: |
| self._num_input += 1 |
| self._input_names.append(i_name) |
| if i_name in self._shape: |
| i_shape = self._shape[i_name] |
| else: |
| if "?" in str(i_shape): |
| warning_msg = ( |
| "Input %s has unknown dimension shapes: %s. " |
| "Specifying static values may improve performance" |
| % (i_name, str(i_shape_name)) |
| ) |
| warnings.warn(warning_msg) |
| if isinstance(self._dtype, dict): |
| dtype = self._dtype[i_name] if i_name in self._dtype else d_type |
| else: |
| dtype = d_type |
| self._nodes[i_name] = new_var(i_name, shape=i_shape, dtype=dtype) |
| self._inputs[i_name] = self._nodes[i_name] |
| |
| def _check_user_inputs_in_outermost_graph_scope(self): |
| """Only check user inputs in the outer-most graph scope.""" |
| if self._old_manager is None: |
| assert all( |
| [name in self._input_names for name in self._shape.keys()] |
| ), "User specified the shape for inputs that weren't found in the graph: " + str( |
| self._shape |
| ) |
| |
| def _check_for_unsupported_ops(self, graph): |
| convert_map = _get_convert_map(self.opset) |
| unsupported_ops = set() |
| for node in graph.node: |
| op_name = node.op_type |
| if ( |
| op_name not in convert_map |
| and op_name != "Constant" |
| and op_name not in _identity_list |
| ): |
| unsupported_ops.add(op_name) |
| if unsupported_ops: |
| msg = "The following operators are not supported for frontend ONNX: " |
| msg += ", ".join(unsupported_ops) |
| raise tvm.error.OpNotImplemented(msg) |
| |
| def _construct_nodes(self, graph): |
| """Nodes are stored as directed acyclic graph.""" |
| for node in graph.node: |
| op_name = node.op_type |
| attr = self._parse_attr(node.attribute) |
| # Create and populate input list. |
| inputs = onnx_input() |
| for i in node.input: |
| if i != "": |
| inputs.append(self._nodes[self._renames.get(i, i)]) |
| else: |
| inputs.append(None) |
| i_name = self._parse_value_proto(node) |
| node_output = self._fix_outputs(op_name, node.output) |
| attr["tvm_custom"] = {} |
| attr["tvm_custom"]["name"] = i_name |
| attr["tvm_custom"]["num_outputs"] = len(node_output) |
| |
| op = self._convert_operator(op_name, inputs, attr, self.opset) |
| if not isinstance(op, _expr.TupleWrapper): |
| outputs_num = 1 |
| else: |
| outputs_num = len(op) |
| |
| if outputs_num == 1: |
| op = fold_constant(op) |
| else: |
| op = _expr.TupleWrapper(fold_constant(op.astuple()), len(op)) |
| |
| if outputs_num > 1: |
| # ONNX supports optional outputs for some nodes. |
| # This block searches for missing outputs in the ONNX graph |
| # and removes any unneeded ops |
| valid_outputs = [False] * outputs_num |
| for i, output in enumerate(node_output): |
| if output != "": |
| valid_outputs[i] = True |
| # If we have outputs ONNX isn't expecting, we need to drop them |
| if not all(valid_outputs): |
| tup = op.astuple() |
| # TupleWrapper can also wrap ops with TupleType outputs |
| if isinstance(tup, _expr.Tuple): |
| # For tuples, we extract the fields instead of using GetTupleItem |
| outputs = [tup.fields[i] for i, valid in enumerate(valid_outputs) if valid] |
| else: |
| # For call nodes, we need to GetTupleItem |
| outputs = [op[i] for i, valid in enumerate(valid_outputs) if valid] |
| # Create the new op with valid outputs |
| if len(outputs) == 1: |
| op = outputs[0] |
| elif len(outputs) != outputs_num: |
| op = _expr.TupleWrapper(_expr.Tuple(outputs), len(outputs)) |
| # Drop invalid outputs for the onnx node |
| outputs_num = len(outputs) |
| node_output = [output for output in node_output if output != ""] |
| assert ( |
| len(node_output) == outputs_num |
| ), "Number of output mismatch {} vs {} in {}.".format( |
| len(node_output), outputs_num, op_name |
| ) |
| |
| if outputs_num == 1: |
| self._nodes[node_output[0]] = op |
| else: |
| for k, i in zip(list(node_output), range(len(node_output))): |
| self._nodes[k] = op[i] |
| |
| def _parse_value_proto(self, value_proto): |
| """Parse ValueProto or raw str.""" |
| try: |
| name = value_proto.name |
| except AttributeError: |
| name = value_proto |
| return name |
| |
| def _parse_array(self, tensor_proto): |
| np_array = get_numpy(tensor_proto).reshape(tuple(tensor_proto.dims)) |
| return _nd.array(np_array) |
| |
| def _parse_attr(self, attr_proto): |
| """Convert a list of AttributeProto to a dict, with names as keys.""" |
| attrs = {} |
| for a in attr_proto: |
| for f in ["f", "i", "s", "g"]: |
| if a.HasField(f): |
| attrs[a.name] = getattr(a, f) |
| for f in ["floats", "ints", "strings"]: |
| if list(getattr(a, f)): |
| assert a.name not in attrs, "Only one type of attr is allowed" |
| attrs[a.name] = tuple(getattr(a, f)) |
| for f in ["t"]: |
| if a.HasField(f): |
| attrs[a.name] = getattr(a, f) |
| for f in ["tensors"]: |
| if list(getattr(a, f)): |
| assert a.name not in attrs, "Only one type of attr is allowed" |
| attrs[a.name] = tuple(getattr(a, f)) |
| for f in ["graphs"]: |
| if list(getattr(a, f)): |
| raise NotImplementedError("Field {} is not supported in relay.".format(f)) |
| if a.name not in attrs: |
| raise ValueError("Cannot parse attribute: \n{}\n.".format(a)) |
| return attrs |
| |
| def _convert_operator(self, op_name, inputs, attrs, opset): |
| """Convert ONNX operator into a Relay operator. |
| The converter must specify conversions explicitly for incompatible name, and |
| apply handlers to operator attributes. |
| |
| Parameters |
| ---------- |
| op_name : str |
| Operator name, such as Convolution, FullyConnected |
| inputs : list of tvm.relay.function.Function |
| List of inputs. |
| attrs : dict |
| Dict of operator attributes |
| opset : int |
| Opset version |
| |
| Returns |
| ------- |
| sym : tvm.relay.function.Function |
| Converted relay function |
| """ |
| convert_map = _get_convert_map(opset) |
| if op_name in _identity_list: |
| sym = get_relay_op(op_name)(*inputs, **attrs) |
| elif op_name in convert_map: |
| sym = convert_map[op_name](inputs, attrs, self._params) |
| else: |
| raise NotImplementedError("Operator {} not implemented.".format(op_name)) |
| return sym |
| |
| def _fix_outputs(self, op_name, outputs): |
| """A hack to handle dropout or similar operator that have more than one out |
| in ONNX. |
| """ |
| if op_name == "Dropout": |
| if len(outputs) == 1: |
| return outputs |
| # TODO(zhreshold): support dropout mask? |
| outputs = outputs[:-1] |
| return outputs |
| |
| |
| def from_onnx( |
| model, shape=None, dtype="float32", opset=None, freeze_params=True, convert_config=None |
| ): |
| """Convert a ONNX model into an equivalent Relay Function. |
| |
| ONNX graphs are represented as Python Protobuf objects. |
| The companion parameters will be handled automatically. |
| However, the input names from onnx graph is vague, mixing inputs and |
| network weights/bias such as "1", "2"... |
| For convenience, we rename the `real` input names to "input_0", |
| "input_1"... And renaming parameters to "param_0", "param_1"... |
| |
| By default, ONNX defines models in terms of dynamic shapes. The ONNX importer |
| retains that dynamism upon import, and the compiler attempts to convert the |
| model into a static shapes at compile time. If this fails, there may still |
| be dynamic operations in the model. Not all TVM kernels currently support |
| dynamic shapes, please file an issue on discuss.tvm.apache.org |
| if you hit an error with dynamic kernels. |
| |
| Parameters |
| ---------- |
| model : protobuf object |
| ONNX ModelProto after ONNX v1.1.0 |
| |
| shape : dict of str to tuple, optional |
| The input shape to the graph |
| |
| dtype : str or dict of str to str |
| The input types to the graph |
| |
| opset : int, optional |
| Override to autodetected opset. |
| This can be helpful for some testing. |
| |
| freeze_params: bool |
| If this parameter is true, the importer will take any provided |
| onnx input values (weights, shapes, etc) and embed them into the relay model |
| as Constants instead of variables. This allows more aggressive optimizations |
| at compile time and helps in making models static if certain inputs represent |
| attributes relay would traditionally consider compile-time constants. |
| |
| convert_config : Optional[Dict[str, Any]] |
| Default config: |
| use_nt_batch_matmul : bool = True |
| True to convert qualified onnx `matmul` to `nn.batch_matmul` strict to NT format |
| (transpose_a=False, transpose_b=True). |
| |
| Returns |
| ------- |
| mod : tvm.IRModule |
| The relay module for compilation |
| |
| params : dict of str to tvm.nd.NDArray |
| The parameter dict to be used by relay |
| """ |
| global ONNX_DEFAULT_CONFIGS |
| if convert_config is not None: |
| ONNX_DEFAULT_CONFIGS.update(convert_config) |
| |
| try: |
| import onnx |
| |
| if hasattr(onnx.checker, "check_model"): |
| # try use onnx's own model checker before converting any model |
| try: |
| onnx.checker.check_model(model) |
| except Exception as e: # pylint: disable=c-extension-no-member, broad-except |
| # the checker is a bit violent about errors, so simply print warnings here |
| warnings.warn(str(e)) |
| except ImportError: |
| pass |
| g = GraphProto(shape, dtype, freeze_params) |
| graph = model.graph |
| |
| try: |
| opset_in_model = model.opset_import[0].version if model.opset_import else 1 |
| except AttributeError: |
| opset_in_model = 1 |
| |
| if opset is None: |
| opset = opset_in_model |
| elif opset < opset_in_model: |
| warnings.warn( |
| "" |
| f"You are overwritting original opset ver = {opset_in_model} by lower ver = {opset}. " |
| f"That might cause model conversion errors." |
| ) |
| |
| # Use the graph proto as a scope so that ops can access other nodes if needed. |
| with g: |
| mod, params = g.from_onnx(graph, opset) |
| |
| if freeze_params: |
| mod = relay.transform.DynamicToStatic()(mod) |
| |
| return mod, params |