ExportOutput

class torch.onnx.ExportOutput(model_proto, input_adapter, output_adapter, diagnostic_context)

An in-memory representation of a PyTorch model that has been exported to ONNX.

adapt_torch_inputs_to_onnx(*model_args, **model_kwargs)

Converts the PyTorch model inputs to exported ONNX model inputs format.

Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., None is allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc.

The actual adapting steps are associated with each individual export. It depends on the PyTorch model, the particular set of model_args and model_kwargs used for the export, and export options.

This method replays the adapting steps recorded during export.

Parameters

• model_args – The PyTorch model inputs.

• model_kwargs – The PyTorch model keyword inputs.

Returns

A sequence of tensors converted from PyTorch model inputs.

Return type

Sequence[Union[Tensor, int, float, bool]]

Example:

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> import torch.onnx
>>> from typing import Dict, Tuple
>>> def func_with_nested_input_structure(
...     x_dict: Dict[str, torch.Tensor],
...     y_tuple: Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
... ):
...     if "a" in x_dict:
...         x = x_dict["a"]
...     elif "b" in x_dict:
...         x = x_dict["b"]
...     else:
...         x = torch.randn(3)
...
...     y1, (y2, y3) = y_tuple
...
...     return x + y1 + y2 + y3
>>> x_dict = {"a": torch.tensor(1.)}
>>> y_tuple = (torch.tensor(2.), (torch.tensor(3.), torch.tensor(4.)))
>>> export_output = torch.onnx.dynamo_export(func_with_nested_input_structure, x_dict, y_tuple)
>>> print(x_dict, y_tuple)
{'a': tensor(1.)} (tensor(2.), (tensor(3.), tensor(4.)))
>>> print(export_output.adapt_torch_inputs_to_onnx(x_dict, y_tuple))
(tensor(1.), tensor(2.), tensor(3.), tensor(4.))

Warning

This API is experimental and is NOT backward-compatible.

adapt_torch_outputs_to_onnx(model_outputs)

Converts the PyTorch model outputs to exported ONNX model outputs format.

Due to design differences, input/output format between PyTorch model and exported ONNX model are often not the same. E.g., None is allowed for PyTorch model, but are not supported by ONNX. Nested constructs of tensors are allowed for PyTorch model, but only flattened tensors are supported by ONNX, etc.

The actual adapting steps are associated with each individual export. It depends on the PyTorch model, the particular set of model_args and model_kwargs used for the export, and export options.

This method replays the adapting steps recorded during export.

Parameters

model_outputs (Any) – The PyTorch model outputs.

Returns

PyTorch model outputs in exported ONNX model outputs format.

Return type

Sequence[Union[Tensor, int, float, bool]]

Example:

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> import torch.onnx
>>> def func_returning_tuples(x, y, z):
...     x = x + y
...     y = y + z
...     z = x + y
...     return (x, (y, z))
>>> x = torch.tensor(1.)
>>> y = torch.tensor(2.)
>>> z = torch.tensor(3.)
>>> export_output = torch.onnx.dynamo_export(func_returning_tuples, x, y, z)
>>> pt_output = func_returning_tuples(x, y, z)
>>> print(pt_output)
(tensor(3.), (tensor(5.), tensor(8.)))
>>> print(export_output.adapt_torch_outputs_to_onnx(pt_output))
[tensor(3.), tensor(5.), tensor(8.)]

Warning

This API is experimental and is NOT backward-compatible.

property diagnostic_context: DiagnosticContext

The diagnostic context associated with the export.

property model_proto: onnx.ModelProto

The exported ONNX model as an onnx.ModelProto.

save(destination, *, model_state_dict=None, serializer=None)

Saves the in-memory ONNX model to destination using specified serializer.

Parameters

• destination (Union[str, BufferedIOBase]) – The destination to save the ONNX model. It can be either a string or a file-like object. When used with model_state_dict, it must be a string with a full path to the destination. In that case, besides saving the ONNX model, a folder with “_initializers” suffix (without extension) will be created to store the each initializer of the ONNX model in a separate file. For example, if the destination is “/path/model.onnx”, the initializers will be saved in “/path/model_initializers/” folder.

• model_state_dict (Optional[Union[Dict[str, Any], str]]) – The state_dict of the PyTorch model containing all weights on it. It can be either a dict as returned by model.state_dict(), or a string with a file name. Required when enable_fake_mode is used but real initializers are needed on the ONNX graph. It can be either a string with the path to a checkpoint or a dictionary with the actual model state.

• serializer (Optional[ExportOutputSerializer]) – The serializer to use. If not specified, the model will be serialized as Protobuf.