| # Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. |
| # |
| # Redistribution and use in source and binary forms, with or without |
| # modification, are permitted provided that the following conditions |
| # are met: |
| # * Redistributions of source code must retain the above copyright |
| # notice, this list of conditions and the following disclaimer. |
| # * Redistributions in binary form must reproduce the above copyright |
| # notice, this list of conditions and the following disclaimer in the |
| # documentation and/or other materials provided with the distribution. |
| # * Neither the name of NVIDIA CORPORATION nor the names of its |
| # contributors may be used to endorse or promote products derived |
| # from this software without specific prior written permission. |
| # |
| # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY |
| # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
| # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR |
| # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR |
| # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, |
| # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, |
| # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR |
| # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY |
| # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| |
| # Based on |
| # https://github.com/NVIDIA/mxnet_to_onnx/blob/master/mx2onnx_converter/mx2onnx_converter.py |
| |
| # coding: utf-8 |
| # pylint: disable=invalid-name,too-many-locals,no-self-use,too-many-arguments, |
| # pylint: disable=maybe-no-member,too-many-nested-blocks,logging-not-lazy |
| # pylint: disable=cell-var-from-loop |
| """MXNet to ONNX graph converter functions""" |
| import logging |
| import json |
| |
| import numpy as np |
| from mxnet import ndarray as nd |
| |
| |
| class MXNetGraph(object): |
| """Class to convert MXNet to ONNX graph""" |
| registry_ = {} |
| input_output_maps_ = {} |
| |
| def __init__(self): |
| # topologically sorted nodes |
| self.nodes = [] |
| self.input_tensors = [] |
| self.output_tensors = [] |
| |
| @staticmethod |
| def register(op_name, opset_version=12): |
| """Register operators""" |
| def wrapper(func): |
| """Helper function to map functions""" |
| try: |
| import onnx as _ |
| op_map = MXNetGraph.registry_.setdefault(opset_version, {}) |
| op_map[op_name] = func |
| except ImportError: |
| pass |
| return func |
| |
| return wrapper |
| |
| @staticmethod |
| def convert_layer(node, **kwargs): |
| """Convert MXNet layer to ONNX""" |
| try: |
| from onnx.defs import onnx_opset_version |
| except ImportError: |
| raise ImportError("Onnx and protobuf need to be installed. " |
| + "Instructions to install - https://github.com/onnx/onnx") |
| |
| op = str(node["op"]) |
| opset_version = kwargs.get("opset_version", onnx_opset_version()) |
| if opset_version < 12: |
| logging.warning('Your ONNX op set version is {}, ' |
| 'which is lower than then lowest tested op set (12), please consider ' |
| 'updating ONNX'.format(str(opset_version))) |
| opset_version = 12 |
| # Fallback to older opset versions if op is not registered in current version |
| convert_func = None |
| for op_version in range(opset_version, 11, -1): |
| if op_version not in MXNetGraph.registry_ or op not in MXNetGraph.registry_[op_version]: |
| continue |
| convert_func = MXNetGraph.registry_[op_version][op] |
| break |
| |
| # The conversion logic is not implemented |
| if convert_func is None: |
| raise AttributeError(f"No conversion function registered for op type {op} yet.") |
| |
| ret = convert_func(node, **kwargs) |
| # in case the conversion function does not specify the returned dtype, we just return None |
| # as the second value |
| if isinstance(ret, list): |
| return ret, None |
| else: |
| return ret |
| |
| @staticmethod |
| def split_params(sym, params): |
| """Helper function to split params dictionary into args and aux params |
| |
| Parameters |
| ---------- |
| sym : :class:`~mxnet.symbol.Symbol` |
| MXNet symbol object |
| params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray` |
| Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format |
| |
| Returns |
| ------- |
| arg_params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray` |
| Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format |
| aux_params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray` |
| Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format |
| """ |
| arg_params = {} |
| aux_params = {} |
| for args in sym.list_arguments(): |
| if args in params: |
| arg_params.update({args: nd.array(params[args])}) |
| for aux in sym.list_auxiliary_states(): |
| if aux in params: |
| aux_params.update({aux: nd.array(params[aux])}) |
| return arg_params, aux_params |
| |
| @staticmethod |
| def get_outputs(sym, params, in_shapes, output_label, in_types, dynamic=False, |
| dynamic_input_shapes=None): |
| """Helper function to collect the output names, types, and shapes |
| |
| Parameters |
| ---------- |
| sym : :class:`~mxnet.symbol.Symbol` |
| MXNet symbol object |
| params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray` |
| Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format |
| in_shapes : list of tuple |
| Input shapes |
| out_label : ``str`` |
| Name of label typically used in loss that may be left in graph. This name is |
| removed from list of inputs required by symbol |
| in_types : list of Int |
| Input ONNX data types |
| dynamic : Boolean |
| If True will allow for dynamic input shapes to the model |
| dynamic_input_shapes: list of tuple |
| Specifies the dynamic input_shapes. If None then all dimensions are set to None |
| |
| Returns |
| in_shapes : list of tuple |
| Updated input shapes |
| graph_outputs : dict ``str`` to dict |
| This maps output name to {'shape':tuple, 'dtype':Int} |
| ------- |
| """ |
| from onnx import mapping |
| import re |
| |
| # Collect graph output names |
| out_names = list() |
| for name in sym.list_outputs(): |
| if name.endswith('_state_output'): # handel special cases for RNN operator |
| out_names.append(name[:-len('_state_output')]+'1') |
| elif name.endswith('_statecell_output'): # handel special cases for RNN operator |
| out_names.append(name[:-len('_statecell_output')]+'2') |
| elif name.endswith('_output'): |
| out_names.append(name[:-len('_output')]) |
| elif name.endswith('_out'): |
| out_names.append(name[:-len('_out')]) |
| elif re.search('.*_output[0-9]$', name): |
| out_names.append(name[:-len('_output0')]+name[-1]) |
| else: |
| logging.info("output '%s' does not end with '_output'", name) |
| out_names.append(name) |
| |
| # Collect graph output shapes |
| # Remove any input listed in params from sym.list_inputs() and bind them to the input shapes provided |
| # by user. Also remove output_label, which is the name of the label symbol that may have been used |
| # as the label for loss during training. |
| inputs = {n: tuple(s) for n, s in |
| zip([n for n in sym.list_inputs() if n not in params and n != output_label], |
| in_shapes)} |
| # Add params and their shape to list of inputs |
| inputs.update({n: v.shape for n, v in params.items() if n in sym.list_inputs()}) |
| # Provide input data as well as input params to infer_shape() |
| _, out_shapes, _ = sym.infer_shape(**inputs) |
| if dynamic: |
| # Keep the dimensionality of the output shapes but change the values to None |
| out_shapes = [tuple(None for _ in i_s) for i_s in out_shapes] |
| |
| if dynamic_input_shapes is None: |
| # Set all dimensions to None |
| in_shapes = [tuple(None for _ in i_s) for i_s in in_shapes] |
| else: |
| assert len(in_shapes) == len(dynamic_input_shapes), "The length of " \ |
| "dynamic_input_shapes must equal to the length of in_shapes." |
| for i_s, d_i_s in zip(in_shapes, dynamic_input_shapes): |
| assert len(i_s) == len(d_i_s), "The dimensionality " \ |
| "of each shape must match." |
| in_shapes = dynamic_input_shapes |
| else: |
| assert dynamic_input_shapes is None, "dynamic_input_shapes is specified. Please " \ |
| "set dynamic_input_shapes=True to enable dynamic input shapes" |
| |
| # Collect graph output types |
| # Remove any input listed in params from sym.list_inputs() and bind them to the input types provided |
| # by user. Also remove output_label |
| in_dtypes = {n: mapping.TENSOR_TYPE_TO_NP_TYPE[t] for n, t in |
| zip([n for n in sym.list_inputs() if n not in params and n != output_label], |
| in_types)} |
| # Add params and their types to list of inputs |
| in_dtypes.update({n: v.dtype for n, v in params.items() if n in sym.list_inputs()}) |
| _, out_type, _ = sym.infer_type(**in_dtypes) |
| out_types = [mapping.NP_TYPE_TO_TENSOR_TYPE[o(0).dtype] for o in out_type] |
| |
| # Make sure the types, names, and shapes all align up |
| assert len(out_types) == len(out_names) == len(out_shapes) |
| |
| # Bind output shapes/types with output names |
| graph_outputs = {n: {'shape': s, 'dtype': d} for n, s, d in zip(out_names, out_shapes, out_types)} |
| |
| return in_shapes, graph_outputs |
| |
| @staticmethod |
| def convert_weights_to_numpy(weights_dict): |
| """Convert weights to numpy""" |
| return dict([(k.replace("arg:", "").replace("aux:", ""), v.asnumpy()) |
| for k, v in weights_dict.items()]) |
| |
| def create_onnx_graph_proto(self, sym, params, in_shapes, in_types, verbose=False, opset_version=None, |
| dynamic=True, dynamic_input_shapes=None): |
| """Convert MXNet graph to ONNX graph |
| |
| Parameters |
| ---------- |
| sym : :class:`~mxnet.symbol.Symbol` |
| MXNet symbol object |
| params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray` |
| Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format |
| in_shapes : List of tuple |
| Input shape of the model e.g [(1,3,224,224)] |
| in_types : List of Int |
| Input ONNX data types |
| verbose : Boolean |
| If true will print logs of the model conversion |
| opset_version : Int |
| ONNX opset version to use for export, defaults to latest supported by onnx package |
| dynamic: Boolean |
| If True will allow for dynamic input shapes to the model |
| dynamic_input_shapes: list of tuple |
| Specifies the dynamic input_shapes. If None then all dimensions are set to None |
| |
| Returns |
| ------- |
| graph : GraphProto |
| ONNX graph |
| """ |
| try: |
| from onnx import (helper, NodeProto, ValueInfoProto, TensorProto) |
| from onnx.helper import make_tensor_value_info |
| from onnx.defs import onnx_opset_version |
| except ImportError: |
| raise ImportError("Onnx and protobuf need to be installed. " |
| + "Instructions to install - https://github.com/onnx/onnx") |
| |
| if opset_version is None: |
| opset_version = onnx_opset_version() |
| |
| # When MXNet model is saved to json file , MXNet adds a node for label. |
| # The name of this node is, name of the last node + "_label" ( i.e if last node |
| # name is "Softmax", this node will have a name "Softmax_label". Also, the new node |
| # will always be second last node in the json graph. |
| # Deriving the output_label name. |
| output_label = sym.get_internals()[len(sym.get_internals()) - 1].name + "_label" |
| |
| weights = MXNetGraph.convert_weights_to_numpy(params) |
| |
| mx_graph = json.loads(sym.tojson())["nodes"] |
| |
| class NodeOutput: |
| def __init__(self, name, dtype): |
| self.name = name |
| self.dtype = np.dtype(dtype) |
| |
| initializer = [] |
| all_processed_nodes = [] |
| onnx_processed_nodes = [] |
| onnx_processed_inputs = [] |
| onnx_processed_outputs = [] |
| outputs_lookup = [] |
| |
| # Determine graph output names, shapes, and dtypes. Also update in_shapes |
| in_shapes, graph_outputs = MXNetGraph.get_outputs(sym, params, in_shapes, output_label, |
| in_types, dynamic, dynamic_input_shapes) |
| appeared_names = set() |
| graph_input_idx = 0 |
| for idx, node in enumerate(mx_graph): |
| op = node["op"] |
| # check if the current node has the same name as nodes before |
| if node["name"] in appeared_names: |
| node["name"] = 'idx_' + str(idx) + '_' + node["name"] |
| else: |
| appeared_names.add(node["name"]) |
| name = node["name"] |
| if verbose: |
| logging.info("Converting idx: %d, op: %s, name: %s", idx, op, name) |
| |
| # A node is an input node if its op_name is "null" and is not |
| # in params dict |
| if op == "null" and name not in params: |
| # Handle graph input |
| |
| # Skip output_label node, as this node is not part of graph |
| # Refer to "output_label" assignment above for more details. |
| if name == output_label: |
| continue |
| |
| converted, dtypes = MXNetGraph.convert_layer( |
| node, |
| is_input=True, |
| mx_graph=mx_graph, |
| weights=weights, |
| in_shape=in_shapes[graph_input_idx], |
| in_type=in_types[graph_input_idx], |
| proc_nodes=all_processed_nodes, |
| initializer=initializer, |
| outputs_lookup=outputs_lookup) |
| graph_input_idx += 1 |
| else: |
| # Handle graph layers |
| converted, dtypes = MXNetGraph.convert_layer( |
| node, |
| is_input=False, |
| mx_graph=mx_graph, |
| weights=weights, |
| proc_nodes=all_processed_nodes, |
| initializer=initializer, |
| outputs_lookup=outputs_lookup, |
| idx=idx, |
| opset_version=opset_version |
| ) |
| if isinstance(converted, list): |
| # Collect all the node's output names |
| node_possible_names = [name] + [name + str(i) for i in range(100)] |
| node_output_names = [] |
| # Collect all the graph's output names |
| graph_output_names = [] |
| # Iterate for all converted nodes |
| for converted_node in converted: |
| # If converted node is ValueInfoProto, add it in inputs |
| if isinstance(converted_node, ValueInfoProto): |
| onnx_processed_inputs.append(converted_node) |
| # If converted node is NodeProto, add it in processed nodes list |
| elif isinstance(converted_node, NodeProto): |
| onnx_processed_nodes.append(converted_node) |
| # some operators have multiple outputs, |
| # therefore, check all output node names |
| node_names = list(converted_node.output) |
| for nodename in node_names: |
| if nodename in node_possible_names: |
| node_output_names.append(nodename) |
| if nodename in graph_outputs: |
| graph_output_names.append(nodename) |
| if verbose: |
| logging.info("Output node is: {}".format(nodename)) |
| elif isinstance(converted_node, TensorProto): |
| raise ValueError("Did not expect TensorProto") |
| else: |
| raise ValueError(f"node is of an unrecognized type: {type(node)}") |
| |
| all_processed_nodes.append(converted_node) |
| |
| # if node_output_names is empty then we use the last returned node as output |
| if not node_output_names: |
| node_output_names = [converted[-1].name] |
| # process node outputs (sort by output index) |
| def str2int(s, l): |
| if len(s) == l: |
| return -1 |
| else: |
| return int(s[l:]) |
| |
| node_output_names = sorted(node_output_names, key=lambda x: str2int(x, len(name))) |
| |
| # match the output names to output dtypes |
| if dtypes is not None: |
| assert len(node_output_names) == len(dtypes) |
| node_outputs = [NodeOutput(node_output_names[i], dtypes[i]) |
| for i in range(len(dtypes))] |
| else: |
| # in case dtypes is None, we just default to the dtype of the first input |
| assert len(node["inputs"]) > 0 |
| first_input = node["inputs"][0] |
| first_input_dtype = outputs_lookup[first_input[0]][first_input[1]].dtype |
| node_outputs = [NodeOutput(n, first_input_dtype) |
| for n in node_output_names] |
| outputs_lookup.append(node_outputs) |
| |
| # process graph outputs (sort by alphabetical order) |
| graph_output_names.sort() |
| for nodename in graph_output_names: |
| onnx_processed_outputs.append( |
| make_tensor_value_info( |
| name=nodename, |
| elem_type=graph_outputs[nodename]['dtype'], |
| shape=graph_outputs[nodename]['shape'] |
| ) |
| ) |
| |
| else: |
| logging.info("Operator converter function should always return a list") |
| |
| # sometimes the graph output can also be in the intializer |
| for i in initializer: |
| if i.name in graph_outputs: |
| onnx_processed_outputs.append( |
| make_tensor_value_info( |
| name=i.name, |
| elem_type=graph_outputs[i.name]['dtype'], |
| shape=graph_outputs[i.name]['shape'] |
| ) |
| ) |
| |
| graph = helper.make_graph( |
| onnx_processed_nodes, |
| "mxnet_converted_model", |
| onnx_processed_inputs, |
| onnx_processed_outputs |
| ) |
| |
| graph.initializer.extend(initializer) |
| |
| return graph |
