mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
expose session option and provider options (#7112)
* expose session option and provider options * merge provider_names and provider_options * integrate into orttrainer options * fix doc string * fix a typo * Update orttraining/orttraining/python/training/orttrainer.py Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com> * Update orttraining/orttraining/python/training/orttrainer.py Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com> * Update orttraining/orttraining/python/training/orttrainer_options.py Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com> * fix the usage of provider_options * Update orttraining/orttraining/python/training/orttrainer.py Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com> * Update orttraining/orttraining/python/training/orttrainer.py Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com> * update expected result in tests * fix default provider options * minor update to trigger rebuild * minor update to trigger rebuild Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com>
This commit is contained in:
Tang, Cheng
GitHub
parent
c4ebc60870
commit
07201bac7a
3 changed files with 90 additions and 18 deletions
|
|
@ -95,7 +95,6 @@ class ORTTrainer(object):
|
|||
Inputs to the combined PyTorch model are concatenation of the :py:attr:`model`'s input and :py:attr:`loss_fn`'s label input.
|
||||
Outputs of the combined PyTorch model are concatenation of :py:attr:`loss_fn`'s loss output and :py:attr:`model`'s outputs.
|
||||
options (ORTTrainerOptions, default is None): options for additional features.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
|
@ -121,7 +120,9 @@ class ORTTrainer(object):
|
|||
ort_trainer = ORTTrainer(model, model_desc, optim_config, loss_fn)
|
||||
"""
|
||||
|
||||
def __init__(self, model, model_desc, optim_config, loss_fn=None, options=None):
|
||||
def __init__(self, model, model_desc, optim_config,
|
||||
loss_fn=None,
|
||||
options=None):
|
||||
assert model is not None, "'model' is required and must be either a 'torch.nn.Module' or ONNX model"
|
||||
assert isinstance(model_desc, dict), "'model_desc' must be a 'dict'"
|
||||
assert isinstance(optim_config, optim._OptimizerConfig),\
|
||||
|
|
@ -205,7 +206,8 @@ class ORTTrainer(object):
|
|||
self._train_step_info = TrainStepInfo(self.optim_config)
|
||||
self._training_session = None
|
||||
self._load_state_dict = None
|
||||
self._init_session()
|
||||
self._init_session(provider_options=self.options._validated_opts['provider_options'],
|
||||
session_options=self.options.session_options)
|
||||
|
||||
def eval_step(self, *args, **kwargs):
|
||||
r"""Evaluation step method
|
||||
|
|
@ -564,7 +566,10 @@ class ORTTrainer(object):
|
|||
|
||||
return onnx_model
|
||||
|
||||
def _create_ort_training_session(self, optimizer_state_dict={}):
|
||||
def _create_ort_training_session(self,
|
||||
optimizer_state_dict={},
|
||||
session_options=None,
|
||||
provider_options=None):
|
||||
# Validating frozen_weights names
|
||||
unused_frozen_weights = [n for n in self.options.utils.frozen_weights\
|
||||
if n not in [i.name for i in self._onnx_model.graph.initializer]]
|
||||
|
|
@ -662,7 +667,7 @@ class ORTTrainer(object):
|
|||
ort_parameters.model_with_training_graph_path = self.options.debug.graph_save_paths.model_with_training_graph_path
|
||||
|
||||
# SessionOptions
|
||||
session_options = ort.SessionOptions()
|
||||
session_options = ort.SessionOptions() if session_options is None else session_options
|
||||
session_options.use_deterministic_compute = self.options.debug.deterministic_compute
|
||||
if (self.options.graph_transformer.attn_dropout_recompute or
|
||||
self.options.graph_transformer.gelu_recompute or
|
||||
|
|
@ -676,10 +681,14 @@ class ORTTrainer(object):
|
|||
del self._training_session
|
||||
|
||||
# Set provider-specific options if needed
|
||||
def get_providers():
|
||||
def get_providers(provider_options):
|
||||
providers = ort.get_available_providers()
|
||||
|
||||
if 'cuda' in self.options.device.id.lower():
|
||||
|
||||
if provider_options:
|
||||
for provider_name in provider_options:
|
||||
providers[providers.index(provider_name)] = (provider_name, provider_options[provider_name])
|
||||
#default: using cuda
|
||||
elif 'cuda' in self.options.device.id.lower():
|
||||
cuda_ep_options = {"device_id": _utils.get_device_index(self.options.device.id)}
|
||||
|
||||
cuda_ep_name = ("ROCMExecutionProvider" if self.is_rocm_pytorch else "CUDAExecutionProvider")
|
||||
|
|
@ -700,7 +709,7 @@ class ORTTrainer(object):
|
|||
|
||||
# TrainingSession
|
||||
self._training_session = ort.TrainingSession(self._onnx_model.SerializeToString(), ort_parameters,
|
||||
session_options, get_providers())
|
||||
session_options, get_providers(provider_options))
|
||||
|
||||
# I/O bindings
|
||||
self._train_io_binding = self._training_session.io_binding()
|
||||
|
|
@ -731,9 +740,13 @@ class ORTTrainer(object):
|
|||
if self._load_state_dict:
|
||||
optimizer_state_dict = self._load_state_dict()
|
||||
|
||||
self._init_session(optimizer_state_dict)
|
||||
self._init_session(optimizer_state_dict,
|
||||
session_options=self.options.session_options,
|
||||
provider_options=self.options._validated_opts['provider_options'])
|
||||
|
||||
def _init_session(self, optimizer_state_dict={}):
|
||||
def _init_session(self, optimizer_state_dict={},
|
||||
session_options=None,
|
||||
provider_options=None):
|
||||
if self._onnx_model is None:
|
||||
return
|
||||
|
||||
|
|
@ -742,7 +755,9 @@ class ORTTrainer(object):
|
|||
|
||||
# Create training session used by train_step
|
||||
# pass all optimizer states to the backend
|
||||
self._create_ort_training_session(optimizer_state_dict)
|
||||
self._create_ort_training_session(optimizer_state_dict,
|
||||
session_options=session_options,
|
||||
provider_options=provider_options)
|
||||
|
||||
# Update model description to update dtype when mixed precision is enabled
|
||||
# C++ backend modifies model's output dtype from float32 to float16 for mixed precision
|
||||
|
|
@ -867,9 +882,20 @@ class ORTTrainer(object):
|
|||
# Keep all finite flag on CPU to match backend implementation
|
||||
# This prevents CPU -> GPU -> CPU copies between frontend and backend
|
||||
target_device = 'cpu'
|
||||
# the self.options.device may be a device that pytorch does not recognize.
|
||||
# in that case, we temporary prefer to leave the input/output on CPU and let ORT session
|
||||
# to move the data between device and host.
|
||||
# so output will be on the same device as input.
|
||||
try:
|
||||
test_pt_device = torch.device(target_device)
|
||||
except:
|
||||
#in this case, input/output must on CPU
|
||||
assert(input.device == 'cpu')
|
||||
target_device = 'cpu'
|
||||
|
||||
torch_tensor = torch.zeros(output_desc.shape, device=target_device,
|
||||
dtype=output_desc.dtype_amp if output_desc.dtype_amp else output_desc.dtype)
|
||||
iobinding.bind_output(output_desc.name, torch_tensor.device.type, _utils.get_device_index(self.options.device.id),
|
||||
iobinding.bind_output(output_desc.name, torch_tensor.device.type, _utils.get_device_index(target_device),
|
||||
_utils.dtype_torch_to_numpy(torch_tensor.dtype),
|
||||
list(torch_tensor.size()), torch_tensor.data_ptr())
|
||||
result[output_desc.name] = torch_tensor
|
||||
|
|
@ -1300,7 +1326,9 @@ class ORTTrainer(object):
|
|||
|
||||
# create a new training session after loading initializer states onto the onnx graph
|
||||
# pass the populated states to the training session to populate the backend graph
|
||||
self._init_session(optimizer_state_dict)
|
||||
self._init_session(optimizer_state_dict,
|
||||
session_options=self.options.session_options,
|
||||
provider_options=self.options._validated_opts['provider_options'])
|
||||
|
||||
def save_checkpoint(self, path, user_dict={}, include_optimizer_states=True):
|
||||
"""Persists ORTTrainer state dictionary on disk along with user_dict.
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import torch
|
|||
|
||||
from .optim import lr_scheduler
|
||||
from .amp import loss_scaler
|
||||
|
||||
import onnxruntime as ort
|
||||
|
||||
class ORTTrainerOptions(object):
|
||||
r"""Settings used by ONNX Runtime training backend
|
||||
|
|
@ -277,7 +277,18 @@ class ORTTrainerOptions(object):
|
|||
'default' : True
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
'provider_options':{
|
||||
'type': 'dict',
|
||||
'default': {},
|
||||
'required': False,
|
||||
'schema': {}
|
||||
},
|
||||
'session_options': {
|
||||
'type': 'SessionOptions',
|
||||
'nullable': True,
|
||||
'default': None
|
||||
},
|
||||
}
|
||||
|
||||
Keyword arguments:
|
||||
|
|
@ -383,6 +394,11 @@ class ORTTrainerOptions(object):
|
|||
_internal_use.enable_onnx_contrib_ops (bool, default is True)
|
||||
enable PyTorch to export nodes as contrib ops in ONNX.
|
||||
This flag may be removed anytime in the future.
|
||||
session_options (onnxruntime.SessionOptions):
|
||||
The SessionOptions instance that TrainingSession will use.
|
||||
provider_options (dict):
|
||||
The provider_options for customized execution providers. it is dict map from EP name to
|
||||
a key-value pairs, like {'EP1' : {'key1' : 'val1'}, ....}
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
|
@ -459,9 +475,13 @@ class ORTTrainerOptionsValidator(cerberus.Validator):
|
|||
_LOSS_SCALER = cerberus.TypeDefinition(
|
||||
'loss_scaler', (loss_scaler.LossScaler,), ())
|
||||
|
||||
_SESSION_OPTIONS = cerberus.TypeDefinition(
|
||||
'session_options', (ort.SessionOptions,),())
|
||||
|
||||
types_mapping = cerberus.Validator.types_mapping.copy()
|
||||
types_mapping['lr_scheduler'] = _LR_SCHEDULER
|
||||
types_mapping['loss_scaler'] = _LOSS_SCALER
|
||||
types_mapping['session_options'] = _SESSION_OPTIONS
|
||||
|
||||
|
||||
def _check_is_callable(field, value, error):
|
||||
|
|
@ -731,5 +751,17 @@ _ORTTRAINER_OPTIONS_SCHEMA = {
|
|||
'default': True
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
'provider_options':{
|
||||
'type': 'dict',
|
||||
'default_setter': lambda _: {},
|
||||
'required': False,
|
||||
'allow_unknown': True,
|
||||
'schema': {}
|
||||
},
|
||||
'session_options': {
|
||||
'type': 'session_options',
|
||||
'nullable': True,
|
||||
'default': None
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from onnxruntime.training import _utils, amp, checkpoint, optim, orttrainer, Tra
|
|||
model_desc_validation as md_val,\
|
||||
orttrainer_options as orttrainer_options
|
||||
import _test_commons,_test_helpers
|
||||
from onnxruntime import SessionOptions
|
||||
|
||||
|
||||
###############################################################################
|
||||
|
|
@ -99,7 +100,9 @@ def testORTTrainerOptionsDefaultValues(test_input):
|
|||
'extra_postprocess': None,
|
||||
'onnx_opset_version' : 12,
|
||||
'enable_onnx_contrib_ops': True,
|
||||
}
|
||||
},
|
||||
'provider_options':{},
|
||||
'session_options': None,
|
||||
}
|
||||
|
||||
actual_values = orttrainer_options.ORTTrainerOptions(test_input)
|
||||
|
|
@ -532,6 +535,15 @@ def testLRSchedulerUpdateImpl(lr_scheduler, expected_values):
|
|||
assert_allclose(lr_list[0],
|
||||
expected_values[optimization_step], rtol=rtol, err_msg="lr mismatch")
|
||||
|
||||
def testInstantiateORTTrainerOptions():
|
||||
session_options = SessionOptions()
|
||||
session_options.enable_mem_pattern = False
|
||||
provider_options = {'EP1': {'key':'val'}}
|
||||
opts = {'session_options' : session_options,
|
||||
'provider_options' : provider_options}
|
||||
opts = orttrainer.ORTTrainerOptions(opts)
|
||||
assert(opts.session_options.enable_mem_pattern is False)
|
||||
assert(opts._validated_opts['provider_options']['EP1']['key'] == 'val')
|
||||
|
||||
@pytest.mark.parametrize("step_fn, lr_scheduler, expected_lr_values, device", [
|
||||
('train_step', None, None, 'cuda'),
|
||||
|
|
|
|||
Loading…
Reference in a new issue