diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 7da8782fe9..aa5e48005e 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -10,7 +10,6 @@ import numpy as np from inspect import signature from torch.utils.dlpack import from_dlpack -from torch._six import container_abcs # Needed to re-implement PyTorch's cpu,cuda,to methods from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict @@ -69,12 +68,6 @@ def _create_iobinding(io_binding, inputs, model, device): io_binding.bind_output(value_info.name, device.type, device_id=_get_device_index(device)) -def _deepcopy_model_input(*inputs, **kwargs): - sample_inputs_copy = [] - for model_input in inputs: - sample_inputs_copy.append(model_input.data if isinstance(model_input, torch.Tensor) else model_input) - sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy)) - return sample_inputs_copy def _onnx_value_info_to_buffer_tensor(value_info, device): '''Create a torch zeroed tensor with the same shape and type of `value_info`''' @@ -83,7 +76,21 @@ def _onnx_value_info_to_buffer_tensor(value_info, device): dtype = _utils.dtype_onnx_to_torch(value_info.type.tensor_type.elem_type) return torch.zeros(shape, device=device, dtype=dtype) -def _parse_inputs_for_onnx_export(module, *inputs, **kwargs): +# TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8, +# and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor +# from forward graph outputs will be converted to torch.uint8 tensor. When this tensor +# is feeded to backward graph as input, it will cause data type mismatch issue during +# inference session running. We cannot change the from_dlpack() in PyTorch side, so we +# have to handle this specially, which will introduce a cast here and there is data copied. +# Always cast from torch.uint8 to torch.bool is not logically right, we need to check the +# real data type of the inputs in the backeard graph, and perform the cast only necessary. +def _ort_output_to_torch_tensor(ort_output): + tensor = from_dlpack(ort_output.to_dlpack()) + return tensor.to(torch.bool) if tensor.dtype == torch.uint8 else tensor + +def _extract_input_information(module, *inputs, **kwargs): + '''Returns the all input names, dynamic_axes information and input names that require gradient''' + # Ignore optional *inputs explicitly specified as None sig = signature(module.forward) all_input_names = sig.parameters.keys() @@ -99,66 +106,10 @@ def _parse_inputs_for_onnx_export(module, *inputs, **kwargs): input_names.append(name) dynamic_axes[name] = {} for dim_idx in range(len(inputs[input_idx].shape)): - dynamic_axes[name].update({dim_idx : 'input{}_dim{}'.format(input_idx, dim_idx)}) + dynamic_axes[name].update({dim_idx : f'input{input_idx}_dim{dim_idx}'}) + return input_names, dynamic_axes, input_names_require_grad -def _parse_outputs_for_onnx_export(module, inputs): - - def _create_output_dim_names(output, output_idx, from_sequence): - if from_sequence and not isinstance(output, torch.Tensor): - raise TypeError('ORTModule does not support the following model output type {} within a Sequence'.format(type(sample_outputs))) - output_names, dynamic_axes = [], {} - name = 'out{}'.format(output_idx) - output_names.append(name) - dynamic_axes[name] = {} - for dim_idx in range(len(output.shape)): - dynamic_axes[name].update({dim_idx : '{}_dim{}'.format(name, dim_idx)}) - return output_names, dynamic_axes - - # Do an inference to grab outputs - is_train_mode = module.training - module.eval() - with torch.no_grad(): - # Deepcopy inputs, since input values may change after model run. - sample_inputs_copy = _deepcopy_model_input(*inputs) - try: - # Deepcopy model, in case model is stateful and changes after model run. - model_copy = copy.deepcopy(module) - except Exception: - model_copy = module - warnings.warn("This model cannot be deep copied (or pickled), which is a required step for stateful models to be properly exported to ONNX." - " Compute will continue, but unexpected results may occur!") - - sample_outputs = model_copy(*sample_inputs_copy) - output_names = [] - output_dynamic_axes = {} - if isinstance(sample_outputs, torch.Tensor): - output_names, output_dynamic_axes = _create_output_dim_names(sample_outputs, 0, False) - elif isinstance(sample_outputs, container_abcs.Mapping): - raise NotImplementedError('Dictionaries are not supported as output yet') - elif isinstance(sample_outputs, container_abcs.Sequence): - for idx, out in enumerate(sample_outputs): - tmp_output_names, tmp_output_dynamic_axes = _create_output_dim_names(out, idx, True) - output_names += tmp_output_names - output_dynamic_axes.update(tmp_output_dynamic_axes) - else: - raise TypeError('ORTModule does not support the following model output type {}'.format(type(sample_outputs))) - if is_train_mode: - module.train() - return output_names, output_dynamic_axes - -# TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8, -# and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor -# from forward graph outputs will be converted to torch.uint8 tensor. When this tensor -# is feeded to backward graph as input, it will cause data type mismatch issue during -# inference session running. We cannot change the from_dlpack() in PyTorch side, so we -# have to handle this specially, which will introduce a cast here and there is data copied. -# Always cast from torch.uint8 to torch.bool is not logically right, we need to check the -# real data type of the inputs in the backeard graph, and perform the cast only necessary. -def _ort_output_to_torch_tensor(ort_output): - tensor = from_dlpack(ort_output.to_dlpack()) - return tensor.to(torch.bool) if tensor.dtype == torch.uint8 else tensor - class ORTModule(torch.nn.Module): def __init__(self, module): @@ -209,7 +160,9 @@ class ORTModule(torch.nn.Module): self._module_gradient_graph_builder.initialize(self._onnx_training.SerializeToString(), grad_builder_config) def _build_training_graph(self, *inputs, **kwargs): - self._onnx_training = self._get_forward_graph(*inputs, **kwargs) + input_names, dynamic_axes, self._input_names_require_grad = \ + _extract_input_information(self._original_module, *inputs, **kwargs) + self._onnx_training = self._get_forward_graph(input_names, dynamic_axes, *inputs, **kwargs) if self._save_onnx: onnx.save(self._onnx_training, self._save_onnx_prefix + '_full_training.onnx') @@ -307,7 +260,7 @@ class ORTModule(torch.nn.Module): if not self._onnx_training: self._build_training_graph(*inputs, **kwargs) - _, _, input_names_require_grad = _parse_inputs_for_onnx_export(self._original_module, *inputs, **kwargs) + _, _, input_names_require_grad = _extract_input_information(self._original_module, *inputs, **kwargs) # If inputs requiring gradient change from one call to forward to the next, the module_gradient_graph_builder # needs to be reinitialized so it can compute the backward output for the new inputs that require_grad if input_names_require_grad != self._input_names_require_grad: @@ -483,36 +436,33 @@ class ORTModule(torch.nn.Module): return result - def _get_forward_graph(self, *inputs, **kwargs): + def _get_forward_graph(self, input_names, dynamic_axes, *inputs, **kwargs): '''Exports PyTorch `module` to ONNX with training flag, using `*inputs` as input TODO: How to support dynamic axes? Dimensions are determined by samples TODO: How to ingest **kwargs in proper order during export? ''' - - # Setup dynamic axes for onnx model - input_names, dynamic_axes, self._input_names_require_grad = _parse_inputs_for_onnx_export(self._original_module, *inputs, **kwargs) - output_names, output_dynamic_axes = _parse_outputs_for_onnx_export(self._original_module, inputs) - dynamic_axes.update(output_dynamic_axes) - - # TODO: Support contrib OPs support? user model has no hint - # from onnxruntime.training import register_custom_ops_pytorch_exporter - # register_custom_ops_pytorch_exporter.register_custom_op() - - # Export torch.nn.Module to ONNX + # Export the model to memory f = io.BytesIO() # Deepcopy inputs, since input values may change after model run. # NOTE: Inputs may contain tensors that have attributes preventing their deepcopy (example grad_fn). # Therefore, deepcopy only the data component of the input tensors for export. - sample_inputs_copy = _deepcopy_model_input(*inputs, **kwargs) + sample_inputs_copy = [] + for model_input in inputs: + sample_inputs_copy.append(model_input.data if isinstance(model_input, torch.Tensor) else model_input) + sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy)) + + # TODO: Support contrib OPs support? user model has no hint + # from onnxruntime.training import register_custom_ops_pytorch_exporter + # register_custom_ops_pytorch_exporter.register_custom_op() with torch.no_grad(): + # Export torch.nn.Module to ONNX torch.onnx.export(self._original_module, sample_inputs_copy, f, input_names=input_names, - output_names=output_names, opset_version=ONNX_OPSET_VERSION, do_constant_folding=False, training=torch.onnx.TrainingMode.TRAINING, diff --git a/orttraining/orttraining/test/python/_test_helpers.py b/orttraining/orttraining/test/python/_test_helpers.py index 45089f2a58..ca31c15cde 100644 --- a/orttraining/orttraining/test/python/_test_helpers.py +++ b/orttraining/orttraining/test/python/_test_helpers.py @@ -96,14 +96,14 @@ def assert_optim_state(expected_state, actual_state, rtol=1e-7, atol=0): "Update_Count": update_tensor # if optimizer is adam, absent otherwise }, ... - "shared_optimizer_state": # if optimizer is shared, absent otherwise. + "shared_optimizer_state": # if optimizer is shared, absent otherwise. So far, only lamb optimizer uses this. { "step": step_tensor # int array of size 1 } Args: - expected_state (dict(dict())): Expected optimizer state + expected_state (dict(dict())): Expected optimizer state actual_state (dict(dict())): Actual optimizer state rtol (float, default is 1e-7): Max relative difference atol (float, default is 0): Max absolute difference @@ -114,24 +114,6 @@ def assert_optim_state(expected_state, actual_state, rtol=1e-7, atol=0): assert_allclose(v, expected_state[param_name][k], rtol=rtol, atol=atol, err_msg=f"Optimizer state mismatch for param {param_name}, key {k}") -def is_dynamic_axes(model): - # Check inputs - for inp in model._onnx_training.graph.input: - shape = inp.type.tensor_type.shape - if shape: - for dim in shape.dim: - if dim.dim_param and not isinstance(dim.dim_param, str): - return False - - # Check outputs - for out in model._onnx_training.graph.output: - shape = out.type.tensor_type.shape - if shape: - for dim in shape.dim: - if dim.dim_param and not isinstance(dim.dim_param, str): - return False - return True - # TODO: thiagofc: Checkpoint related for redesign def _get_name(name): if os.path.exists(name): diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 3c5f788d4c..eb75952ace 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -9,7 +9,6 @@ from unittest.mock import patch import onnxruntime from onnxruntime.training import ORTModule -import _test_helpers # PyTorch model definitions for tests @@ -123,9 +122,10 @@ def test_forward_call_single_positional_argument(): model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) model = ORTModule(model) x = torch.randn(N, D_in, device=device) - # Make sure model runs without any exception - output = model(x) - assert output is not None + try: + model(x) + except Exception as exception: + raise exception def test_forward_call_multiple_positional_arguments(): device = 'cuda' @@ -135,10 +135,10 @@ def test_forward_call_multiple_positional_arguments(): model = ORTModule(model) x = torch.randn(N, D_in, device=device) y = torch.randn(N, D_in, device=device) - - # Make sure model runs without any exception - output = model(x, y) - assert output is not None + try: + model(x, y) + except Exception as exception: + raise exception def test_forward_call_positional_arguments(): device = 'cuda' @@ -147,10 +147,10 @@ def test_forward_call_positional_arguments(): model = NeuralNetPositionalArguments(input_size=D_in, hidden_size=H, num_classes=D_out).to(device) model = ORTModule(model) args = [torch.randn(N, D_in, device=device), torch.randn(N, D_in, device=device), torch.randn(N, D_in, device=device)] - - # Make sure model runs without any exception - output = model(*args) - assert output is not None + try: + model(*args) + except Exception as exception: + raise exception def test_forward_call_keyword_arguments(): device = 'cuda' @@ -161,10 +161,10 @@ def test_forward_call_keyword_arguments(): x = torch.randn(N, D_in, device=device) y = torch.randn(N, D_in, device=device) z = torch.randn(N, D_in, device=device) - - # Make sure model runs without any exception - output = model(x, y, z) - assert output is not None + try: + model(x, y, z) + except Exception as exception: + raise exception def test_forward_call_positional_and_keyword_arguments(): device = 'cuda' @@ -176,10 +176,10 @@ def test_forward_call_positional_and_keyword_arguments(): x = torch.randn(N, D_in, device=device) y = torch.randn(N, D_in, device=device) z = torch.randn(N, D_in, device=device) - - # Make sure model runs without any exception - output = model(a, x, y, z) - assert output is not None + try: + model(a, x, y, z) + except Exception as exception: + raise exception def test_model_cuda(): original_device = 'cpu' @@ -345,99 +345,3 @@ def test_gpu_reserved_memory_with_torch_no_grad(): assert mem_reserved_after_export_with_torch_no_grad < mem_reserved_after_export_without_torch_no_grad assert mem_reserved_before_export < mem_reserved_after_export_with_torch_no_grad - -@pytest.mark.parametrize("device", ['cpu', 'cuda']) -def test_exception_raised_for_dict_return_value_module(device): - class NeuralNetDictOutput(torch.nn.Module): - def __init__(self, input_size, hidden_size, num_classes): - super(NeuralNetDictOutput, self).__init__() - - self.fc1_1 = torch.nn.Linear(input_size, hidden_size) - self.relu1 = torch.nn.ReLU() - self.fc1_2 = torch.nn.Linear(hidden_size, num_classes) - - self.fc2_1 = torch.nn.Linear(input_size, hidden_size) - self.relu2 = torch.nn.ReLU() - self.fc2_2 = torch.nn.Linear(hidden_size, num_classes) - - self.fc3_1 = torch.nn.Linear(input_size, hidden_size) - self.relu3 = torch.nn.ReLU() - self.fc3_2 = torch.nn.Linear(hidden_size, num_classes) - - def forward(self, input1, input2, input3): - out1 = self.fc1_2(self.relu1(self.fc1_1(input1))) - out2 = self.fc2_2(self.relu2(self.fc2_1(input2))) - out3 = self.fc3_2(self.relu3(self.fc3_1(input2))) - return {'a': out1, 'b': out2, 'c': out3} - - N, D_in, H, D_out = 64, 784, 500, 10 - model = NeuralNetDictOutput(D_in, H, D_out).to(device) - model = ORTModule(model) - x = torch.randn(N, D_in, device=device) - y = torch.randn(N, D_in, device=device) - z = torch.randn(N, D_in, device=device) - - with pytest.raises(NotImplementedError) as not_implemented_error: - model(x, y, z) - assert str(not_implemented_error.value) == 'Dictionaries are not supported as output yet' - -@pytest.mark.parametrize("device", ['cpu', 'cuda']) -def test_exception_raised_for_custom_class_return_value_module(device): - class CustomClass(object): - def __init__(self, out1, out2, out3): - self.out1 = out1 - self.out2 = out2 - self.out3 = out3 - - class NeuralNetCustomClassOutput(torch.nn.Module): - def __init__(self, input_size, hidden_size, num_classes): - super(NeuralNetCustomClassOutput, self).__init__() - - self.fc1_1 = torch.nn.Linear(input_size, hidden_size) - self.relu1 = torch.nn.ReLU() - self.fc1_2 = torch.nn.Linear(hidden_size, num_classes) - - self.fc2_1 = torch.nn.Linear(input_size, hidden_size) - self.relu2 = torch.nn.ReLU() - self.fc2_2 = torch.nn.Linear(hidden_size, num_classes) - - self.fc3_1 = torch.nn.Linear(input_size, hidden_size) - self.relu3 = torch.nn.ReLU() - self.fc3_2 = torch.nn.Linear(hidden_size, num_classes) - - def forward(self, input1, input2, input3): - out1 = self.fc1_2(self.relu1(self.fc1_1(input1))) - out2 = self.fc2_2(self.relu2(self.fc2_1(input2))) - out3 = self.fc3_2(self.relu3(self.fc3_1(input2))) - return CustomClass(out1, out2, out3) - - N, D_in, H, D_out = 64, 784, 500, 10 - model = NeuralNetCustomClassOutput(D_in, H, D_out).to(device) - model = ORTModule(model) - x = torch.randn(N, D_in, device=device) - y = torch.randn(N, D_in, device=device) - z = torch.randn(N, D_in, device=device) - - with pytest.raises(TypeError) as runtime_error: - model(x, y, z) - assert 'ORTModule does not support the following model output type' in str(runtime_error.value) - -def test_dynamic_axes_config_NeuralNetSinglePositionalArgument(device = 'cuda'): - N, D_in, H, D_out = 64, 784, 500, 10 - model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) - model = ORTModule(model) - x = torch.randn(N, D_in, device=device) - - output = model(x) - assert output is not None - assert _test_helpers.is_dynamic_axes(model) - del model, output - -def test_dynamic_axes_config_BertForSequenceClassification(device = 'cuda'): - model = _get_bert_for_sequence_classification_model(device) - model = ORTModule(model).to(device) - x, y, z = _get_bert_for_sequence_classification_sample_data(device) - - output = model(x, y, None, None, None, None, z) - assert output is not None - assert _test_helpers.is_dynamic_axes(model)