saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Allow custom exporter args + bug fix (#9242)

Browse source
This commit is contained in:
Thiago Crepaldi 2021-10-04 11:32:42 -04:00 committed by GitHub
parent 67ff339df7
commit 6e2f66ee9c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 97 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -125,6 +125,10 @@ class GraphExecutionManager(GraphExecutionInterface):
# To be instantiated in the concrete implementation of GraphExecutionManager
self._export_mode = None
# Exporter can take extra arguments for ORTModule extensions
# It cannot overlap with required/immutable arguments (validated in runtime)
self._export_extra_kwargs = {}
# Related to training graph shape inference
self._current_input_shape = None
# default execution order is priority-based for both dynamic/static shape input for now
@ -353,18 +357,23 @@ class GraphExecutionManager(GraphExecutionInterface):
try:
with torch.set_grad_enabled(self._enable_custom_autograd_function), \
_logger.suppress_os_stream_output(log_level=self._debug_options.logging.log_level):
required_export_kwargs = {'input_names': self._input_info.names,
'output_names': output_names,
'opset_version': ONNX_OPSET_VERSION,
'do_constant_folding': False,
'training': self._export_mode,
'dynamic_axes': self._input_info.dynamic_axes,
'verbose': self._debug_options.logging.log_level < LogLevel.WARNING,
'export_params': False,
'keep_initializers_as_inputs': True}
invalid_args = self._export_extra_kwargs.keys() & required_export_kwargs.keys()
assert len(invalid_args) == 0,\
f"The following PyTorch exporter arguments cannot be specified: '{invalid_args}'."
torch.onnx.export(self._flattened_module,
sample_inputs_as_tuple,
f,
input_names=self._input_info.names,
output_names=output_names,
opset_version=ONNX_OPSET_VERSION,
do_constant_folding=False,
training=self._export_mode,
dynamic_axes=self._input_info.dynamic_axes,
verbose=self._debug_options.logging.log_level < LogLevel.WARNING,
export_params=False,
keep_initializers_as_inputs=True)
**required_export_kwargs,
**self._export_extra_kwargs)
except Exception as e:
raise wrap_exception(ORTModuleONNXModelException,
RuntimeError(f'There was an error while exporting the PyTorch model to ONNX: '

81
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -23,7 +23,13 @@ from distutils.version import LooseVersion
from onnxruntime.training.ortmodule._custom_gradient_registry import register_gradient
from onnxruntime.training.ortmodule import ORTModule, _utils, _io, DebugOptions, LogLevel, _fallback, _graph_execution_manager
from onnxruntime.training.ortmodule import (ORTModule,
_utils,
_io,
DebugOptions,
LogLevel,
_fallback,
_graph_execution_manager)
import _test_helpers
# Import autocasting libs
@ -3862,7 +3868,7 @@ def test_ortmodule_determinism_flag(is_training,deterministic):
model = ORTModule(model)
model.train(is_training)
for i in range(5):
for _ in range(5):
x = torch.randn(N, D_in)
_ = model(x)
@ -3900,3 +3906,74 @@ def test_ortmodule_gradient_builder():
ort_prediction = run_step(ort_model, ort_x)
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
_test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad )
def test_override_pytorch_exporter_kwargs():
device = 'cuda'
N, D_in, H, D_out = 64, 784, 500, 10
x = torch.randn(N, D_in, device=device)
model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device)
ort_model = ORTModule(model)
ort_model._torch_module._execution_manager(True)._export_extra_kwargs = {'custom_opsets': None}
# Make sure model runs without any exception
prediction = ort_model(x)
assert prediction is not None
prediction = prediction.sum()
prediction.backward()
def test_override_pytorch_exporter_kwargs__invalid():
device = 'cuda'
N, D_in, H, D_out = 64, 784, 500, 10
x = torch.randn(N, D_in, device=device)
model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device)
ort_model = ORTModule(model)
ort_model._torch_module._execution_manager(True)._export_extra_kwargs = {'verbose': False}
with pytest.raises(_fallback.ORTModuleONNXModelException) as type_error:
_ = ort_model(x)
assert "The following PyTorch exporter arguments cannot be specified: '{'verbose'}'." in str(type_error.value)
def test_override_pytorch_exporter_kwargs_using_ortmodule_extension__invalid():
device = 'cuda'
class ORTModuleExtension(ORTModule):
def __init__(self, module, debug_options=None):
super().__init__(module, debug_options)
for training_mode in [False, True]:
self._torch_module._execution_manager(training_mode)._export_extra_kwargs = {'verbose': None}
N, D_in, H, D_out = 64, 784, 500, 10
x = torch.randn(N, D_in, device=device)
model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device)
ort_model = ORTModuleExtension(model)
with pytest.raises(_fallback.ORTModuleONNXModelException) as type_error:
_ = ort_model(x)
assert "The following PyTorch exporter arguments cannot be specified: '{'verbose'}'." in str(type_error.value)
def test_override_pytorch_exporter_kwargs_using_ortmodule_extension():
device = 'cuda'
class ORTModuleExtension(ORTModule):
def __init__(self, module, debug_options=None):
super().__init__(module, debug_options)
# modify GraphExecutionManager internally
for training_mode in [False, True]:
self._torch_module._execution_manager(training_mode)._export_extra_kwargs = {'custom_opsets': None}
N, D_in, H, D_out = 64, 784, 500, 10
x = torch.randn(N, D_in, device=device)
model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device)
ort_model = ORTModuleExtension(model)
# Make sure model runs without any exception
prediction = ort_model(x)
assert prediction is not None
prediction = prediction.sum()
prediction.backward()