This document describes ORTModule PyTorch frontend API for the ONNX Runtime (aka ORT) training acelerator.
+What is new¶
+Version 0.1¶
+-
+
Initial version
+
Overview¶
+The aim of ORTModule is to provide a drop-in replacement for one or more torch.nn.Module objects in a user’s PyTorch program, +and execute the forward and backward passes of those modules using ORT.
+As a result, the user will be able to accelerate their training script gradually using ORT, +without having to modify their training loop.
+Users will be able to use standard PyTorch debugging techniques for convergence issues, +e.g. by probing the computed gradients on the model’s parameters.
+The following code example illustrates how ORTModule would be used in a user’s training script, +in the simple case where the entire model can be offloaded to ONNX Runtime:
+# Original PyTorch model
+class NeuralNet(torch.nn.Module):
+ def __init__(self, input_size, hidden_size, num_classes):
+ ...
+ def forward(self, x):
+ ...
+
+model = NeuralNet(input_size=784, hidden_size=500, num_classes=10)
+model = ORTModule(model) # Only change to original PyTorch script
+criterion = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
+
+# Training Loop is unchanged
+for data, target in data_loader:
+ optimizer.zero_grad()
+ output = model(data)
+ loss = criterion(output, target)
+ loss.backward()
+ optimizer.step()
+API¶
+-
+
- +class onnxruntime.training.ortmodule.ortmodule.ORTModule(module)¶ +
Bases:
+torch.nn.modules.module.ModuleExtends user’s
+torch.nn.Modulemodel to leverage ONNX Runtime super fast training engine.ORTModule specializes the user’s
+torch.nn.Modulemodel, providingforward(), +backward()along with all otherstorch.nn.Module’s APIs.-
+
- +forward(*inputs, **kwargs)¶ +
Dummy documentation for forward method
+
-
+
- +apply(fn: Callable[[Module], None]) → onnxruntime.training.ortmodule.ortmodule.T¶ +
Override original method to delegate execution to the flattened PyTorch user module
+
-
+
- +train(mode: bool = True) → onnxruntime.training.ortmodule.ortmodule.T¶ +
Override original method to delegate execution to the flattened PyTorch user module
+
-
+
- +state_dict(destination=None, prefix='', keep_vars=False)¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +load_state_dict(state_dict: OrderedDict[str, Tensor], strict: bool = True)¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +register_buffer(name: str, tensor: Optional[torch.Tensor], persistent: bool = True) → None¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +register_parameter(name: str, param: Optional[torch.nn.parameter.Parameter]) → None¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +get_parameter(target: str) → torch.nn.parameter.Parameter¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +get_buffer(target: str) → torch.Tensor¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +parameters(recurse: bool = True) → Iterator[torch.nn.parameter.Parameter]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +named_parameters(prefix: str = '', recurse: bool = True) → Iterator[Tuple[str, torch.nn.parameter.Parameter]]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +buffers(recurse: bool = True) → Iterator[torch.Tensor]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +named_buffers(prefix: str = '', recurse: bool = True) → Iterator[Tuple[str, torch.Tensor]]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +named_children() → Iterator[Tuple[str, Module]]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +modules() → Iterator[Module]¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +named_modules(memo: Optional[Set[Module]] = None, prefix: str = '')¶ +
Override original method to delegate execution to the original PyTorch user module
+
-
+
- +add_module(name: str, module: Optional[Module]) → None¶ +
Raises a NotImplementedError exception since ORTModule does not support adding modules to it
+
-
+
- +bfloat16() → torch.nn.modules.module.T¶ +
Casts all floating point parameters and buffers to
+bfloat16datatype.-
+
- Returns:
Module: self
+
+
-
+
- +children() → Iterator[torch.nn.modules.module.Module]¶ +
Returns an iterator over immediate children modules.
+-
+
- Yields:
Module: a child module
+
+
-
+
- +cpu() → torch.nn.modules.module.T¶ +
Moves all model parameters and buffers to the CPU.
+-
+
- Returns:
Module: self
+
+
-
+
- +cuda(device: Optional[Union[int, torch.device]] = None) → torch.nn.modules.module.T¶ +
Moves all model parameters and buffers to the GPU.
+This also makes associated parameters and buffers different objects. So +it should be called before constructing optimizer if the module will +live on GPU while being optimized.
+-
+
- Args:
-
+
- device (int, optional): if specified, all parameters will be
copied to that device
+
+
+- Returns:
Module: self
+
+
-
+
- +double() → torch.nn.modules.module.T¶ +
Casts all floating point parameters and buffers to
+doubledatatype.-
+
- Returns:
Module: self
+
+
-
+
- +eval() → torch.nn.modules.module.T¶ +
Sets the module in evaluation mode.
+This has any effect only on certain modules. See documentations of +particular modules for details of their behaviors in training/evaluation +mode, if they are affected, e.g.
+Dropout,BatchNorm, +etc.This is equivalent with
+self.train(False).-
+
- Returns:
Module: self
+
+
-
+
- +extra_repr() → str¶ +
Set the extra representation of the module
+To print customized extra information, you should re-implement +this method in your own modules. Both single-line and multi-line +strings are acceptable.
+
-
+
- +float() → torch.nn.modules.module.T¶ +
Casts all floating point parameters and buffers to float datatype.
+-
+
- Returns:
Module: self
+
+
-
+
- +half() → torch.nn.modules.module.T¶ +
Casts all floating point parameters and buffers to
+halfdatatype.-
+
- Returns:
Module: self
+
+
-
+
- +register_backward_hook(hook: Callable[[torch.nn.modules.module.Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[None, torch.Tensor]]) → torch.utils.hooks.RemovableHandle¶ +
Registers a backward hook on the module.
+This function is deprecated in favor of
+nn.Module.register_full_backward_hook()and +the behavior of this function will change in future versions.-
+
- Returns:
-
+
torch.utils.hooks.RemovableHandle:a handle that can be used to remove the added hook by calling +
+handle.remove()
+
+
-
+
- +register_forward_hook(hook: Callable[[...], None]) → torch.utils.hooks.RemovableHandle¶ +
Registers a forward hook on the module.
+The hook will be called every time after
+forward()has computed an output. +It should have the following signature:++hook(module, input, output) -> None or modified output +
The input contains only the positional arguments given to the module. +Keyword arguments won’t be passed to the hooks and only to the
+forward. +The hook can modify the output. It can modify the input inplace but +it will not have effect on forward since this is called after +forward()is called.-
+
- Returns:
-
+
torch.utils.hooks.RemovableHandle:a handle that can be used to remove the added hook by calling +
+handle.remove()
+
+
-
+
- +register_forward_pre_hook(hook: Callable[[...], None]) → torch.utils.hooks.RemovableHandle¶ +
Registers a forward pre-hook on the module.
+The hook will be called every time before
+forward()is invoked. +It should have the following signature:++hook(module, input) -> None or modified input +
The input contains only the positional arguments given to the module. +Keyword arguments won’t be passed to the hooks and only to the
+forward. +The hook can modify the input. User can either return a tuple or a +single modified value in the hook. We will wrap the value into a tuple +if a single value is returned(unless that value is already a tuple).-
+
- Returns:
-
+
torch.utils.hooks.RemovableHandle:a handle that can be used to remove the added hook by calling +
+handle.remove()
+
+
-
+
- +register_full_backward_hook(hook: Callable[[torch.nn.modules.module.Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[None, torch.Tensor]]) → torch.utils.hooks.RemovableHandle¶ +
Registers a backward hook on the module.
+The hook will be called every time the gradients with respect to module +inputs are computed. The hook should have the following signature:
+++hook(module, grad_input, grad_output) -> tuple(Tensor) or None +
The
+grad_inputandgrad_outputare tuples that contain the gradients +with respect to the inputs and outputs respectively. The hook should +not modify its arguments, but it can optionally return a new gradient with +respect to the input that will be used in place ofgrad_inputin +subsequent computations.grad_inputwill only correspond to the inputs given +as positional arguments and all kwarg arguments are ignored. Entries +ingrad_inputandgrad_outputwill beNonefor all non-Tensor +arguments.++Warning
+Modifying inputs or outputs inplace is not allowed when using backward hooks and +will raise an error.
+-
+
- Returns:
-
+
torch.utils.hooks.RemovableHandle:a handle that can be used to remove the added hook by calling +
+handle.remove()
+
+
-
+
- +requires_grad_(requires_grad: bool = True) → torch.nn.modules.module.T¶ +
Change if autograd should record operations on parameters in this +module.
+This method sets the parameters’
+requires_gradattributes +in-place.This method is helpful for freezing part of the module for finetuning +or training parts of a model individually (e.g., GAN training).
+-
+
- Args:
-
+
- requires_grad (bool): whether autograd should record operations on
parameters in this module. Default:
+True.
+
+- Returns:
Module: self
+
+
-
+
- +to(*args, **kwargs)¶ +
Moves and/or casts the parameters and buffers.
+This can be called as
+-
+
- +to(device=None, dtype=None, non_blocking=False)¶ +
-
+
- +to(dtype, non_blocking=False)¶ +
-
+
- +to(tensor, non_blocking=False)¶ +
-
+
- +to(memory_format=torch.channels_last)¶ +
Its signature is similar to
+torch.Tensor.to(), but only accepts +floating point or complexdtype`s. In addition, this method will +only cast the floating point or complex parameters and buffers to :attr:`dtype+(if given). The integral parameters and buffers will be moved +device, if that is given, but with dtypes unchanged. When +non_blockingis set, it tries to convert/move asynchronously +with respect to the host if possible, e.g., moving CPU Tensors with +pinned memory to CUDA devices.See below for examples.
+++Note
+This method modifies the module in-place.
+-
+
- Args:
-
+
- device (
torch.device): the desired device of the parameters and buffers in this module
+
+- dtype (
torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module
+
+- tensor (torch.Tensor): Tensor whose dtype and device are the desired
dtype and device for all parameters and buffers in this module
+
+- memory_format (
torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword +only argument)
+
+
+- device (
- Returns:
Module: self
+
+
Examples:
+++>>> linear = nn.Linear(2, 2) +>>> linear.weight +Parameter containing: +tensor([[ 0.1913, -0.3420], + [-0.5113, -0.2325]]) +>>> linear.to(torch.double) +Linear(in_features=2, out_features=2, bias=True) +>>> linear.weight +Parameter containing: +tensor([[ 0.1913, -0.3420], + [-0.5113, -0.2325]], dtype=torch.float64) +>>> gpu1 = torch.device("cuda:1") +>>> linear.to(gpu1, dtype=torch.half, non_blocking=True) +Linear(in_features=2, out_features=2, bias=True) +>>> linear.weight +Parameter containing: +tensor([[ 0.1914, -0.3420], + [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1') +>>> cpu = torch.device("cpu") +>>> linear.to(cpu) +Linear(in_features=2, out_features=2, bias=True) +>>> linear.weight +Parameter containing: +tensor([[ 0.1914, -0.3420], + [-0.5112, -0.2324]], dtype=torch.float16) + +>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble) +>>> linear.weight +Parameter containing: +tensor([[ 0.3741+0.j, 0.2382+0.j], + [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128) +>>> linear(torch.ones(3, 2, dtype=torch.cdouble)) +tensor([[0.6122+0.j, 0.1150+0.j], + [0.6122+0.j, 0.1150+0.j], + [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128) +
-
+
- +type(dst_type: Union[torch.dtype, str]) → torch.nn.modules.module.T¶ +
Casts all parameters and buffers to
+dst_type.-
+
- Args:
dst_type (type or string): the desired type
+
+- Returns:
Module: self
+
+
-
+
- +xpu(device: Optional[Union[int, torch.device]] = None) → torch.nn.modules.module.T¶ +
Moves all model parameters and buffers to the XPU.
+This also makes associated parameters and buffers different objects. So +it should be called before constructing optimizer if the module will +live on XPU while being optimized.
+-
+
- Arguments:
-
+
- device (int, optional): if specified, all parameters will be
copied to that device
+
+
+- Returns:
Module: self
+
+
-
+
- +zero_grad(set_to_none: bool = False) → None¶ +
Sets gradients of all model parameters to zero. See similar function +under
+torch.optim.Optimizerfor more context.-
+
- Args:
-
+
- set_to_none (bool): instead of setting to zero, set the grads to None.
See
+torch.optim.Optimizer.zero_grad()for details.
+
+