diff --git a/orttraining/orttraining/python/training/_ortmodule_output_transformation.py b/orttraining/orttraining/python/training/_ortmodule_output_transformation.py new file mode 100644 index 0000000000..9644a382e9 --- /dev/null +++ b/orttraining/orttraining/python/training/_ortmodule_output_transformation.py @@ -0,0 +1,220 @@ +from collections import abc +import copy +import functools +import torch +import warnings + +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)) + + sample_kwargs_copy = {} + for name, model_input in kwargs.items(): + sample_kwargs_copy[name] = model_input.data if isinstance(model_input, torch.Tensor) else model_input + sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy) + + return sample_inputs_copy, sample_kwargs_copy + +class _TensorStub: + # Stub for a torch.Tensor value to be used to formulate the output schema + pass + +def populate_user_output_from_schema_and_outputs(output_schema, output_names, outputs): + """Follows the schema to generate an output that is expected by the user""" + + def _replace_stub_with_tensor_value(user_output, outputs, output_idx): + # Recursively traverse across user_output and replace all _TensorStub + # with torch.Tensor values from outputs following output_idx + + if isinstance(user_output, _TensorStub): + output_idx[0] += 1 + return outputs[output_idx[0]-1] + + if isinstance(user_output, abc.Sequence): + sequence_type = type(user_output) + user_output = list(user_output) + for idx in range(len(user_output)): + user_output[idx] = _replace_stub_with_tensor_value(user_output[idx], outputs, output_idx) + try: + # namedtuple can be created by passing the list sequence to method _make + user_output = sequence_type._make(user_output) + except AttributeError: + # If attribute error encountered, create the sequence directly + user_output = sequence_type(user_output) + elif isinstance(user_output, abc.Mapping): + for key in sorted(user_output): + user_output[key] = _replace_stub_with_tensor_value(user_output[key], outputs, output_idx) + else: + raise TypeError(f'ORTModule does not support the following model output type {type(user_output)}.') + + return user_output + + # Order the outputs according to the names so that the traversal order is consistent + outputs = [x for _, x in sorted(zip(output_names, outputs))] + + # Replace every _TensorStub value in the schema with the torch.Tensor outputs calculated + output_schema_copy = copy.deepcopy(output_schema) + output_idx = [0] + user_output = _replace_stub_with_tensor_value(output_schema_copy, outputs, output_idx) + + return user_output + +def _extract_output_schema(output): + """Extract the output schema by replacing every torch.Tensor value with _TensorStub""" + + # Depth first traversal to iterate over the output to replace every tensor with a stub + if isinstance(output, torch.Tensor): + return _TensorStub() + + if isinstance(output, abc.Sequence): + sequence_type = type(output) + output = list(output) + for idx in range(len(output)): + output[idx] = _extract_output_schema(output[idx]) + try: + # namedtuple can be created by passing the list sequence to method _make + output = sequence_type._make(output) + except AttributeError: + # If attribute error encountered, create the sequence directly + output = sequence_type(output) + elif isinstance(output, abc.Mapping): + for key in sorted(output): + output[key] = _extract_output_schema(output[key]) + else: + raise TypeError(f'ORTModule does not support the following model output type {type(output)}') + + return output + +def _parse_outputs_and_extract_names_and_dynamic_axes(module_output): + """Parses through the module output and returns output names and dynamic axes""" + + def _populate_output_names_and_dynamic_axes(output, output_names, output_dynamic_axes, output_idx): + # Depth first traversal to traverse through the entire output collecting output names and dynamic axes + + if isinstance(output, torch.Tensor): + output_name = f'output{output_idx[0]}' + output_idx[0] += 1 + output_names.append(output_name) + output_dynamic_axes[output_name] = {} + for dim_idx in range(len(output.shape)): + output_dynamic_axes[output_name].update({dim_idx: f'{output_name}_dim{dim_idx}'}) + return + + if isinstance(output, abc.Sequence): + for value in output: + _populate_output_names_and_dynamic_axes(value, output_names, output_dynamic_axes, output_idx) + elif isinstance(output, abc.Mapping): + for _, value in sorted(output.items()): + _populate_output_names_and_dynamic_axes(value, output_names, output_dynamic_axes, output_idx) + else: + raise TypeError(f'ORTModule does not support the following model output type {type(output)}') + + output_names = [] + output_dynamic_axes = {} + output_idx = [0] + _populate_output_names_and_dynamic_axes(module_output, output_names, output_dynamic_axes, output_idx) + + return output_names, output_dynamic_axes + +def get_flattened_output_module(original_module): + """Returns a torch.nn.Module that flattens the output of the original module in its forward method""" + + def _transform_output_to_flat_tuple(output): + """Converts the output to a flat tuple by iterating over the entire output structure""" + + def _flatten_output(output, flat_output): + # Recursively traverse over the output and populate the flat_output with torch.Tensors + + if isinstance(output, torch.Tensor): + flat_output.append(output) + elif isinstance(output, abc.Sequence): + for value in output: + _flatten_output(value, flat_output) + elif isinstance(output, abc.Mapping): + for _, value in sorted(output.items()): + _flatten_output(value, flat_output) + else: + raise TypeError(f'ORTModule does not support the following output type {type(output)}.') + + flat_output = [] + _flatten_output(output, flat_output) + return tuple(flat_output) + + class FlattenedOutputModule(torch.nn.Module): + def __init__(self, module): + super(FlattenedOutputModule, self).__init__() + self._base_module = module + + def _forward(self, *args, **kwargs): + return _transform_output_to_flat_tuple(self._base_module(*args, **kwargs)) + + # Exporter does not support use of **kwargs in the forward method. + # Work around it by making the signature of the forward method to resemble that of the + # original model + # Copy the forward signature from the original PyTorch module. + self.forward = _forward.__get__(self) + functools.update_wrapper(self.forward.__func__, module.forward.__func__) + + return FlattenedOutputModule(original_module) + +def parse_inputs_for_onnx_export(all_input_names, onnx_graph, *inputs, **kwargs): + # Ignore optional inputs explicitly specified as None + # ONNX exporter may remove unused inputs + onnx_graph_input_names = [] + if onnx_graph is not None: + onnx_graph_input_names = set([inp.name for inp in onnx_graph.graph.input]) + + input_names = [] + dynamic_axes = {} + input_names_require_grad = [] + input_shape = [] + + for input_idx, name in enumerate(all_input_names): + inp = None + if input_idx < len(inputs) and inputs[input_idx] is not None: + inp = inputs[input_idx] + elif name in kwargs and kwargs[name] is not None: + inp = kwargs[name] + if inp is not None and (onnx_graph is None or name in onnx_graph_input_names): + if inp.requires_grad: + # input_names_require_grad holds all input tensors that have requires_grad + input_names_require_grad.append(name) + + input_names.append(name) + dynamic_axes[name] = {} + for dim_idx in range(len(inp.shape)): + dynamic_axes[name].update({dim_idx : f'input{input_idx}_dim{dim_idx}'}) + + input_shape.append(list(inp.size())) + return input_names, dynamic_axes, input_names_require_grad, input_shape + +def parse_outputs_for_onnx_export_and_extract_output_schema(module, inputs, kwargs): + + # Do an inference to grab outputs + is_train_mode = module.training + module.eval() + output_names = None + output_dynamic_axes = None + with torch.no_grad(): + # Deepcopy inputs, since input values may change after model run. + sample_inputs_copy, sample_kwargs_copy = deepcopy_model_input(*inputs, **kwargs) + 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, **sample_kwargs_copy) + + # Parse the output and extract the output_names and output_dynamic_axes to be used for onnx export + output_names, output_dynamic_axes = \ + _parse_outputs_and_extract_names_and_dynamic_axes(sample_outputs) + if is_train_mode: + module.train() + + # Return output names, output dynamic axes and output schema + return output_names, output_dynamic_axes, _extract_output_schema(sample_outputs) diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 5ce580e758..cbbfc9d4d3 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -1,24 +1,19 @@ -import copy import io import logging import onnx import onnxruntime -import os import torch -import warnings -import numpy as np from inspect import signature from torch.utils.dlpack import from_dlpack from torch.utils.cpp_extension import load_inline -from collections import abc # Needed to re-implement PyTorch's cpu,cuda,to methods from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict from onnxruntime.capi import _pybind_state as C from onnxruntime.training import register_custom_ops_pytorch_exporter -from . import _utils +from . import _utils, _ortmodule_output_transformation ONNX_OPSET_VERSION = 12 @@ -41,19 +36,6 @@ def _create_iobinding(io_binding, inputs, model, device): io_binding.bind_output(value_info.name, device.type, device_id=_utils.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)) - - sample_kwargs_copy = {} - for name, model_input in kwargs.items(): - sample_kwargs_copy[name] = model_input.data if isinstance(model_input, torch.Tensor) else model_input - sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy) - - return sample_inputs_copy, sample_kwargs_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`''' @@ -61,109 +43,6 @@ 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(all_input_names, onnx_graph, *inputs, **kwargs): - # Ignore optional inputs explicitly specified as None - # ONNX exporter may remove unused inputs - onnx_graph_input_names = [] - if onnx_graph is not None: - onnx_graph_input_names = set([inp.name for inp in onnx_graph.graph.input]) - - input_names = [] - dynamic_axes = {} - input_names_require_grad = [] - input_shape = [] - - for input_idx, name in enumerate(all_input_names): - inp = None - if input_idx < len(inputs) and inputs[input_idx] is not None: - inp = inputs[input_idx] - elif name in kwargs and kwargs[name] is not None: - inp = kwargs[name] - if inp is not None and (onnx_graph is None or name in onnx_graph_input_names): - if inp.requires_grad: - # input_names_require_grad holds all input tensors that have requires_grad - input_names_require_grad.append(name) - - input_names.append(name) - dynamic_axes[name] = {} - for dim_idx in range(len(inp.shape)): - dynamic_axes[name].update({dim_idx : f'input{input_idx}_dim{dim_idx}'}) - - input_shape.append(list(inp.size())) - return input_names, dynamic_axes, input_names_require_grad, input_shape - -def _parse_outputs_for_onnx_export(module, inputs, kwargs): - - def _create_output_dim_names_from_mapping(output): - output_names, dynamic_axes = [], {} - for name, value in output.items(): - if not isinstance(value, torch.Tensor): - raise TypeError('ORTModule does not support the following model output type {} within a Mapping'.format(type(value))) - output_names.append(name) - dynamic_axes[name] = {} - for dim_idx in range(len(value.shape)): - dynamic_axes[name].update({dim_idx: '{}_dim{}'.format(name, dim_idx)}) - return output_names, dynamic_axes - - 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(output))) - output_names, dynamic_axes = [], {} - name = 'output{}'.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() - output_names = [] - output_dynamic_axes = {} - sample_output_type = None - with torch.no_grad(): - # Deepcopy inputs, since input values may change after model run. - sample_inputs_copy, sample_kwargs_copy = _deepcopy_model_input(*inputs, **kwargs) - 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, **sample_kwargs_copy) - sample_output_type = type(sample_outputs) - if isinstance(sample_outputs, torch.Tensor): - output_names, output_dynamic_axes = _create_output_dim_names(sample_outputs, 0, False) - elif isinstance(sample_outputs, abc.Mapping): - output_names, output_dynamic_axes = _create_output_dim_names_from_mapping(sample_outputs) - elif isinstance(sample_outputs, abc.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, sample_output_type - -def _populate_user_output(user_output_type, user_output_names, user_outputs): - if issubclass(user_output_type, Mapping): - key_value_pairs = [(user_output_names[i], user_outputs[i]) for i in range(len(user_output_names))] - return user_output_type(key_value_pairs) - elif issubclass(user_output_type, tuple): - try: - # Try constructing the user named tuple from the output tuple - return user_output_type(*user_outputs) - except TypeError: - # The expected output type is not a namedtuple, but is a regular tuple type - pass - - return user_outputs - # 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 @@ -206,6 +85,9 @@ class ORTModule(torch.nn.Module): # User module is wrapped to use its initializers and save computed gradients self._original_module = module + # Get the module that flattens the output from the original module into a tuple + self._flattened_output_module = \ + _ortmodule_output_transformation.get_flattened_output_module(self._original_module) sig = signature(self._original_module.forward) self._original_module_input_names = sig.parameters.keys() self._onnx_inference = None @@ -215,7 +97,7 @@ class ORTModule(torch.nn.Module): self._current_input_shape = None self._module_gradient_graph_builder = None self._input_names_require_grad = None - self._original_module_output_type = None + self._original_module_output_schema = None # Training model self._onnx_training = None @@ -239,7 +121,7 @@ class ORTModule(torch.nn.Module): def _initialize_module_gradient_graph_builder(self): # TODO: PyTorch exporter bug: changes the initializer order in ONNX model - initializer_names = [p[0] for p in self._original_module.named_parameters()] + initializer_names = [p[0] for p in self._flattened_output_module.named_parameters()] onnx_initializer_names = [p.name for p in self._onnx_inference.graph.initializer] initializer_names = [p for p in initializer_names if p in onnx_initializer_names] @@ -298,11 +180,11 @@ class ORTModule(torch.nn.Module): def eval(self: T) -> T: self._is_training = False - self._original_module.eval() + self._flattened_output_module.eval() def train(self: T, mode: bool = True) -> T: self._is_training = mode - self._original_module.train(mode) + self._flattened_output_module.train(mode) def forward(self, *inputs, **kwargs): '''Forward pass starts here and continues at `_ORTModuleFunction.forward` @@ -324,7 +206,9 @@ class ORTModule(torch.nn.Module): raise RuntimeError('A device must be specified in the model or data!') self._get_inference_graph_and_init_gradient_graph_builder(*inputs, **kwargs) - _, _, input_names_require_grad, new_input_shape = _parse_inputs_for_onnx_export(self._original_module_input_names, self._onnx_inference, *inputs, **kwargs) + _, _, input_names_require_grad, new_input_shape = \ + _ortmodule_output_transformation.parse_inputs_for_onnx_export( + self._original_module_input_names, self._onnx_inference, *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: @@ -362,7 +246,7 @@ class ORTModule(torch.nn.Module): # Run and return module outputs. user_outputs = tuple(_ort_output_to_torch_tensor(forward_output) \ for forward_output in self._training_session.run_forward(self._training_io_binding, self._run_options)) - return user_outputs[0] if len(user_outputs) == 1 else user_outputs + return user_outputs @staticmethod def backward(ctx, *grad_output): @@ -398,7 +282,8 @@ class ORTModule(torch.nn.Module): for backward_output in backward_outputs[num_user_input_grads:]] return tuple(results) - return _populate_user_output(self._original_module_output_type, self._onnx_graphs_info.user_output_names, + return _ortmodule_output_transformation.populate_user_output_from_schema_and_outputs(self._original_module_output_schema, + self._onnx_graphs_info.user_output_names, _ORTModuleFunction.apply(*self._convert_training_graph_input_to_list(*inputs, **kwargs))) @_utils.timeit(enabled=__TEMP_ENABLE_METHOD_TIMING__) @@ -423,7 +308,7 @@ class ORTModule(torch.nn.Module): result.append(inp) # Initializers - for param in self._original_module.named_parameters(): + for param in self._flattened_output_module.named_parameters(): result.append(param[1]) return result @@ -436,8 +321,12 @@ class ORTModule(torch.nn.Module): ''' # Setup dynamic axes for onnx model - input_names, dynamic_axes, self._input_names_require_grad, _ = _parse_inputs_for_onnx_export(self._original_module_input_names, None, *inputs, **kwargs) - output_names, output_dynamic_axes, self._original_module_output_type = _parse_outputs_for_onnx_export(self._original_module, inputs, kwargs) + input_names, dynamic_axes, self._input_names_require_grad, _ = \ + _ortmodule_output_transformation.parse_inputs_for_onnx_export( + self._original_module_input_names, None, *inputs, **kwargs) + output_names, output_dynamic_axes, self._original_module_output_schema = \ + _ortmodule_output_transformation.parse_outputs_for_onnx_export_and_extract_output_schema( + self._original_module, inputs, kwargs) dynamic_axes.update(output_dynamic_axes) # Export torch.nn.Module to ONNX @@ -446,11 +335,12 @@ class ORTModule(torch.nn.Module): # 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, sample_kwargs_copy = _deepcopy_model_input(*inputs, **kwargs) + sample_inputs_copy, sample_kwargs_copy = \ + _ortmodule_output_transformation.deepcopy_model_input(*inputs, **kwargs) try: with torch.no_grad(): - torch.onnx.export(self._original_module, + torch.onnx.export(self._flattened_output_module, sample_inputs_copy + (sample_kwargs_copy, ), f, input_names=input_names, diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index ce0279dd6c..5e9bf4520a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -103,16 +103,18 @@ class NeuralNetSimplePositionalAndKeywordArguments(torch.nn.Module): return torch.mean(self.a) + 3 * y return torch.mean(self.a) + x -def _get_bert_for_sequence_classification_model(device): +def _get_bert_for_sequence_classification_model(device, output_attentions = False, \ + output_hidden_states = False, return_dict = True): """Returns the BertForSequenceClassification pretrained model""" config = AutoConfig.from_pretrained( "bert-base-uncased", num_labels=2, num_hidden_layers=1, - output_attentions = False, - output_hidden_states = False, + output_attentions = output_attentions, + output_hidden_states = output_hidden_states, ) + config.return_dict = return_dict model = BertForSequenceClassification.from_pretrained( "bert-base-uncased", @@ -383,11 +385,9 @@ def test_gpu_reserved_memory_with_torch_no_grad(): torch.cuda.empty_cache() model_with_no_grad = ORTModule(model_with_no_grad) - mem_reserved_before_export = torch.cuda.memory_reserved(device) - model_with_no_grad(x, y, None, None, None, None, z) + model_with_no_grad(x, attention_mask=y, labels=z) mem_reserved_after_export_with_torch_no_grad = torch.cuda.memory_reserved(device) del model_with_no_grad - mem_reserved_after_cache_empty = torch.cuda.memory_reserved(device) # Create another model and get the memory_reserved when torch.no_grad has not been enabled after export. model_without_no_grad = _get_bert_for_sequence_classification_model(device) @@ -395,7 +395,7 @@ def test_gpu_reserved_memory_with_torch_no_grad(): mem_reserved_after_export_without_torch_no_grad = 0 with patch('torch.no_grad'): - model_without_no_grad(x, y, None, None, None, None, z) + model_without_no_grad(x, attention_mask=y, labels=z) mem_reserved_after_export_without_torch_no_grad = torch.cuda.memory_reserved(device) assert mem_reserved_after_export_with_torch_no_grad < mem_reserved_after_export_without_torch_no_grad @@ -473,9 +473,46 @@ def test_dict_of_tuple_return_value_module(device): y = torch.randn(N, D_in, device=device) z = torch.randn(N, D_in, device=device) - with pytest.raises(TypeError) as type_error: - model(x, y, z) - assert 'ORTModule does not support the following model output type' in str(type_error.value) + output = model(x, y, z) + assert 'loss' in output + assert len(output['loss']) == 3 + +@pytest.mark.parametrize("device", ['cuda', 'cpu']) +def test_tuple_of_tuple_return_value_module(device): + class NeuralNetTupleOfTuplesOutput(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetTupleOfTuplesOutput, 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(input3))) + return ((out1, out2), out3) + + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetTupleOfTuplesOutput(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) + + output = model(x, y, z) + assert len(output) == 2 + assert isinstance(output[0], tuple) + assert len(output[0]) == 2 + assert isinstance(output[1], torch.Tensor) @pytest.mark.parametrize("device", ['cpu', 'cuda']) def test_named_tuple_return_value_module(device): @@ -572,7 +609,7 @@ def test_dynamic_axes_config(): model_with_no_grad = _get_bert_for_sequence_classification_model(device) model_with_no_grad = ORTModule(model_with_no_grad) x, y, z = _get_bert_for_sequence_classification_sample_data(device) - output = model_with_no_grad(x, y, None, None, None, None, z) + output = model_with_no_grad(x, attention_mask=y, labels=z) assert output is not None assert _test_helpers.is_dynamic_axes(model_with_no_grad) @@ -730,3 +767,78 @@ def test_wrap_ortmodule_and_change_device(): # Checking training finished normally assert y_pred is not None and loss is not None + +@pytest.mark.parametrize("return_dict", [True, False]) +def test_hf_model_output_with_tuples(return_dict): + device = 'cuda' + + model = _get_bert_for_sequence_classification_model(device, output_attentions=True, + output_hidden_states=True, return_dict=return_dict) + x, y, z = _get_bert_for_sequence_classification_sample_data(device) + + model = ORTModule(model) + output = model(x, attention_mask=y, labels=z) + + if return_dict: + assert isinstance(output, SequenceClassifierOutput) + assert 'loss' in output and 'logits' in output and \ + 'attentions' in output and 'hidden_states' in output + assert isinstance(output['loss'], torch.Tensor) + assert isinstance(output['logits'], torch.Tensor) + assert isinstance(output['attentions'], tuple) + assert isinstance(output['hidden_states'], tuple) + else: + assert isinstance(output, tuple) + assert isinstance(output[0], torch.Tensor) + assert isinstance(output[1], torch.Tensor) + assert isinstance(output[2], tuple) + assert isinstance(output[3], tuple) + +@pytest.mark.parametrize("device", ['cuda', 'cpu']) +def test_nested_return_value_module(device): + class NeuralNetNestedOutput(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetNestedOutput, self).__init__() + + self.fc1_1 = torch.nn.Linear(input_size, hidden_size) + self.relu1 = torch.nn.ReLU() + self.relu = 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.relu(self.relu3(self.fc3_1(input3)))) + return { + 'a': { + 'b': { + 'c': out1 + }, + 'd': (out2, out3) + } + } + + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetNestedOutput(D_in, H, D_out).to(device) + model = ORTModule(model) + model._save_onnx = True + model._save_onnx_prefix = 'nested_model_output' + x = torch.randn(N, D_in, device=device) + y = torch.randn(N, D_in, device=device) + z = torch.randn(N, D_in, device=device) + + output = model(x, y, z) + assert 'a' in output and 'b' in output['a'] and 'c' in output['a']['b'] + assert isinstance(output['a']['b']['c'], torch.Tensor) + + assert 'd' in output['a'] + assert isinstance(output['a']['d'], tuple) + assert len(output['a']['d']) == 2