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

Revert "Add support for dynamic axes for outputs + check model output type before export (#6491)" (#6566)

Browse source 
This reverts commit c983b84316.
This commit is contained in:
Sherlock 2021-02-04 00:43:13 -08:00 committed by GitHub
parent c983b84316
commit bc0d04bf07
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 218 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

114
orttraining/orttraining/python/training/ortmodule.py
View file

@ -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,

22
orttraining/orttraining/test/python/_test_helpers.py
View file

@ -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):

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

@ -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)