mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
0c491702c4
Define the `TORCH_ONNX_USE_EXPERIMENTAL_LOGIC` flag to allow for enabling the new torch.onnx logic and hiding them during migration and testing. The actual logic migration will happen after. Pull Request resolved: https://github.com/pytorch/pytorch/pull/132299 Approved by: https://github.com/titaiwangms
28 lines
1017 B
Python
28 lines
1017 B
Python
"""Experimental classes and functions used by ONNX export."""
|
|
|
|
import dataclasses
|
|
from typing import Mapping, Optional, Sequence, Set, Type, Union
|
|
|
|
import torch
|
|
import torch._C._onnx as _C_onnx
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ExportOptions:
|
|
"""Arguments used by :func:`torch.onnx.export`."""
|
|
|
|
# TODO(justinchuby): Deprecate and remove this class.
|
|
|
|
export_params: bool = True
|
|
verbose: bool = False
|
|
training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL
|
|
input_names: Optional[Sequence[str]] = None
|
|
output_names: Optional[Sequence[str]] = None
|
|
operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX
|
|
opset_version: Optional[int] = None
|
|
do_constant_folding: bool = True
|
|
dynamic_axes: Optional[Mapping[str, Union[Mapping[int, str], Sequence[int]]]] = None
|
|
keep_initializers_as_inputs: Optional[bool] = None
|
|
custom_opsets: Optional[Mapping[str, int]] = None
|
|
export_modules_as_functions: Union[bool, Set[Type[torch.nn.Module]]] = False
|