mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
0d76299ff7
Cleaning up onnx module imports to prepare for updating `__init__`.
- Simplify importing the `_C` and `_C._onnx` name spaces
- Remove alias of the symbolic_helper module in imports
- Remove any module level function imports. Import modules instead
- Alias `symbilic_opsetx` as `opsetx`
- Fix some docstrings
Requires:
- https://github.com/pytorch/pytorch/pull/77448
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77423
Approved by: https://github.com/BowenBao
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
"""Globals used internally by the ONNX exporter.
|
|
|
|
Do not use this module outside of `torch.onnx` and its tests.
|
|
|
|
Be very judicious when adding any new global variables. Do not create new global
|
|
variables unless they are absolutely necessary.
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
import torch._C._onnx as _C_onnx
|
|
|
|
# This module should only depend on _constants and nothing else in torch.onnx to keep
|
|
# dependency direction clean.
|
|
from torch.onnx import _constants
|
|
|
|
|
|
class _InternalGlobals:
|
|
"""Globals used internally by ONNX exporter.
|
|
|
|
NOTE: Be very judicious when adding any new variables. Do not create new
|
|
global variables unless they are absolutely necessary.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._export_onnx_opset_version = _constants.onnx_default_opset
|
|
self.operator_export_type: Optional[_C_onnx.OperatorExportTypes] = None
|
|
self.training_mode: Optional[_C_onnx.TrainingMode] = None
|
|
self.onnx_shape_inference: bool = False
|
|
|
|
@property
|
|
def export_onnx_opset_version(self):
|
|
return self._export_onnx_opset_version
|
|
|
|
@export_onnx_opset_version.setter
|
|
def export_onnx_opset_version(self, value: int):
|
|
supported_versions = [_constants.onnx_main_opset]
|
|
supported_versions.extend(_constants.onnx_stable_opsets)
|
|
if value not in supported_versions:
|
|
raise ValueError(f"Unsupported ONNX opset version: {value}")
|
|
self._export_onnx_opset_version = value
|
|
|
|
|
|
GLOBALS = _InternalGlobals()
|