mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Type hint for ORTModule (#15938)
### Type hint for ORTModule
Add Type hint for ORTModule
Refine comments.
The reason of removing theinterface execution_session_run_forward from
`orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py`:
PR
cc275e7529 (diff-497e18dc8878818205b81fd80f85942548d8aa15d0f1204ce3e3d9795e3dd195)
and some commit before it breaks the function interface contracts
between parent calss _graph_execution_manager.py and its children
_training_manager.py and _inference_manager.py. So there is no need to
have this interface.
### Other EE work opportunities
1. Use logger correctly.
2. Remove few duplication logic parsing input/output recursively.
3. Clean up environment variable usage.
This commit is contained in:
pengwa
GitHub
parent
70d2dc8209
commit
34fe8fb069
8 changed files with 349 additions and 159 deletions
|
|
@ -12,7 +12,7 @@ import warnings
|
|||
from abc import ABC, abstractmethod # noqa: F401
|
||||
from enum import IntFlag
|
||||
from functools import reduce
|
||||
from typing import Dict, Tuple
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
|
@ -34,6 +34,7 @@ from ._fallback import (
|
|||
)
|
||||
from ._gradient_accumulation_manager import GradientAccumulationManager
|
||||
from ._graph_execution_interface import GraphExecutionInterface
|
||||
from ._io import _FlattenedModule, _InputInfo, _ModelInputOutputSchemaType
|
||||
from .debug_options import DebugOptions, LogLevel
|
||||
from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension
|
||||
|
||||
|
|
@ -41,7 +42,7 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
|
||||
class _RunStateInfo:
|
||||
def __init__(self, state, output_info):
|
||||
def __init__(self, state, output_info: List[Tuple[torch.Size, torch.device, torch.dtype]]):
|
||||
"""
|
||||
:param state: State of partial run that contains intermediate tensors needed to resume the run later.
|
||||
:param output_info: Output info.
|
||||
|
|
@ -73,7 +74,7 @@ class _SkipCheck(IntFlag):
|
|||
|
||||
|
||||
class GraphExecutionManager(GraphExecutionInterface):
|
||||
def __init__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
"""Manages construction and execution of ONNX graphs"""
|
||||
|
||||
super().__init__(module._original_module)
|
||||
|
|
@ -92,9 +93,9 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# Model after inference optimization or gradient building.
|
||||
self._graph_builder = None
|
||||
self._graph_info = None
|
||||
self._graph_initializer_names = None
|
||||
self._graph_initializer_names_to_train = None
|
||||
self._graph_initializers = None
|
||||
self._graph_initializer_names = set()
|
||||
self._graph_initializer_names_to_train = set()
|
||||
self._graph_initializers: List[torch.nn.parameter.Parameter] = []
|
||||
|
||||
# Update constant ONNX_OPSET_VERSION with env var ORTMODULE_ONNX_OPSET_VERSION
|
||||
# if defined.
|
||||
|
|
@ -105,7 +106,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# TrainingAgent or InferenceAgent
|
||||
self._execution_agent = None
|
||||
|
||||
# indicators of some logic have been executed previously and thus could be skipped for faster training
|
||||
# Indicators of some logic have been executed previously and thus could be skipped for faster training
|
||||
# default is enabled, if not defined in os env
|
||||
self._skip_check = _SkipCheck(
|
||||
_SkipCheck.SKIP_CHECK_DEVICE | _SkipCheck.SKIP_CHECK_BUILD_GRADIENT | _SkipCheck.SKIP_CHECK_EXECUTION_AGENT
|
||||
|
|
@ -144,8 +145,6 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
# It cannot overlap with required/immutable arguments (validated in runtime)
|
||||
self._export_extra_kwargs = {}
|
||||
|
||||
# Related to training graph shape inference
|
||||
self._current_input_shape = None
|
||||
# default execution order is priority-based for both dynamic/static shape input for now
|
||||
# if we observe the benefit of static shape, we can expose this flag to the user
|
||||
self._use_static_shape = False
|
||||
|
|
@ -158,11 +157,17 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
self._enable_custom_autograd_function = custom_autograd_function_enabler.state
|
||||
|
||||
self._input_info = None
|
||||
self._module_output_schema = None
|
||||
self._device = _utils.get_device_from_module(module)
|
||||
# Input and output infos (including schema) for exported model.
|
||||
self._input_info: Optional[_InputInfo] = None
|
||||
self._module_output_schema: Optional[_ModelInputOutputSchemaType] = None
|
||||
|
||||
self._module_parameters = list(inspect.signature(self._original_module.forward).parameters.values())
|
||||
# Device where the model is placed.
|
||||
self._device: Optional[torch.device] = _utils.get_device_from_module(module)
|
||||
|
||||
# Forward function input parameters of the original module.
|
||||
self._module_parameters: List[inspect.Parameter] = list(
|
||||
inspect.signature(self._original_module.forward).parameters.values()
|
||||
)
|
||||
|
||||
# TODO: remove after PyTorch ONNX exporter supports VAR_KEYWORD parameters.
|
||||
for input_parameter in self._module_parameters:
|
||||
|
|
@ -230,26 +235,6 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
|
||||
"""Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs`
|
||||
|
||||
This is a helper that can be called by the actual `GraphExecutionManager.forward` method
|
||||
|
||||
Args:
|
||||
execution_session (InferenceAgent or InferenceAgent): Agent which runs either inference or train
|
||||
onnx_model (onnx.ModelProto): ONNX model
|
||||
device (torch.device): PyTorch device
|
||||
inputs: (torch.Tensor or a container of): User input
|
||||
|
||||
Returns:
|
||||
Returns a tuple (user_outputs, run_info):
|
||||
user_outputs: The model output (either torch.Tensor or a container of torch.Tensor)
|
||||
run_info: A _RunStateInfo which contains extra information about the execution of the graph
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def forward(self):
|
||||
"""Executes the forward method for ORTModule
|
||||
|
|
@ -330,7 +315,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
return session_options, providers, provider_options
|
||||
|
||||
def _export_model(self, *inputs, **kwargs):
|
||||
def _export_model(self, *inputs, **kwargs) -> bool:
|
||||
# 1. Set the self._device from the user module
|
||||
# 2. Verify input schema matches the schema used on the previous model export
|
||||
# 3. Export the user model under self._export_training_flag mode
|
||||
|
|
@ -374,8 +359,9 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
|
||||
return True
|
||||
|
||||
def _get_exported_model(self, input_schema, *inputs, **kwargs):
|
||||
"""Exports PyTorch `self._flattened_module` to ONNX for inferencing or training, using `*inputs` and `**kwargs` as input
|
||||
def _get_exported_model(self, input_schema: _ModelInputOutputSchemaType, *inputs, **kwargs) -> onnx.ModelProto:
|
||||
"""Exports PyTorch `self._flattened_module` to ONNX for inferencing or training,
|
||||
using `*inputs` and `**kwargs` as input
|
||||
|
||||
TODO: How to support dynamic axes? Dimensions are determined by samples
|
||||
"""
|
||||
|
|
@ -459,7 +445,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
ORTModuleDeviceException, RuntimeError("A device must be specified in the model or inputs!")
|
||||
)
|
||||
|
||||
def _get_graph_transformer_config(self):
|
||||
def _get_graph_transformer_config(self) -> C.TrainingGraphTransformerConfiguration:
|
||||
graph_transformer_config = C.TrainingGraphTransformerConfiguration()
|
||||
graph_transformer_config.propagate_cast_ops_config = C.PropagateCastOpsConfiguration()
|
||||
graph_transformer_config.propagate_cast_ops_config.level = self._propagate_cast_ops_level
|
||||
|
|
|
|||
|
|
@ -3,18 +3,21 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from typing import Union
|
||||
|
||||
from ._fallback import _FallbackManager
|
||||
from ._inference_manager import InferenceManager
|
||||
from ._io import _FlattenedModule
|
||||
from ._training_manager import TrainingManager
|
||||
from .debug_options import DebugOptions
|
||||
|
||||
|
||||
class GraphExecutionManagerFactory:
|
||||
def __init__(self, module, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __init__(self, module: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
self._training_manager = TrainingManager(module, debug_options, fallback_manager)
|
||||
self._inference_manager = InferenceManager(module, debug_options, fallback_manager)
|
||||
|
||||
def __call__(self, is_training):
|
||||
def __call__(self, is_training) -> Union[InferenceManager, TrainingManager]:
|
||||
if is_training:
|
||||
return self._training_manager
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
# --------------------------------------------------------------------------
|
||||
|
||||
import warnings
|
||||
from typing import Tuple
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
|
@ -20,7 +21,7 @@ from .debug_options import DebugOptions
|
|||
class InferenceManager(GraphExecutionManager):
|
||||
"""Concrete instance of GraphExecutionManager that is able to manage the inference model
|
||||
|
||||
InferenceManager is resposible for building and running the forward graph of the inference model
|
||||
InferenceManager is responsible for building and running the forward graph of the inference model
|
||||
"""
|
||||
|
||||
def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
|
|
@ -28,9 +29,25 @@ class InferenceManager(GraphExecutionManager):
|
|||
self._export_mode = torch.onnx.TrainingMode.EVAL
|
||||
|
||||
@staticmethod
|
||||
def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
|
||||
"""Runs the forward graph on execution_session with given model inputs and device"""
|
||||
def execution_session_run_forward(
|
||||
execution_session,
|
||||
onnx_model: onnx.ModelProto,
|
||||
device: torch.device,
|
||||
*inputs,
|
||||
) -> Tuple[Tuple[torch.Tensor, ...], _RunStateInfo]:
|
||||
"""Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs`
|
||||
|
||||
Args:
|
||||
execution_session InferenceAgent: Agent which runs inference
|
||||
onnx_model (onnx.ModelProto): ONNX model
|
||||
device (torch.device): PyTorch device
|
||||
inputs: (torch.Tensor or a container of): User inputs passed from ORTModule.forward().
|
||||
|
||||
Returns:
|
||||
Returns a tuple (user_outputs, run_info):
|
||||
user_outputs: The model output (either torch.Tensor or a container of torch.Tensor)
|
||||
run_info: A _RunStateInfo which contains extra information about the execution of the graph
|
||||
"""
|
||||
# TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
|
||||
# especially the backward graph outputs.
|
||||
# REVIEW(codemzs): Consolidate Training Agent with InferenceAgent on C++ side to not
|
||||
|
|
@ -58,6 +75,14 @@ class InferenceManager(GraphExecutionManager):
|
|||
ONNX model is exported the first time this method is executed.
|
||||
Next, we build an optimized inference graph with module_graph_builder.
|
||||
Finally, we instantiate the ONNX Runtime InferenceSession through the InferenceAgent.
|
||||
|
||||
The call stack is as follows:
|
||||
ORTModule.forward(*inputs, **kwargs) ->
|
||||
ORTModule._torch_module.forward(*inputs, **kwargs) where _torch_module is a TorchModuleORT instance ->
|
||||
ORTModule._torch_module._execution_manager(is_training()).forward(*inputs, **kwargs) where:
|
||||
TorchModuleORT._execution_manager(true) is a TrainingManager instance;
|
||||
and TorchModuleORT._execution_manager(false) is an InferenceManager instance.
|
||||
|
||||
"""
|
||||
|
||||
# Fallback to PyTorch due to failures *external* to forward(),
|
||||
|
|
|
|||
|
|
@ -8,14 +8,84 @@ import gc
|
|||
import inspect
|
||||
import warnings
|
||||
from collections import abc
|
||||
from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
||||
from ._fallback import ORTModuleIOError, ORTModuleONNXModelException, wrap_exception
|
||||
from ._logger import suppress_os_stream_output
|
||||
from ._logger import LogLevel, suppress_os_stream_output
|
||||
from ._runtime_inspector import RuntimeInspector
|
||||
from ._utils import warn_of_constant_inputs
|
||||
|
||||
_ModelInputOutputType = Union[
|
||||
None,
|
||||
str,
|
||||
int,
|
||||
bool,
|
||||
float,
|
||||
torch.Tensor,
|
||||
Sequence["_ModelInputOutputType"],
|
||||
Mapping[str, "_ModelInputOutputType"],
|
||||
]
|
||||
|
||||
|
||||
class _TensorStub:
|
||||
"""Tensor stub class used to represent model's input or output"""
|
||||
|
||||
__slots__ = ["name", "dtype", "shape", "shape_dims"]
|
||||
|
||||
def __init__(
|
||||
self, name: Optional[str] = None, dtype: Optional[str] = None, shape=None, shape_dims: Optional[int] = None
|
||||
):
|
||||
self.name: Optional[str] = name
|
||||
self.dtype: Optional[str] = dtype
|
||||
self.shape = shape
|
||||
self.shape_dims: Optional[int] = shape_dims # r.g. rank.
|
||||
|
||||
def __repr__(self) -> str:
|
||||
result = "_TensorStub("
|
||||
if self.name is not None:
|
||||
result += f"name={self.name}"
|
||||
if self.dtype is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"dtype={self.dtype}"
|
||||
if self.shape is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"shape={self.shape}"
|
||||
if self.shape_dims is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"shape_dims={self.shape_dims}"
|
||||
result += ")"
|
||||
return result
|
||||
|
||||
def __eq__(self, other):
|
||||
if not other:
|
||||
return False
|
||||
elif not isinstance(other, _TensorStub):
|
||||
raise NotImplementedError("_TensorStub must only be compared to another _TensorStub instance!")
|
||||
elif self.name != other.name:
|
||||
return False
|
||||
elif self.dtype != other.dtype:
|
||||
return False
|
||||
elif self.shape != other.shape:
|
||||
return False
|
||||
elif self.shape_dims != other.shape_dims:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
_ModelInputOutputSchemaType = Union[
|
||||
None,
|
||||
str,
|
||||
_TensorStub,
|
||||
Sequence["_ModelInputOutputSchemaType"],
|
||||
Mapping[str, "_ModelInputOutputSchemaType"],
|
||||
]
|
||||
|
||||
|
||||
class _OutputIdentityOp(torch.autograd.Function):
|
||||
"""Internal class used to prepend Identity ops in model's outputs
|
||||
|
|
@ -76,7 +146,7 @@ class _PrimitiveType:
|
|||
return type(value) in _PrimitiveType._primitive_types
|
||||
|
||||
@staticmethod
|
||||
def get_tensor(value, device):
|
||||
def get_tensor(value, device) -> torch.Tensor:
|
||||
return torch.tensor(value, device=device)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -115,18 +185,18 @@ def flatten_kwargs(kwargs, device):
|
|||
class _InputInfo:
|
||||
def __init__(
|
||||
self,
|
||||
names,
|
||||
shape,
|
||||
require_grad_names=None,
|
||||
dynamic_axes=None,
|
||||
schema=None,
|
||||
names: List[str],
|
||||
shape: List[List[int]],
|
||||
require_grad_names: Optional[List[str]] = None,
|
||||
dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
|
||||
schema: Optional[_ModelInputOutputSchemaType] = None,
|
||||
num_positionals=0,
|
||||
):
|
||||
self.names = names
|
||||
self.shape = shape
|
||||
self.require_grad_names = require_grad_names if require_grad_names else []
|
||||
self.dynamic_axes = dynamic_axes if dynamic_axes else {}
|
||||
self.schema = schema if schema else []
|
||||
self.names: List[str] = names
|
||||
self.shape: List[List[int]] = shape
|
||||
self.require_grad_names: List[str] = require_grad_names if require_grad_names else []
|
||||
self.dynamic_axes: Dict[str, Dict[int, str]] = dynamic_axes if dynamic_axes else {}
|
||||
self.schema: _ModelInputOutputSchemaType = schema if schema else []
|
||||
self.num_positionals = num_positionals
|
||||
self.kwargs = None
|
||||
|
||||
|
|
@ -139,7 +209,9 @@ class _InputInfo:
|
|||
\tSchema: {self.schema}
|
||||
\t#Positionals (total): {self.num_positionals}"""
|
||||
|
||||
def flatten(self, args, kwargs, device):
|
||||
def flatten(
|
||||
self, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], device: torch.device
|
||||
) -> Sequence[_ModelInputOutputType]:
|
||||
"""Flatten args and kwargs in a single tuple of tensors with strict ordering"""
|
||||
|
||||
ret = [_PrimitiveType.get_tensor(arg, device) if _PrimitiveType.is_primitive_type(arg) else arg for arg in args]
|
||||
|
|
@ -154,7 +226,9 @@ class _InputInfo:
|
|||
|
||||
return ret
|
||||
|
||||
def unflatten(self, flat_args):
|
||||
def unflatten(
|
||||
self, flat_args: Sequence[_ModelInputOutputType]
|
||||
) -> Tuple[Sequence[_ModelInputOutputType], Mapping[str, _ModelInputOutputType]]:
|
||||
"""Unflatten tuple of tensors into args and kwargs"""
|
||||
|
||||
args = tuple(flat_args[: self.num_positionals])
|
||||
|
|
@ -164,13 +238,20 @@ class _InputInfo:
|
|||
|
||||
|
||||
def _combine_input_buffers_initializers(
|
||||
params, onnx_input_names, input_info, buffer_names, inputs, kwargs, device, rt_inspector: RuntimeInspector
|
||||
params: List[torch.nn.parameter.Parameter],
|
||||
onnx_input_names: List[str],
|
||||
input_info: Optional[_InputInfo],
|
||||
named_buffer: Iterator[Tuple[str, torch.Tensor]],
|
||||
inputs: Sequence[_ModelInputOutputType],
|
||||
kwargs: Mapping[str, _ModelInputOutputType],
|
||||
device: torch.device,
|
||||
rt_inspector: RuntimeInspector,
|
||||
):
|
||||
"""Creates forward `*inputs` list from user input and PyTorch initializers
|
||||
|
||||
ONNX Runtime forward requires an ordered list of:
|
||||
* User input: computed from forward InferenceSession
|
||||
* Initializers: computed from original PyTorch model parameters
|
||||
* Initializers: computed from original PyTorch model parameters.
|
||||
"""
|
||||
|
||||
def _expand_inputs(current_input, non_none_inputs, name=""):
|
||||
|
|
@ -202,7 +283,7 @@ def _combine_input_buffers_initializers(
|
|||
_expand_inputs(inputs, non_none_inputs)
|
||||
flattened_kwargs_inputs = {}
|
||||
_expand_inputs(kwargs, flattened_kwargs_inputs)
|
||||
buffer_names_dict = {buffer_name: inp for buffer_name, inp in buffer_names}
|
||||
buffer_names_dict = {buffer_name: inp for buffer_name, inp in named_buffer}
|
||||
result = []
|
||||
embed_sparsity_results = []
|
||||
label_sparsity_results = []
|
||||
|
|
@ -256,7 +337,9 @@ def _combine_input_buffers_initializers(
|
|||
return result, embed_sparsity_results, label_sparsity_results
|
||||
|
||||
|
||||
def deepcopy_model_input(*inputs, **kwargs):
|
||||
def deepcopy_model_input(
|
||||
*args, **kwargs
|
||||
) -> Tuple[Sequence[_ModelInputOutputType], Mapping[str, _ModelInputOutputType]]:
|
||||
def extract_tensor(value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.requires_grad:
|
||||
|
|
@ -266,64 +349,18 @@ def deepcopy_model_input(*inputs, **kwargs):
|
|||
else:
|
||||
return value
|
||||
|
||||
sample_inputs_copy = [extract_tensor(value) for value in inputs]
|
||||
sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy))
|
||||
sample_args_copy: Sequence[_ModelInputOutputType] = [extract_tensor(value) for value in args]
|
||||
sample_args_copy = copy.deepcopy(tuple(sample_args_copy))
|
||||
|
||||
sample_kwargs_copy = {}
|
||||
sample_kwargs_copy: Mapping[str, _ModelInputOutputType] = {}
|
||||
for name, value in kwargs.items():
|
||||
sample_kwargs_copy[name] = extract_tensor(value)
|
||||
sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy)
|
||||
|
||||
return sample_inputs_copy, sample_kwargs_copy
|
||||
return sample_args_copy, sample_kwargs_copy
|
||||
|
||||
|
||||
class _TensorStub:
|
||||
"""Tensor stub class used to represent model's input or output"""
|
||||
|
||||
__slots__ = ["name", "dtype", "shape", "shape_dims"]
|
||||
|
||||
def __init__(self, name=None, dtype=None, shape=None, shape_dims=None):
|
||||
self.name = name
|
||||
self.dtype = dtype
|
||||
self.shape = shape
|
||||
self.shape_dims = shape_dims
|
||||
|
||||
def __repr__(self) -> str:
|
||||
result = "_TensorStub("
|
||||
if self.name is not None:
|
||||
result += f"name={self.name}"
|
||||
if self.dtype is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"dtype={self.dtype}"
|
||||
if self.shape is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"shape={self.shape}"
|
||||
if self.shape_dims is not None:
|
||||
if result[-1] != "(":
|
||||
result += ", "
|
||||
result += f"shape_dims={self.shape_dims}"
|
||||
result += ")"
|
||||
return result
|
||||
|
||||
def __eq__(self, other):
|
||||
if not other:
|
||||
return False
|
||||
elif not isinstance(other, _TensorStub):
|
||||
raise NotImplementedError("_TensorStub must only be compared to another _TensorStub instance!")
|
||||
elif self.name != other.name:
|
||||
return False
|
||||
elif self.dtype != other.dtype:
|
||||
return False
|
||||
elif self.shape != other.shape:
|
||||
return False
|
||||
elif self.shape_dims != other.shape_dims:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def unflatten_user_output(output_schema, outputs):
|
||||
def unflatten_user_output(output_schema: Optional[_ModelInputOutputSchemaType], 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):
|
||||
|
|
@ -368,8 +405,49 @@ def unflatten_user_output(output_schema, outputs):
|
|||
return user_output
|
||||
|
||||
|
||||
def _extract_schema(data):
|
||||
"""Extract the data schema by replacing every torch.Tensor value with _TensorStub"""
|
||||
def _extract_schema(data: _ModelInputOutputType) -> _ModelInputOutputSchemaType:
|
||||
"""Extract the data schema by replacing every torch.Tensor value with _TensorStub.
|
||||
|
||||
Depth first traversal to iterate over the data:
|
||||
> Replace every tensor with a stub
|
||||
> Replace None/str typed data with itself
|
||||
> Recreate tensor from data for other primitive types, and replace them with a stub.
|
||||
|
||||
TODO(pengwa): use a unified logger and log level control.
|
||||
Examples:
|
||||
Example 1, list:
|
||||
data = [torch.tensor(1), torch.tensor(2)]
|
||||
schema = [_TensorStub(shape=()), _TensorStub(shape=())]
|
||||
|
||||
Example 2, dict:
|
||||
data = {"a": torch.tensor(1), "b": torch.tensor(2)}
|
||||
schema = {"a": _TensorStub(shape=()), "b": _TensorStub(shape=())}
|
||||
|
||||
Example 3, dict of list:
|
||||
data = {"a": [torch.tensor(1), torch.tensor(2)], "b": [torch.tensor(3), torch.tensor(4)]}
|
||||
schema = {"a": [_TensorStub(shape=()), _TensorStub(shape=())],
|
||||
"b": [_TensorStub(shape=()), _TensorStub(shape=())]}
|
||||
|
||||
Example 4, nested dict:
|
||||
data = {"a": {"b": torch.tensor(1), "c": torch.tensor(2)},
|
||||
"d": {"e": torch.tensor(3), "f": torch.tensor(4)}}
|
||||
schema = {"a": {"b": _TensorStub(shape=()), "c": _TensorStub(shape=())},
|
||||
"d": {"e": _TensorStub(shape=()), "f": _TensorStub(shape=())}}
|
||||
|
||||
Example 5, dict of mixed list and dict:
|
||||
data = {"a": [torch.tensor(1), torch.tensor(2)], "b": {"c": torch.tensor(3), "d": torch.tensor(4)}}
|
||||
schema = {"a": [_TensorStub(shape=()), _TensorStub(shape=())],
|
||||
"b": {"c": _TensorStub(shape=()), "d": _TensorStub(shape=())}}
|
||||
|
||||
Args:
|
||||
data: The data to extract the schema from, which can be in any kind of nested structure,
|
||||
including Sequence and Mapping.
|
||||
|
||||
|
||||
Returns:
|
||||
The schema of the data, which has the same structure as the data.
|
||||
|
||||
"""
|
||||
|
||||
if data is None:
|
||||
return data
|
||||
|
|
@ -386,7 +464,7 @@ def _extract_schema(data):
|
|||
|
||||
# Instead of replacing the tensor with a stub in the original user input, build the stubbed_schema
|
||||
# from scratch from the user input.
|
||||
stubbed_schema = None
|
||||
stubbed_schema: Optional[_ModelInputOutputSchemaType] = None
|
||||
if isinstance(data, abc.Sequence):
|
||||
sequence_type = type(data)
|
||||
stubbed_schema = [_extract_schema(val) for val in data]
|
||||
|
|
@ -407,10 +485,12 @@ def _extract_schema(data):
|
|||
return stubbed_schema
|
||||
|
||||
|
||||
def _parse_outputs_and_extract_names_and_dynamic_axes(module_output):
|
||||
def _parse_outputs_and_extract_names_and_dynamic_axes(module_output) -> Tuple[List[str], Dict[str, Dict[int, str]]]:
|
||||
"""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):
|
||||
def _populate_output_names_and_dynamic_axes(
|
||||
output, output_names: List[str], output_dynamic_axes: Dict[str, Dict[int, str]], output_idx: List[int]
|
||||
):
|
||||
# Depth first traversal to traverse through the entire output collecting output names and dynamic axes
|
||||
|
||||
if output is None:
|
||||
|
|
@ -438,9 +518,9 @@ def _parse_outputs_and_extract_names_and_dynamic_axes(module_output):
|
|||
TypeError(f"ORTModule does not support the following model output type {type(output)}"),
|
||||
)
|
||||
|
||||
output_names = []
|
||||
output_dynamic_axes = {}
|
||||
output_idx = [0]
|
||||
output_names: List[str] = []
|
||||
output_dynamic_axes: Dict[str, Dict[int, str]] = {}
|
||||
output_idx: List[int] = [0]
|
||||
_populate_output_names_and_dynamic_axes(module_output, output_names, output_dynamic_axes, output_idx)
|
||||
|
||||
return output_names, output_dynamic_axes
|
||||
|
|
@ -474,21 +554,53 @@ def _transform_output_to_flat_tuple(data):
|
|||
|
||||
|
||||
class _FlattenedModule(torch.nn.Module):
|
||||
def __init__(self, original_module):
|
||||
def __init__(self, original_module: torch.nn.Module):
|
||||
super().__init__()
|
||||
self._original_module = original_module
|
||||
self._original_module: torch.nn.Module = original_module
|
||||
|
||||
# Before `forward` is called, _ort_module must be assigned
|
||||
# Updated input info is needed to expand args into *args, **kwargs
|
||||
self._input_info = None
|
||||
self._input_info: Optional[_InputInfo] = None
|
||||
|
||||
def forward(self, *args):
|
||||
new_args, new_kwargs = self._input_info.unflatten(args)
|
||||
return _transform_output_to_flat_tuple(self._original_module(*new_args, **new_kwargs))
|
||||
|
||||
|
||||
def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, inputs, kwargs):
|
||||
def _add_dynamic_shape(name, input):
|
||||
def parse_inputs_for_onnx_export(
|
||||
all_input_parameters: List[inspect.Parameter],
|
||||
onnx_graph: Optional[onnx.ModelProto],
|
||||
schema: _ModelInputOutputSchemaType,
|
||||
args: Sequence[_ModelInputOutputType],
|
||||
kwargs: Mapping[str, _ModelInputOutputType],
|
||||
) -> _InputInfo:
|
||||
"""Parses through the model inputs and returns _InputInfo.
|
||||
|
||||
Loop through all input parameters, try to flatten them into a 1-D list of inputs. For nested data in the inputs,
|
||||
construct the name in hierarchical order.
|
||||
|
||||
Example 1, arg is a list, kwarg is a dict:
|
||||
args = [arg1, arg2], kwargs = {"a": 4, "b": 5},
|
||||
input_names = ["arg1", "arg2", "a", "b"].
|
||||
|
||||
Example 2, arg is a list, kwarg is a dict of mixed list and scalar:
|
||||
args = [arg1, arg2], kwargs = {"a": [4, 5], "b": 6},
|
||||
input_names = ["arg1", "arg2", "a_0", "a_1", "b"].
|
||||
|
||||
Example 3, arg is a list, kwarg is a dict of mixed dict and scalar:
|
||||
args = [arg1, arg2], kwargs = {"a": {"c": 4, "d": 5}, "b": 6},
|
||||
input_names = ["arg1", "arg2", "a_c", "a_d", "b"].
|
||||
|
||||
Args:
|
||||
all_input_parameters: All inspected input parameters from the original model forward function signature.
|
||||
onnx_graph: (optional) The onnx graph.
|
||||
schema: The schema extracted from the positional arguments and keyword arguments of the model.
|
||||
args: The positional arguments of the model.
|
||||
kwargs: The keyword arguments of the model.
|
||||
|
||||
"""
|
||||
|
||||
def _add_dynamic_shape(name, input) -> Dict[str, Dict[int, str]]:
|
||||
dynamic_axes[name] = {}
|
||||
for dim_idx in range(len(input.shape)):
|
||||
dynamic_axes[name].update({dim_idx: f"{name}_dim{dim_idx}"})
|
||||
|
|
@ -534,24 +646,43 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input
|
|||
|
||||
# Ignore optional inputs explicitly specified as None
|
||||
# ONNX exporter may remove unused inputs
|
||||
onnx_graph_input_names = []
|
||||
onnx_graph_input_names: List[str] = []
|
||||
if onnx_graph is not None:
|
||||
onnx_graph_input_names = {inp.name for inp in onnx_graph.graph.input}
|
||||
|
||||
input_names = []
|
||||
dynamic_axes = {}
|
||||
input_names_require_grad = []
|
||||
input_shape = []
|
||||
input_names: List[str] = []
|
||||
dynamic_axes: Dict[str, Dict[int, str]] = {}
|
||||
input_names_require_grad: List[str] = []
|
||||
input_shape: List[List[int]] = []
|
||||
var_positional_idx = 0
|
||||
|
||||
# Be noted, all_input_parameters is a list of inspect.Parameters parsed from the original module's forward method.
|
||||
# While the execution manger's forward function will map all given model inputs to *args and **kwargs, so it is
|
||||
# possible the input parameter list cannot represent the real model inputs given here (e.g., *args, **kwargs).
|
||||
# But it is still fine to use all_input_parameters to make sure all model inputs are covered.
|
||||
#
|
||||
# Here is an example caused by the mismatch between all_input_parameters and real model inputs.
|
||||
# def foo(*args, named_kwargs, **kwargs):
|
||||
# ... print("foo")
|
||||
# From inspection,
|
||||
# > ('args', <_ParameterKind.VAR_POSITIONAL: 2>)
|
||||
# > ('named_kwargs', <_ParameterKind.KEYWORD_ONLY: 3>)
|
||||
# > ('kwargs', <_ParameterKind.VAR_KEYWORD: 4>)
|
||||
#
|
||||
# At this point, 'named_kwargs' exists in **kwargs as a result of ORTModule's forward parse all original
|
||||
# model inputs in to *args and **kwargs.
|
||||
# When we loop `all_input_parameters``, for the `named_kwargs`, we will try to handle it in KEYWORD_ONLY branch.
|
||||
# Additionally in VAR_KEYWORD branch, we will get the `named_kwargs` value again, because its name exists in the
|
||||
# `kwargs`. So _add_input avoids handling the `named_kwargs` twice, check test case `test_named_kwargs_dict_input`
|
||||
# for the details.
|
||||
for input_idx, input_parameter in enumerate(all_input_parameters):
|
||||
if input_parameter.kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
# VAR_POSITIONAL parameter carries all *args parameters from original forward method
|
||||
|
||||
for args_i in range(input_idx, len(inputs)):
|
||||
for args_i in range(input_idx, len(args)):
|
||||
name = f"{input_parameter.name}_{var_positional_idx}"
|
||||
var_positional_idx += 1
|
||||
inp = inputs[args_i]
|
||||
inp = args[args_i]
|
||||
_add_input(name, inp, onnx_graph, onnx_graph_input_names)
|
||||
elif (
|
||||
input_parameter.kind == inspect.Parameter.POSITIONAL_ONLY
|
||||
|
|
@ -562,8 +693,8 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input
|
|||
name = input_parameter.name
|
||||
inp = None
|
||||
input_idx += var_positional_idx # noqa: PLW2901
|
||||
if input_idx < len(inputs) and inputs[input_idx] is not None:
|
||||
inp = inputs[input_idx]
|
||||
if input_idx < len(args) and args[input_idx] is not None:
|
||||
inp = args[input_idx]
|
||||
elif name in kwargs and kwargs[name] is not None:
|
||||
inp = kwargs[name]
|
||||
_add_input(name, inp, onnx_graph, onnx_graph_input_names)
|
||||
|
|
@ -578,18 +709,20 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, schema, input
|
|||
require_grad_names=input_names_require_grad,
|
||||
dynamic_axes=dynamic_axes,
|
||||
schema=schema,
|
||||
num_positionals=len(inputs),
|
||||
num_positionals=len(args),
|
||||
)
|
||||
|
||||
|
||||
def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs, log_level):
|
||||
def parse_outputs_for_onnx_export_and_extract_schema(
|
||||
module, args: Sequence[_ModelInputOutputType], kwargs: Mapping[str, _ModelInputOutputType], log_level: LogLevel
|
||||
):
|
||||
# Perform a forward call to grab outputs
|
||||
output_names = None
|
||||
output_dynamic_axes = None
|
||||
is_deepcopy = False
|
||||
with torch.no_grad(), suppress_os_stream_output(log_level=log_level):
|
||||
# Deepcopy inputs, since input values may change after model run.
|
||||
sample_inputs_copy, sample_kwargs_copy = deepcopy_model_input(*inputs, **kwargs)
|
||||
sample_args_copy, sample_kwargs_copy = deepcopy_model_input(*args, **kwargs)
|
||||
try:
|
||||
# Deepcopy model, in case model is stateful and changes after model run.
|
||||
model_copy = copy.deepcopy(module)
|
||||
|
|
@ -602,7 +735,7 @@ def parse_outputs_for_onnx_export_and_extract_schema(module, inputs, kwargs, log
|
|||
" Compute will continue, but unexpected results may occur!"
|
||||
)
|
||||
|
||||
sample_outputs = model_copy(*sample_inputs_copy, **sample_kwargs_copy)
|
||||
sample_outputs = model_copy(*sample_args_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)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
|
@ -31,9 +32,9 @@ class ONNXModels:
|
|||
InferenceSession and is saved by the InferenceSession.
|
||||
"""
|
||||
|
||||
exported_model: onnx.ModelProto = None
|
||||
optimized_model: onnx.ModelProto = None
|
||||
optimized_pre_grad_model: onnx.ModelProto = None
|
||||
exported_model: Optional[onnx.ModelProto] = None
|
||||
optimized_model: Optional[onnx.ModelProto] = None
|
||||
optimized_pre_grad_model: Optional[onnx.ModelProto] = None
|
||||
|
||||
def save_exported_model(self, path, name_prefix, export_mode):
|
||||
# save the ortmodule exported model
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
# --------------------------------------------------------------------------
|
||||
|
||||
import warnings
|
||||
from typing import Tuple
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
|
|
@ -14,24 +15,46 @@ from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_typ
|
|||
from . import _are_deterministic_algorithms_enabled, _io, _logger, _use_deterministic_algorithms, _utils
|
||||
from ._execution_agent import TrainingAgent
|
||||
from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy
|
||||
from ._gradient_accumulation_manager import GradientAccumulationManager
|
||||
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck
|
||||
from ._io import _FlattenedModule, _InputInfo
|
||||
from .debug_options import DebugOptions
|
||||
|
||||
|
||||
class TrainingManager(GraphExecutionManager):
|
||||
"""Concrete instance of GraphExecutionManager that is able to manage the training model
|
||||
|
||||
TrainingManager is responsible for building and running the forward and backward graph of the training model
|
||||
TrainingManager is responsible for building and running the forward and backward graph of the training model.
|
||||
"""
|
||||
|
||||
def __init__(self, model, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
def __init__(self, model: _FlattenedModule, debug_options: DebugOptions, fallback_manager: _FallbackManager):
|
||||
super().__init__(model, debug_options, fallback_manager)
|
||||
|
||||
self._export_mode = torch.onnx.TrainingMode.TRAINING
|
||||
self._forward_class = self._create_autofunction_class()
|
||||
|
||||
@staticmethod
|
||||
def execution_session_run_forward(execution_session, onnx_model, device, gradient_accumulation_manager, *inputs):
|
||||
"""Runs the forward graph on execution_session with given model inputs and device"""
|
||||
def execution_session_run_forward(
|
||||
execution_session,
|
||||
onnx_model: onnx.ModelProto,
|
||||
device: torch.device,
|
||||
gradient_accumulation_manager: GradientAccumulationManager,
|
||||
*inputs,
|
||||
) -> Tuple[Tuple[torch.Tensor, ...], _RunStateInfo]:
|
||||
"""Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs`
|
||||
|
||||
Args:
|
||||
execution_session (InferenceAgent or TrainingAgent): Agent which runs training.
|
||||
onnx_model (onnx.ModelProto): ONNX model
|
||||
device (torch.device): PyTorch device
|
||||
gradient_accumulation_manager (GradientAccumulationManager): Gradient accumulation manager
|
||||
inputs: (torch.Tensor or a container of): User inputs passed from ORTModule.forward().
|
||||
|
||||
Returns:
|
||||
Returns a tuple (user_outputs, run_info):
|
||||
user_outputs: The model output (either torch.Tensor or a container of torch.Tensor)
|
||||
run_info: A _RunStateInfo which contains extra information about the execution of the graph
|
||||
"""
|
||||
|
||||
# TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
|
||||
# especially the backward graph outputs.
|
||||
|
|
@ -54,7 +77,9 @@ class TrainingManager(GraphExecutionManager):
|
|||
# Run and return module outputs.
|
||||
execution_session.run_forward(forward_inputs, forward_outputs, state, gradient_accumulation_manager.cache)
|
||||
|
||||
user_outputs = gradient_accumulation_manager.extract_outputs_and_maybe_update_cache(forward_outputs, device)
|
||||
user_outputs: Tuple[torch.Tensor, ...] = gradient_accumulation_manager.extract_outputs_and_maybe_update_cache(
|
||||
forward_outputs, device
|
||||
)
|
||||
|
||||
output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
|
||||
run_info = _RunStateInfo(state, output_info)
|
||||
|
|
@ -174,6 +199,14 @@ class TrainingManager(GraphExecutionManager):
|
|||
ONNX model is exported the first time this method is executed.
|
||||
Next, we build a full training graph with module_graph_builder.
|
||||
Finally, we instantiate the ONNX Runtime InferenceSession.
|
||||
|
||||
The call stack is as follows:
|
||||
ORTModule.forward(*inputs, **kwargs) ->
|
||||
ORTModule._torch_module.forward(*inputs, **kwargs) where _torch_module is a TorchModuleORT instance ->
|
||||
ORTModule._torch_module._execution_manager(is_training()).forward(*inputs, **kwargs) where:
|
||||
TorchModuleORT._execution_manager(true) is a TrainingManager instance;
|
||||
and TorchModuleORT._execution_manager(false) is an InferenceManager instance.
|
||||
|
||||
"""
|
||||
|
||||
# Fallback to PyTorch due to failures *external* to forward(),
|
||||
|
|
@ -209,7 +242,7 @@ class TrainingManager(GraphExecutionManager):
|
|||
# If model was exported, then initialize the graph builder
|
||||
self._initialize_graph_builder()
|
||||
|
||||
# since the schema was just extracted while trying to export the model and it was either
|
||||
# Since the schema was just extracted while trying to export the model and it was either
|
||||
# saved to self._input_info.schema or checked for equality with the self._input_info.schema
|
||||
# it should not need to be updated again. Pass it inside parse_inputs_for_onnx_export.
|
||||
input_info = _io.parse_inputs_for_onnx_export(
|
||||
|
|
@ -366,7 +399,7 @@ class TrainingManager(GraphExecutionManager):
|
|||
local_device_rank,
|
||||
)
|
||||
|
||||
def _reinitialize_graph_builder(self, input_info):
|
||||
def _reinitialize_graph_builder(self, input_info: _InputInfo):
|
||||
"""Return true if the module graph builder was reinitialized"""
|
||||
|
||||
# Model may have unused params dropped after export and not part of self._graph_initializer_names_to_train
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import random
|
|||
import traceback
|
||||
import types
|
||||
import warnings
|
||||
from typing import List # noqa: F401
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
|
@ -49,7 +49,8 @@ def set_random_states(states):
|
|||
torch.cuda.set_rng_state(torch_cuda_state)
|
||||
|
||||
|
||||
def _ortvalue_from_torch_tensor(torch_tensor):
|
||||
def _ortvalue_from_torch_tensor(torch_tensor: torch.Tensor) -> C.OrtValue:
|
||||
"""Converts a torch tensor to an OrtValue."""
|
||||
# TODO: Current DLPack doesn't support bool and PyTorch disables converting bool tensor to DLPack in recent commit.
|
||||
# https://github.com/pytorch/pytorch/blob/7e7be526c9d9179f35084e9cca5b5c5ad5172100/aten/src/ATen/DLConvertor.cpp#L41
|
||||
# We need to convert bool tensor to unit8 tensor to workaround this.
|
||||
|
|
@ -63,7 +64,9 @@ def _ortvalue_from_torch_tensor(torch_tensor):
|
|||
return C.OrtValue.from_dlpack(to_dlpack(torch_tensor), is_bool_tensor)
|
||||
|
||||
|
||||
def _ortvalues_to_torch_tensor(ortvalues, device=None):
|
||||
def _ortvalues_to_torch_tensor(
|
||||
ortvalues: C.OrtValueVector, device: Optional[torch.device] = None
|
||||
) -> Tuple[torch.Tensor, ...]:
|
||||
if len(ortvalues) == 0:
|
||||
return tuple()
|
||||
|
||||
|
|
@ -75,7 +78,7 @@ def _ortvalues_to_torch_tensor(ortvalues, device=None):
|
|||
if not isinstance(ortvalues, C.OrtValueVector):
|
||||
raise TypeError("ortvalues must be an instance of OrtValueVector not %r." % type(ortvalues))
|
||||
|
||||
res = ortvalues.to_dlpacks(_from_dlpack)
|
||||
res: List[torch.Tensor] = ortvalues.to_dlpacks(_from_dlpack)
|
||||
bool_indices = ortvalues.bool_tensor_indices()
|
||||
if len(bool_indices):
|
||||
# DLPack structure does not know for sure if it stores boolean
|
||||
|
|
@ -97,7 +100,7 @@ def _ortvalues_to_torch_tensor(ortvalues, device=None):
|
|||
return tuple(res)
|
||||
|
||||
|
||||
def _torch_tensor_to_dlpack(tensor):
|
||||
def _torch_tensor_to_dlpack(tensor: torch.Tensor):
|
||||
# TODO: Current DLPack doesn't support bool and PyTorch disables converting bool tensor to DLPack in recent commit.
|
||||
# https://github.com/pytorch/pytorch/blob/7e7be526c9d9179f35084e9cca5b5c5ad5172100/aten/src/ATen/DLConvertor.cpp#L41
|
||||
# We need to convert bool tensor to unit8 tensor to workaround this.
|
||||
|
|
@ -110,7 +113,7 @@ def _torch_tensor_to_dlpack(tensor):
|
|||
return to_dlpack(tensor)
|
||||
|
||||
|
||||
def _check_same_device(device, argument_str, *args):
|
||||
def _check_same_device(device: torch.device, argument_str: str, *args):
|
||||
"""Check that all tensor arguments in *args reside on the same device as the input device"""
|
||||
|
||||
assert isinstance(device, torch.device), "`device` must be a valid `torch.device` object"
|
||||
|
|
@ -126,7 +129,7 @@ def _check_same_device(device, argument_str, *args):
|
|||
)
|
||||
|
||||
|
||||
def get_device_index(device):
|
||||
def get_device_index(device: Union[str, int, torch.device]) -> int:
|
||||
if isinstance(device, str):
|
||||
# could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
|
||||
device = torch.device(device)
|
||||
|
|
@ -135,7 +138,7 @@ def get_device_index(device):
|
|||
return 0 if device.index is None else device.index
|
||||
|
||||
|
||||
def get_device_str(device):
|
||||
def get_device_str(device: Union[str, int, torch.device]) -> str:
|
||||
if isinstance(device, str):
|
||||
# could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
|
||||
if device.find(":") == -1:
|
||||
|
|
@ -152,12 +155,15 @@ def get_device_str(device):
|
|||
return device
|
||||
|
||||
|
||||
def get_device_from_module(module):
|
||||
def get_device_from_module(module) -> Optional[torch.device]:
|
||||
"""Returns the first device found in the `module`'s parameters or None
|
||||
|
||||
Args:
|
||||
module (torch.nn.Module): PyTorch model to extract device from
|
||||
|
||||
Returns:
|
||||
torch.device: Device extracted from `module`'s parameters or None
|
||||
|
||||
Raises:
|
||||
ORTModuleFallbackException: When more than one device is found at `module`
|
||||
"""
|
||||
|
|
@ -175,12 +181,15 @@ def get_device_from_module(module):
|
|||
return device
|
||||
|
||||
|
||||
def get_device_from_inputs(args, kwargs):
|
||||
def get_device_from_inputs(args, kwargs) -> Optional[torch.device]:
|
||||
"""Returns device from first PyTorch Tensor within args or kwargs
|
||||
|
||||
Args:
|
||||
args: List with inputs
|
||||
kwargs: Dictionary with inputs
|
||||
|
||||
Returns:
|
||||
torch.device: Device extracted from `args` or `kwargs` or None
|
||||
"""
|
||||
|
||||
device = None
|
||||
|
|
@ -191,7 +200,7 @@ def get_device_from_inputs(args, kwargs):
|
|||
return device
|
||||
|
||||
|
||||
def _create_iobinding(io_binding, inputs, model, device):
|
||||
def _create_iobinding(io_binding, inputs, model, device: torch.device):
|
||||
"""Creates IO binding for a `model` inputs and output"""
|
||||
for idx, value_info in enumerate(model.graph.input):
|
||||
io_binding.bind_ortvalue_input(
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ class ORTModule(torch.nn.Module):
|
|||
debug_options (:obj:`DebugOptions`, optional): debugging options for ORTModule.
|
||||
"""
|
||||
|
||||
def __init__(self, module, debug_options=None):
|
||||
def __init__(self, module: torch.nn.Module, debug_options: Optional[DebugOptions] = None):
|
||||
# NOTE: torch.nn.Modules that call setattr on their internal attributes regularly
|
||||
# (for example PyTorch Lightning), will trigger regular re-exports. This is
|
||||
# because ORTModule auto detects such setattrs on the original module and
|
||||
|
|
|
|||
Loading…
Reference in a new issue