ExportOutput

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

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[Tensor]

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[Tensor]

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 model_proto: onnx.ModelProto

The exported ONNX model as an onnx.ModelProto.

save(destination, *, serializer=None)

Saves the in-memory ONNX model to destination using specified serializer. If no serializer is specified, the model will be serialized as Protobuf.