def multinomial_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) helper.add_attribute(node, 'dtype', helper.tensor_type('int64')) for arg in op_def.arg: if arg.name == 'sample_size': helper.add_attribute(node, 'sample_size', arg.i) return node, const_tensors
def cast_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'Cast' if len(node.input) == 0: raise ValueError('ONNX does not support in-place cast.') for arg in op_def.arg: if arg.name == 'dtype': helper.add_attribute(node, 'to', helper.tensor_type(arg.s)) return node, const_tensors
def eye_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) if len(op_def.input) > 0: node.op_type += 'Like' else: output_shape = list(context.blob_shapes[op_def.output[0]]) helper.add_attribute(node, 'shape', output_shape) for arg in op_def.arg: if arg.name == 'k': helper.add_attribute(node, 'k', arg.i) elif arg.name == 'dtype': helper.add_attribute(node, 'dtype', helper.tensor_type(arg.s)) return node, const_tensors
def export( inputs, outputs, f, input_names=None, output_names=None, input_shapes=None, opset_version=None, verbose=False, enable_onnx_checker=True, ): """Export the recorded graph to an onnx model. Enter into the record mode to export operators into an onnx model: ```python x = dragon.constant([1, 2, 3]) with dragon.onnx.record(): y = x * x dragon.onnx.export(inputs=[x], outputs=[y], f='model.onnx') ``` Parameters ---------- inputs : Union[Sequence, Dict] The model inputs. outputs : Union[Sequence, Dict] The model outputs. f : str The filename for exporting model. input_names : Sequence[str], optional The name to the inputs. output_names : Sequence[str], optional The name to the outputs. input_shapes : Union[Sequence, Dict], optional The optional rewritten for input shapes. opset_version : int, optional The version of operator set. verbose : bool, optional, default=False Whether to print the debug string of graph. enable_onnx_checker : bool, optional, default=True Whether to check if model is valid. """ # Process the inputs. if isinstance(inputs, dict): if input_names is not None: raise ValueError( 'Excepted the input names from <inputs>.\n' 'You should set the <input_names> to None.') inputs, input_names = list(inputs.values()), list(inputs.keys()) else: inputs = nest.flatten(inputs) # Process the outputs. if isinstance(outputs, dict): if output_names is not None: raise ValueError( 'Excepted the output names from <outputs>.\n' 'You should set the <output_names> to None.') outputs, output_names = list(outputs.values()), list(outputs.keys()) else: outputs = nest.flatten(outputs) if eager_context.executing_eagerly(): op_defs = [] tape = backprop.get_default_tape() if tape is None: raise RuntimeError('Please enter with ``onnx.frontend.record()``.') for op_def in tape._defs: op_defs.append(dragon_pb2.OperatorDef()) op_defs[-1].ParseFromString(op_def.SerializeAs()) graph_def = dragon_pb2.GraphDef(op=op_defs) else: symbolic_outputs = [] for output in outputs: if types.is_symbolic_tensor(output): symbolic_outputs.append(output) graph_func = function_lib.create_function(outputs=symbolic_outputs) graph_func.callback() graph_def = graph_func.graph_def graph_def.name = '' # Add inputs and outputs. for i, input in enumerate(inputs): if hasattr(input, 'id'): graph_def.input.extend([input.id]) elif input_names is not None: graph_def.input.extend([input_names[i]]) for i, output in enumerate(outputs): if hasattr(output, 'id'): graph_def.output.extend([output.id]) elif output_names is not None: graph_def.output.extend([output_names[i]]) # Make value info from inputs and outputs. value_names = graph_def.input[:] + graph_def.output[:] value_info = dict([(k, (helper.tensor_type(v.dtype), v.shape)) for k, v in zip(value_names, inputs + outputs)]) # Extract the constants from inputs and outputs. constants = collections.OrderedDict() for k, v in zip(value_names, inputs + outputs): if isinstance(v, numpy.ndarray): constants[k] = v # Export. model = graph_def_to_onnx_model( graph_def=graph_def, input_names=input_names, output_names=output_names, input_shapes=input_shapes, constants=constants, value_info=value_info, opset_version=opset_version, workspace=workspace_util.get_workspace(), verbose=verbose, enable_onnx_checker=enable_onnx_checker, ) serialization.save_bytes(serialization.serialize_proto(model), f)
def export( model, args, f, input_names=None, output_names=None, input_shapes=None, opset_version=None, verbose=False, enable_onnx_checker=True, ): """Export the recorded graph to an onnx model. The outputs will be obtained by calling ``model(*args)``, both the tensor or numpy array are allowed: ```python class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() self.fc = torch.nn.Linear(3, 3) def forward(self, x): y = self.fc(x) return y, np.ones((2, 3)) m = MyModule() x = torch.zeros(2, 3) torch.onnx.export( m, args=(x,), f='my_module.onnx', input_names=('x',), output_names=('y', 'ones'), ) ``` You can either specify the ``input_names``, or pass a *dict* to the ``args``. In the same way, ``model`` could return a *dict* to specify the ``output_names``: ```python class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() self.fc = torch.nn.Linear(3, 3) def forward(self, inputs): y = self.fc(inputs['x']) return {'y': y, 'ones': np.ones((2, 3))} m = MyModule() x = torch.zeros(2, 3) torch.onnx.export( m, args={'x': x}, f='my_module.onnx', ) ``` Also note that if a numpy array is given or returned, it's name is definitely required. Otherwise, ONNX can't export this value due to the lacking of *id*. Parameters ---------- model : dragon.vm.torch.nn.Module The module to export. args : Union[Sequence, Dict] The model inputs. f : str The filename for exporting model. input_names : Sequence[str], optional The name to the inputs. output_names : Sequence[str], optional The name to the outputs. input_shapes : Union[Sequence, Dict], optional The optional rewritten for input shapes. opset_version : int, optional The version of operator set. verbose : bool, optional, default=False Whether to print the debug string of graph. enable_onnx_checker : bool, optional, default=True Whether to check if model is valid. """ # Process the inputs. if isinstance(args, dict): if input_names is not None: raise ValueError('Excepted the input names from <args>.\n' 'You should set the <input_names> to None.') inputs, input_names, args = \ list(args.values()), list(args.keys()), [args] else: inputs = args = nest.flatten(args) # Run the model to get the outputs. execute_ws = workspace.Workspace() execute_ws.merge_from(workspace.get_workspace()) with execute_ws.as_default(): with tapes.Tape() as model_tape: model_tape._exporting = True outputs = model(*args) # Process the outputs if isinstance(outputs, dict): if output_names is not None: raise ValueError('Excepted the output names from <outputs>.\n' 'You should set the <output_names> to None.') outputs, output_names = list(outputs.values()), list(outputs.keys()) else: outputs = nest.flatten(outputs) # Make graph def. ops_def, graph_def = [], dragon_pb2.GraphDef() # Add inputs and outputs. for i, input in enumerate(inputs): if hasattr(input, 'id'): graph_def.input.extend([input.id]) elif input_names is not None: graph_def.input.extend([input_names[i]]) for i, output in enumerate(outputs): if hasattr(output, 'id'): graph_def.output.extend([output.id]) elif output_names is not None: graph_def.output.extend([output_names[i]]) # Add operators. for op_def in model_tape.get_elements(): ops_def.append(dragon_pb2.OperatorDef()) ops_def[-1].ParseFromString(op_def.SerializeAs()) graph_def.op.extend(ops_def) # Make value info from inputs and outputs. value_names = graph_def.input[:] + graph_def.output[:] value_info = dict([(k, (helper.tensor_type(v.dtype), v.shape)) for k, v in zip(value_names, inputs + outputs)]) # Extract the constants from inputs and outputs. constants = collections.OrderedDict() for k, v in zip(value_names, inputs + outputs): if isinstance(v, numpy.ndarray): constants[k] = v # Export. with execute_ws.as_default(): model = graph_def_to_onnx_model( graph_def=graph_def, input_names=input_names, output_names=output_names, input_shapes=input_shapes, constants=constants, value_info=value_info, opset_version=opset_version, workspace=execute_ws, verbose=verbose, enable_onnx_checker=enable_onnx_checker, ) serialization.save_bytes(serialization.serialize_proto(model), f)