mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
b319fa3fd9
Add ONNX directory to use ruff format. Pull Request resolved: https://github.com/pytorch/pytorch/pull/134120 Approved by: https://github.com/XuehaiPan, https://github.com/Skylion007
431 lines
15 KiB
Python
431 lines
15 KiB
Python
# mypy: allow-untyped-defs
|
|
from __future__ import annotations
|
|
|
|
|
|
__all__ = [
|
|
# Modules
|
|
"symbolic_helper",
|
|
"utils",
|
|
"errors",
|
|
# All opsets
|
|
"symbolic_caffe2",
|
|
"symbolic_opset7",
|
|
"symbolic_opset8",
|
|
"symbolic_opset9",
|
|
"symbolic_opset10",
|
|
"symbolic_opset11",
|
|
"symbolic_opset12",
|
|
"symbolic_opset13",
|
|
"symbolic_opset14",
|
|
"symbolic_opset15",
|
|
"symbolic_opset16",
|
|
"symbolic_opset17",
|
|
"symbolic_opset18",
|
|
"symbolic_opset19",
|
|
"symbolic_opset20",
|
|
# Enums
|
|
"ExportTypes",
|
|
"OperatorExportTypes",
|
|
"TrainingMode",
|
|
"TensorProtoDataType",
|
|
"JitScalarType",
|
|
# Public functions
|
|
"export",
|
|
"export_to_pretty_string",
|
|
"is_in_onnx_export",
|
|
"select_model_mode_for_export",
|
|
"register_custom_op_symbolic",
|
|
"unregister_custom_op_symbolic",
|
|
"disable_log",
|
|
"enable_log",
|
|
# Errors
|
|
"CheckerError", # Backwards compatibility
|
|
# Dynamo Exporter
|
|
"DiagnosticOptions",
|
|
"ExportOptions",
|
|
"ONNXProgram",
|
|
"ONNXProgramSerializer",
|
|
"ONNXRuntimeOptions",
|
|
"InvalidExportOptionsError",
|
|
"OnnxExporterError",
|
|
"OnnxRegistry",
|
|
"dynamo_export",
|
|
"enable_fake_mode",
|
|
# DORT / torch.compile
|
|
"is_onnxrt_backend_supported",
|
|
]
|
|
|
|
from typing import Any, Collection, Mapping, Sequence, TYPE_CHECKING
|
|
|
|
import torch
|
|
from torch import _C
|
|
from torch._C import _onnx as _C_onnx
|
|
from torch._C._onnx import OperatorExportTypes, TensorProtoDataType, TrainingMode
|
|
|
|
from ._exporter_states import ExportTypes
|
|
from ._internal.onnxruntime import (
|
|
is_onnxrt_backend_supported,
|
|
OrtBackend as _OrtBackend,
|
|
OrtBackendOptions as _OrtBackendOptions,
|
|
OrtExecutionProvider as _OrtExecutionProvider,
|
|
)
|
|
from ._type_utils import JitScalarType
|
|
from .errors import CheckerError # Backwards compatibility
|
|
from .utils import (
|
|
_optimize_graph,
|
|
_run_symbolic_function,
|
|
_run_symbolic_method,
|
|
export_to_pretty_string,
|
|
is_in_onnx_export,
|
|
register_custom_op_symbolic,
|
|
select_model_mode_for_export,
|
|
unregister_custom_op_symbolic,
|
|
)
|
|
|
|
|
|
from . import ( # usort: skip. Keep the order instead of sorting lexicographically
|
|
errors,
|
|
symbolic_caffe2,
|
|
symbolic_helper,
|
|
symbolic_opset7,
|
|
symbolic_opset8,
|
|
symbolic_opset9,
|
|
symbolic_opset10,
|
|
symbolic_opset11,
|
|
symbolic_opset12,
|
|
symbolic_opset13,
|
|
symbolic_opset14,
|
|
symbolic_opset15,
|
|
symbolic_opset16,
|
|
symbolic_opset17,
|
|
symbolic_opset18,
|
|
symbolic_opset19,
|
|
symbolic_opset20,
|
|
utils,
|
|
)
|
|
|
|
|
|
from ._internal._exporter_legacy import ( # usort: skip. needs to be last to avoid circular import
|
|
DiagnosticOptions,
|
|
ExportOptions,
|
|
ONNXProgram,
|
|
ONNXProgramSerializer,
|
|
ONNXRuntimeOptions,
|
|
InvalidExportOptionsError,
|
|
OnnxExporterError,
|
|
OnnxRegistry,
|
|
dynamo_export,
|
|
enable_fake_mode,
|
|
)
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
import os
|
|
|
|
# Set namespace for exposed private names
|
|
ExportTypes.__module__ = "torch.onnx"
|
|
JitScalarType.__module__ = "torch.onnx"
|
|
ExportOptions.__module__ = "torch.onnx"
|
|
ONNXProgram.__module__ = "torch.onnx"
|
|
ONNXProgramSerializer.__module__ = "torch.onnx"
|
|
ONNXRuntimeOptions.__module__ = "torch.onnx"
|
|
dynamo_export.__module__ = "torch.onnx"
|
|
InvalidExportOptionsError.__module__ = "torch.onnx"
|
|
OnnxExporterError.__module__ = "torch.onnx"
|
|
enable_fake_mode.__module__ = "torch.onnx"
|
|
OnnxRegistry.__module__ = "torch.onnx"
|
|
DiagnosticOptions.__module__ = "torch.onnx"
|
|
is_onnxrt_backend_supported.__module__ = "torch.onnx"
|
|
_OrtExecutionProvider.__module__ = "torch.onnx"
|
|
_OrtBackendOptions.__module__ = "torch.onnx"
|
|
_OrtBackend.__module__ = "torch.onnx"
|
|
|
|
producer_name = "pytorch"
|
|
producer_version = _C_onnx.PRODUCER_VERSION
|
|
|
|
|
|
def export(
|
|
model: torch.nn.Module
|
|
| torch.export.ExportedProgram
|
|
| torch.jit.ScriptModule
|
|
| torch.jit.ScriptFunction,
|
|
args: tuple[Any, ...],
|
|
f: str | os.PathLike | None = None,
|
|
*,
|
|
kwargs: dict[str, Any] | None = None,
|
|
export_params: bool = True,
|
|
verbose: bool | None = None,
|
|
input_names: Sequence[str] | None = None,
|
|
output_names: Sequence[str] | None = None,
|
|
opset_version: int | None = None,
|
|
dynamic_axes: Mapping[str, Mapping[int, str]]
|
|
| Mapping[str, Sequence[int]]
|
|
| None = None,
|
|
keep_initializers_as_inputs: bool = False,
|
|
dynamo: bool = False,
|
|
# Dynamo only options
|
|
external_data: bool = True,
|
|
dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,
|
|
report: bool = False,
|
|
verify: bool = False,
|
|
profile: bool = False,
|
|
dump_exported_program: bool = False,
|
|
artifacts_dir: str | os.PathLike = ".",
|
|
fallback: bool = False,
|
|
# Deprecated options
|
|
training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL,
|
|
operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX,
|
|
do_constant_folding: bool = True,
|
|
custom_opsets: Mapping[str, int] | None = None,
|
|
export_modules_as_functions: bool | Collection[type[torch.nn.Module]] = False,
|
|
autograd_inlining: bool = True,
|
|
**_: Any, # ignored options
|
|
) -> Any | None:
|
|
r"""Exports a model into ONNX format.
|
|
|
|
Args:
|
|
model: The model to be exported.
|
|
args: Example positional inputs. Any non-Tensor arguments will be hard-coded into the
|
|
exported model; any Tensor arguments will become inputs of the exported model,
|
|
in the order they occur in the tuple.
|
|
f: Path to the output ONNX model file. E.g. "model.onnx".
|
|
kwargs: Optional example keyword inputs.
|
|
export_params: If false, parameters (weights) will not be exported.
|
|
verbose: Whether to enable verbose logging.
|
|
input_names: names to assign to the input nodes of the graph, in order.
|
|
output_names: names to assign to the output nodes of the graph, in order.
|
|
opset_version: The version of the
|
|
`default (ai.onnx) opset <https://github.com/onnx/onnx/blob/master/docs/Operators.md>`_
|
|
to target. Must be >= 7.
|
|
dynamic_axes:
|
|
|
|
By default the exported model will have the shapes of all input and output tensors
|
|
set to exactly match those given in ``args``. To specify axes of tensors as
|
|
dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema:
|
|
|
|
* KEY (str): an input or output name. Each name must also be provided in ``input_names`` or
|
|
``output_names``.
|
|
* VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a
|
|
list, each element is an axis index.
|
|
|
|
For example::
|
|
|
|
class SumModule(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.sum(x, dim=1)
|
|
|
|
|
|
torch.onnx.export(
|
|
SumModule(),
|
|
(torch.ones(2, 2),),
|
|
"onnx.pb",
|
|
input_names=["x"],
|
|
output_names=["sum"],
|
|
)
|
|
|
|
Produces::
|
|
|
|
input {
|
|
name: "x"
|
|
...
|
|
shape {
|
|
dim {
|
|
dim_value: 2 # axis 0
|
|
}
|
|
dim {
|
|
dim_value: 2 # axis 1
|
|
...
|
|
output {
|
|
name: "sum"
|
|
...
|
|
shape {
|
|
dim {
|
|
dim_value: 2 # axis 0
|
|
...
|
|
|
|
While::
|
|
|
|
torch.onnx.export(
|
|
SumModule(),
|
|
(torch.ones(2, 2),),
|
|
"onnx.pb",
|
|
input_names=["x"],
|
|
output_names=["sum"],
|
|
dynamic_axes={
|
|
# dict value: manually named axes
|
|
"x": {0: "my_custom_axis_name"},
|
|
# list value: automatic names
|
|
"sum": [0],
|
|
},
|
|
)
|
|
|
|
Produces::
|
|
|
|
input {
|
|
name: "x"
|
|
...
|
|
shape {
|
|
dim {
|
|
dim_param: "my_custom_axis_name" # axis 0
|
|
}
|
|
dim {
|
|
dim_value: 2 # axis 1
|
|
...
|
|
output {
|
|
name: "sum"
|
|
...
|
|
shape {
|
|
dim {
|
|
dim_param: "sum_dynamic_axes_1" # axis 0
|
|
...
|
|
|
|
keep_initializers_as_inputs: If True, all the
|
|
initializers (typically corresponding to model weights) in the
|
|
exported graph will also be added as inputs to the graph. If False,
|
|
then initializers are not added as inputs to the graph, and only
|
|
the user inputs are added as inputs.
|
|
|
|
Set this to True if you intend to supply model weights at runtime.
|
|
Set it to False if the weights are static to allow for better optimizations
|
|
(e.g. constant folding) by backends/runtimes.
|
|
|
|
dynamo: Whether to export the model with ``torch.export`` ExportedProgram instead of TorchScript.
|
|
external_data: Whether to save the model weights as an external data file.
|
|
This is required for models with large weights that exceed the ONNX file size limit (2GB).
|
|
When False, the weights are saved in the ONNX file with the model architecture.
|
|
dynamic_shapes: A dictionary of dynamic shapes for the model inputs. Refer to
|
|
:func:`torch.export.export` for more details.
|
|
report: Whether to generate a markdown report for the export process.
|
|
verify: Whether to verify the exported model using ONNX Runtime.
|
|
profile: Whether to profile the export process.
|
|
dump_exported_program: Whether to dump the :class:`torch.export.ExportedProgram` to a file.
|
|
This is useful for debugging the exporter.
|
|
artifacts_dir: The directory to save the debugging artifacts like the report and the serialized
|
|
exported program.
|
|
fallback: Whether to fallback to the TorchScript exporter if the dynamo exporter fails.
|
|
|
|
training: Deprecated option. Instead, set the training mode of the model before exporting.
|
|
operator_export_type: Deprecated option. Only ONNX is supported.
|
|
do_constant_folding: Deprecated option. The exported graph is always optimized.
|
|
custom_opsets: Deprecated.
|
|
A dictionary:
|
|
|
|
* KEY (str): opset domain name
|
|
* VALUE (int): opset version
|
|
|
|
If a custom opset is referenced by ``model`` but not mentioned in this dictionary,
|
|
the opset version is set to 1. Only custom opset domain name and version should be
|
|
indicated through this argument.
|
|
export_modules_as_functions: Deprecated option.
|
|
|
|
Flag to enable
|
|
exporting all ``nn.Module`` forward calls as local functions in ONNX. Or a set to indicate the
|
|
particular types of modules to export as local functions in ONNX.
|
|
This feature requires ``opset_version`` >= 15, otherwise the export will fail. This is because
|
|
``opset_version`` < 15 implies IR version < 8, which means no local function support.
|
|
Module variables will be exported as function attributes. There are two categories of function
|
|
attributes.
|
|
|
|
1. Annotated attributes: class variables that have type annotations via
|
|
`PEP 526-style <https://www.python.org/dev/peps/pep-0526/#class-and-instance-variable-annotations>`_
|
|
will be exported as attributes.
|
|
Annotated attributes are not used inside the subgraph of ONNX local function because
|
|
they are not created by PyTorch JIT tracing, but they may be used by consumers
|
|
to determine whether or not to replace the function with a particular fused kernel.
|
|
|
|
2. Inferred attributes: variables that are used by operators inside the module. Attribute names
|
|
will have prefix "inferred::". This is to differentiate from predefined attributes retrieved from
|
|
python module annotations. Inferred attributes are used inside the subgraph of ONNX local function.
|
|
|
|
* ``False`` (default): export ``nn.Module`` forward calls as fine grained nodes.
|
|
* ``True``: export all ``nn.Module`` forward calls as local function nodes.
|
|
* Set of type of nn.Module: export ``nn.Module`` forward calls as local function nodes,
|
|
only if the type of the ``nn.Module`` is found in the set.
|
|
autograd_inlining: Deprecated.
|
|
Flag used to control whether to inline autograd functions.
|
|
Refer to https://github.com/pytorch/pytorch/pull/74765 for more details.
|
|
"""
|
|
if dynamo is True or isinstance(model, torch.export.ExportedProgram):
|
|
from torch.onnx._internal import exporter
|
|
|
|
if isinstance(args, torch.Tensor):
|
|
args = (args,)
|
|
return exporter.export_compat(
|
|
model,
|
|
args,
|
|
f,
|
|
kwargs=kwargs,
|
|
export_params=export_params,
|
|
verbose=verbose,
|
|
input_names=input_names,
|
|
output_names=output_names,
|
|
opset_version=opset_version,
|
|
dynamic_axes=dynamic_axes,
|
|
keep_initializers_as_inputs=keep_initializers_as_inputs,
|
|
external_data=external_data,
|
|
dynamic_shapes=dynamic_shapes,
|
|
report=report,
|
|
verify=verify,
|
|
profile=profile,
|
|
dump_exported_program=dump_exported_program,
|
|
artifacts_dir=artifacts_dir,
|
|
fallback=fallback,
|
|
)
|
|
else:
|
|
from torch.onnx.utils import export
|
|
|
|
export(
|
|
model,
|
|
args,
|
|
f, # type: ignore[arg-type]
|
|
kwargs=kwargs,
|
|
export_params=export_params,
|
|
verbose=verbose is True,
|
|
input_names=input_names,
|
|
output_names=output_names,
|
|
opset_version=opset_version,
|
|
dynamic_axes=dynamic_axes,
|
|
keep_initializers_as_inputs=keep_initializers_as_inputs,
|
|
training=training,
|
|
operator_export_type=operator_export_type,
|
|
do_constant_folding=do_constant_folding,
|
|
custom_opsets=custom_opsets,
|
|
export_modules_as_functions=export_modules_as_functions,
|
|
autograd_inlining=autograd_inlining,
|
|
)
|
|
return None
|
|
|
|
|
|
# TODO(justinchuby): Deprecate these logging functions in favor of the new diagnostic module.
|
|
|
|
# Returns True iff ONNX logging is turned on.
|
|
is_onnx_log_enabled = _C._jit_is_onnx_log_enabled
|
|
|
|
|
|
def enable_log() -> None:
|
|
r"""Enables ONNX logging."""
|
|
_C._jit_set_onnx_log_enabled(True)
|
|
|
|
|
|
def disable_log() -> None:
|
|
r"""Disables ONNX logging."""
|
|
_C._jit_set_onnx_log_enabled(False)
|
|
|
|
|
|
"""Sets output stream for ONNX logging.
|
|
|
|
Args:
|
|
stream_name (str, default "stdout"): Only 'stdout' and 'stderr' are supported
|
|
as ``stream_name``.
|
|
"""
|
|
set_log_stream = _C._jit_set_onnx_log_output_stream
|
|
|
|
|
|
"""A simple logging facility for ONNX exporter.
|
|
|
|
Args:
|
|
args: Arguments are converted to string, concatenated together with a newline
|
|
character appended to the end, and flushed to output stream.
|
|
"""
|
|
log = _C._jit_onnx_log
|