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)
