def _convert(self) -> None:
prev_opset_version = None
prev_export_type = None
prev_shape_inference = None
try:
assert not to_utils.is_in_onnx_export() # type: ignore[no-untyped-call]
with to_utils.select_model_mode_for_export(self.original_model, self.training):
to_utils.__IN_ONNX_EXPORT = True
prev_opset_version = sym_hel._export_onnx_opset_version
sym_hel._set_opset_version(self.opset_version) # type: ignore[no-untyped-call]
prev_export_type = sym_hel._operator_export_type
sym_hel._set_operator_export_type(self.operator_export_type) # type: ignore[no-untyped-call]
prev_shape_inference = sym_hel._onnx_shape_inference
sym_hel._set_onnx_shape_inference( # type: ignore[no-untyped-call]
False # TODO(twata): Use `self.onnx_shape_inference`
)
self._run_trace()
self.model: onnx.ModelProto = self.generate_onnx()
finally:
to_utils.__IN_ONNX_EXPORT = False
if prev_opset_version is not None:
sym_hel._set_opset_version(prev_opset_version) # type: ignore[no-untyped-call]
if prev_shape_inference is not None:
sym_hel._set_operator_export_type(prev_export_type) # type: ignore[no-untyped-call]
if prev_shape_inference is not None:
sym_hel._set_onnx_shape_inference(prev_shape_inference) # type: ignore[no-untyped-call]
def select_model_mode_for_export(model, mode):
r"""
A context manager to temporarily set the training mode of 'model'
to 'mode', resetting it when we exit the with-block. A no-op if
mode is None.
"""
from torch.onnx import utils
return utils.select_model_mode_for_export(model, mode)
def select_model_mode_for_export(model, mode):
r"""
A context manager to temporarily set the training mode of "model"
to "mode", resetting it when we exit the with-block. A no-op if
mode is None.
In version 1.6 changed to this from set_training
"""
from torch.onnx import utils
return utils.select_model_mode_for_export(model, mode)
def select_model_mode_for_export(model, mode):
r"""
A context manager to temporarily set the training mode of ``model``
to ``mode``, resetting it when we exit the with-block. A no-op if
mode is None.
Args:
model: Same type and meaning as ``model`` arg to :func:`export`.
mode: Same type and meaning as ``training`` arg to :func:`export`.
"""
from torch.onnx import utils
return utils.select_model_mode_for_export(model, mode)
