From 88c95ef06bf81b829fb2e194131af01b34778d62 Mon Sep 17 00:00:00 2001 From: baijumeswani Date: Mon, 10 May 2021 10:59:47 -0700 Subject: [PATCH] Support for primitive types in ortmodule (#7588) --- .../ortmodule/_graph_execution_manager.py | 2 +- .../training/ortmodule/_inference_manager.py | 3 +- .../python/training/ortmodule/_io.py | 41 ++++++++-- .../training/ortmodule/_training_manager.py | 7 +- .../python/orttraining_test_ortmodule_api.py | 76 +++++++++++++++++++ 5 files changed, 119 insertions(+), 10 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 88d8dacefa..37a7e70e1b 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -229,7 +229,7 @@ class GraphExecutionManager(ABC): # Therefore, deepcopy only the data component of the input tensors for export. sample_inputs_copy, sample_kwargs_copy = _io.deepcopy_model_input(*inputs, **kwargs) # NOTE: Flattening the input will change the 'input schema', resulting in a re-export - sample_inputs_as_tuple = tuple(self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy)) + sample_inputs_as_tuple = tuple(self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy, self._device)) # Ops behaving differently under train/eval mode need to exported with the # correct training flag to reflect the expected behavior. # For example, the Dropout node in a model is dropped under eval mode. diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 1f9aea92f2..57f0ca579b 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -94,7 +94,8 @@ class InferenceManager(GraphExecutionManager): self._input_info, self._flattened_module.named_buffers(), inputs, - kwargs)) + kwargs, + self._device)) return _io.unflatten_user_output(self._module_output_schema, self._graph_info.user_output_names, diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py index d47beae74d..7c02082937 100644 --- a/orttraining/orttraining/python/training/ortmodule/_io.py +++ b/orttraining/orttraining/python/training/ortmodule/_io.py @@ -9,6 +9,22 @@ import inspect import torch import warnings +class _PrimitiveType(object): + _primitive_types = {int, bool, float} + @staticmethod + def is_primitive_type(value): + return type(value) in _PrimitiveType._primitive_types + + @staticmethod + def get_tensor(value, device): + return torch.tensor(value, device=device) + + @staticmethod + def get_primitive_dtype(value): + # If `value` is a boolean, save the value of the boolean in dtype. + # This way, if the value changes from one forward call to the next, the schema will mismatch, + # and the model will be re-exported. + return f"{str(type(value))}_{value}" if isinstance(value, bool) else str(type(value)) class _InputInfo(object): def __init__(self, @@ -40,11 +56,13 @@ class _InputInfo(object): \t#Positionals (non-None): {self.num_positionals_non_none} \tKeyword names: {self.keyword_names}''' - def flatten(self, args, kwargs): + def flatten(self, args, kwargs, device): '''Flatten args and kwargs in a single tuple of tensors with strict ordering''' - ret = list(args) - ret += [kwargs[name] for name in self.names if name in kwargs] + ret = [_PrimitiveType.get_tensor(arg, device) if _PrimitiveType.is_primitive_type(arg) else arg for arg in args] + ret += [_PrimitiveType.get_tensor(kwargs[name], device) if _PrimitiveType.is_primitive_type(kwargs[name]) + else kwargs[name] for name in self.names if name in kwargs] + return ret def unflatten(self, flat_args): @@ -55,7 +73,7 @@ class _InputInfo(object): if name in self.keyword_names} return args, kwargs -def _combine_input_buffers_initializers(param_names, onnx_input_names, input_info, buffer_names, inputs, kwargs): +def _combine_input_buffers_initializers(param_names, onnx_input_names, input_info, buffer_names, inputs, kwargs, device): '''Creates forward `*inputs` list from user input and PyTorch initializers ONNX Runtime forward requires an ordered list of: @@ -88,6 +106,8 @@ def _combine_input_buffers_initializers(param_names, onnx_input_names, input_inf raise KeyError(f'Registered buffer name {name} not found.') if inp is not None: + if _PrimitiveType.is_primitive_type(inp): + inp = _PrimitiveType.get_tensor(inp, device) result.append(inp) else: raise RuntimeError(f'Input is present in ONNX graph but not provided: {name}.') @@ -201,6 +221,8 @@ def _extract_schema(data): if data is None: return None + elif _PrimitiveType.is_primitive_type(data): + return _TensorStub(dtype=_PrimitiveType.get_primitive_dtype(data), shape_dims=0) # Depth first traversal to iterate over the data to replace every tensor with a stub elif isinstance(data, torch.Tensor): return _TensorStub(dtype=str(data.dtype), shape_dims=len(data.size())) @@ -305,10 +327,17 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, inputs, kwarg return dynamic_axes def _add_input(name, input, onnx_graph, onnx_graph_input_names): - if input is not None and (onnx_graph is None or name in onnx_graph_input_names): + if input is None: + # Drop all None inputs. + return + + # InputInfo should contain all the names irrespective of whether they are + # a part of the onnx graph or not. + input_names.append(name) + + if (onnx_graph is None or name in onnx_graph_input_names) and isinstance(input, torch.Tensor): if input.requires_grad: input_names_require_grad.append(name) - input_names.append(name) dynamic_axes.update(_add_dynamic_shape(name, input)) input_shape.append(list(input.size())) diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 9034bab74e..00f19e1ea5 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -73,7 +73,9 @@ class TrainingManager(GraphExecutionManager): kwargs) # Reinitialize graph builder if the inputs or initializers requiring gradient have changed. - build_gradient_graph = build_gradient_graph or self._reinitialize_graph_builder(input_info) + # Order of or operation is important here because we always need to call + # _reinitialize_graph_builder irrespective of the value of build_gradient_graph. + build_gradient_graph = self._reinitialize_graph_builder(input_info) or build_gradient_graph # Build the gradient graph if build_gradient_graph: @@ -191,7 +193,8 @@ class TrainingManager(GraphExecutionManager): self._input_info, self._flattened_module.named_buffers(), inputs, - kwargs))) + kwargs, + self._device))) def _build_graph(self): """Build an optimized gradient graph using the module_graph_builder""" diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index ec4ee18112..5d393094e2 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -2252,3 +2252,79 @@ def test_forward_call_lots_None(): a, b, c, d, e, f, y, z) run_step(a.item() + b.item() + c.item() + d.item() + e.item() + f.item() + y.item() + z.item(), **{'a': a, 'b': b, 'c': c, 'd': d, 'e': e, 'f': f, 'y': y, 'z': z}) + +@pytest.mark.parametrize("bool_argument", [True, False]) +@pytest.mark.parametrize("int_argument", [100, 100000, 100000000, -100, -100000, -100000000]) +@pytest.mark.parametrize("float_argument", [1.23, 11209123.12452, 12093702935.1249863, -1.23, -11209123.12452, -12093702935.1249863]) +def test_primitive_inputs(bool_argument, int_argument, float_argument): + class PrimitiveTypesInputNet(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(PrimitiveTypesInputNet, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(hidden_size, num_classes) + + def forward(self, input1, bool_argument, int_argument, float_argument): + input1 = input1 + int_argument + float_argument + if bool_argument: + out = self.fc1(input1) + out = self.relu(out) + out = self.fc2(out) + else: + out = self.fc1(input1) + out = self.fc2(out) + out = self.relu(out) + return out + + assert type(bool_argument) is bool + assert type(int_argument) is int + assert type(float_argument) is float + + device = 'cuda' + N, D_in, H, D_out = 32, 784, 500, 10 + pt_model = PrimitiveTypesInputNet(D_in, H, D_out).to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + input1 = torch.randn(N, D_in, device=device) + pt_out = pt_model(input1, bool_argument, int_argument, float_argument) + ort_out = ort_model(input1, bool_argument, int_argument, float_argument) + assert torch.equal(pt_out, ort_out) + +@pytest.mark.parametrize("bool_arguments", [(True, False), (False, True)]) +def test_changing_bool_input_re_exports_model(bool_arguments): + class PrimitiveTypesInputNet(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(PrimitiveTypesInputNet, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(hidden_size, num_classes) + + def forward(self, input1, bool_argument): + if bool_argument: + out = self.fc1(input1) + out = self.relu(out) + out = self.fc2(out) + else: + out = self.fc1(input1) + out = self.fc2(out) + out = self.relu(out) + return out + + assert type(bool_arguments[0]) is bool + assert type(bool_arguments[1]) is bool + + device = 'cuda' + N, D_in, H, D_out = 32, 784, 500, 10 + pt_model = PrimitiveTypesInputNet(D_in, H, D_out).to(device) + ort_model = ORTModule(pt_model) + + input1 = torch.randn(N, D_in, device=device) + ort_model(input1, bool_arguments[0]) + exported_model1 = ort_model._execution_manager(ort_model._is_training())._onnx_model + + ort_model(input1, bool_arguments[1]) + exported_model2 = ort_model._execution_manager(ort_model._is_training())._onnx_model + + assert exported_model1 != exported_model2