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

Add pytorch version check before loading Python ONNX Runtime training module (#7377)

Browse source
This commit is contained in:
Thiago Crepaldi 2021-04-26 14:53:50 -07:00 committed by GitHub
parent 4804ede501
commit 0702a14ee7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
21 changed files with 272 additions and 183 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cmake/onnxruntime_python.cmake
View file

@ -223,6 +223,9 @@ if (onnxruntime_ENABLE_TRAINING)
file(GLOB onnxruntime_python_optim_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/optim/*.py"
)
file(GLOB onnxruntime_python_ortmodule_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/*.py"
)
file(GLOB onnxruntime_python_train_tools_srcs CONFIGURE_DEPENDS
"${REPO_ROOT}/tools/python/register_custom_ops_pytorch_exporter.py"
)
@ -388,6 +391,7 @@ if (onnxruntime_ENABLE_TRAINING)
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/amp
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/optim
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_capi_training_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/capi/training/
@ -400,6 +404,9 @@ if (onnxruntime_ENABLE_TRAINING)
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_optim_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/optim/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_ortmodule_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_train_tools_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/

5
orttraining/orttraining/python/training/__init__.py
View file

@ -2,12 +2,11 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from onnxruntime.capi._pybind_state import TrainingParameters
from onnxruntime.capi.training.training_session import TrainingSession
from .orttrainer_options import ORTTrainerOptions
from .orttrainer import ORTTrainer, TrainStepInfo
from . import amp, checkpoint, optim, model_desc_validation
from .execution_agent import InferenceAgent, TrainingAgent
from .ortmodule import ORTModule
from .runstateinfo import RunStateInfo

85
orttraining/orttraining/python/training/_ortmodule_inference_manager.py
View file

@ -1,85 +0,0 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _ortmodule_utils as _utils, _ortmodule_io as _io
from ._ortmodule_graph_execution_manager import GraphExecutionManager, _run_forward
import copy
import onnx
import onnxruntime
import torch
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
"""
def __init__(self, model):
super().__init__(model)
self._export_mode = torch.onnx.TrainingMode.EVAL
def forward(self, *inputs, **kwargs):
'''Forward pass of the inference model
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.
'''
# Exporting module to ONNX for the first time
build_graph = self._export_model(*inputs, **kwargs)
if build_graph:
# If model was exported, then initialize the graph builder
self._initialize_graph_builder(training=False)
# Save the onnx model if the model was exported
if self._save_onnx:
onnx.save(self._onnx_model, self._save_onnx_prefix + '_exported_inference_model.onnx')
# Build the inference graph
if build_graph:
self._build_graph()
module_device = _utils.get_device_from_module(self._original_module)
# The inference session should be created every time
# the graph was built or if the device changed between calls to forward
create_execution_session = build_graph or self._device != module_device
if self._device != module_device:
self._device = module_device
if create_execution_session:
# Create execution session creates the inference_session
self._create_execution_agent()
user_outputs, _ = _run_forward(self._execution_agent,
self._optimized_onnx_model,
self._device,
*_io._combine_input_buffers_initializers(
self._flattened_module.named_parameters(),
self._graph_info.user_input_names,
self._input_info,
self._flattened_module.named_buffers(),
inputs,
kwargs))
return _io.unflatten_user_output(self._module_output_schema,
self._graph_info.user_output_names,
user_outputs)
def _build_graph(self):
"""Build an optimized inference graph using the module_graph_builder"""
super()._build_graph()
if self._save_onnx:
onnx.save(self._optimized_onnx_model, self._save_onnx_prefix + '_inference.onnx')
def _create_execution_agent(self):
"""Creates an InferenceAgent that can run forward graph on an inference model"""
session_options, providers, provider_options = self._get_session_config()
self._execution_agent = onnxruntime.training.InferenceAgent(self._optimized_onnx_model.SerializeToString(),
session_options, providers, provider_options)

25
orttraining/orttraining/python/training/ortmodule/__init__.py Normal file
View file

@ -0,0 +1,25 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from packaging import version
# All global constant goes here, before ORTModule is imported
ONNX_OPSET_VERSION = 12
MINIMUM_TORCH_VERSION_STR = '1.8.1'
from .ortmodule import ORTModule
# Verify proper PyTorch is installed before proceding to ONNX Runtime initializetion
try:
import torch
torch_version = version.parse(torch.__version__.split('+')[0])
minimum_torch_version = version.parse(MINIMUM_TORCH_VERSION_STR)
if torch_version < minimum_torch_version:
raise RuntimeError(
f'ONNXRuntime ORTModule frontend requires PyTorch version greater or equal to {MINIMUM_TORCH_VERSION_STR}, '
f'but version {torch.__version__} was found instead.')
except:
raise(f'PyTorch {MINIMUM_TORCH_VERSION_STR} must be installed in order to run ONNXRuntime ORTModule frontend!')

0
orttraining/orttraining/python/training/execution_agent.py → orttraining/orttraining/python/training/ortmodule/_execution_agent.py
View file

76
orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py → orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -2,11 +2,11 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _ortmodule_utils as _utils, _ortmodule_io as _io
from . import _ortmodule_logger as _logger
from . import _utils, _io, _logger
from onnxruntime.training.ortmodule import ONNX_OPSET_VERSION
from onnxruntime.capi import _pybind_state as C
from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
from abc import ABC, abstractmethod
@ -20,49 +20,15 @@ import warnings
from torch.utils.cpp_extension import ROCM_HOME
ONNX_OPSET_VERSION = 12
def _run_forward(execution_session, onnx_model, device, *inputs, **kwargs):
"""Runs the forward graph on execution_session with given model inputs and device"""
# Assert that the input and model device match
_utils._check_same_device(device, "Input argument to forward", *inputs)
# 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
# have the need for passing IOBinding.
if isinstance(execution_session, onnxruntime.training.InferenceAgent):
io_binding = execution_session.io_binding()
run_options = C.RunOptions()
# Use IO binding
_utils._create_iobinding(io_binding, inputs, onnx_model, device)
# Run and return module outputs.
ort_output = execution_session.run_forward(io_binding, run_options)
forward_outputs, run_id = ort_output.ortvalues, ort_output.run_id
user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output._ortvalue) for forward_output in forward_outputs)
state = None
else:
state = C.PartialGraphExecutionState()
forward_inputs = C.OrtValueVector()
for input in inputs:
forward_inputs.append(_utils._ortvalue_from_torch_tensor(input))
forward_outputs = C.OrtValueVector()
# Run and return module outputs.
execution_session.run_forward(forward_inputs, forward_outputs, state)
user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs)
# Assert that the outputs and model device match
_utils._check_same_device(device, "Output argument from forward", *user_outputs)
output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
run_info = onnxruntime.training.RunStateInfo(state, output_info)
# Return user outputs and forward run information
return user_outputs, run_info
class RunStateInfo(object):
def __init__(self, state, output_info):
"""
:param state: State of partial run that contains intermediate tensors needed to resume the run later.
:param output_info: Output info.
"""
self.state = state
self.output_info = output_info
class GraphExecutionManager(ABC):
def __init__(self, module):
@ -145,6 +111,26 @@ class GraphExecutionManager(ABC):
self._torch_alloc = self._torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address()
self._torch_free = self._torch_gpu_allocator.gpu_caching_allocator_raw_delete_address()
@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 NotImplemented
@abstractmethod
def forward(self):
"""Executes the forward method for ORTModule

4
orttraining/orttraining/python/training/_ortmodule_graph_execution_manager_factory.py → orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py
View file

@ -3,8 +3,8 @@
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from ._ortmodule_training_manager import TrainingManager
from ._ortmodule_inference_manager import InferenceManager
from ._training_manager import TrainingManager
from ._inference_manager import InferenceManager
class GraphExecutionManagerFactory(object):

115
orttraining/orttraining/python/training/ortmodule/_inference_manager.py Normal file
View file

@ -0,0 +1,115 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _utils, _io
from ._graph_execution_manager import GraphExecutionManager, RunStateInfo
from ._execution_agent import InferenceAgent
from onnxruntime.capi import _pybind_state as C
import onnx
import torch
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
"""
def __init__(self, model):
super().__init__(model)
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"""
# Assert that the input and model device match
_utils._check_same_device(device, "Input argument to forward", *inputs)
# 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
# have the need for passing IOBinding.
io_binding = execution_session.io_binding()
run_options = C.RunOptions()
# Use IO binding
_utils._create_iobinding(io_binding, inputs, onnx_model, device)
# Run and return module outputs.
ort_output = execution_session.run_forward(io_binding, run_options)
forward_outputs, run_id = ort_output.ortvalues, ort_output.run_id
user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output._ortvalue) for forward_output in forward_outputs)
state = None
# Assert that the outputs and model device match
_utils._check_same_device(device, "Output argument from forward", *user_outputs)
output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
run_info = RunStateInfo(state, output_info)
# Return user outputs and forward run information
return user_outputs, run_info
def forward(self, *inputs, **kwargs):
'''Forward pass of the inference model
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.
'''
# Exporting module to ONNX for the first time
build_graph = self._export_model(*inputs, **kwargs)
if build_graph:
# If model was exported, then initialize the graph builder
self._initialize_graph_builder(training=False)
# Save the onnx model if the model was exported
if self._save_onnx:
onnx.save(self._onnx_model, self._save_onnx_prefix + '_exported_inference_model.onnx')
# Build the inference graph
if build_graph:
self._build_graph()
module_device = _utils.get_device_from_module(self._original_module)
# The inference session should be created every time
# the graph was built or if the device changed between calls to forward
create_execution_session = build_graph or self._device != module_device
if self._device != module_device:
self._device = module_device
if create_execution_session:
# Create execution session creates the inference_session
self._create_execution_agent()
user_outputs, _ = InferenceManager.execution_session_run_forward(self._execution_agent,
self._optimized_onnx_model,
self._device,
*_io._combine_input_buffers_initializers(
self._flattened_module.named_parameters(),
self._graph_info.user_input_names,
self._input_info,
self._flattened_module.named_buffers(),
inputs,
kwargs))
return _io.unflatten_user_output(self._module_output_schema,
self._graph_info.user_output_names,
user_outputs)
def _build_graph(self):
"""Build an optimized inference graph using the module_graph_builder"""
super()._build_graph()
if self._save_onnx:
onnx.save(self._optimized_onnx_model, self._save_onnx_prefix + '_inference.onnx')
def _create_execution_agent(self):
"""Creates an InferenceAgent that can run forward graph on an inference model"""
session_options, providers, provider_options = self._get_session_config()
self._execution_agent = InferenceAgent(self._optimized_onnx_model.SerializeToString(),
session_options, providers, provider_options)

0
orttraining/orttraining/python/training/_ortmodule_io.py → orttraining/orttraining/python/training/ortmodule/_io.py
View file

0
orttraining/orttraining/python/training/_ortmodule_logger.py → orttraining/orttraining/python/training/ortmodule/_logger.py
View file

74
orttraining/orttraining/python/training/_ortmodule_training_manager.py → orttraining/orttraining/python/training/ortmodule/_training_manager.py
View file

@ -3,14 +3,14 @@
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _ortmodule_utils as _utils, _ortmodule_io as _io
from ._ortmodule_graph_execution_manager import GraphExecutionManager, _run_forward
from . import _utils, _io
from ._graph_execution_manager import GraphExecutionManager, RunStateInfo
from ._execution_agent import TrainingAgent
from onnxruntime.capi import _pybind_state as C
from . import _utils as _utils_ort
from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_type
import onnx
import onnxruntime
import torch
@ -24,6 +24,35 @@ class TrainingManager(GraphExecutionManager):
super().__init__(model)
self._export_mode = torch.onnx.TrainingMode.TRAINING
@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"""
# Assert that the input and model device match
_utils._check_same_device(device, "Input argument to forward", *inputs)
# 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
# have the need for passing IOBinding.
state = C.PartialGraphExecutionState()
forward_inputs = C.OrtValueVector()
for input in inputs:
forward_inputs.append(_utils._ortvalue_from_torch_tensor(input))
forward_outputs = C.OrtValueVector()
# Run and return module outputs.
execution_session.run_forward(forward_inputs, forward_outputs, state)
user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs)
# Assert that the outputs and model device match
_utils._check_same_device(device, "Output argument from forward", *user_outputs)
output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
run_info = RunStateInfo(state, output_info)
# Return user outputs and forward run information
return user_outputs, run_info
def forward(self, *inputs, **kwargs):
'''Forward pass starts here and continues at `_ORTModuleFunction.forward`
@ -65,24 +94,24 @@ class TrainingManager(GraphExecutionManager):
gradient implementation for self.forward_graph.'''
@staticmethod
def forward(ctx, *inputs, **kwargs):
def forward(ctx, *inputs):
'''Performs forward pass based on user input and PyTorch initializer
Autograd Function's apply() doesn't support keyword arguments,
so `*inputs` has all the arguments - keyword arguments converted
to positional by the caller.
to positional/keywords during `TrainingManager.forward`.
Module outputs are returned to the user
'''
user_outputs, ctx.run_info = _run_forward(self._execution_agent,
self._optimized_onnx_model,
self._device,
*inputs,
**kwargs)
user_outputs, ctx.run_info = TrainingManager.execution_session_run_forward(self._execution_agent,
self._optimized_onnx_model,
self._device,
*inputs)
# Disable materializing grads then None object will not be converted to a tensor filled with zeros prior to calling backward.
# Also save shape, device and type info to ctx for materializing tensor in backward if output grad is None.
# Disable materializing grads then None object will not be
# converted to a tensor filled with zeros prior to calling backward.
# Save shape, device and type info to ctx for materializing tensor in backward if output grad is None.
ctx.set_materialize_grads(False)
return user_outputs
@ -179,19 +208,24 @@ class TrainingManager(GraphExecutionManager):
fw_outputs_device_info = []
for idx in range(len(self._graph_info.user_output_names)):
fw_outputs_device_info.append(C.OrtDevice(get_ort_device_type(self._device.type),
C.OrtDevice.default_memory(), _utils_ort.get_device_index(self._device)))
C.OrtDevice.default_memory(), _utils.get_device_index(self._device)))
bw_fetches_names = [output.name for output in self._optimized_onnx_model.graph.output]
bw_outputs_device_info = []
for idx in range(len(bw_fetches_names)):
bw_outputs_device_info.append(C.OrtDevice(get_ort_device_type(self._device.type),
C.OrtDevice.default_memory(), _utils_ort.get_device_index(self._device)))
C.OrtDevice.default_memory(), _utils.get_device_index(self._device)))
self._execution_agent = onnxruntime.training.TrainingAgent(self._optimized_onnx_model.SerializeToString(),
fw_feed_names, self._graph_info.user_output_names,
fw_outputs_device_info, self._graph_info.module_output_gradient_name,
bw_fetches_names, bw_outputs_device_info,
session_options, providers, provider_options)
self._execution_agent = TrainingAgent(self._optimized_onnx_model.SerializeToString(),
fw_feed_names,
self._graph_info.user_output_names,
fw_outputs_device_info,
self._graph_info.module_output_gradient_name,
bw_fetches_names,
bw_outputs_device_info,
session_options,
providers,
provider_options)
def _reinitialize_graph_builder(self, input_info):
"""Return true if the module graph builder was reinitialized"""

30
orttraining/orttraining/python/training/_ortmodule_utils.py → orttraining/orttraining/python/training/ortmodule/_utils.py
View file

@ -3,8 +3,6 @@
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _utils
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
from onnxruntime.capi import _pybind_state as C
@ -60,6 +58,32 @@ def _check_same_device(device, argument_str, *args):
f"{argument_str} found on device {arg_device}, but expected it to be on module device {device}.")
def get_device_index(device):
if isinstance(device, str):
# could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
device = torch.device(device)
elif isinstance(device, int):
return device
return 0 if device.index is None else device.index
def get_device_str(device):
if isinstance(device, str):
# could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
if device.find(':') == -1:
device += ':' + str(torch.cuda.current_device())
elif isinstance(device, int):
device = 'cuda:' + str(device)
elif isinstance(device, torch.device):
if device.index is None:
device = device.type + ':' + str(torch.cuda.current_device())
else:
device = device.type + ':' + str(device.index)
else:
raise RuntimeError('Unsupported device type')
return device
def get_device_from_module(module):
'''Returns the first device found in the `module`'s parameters or None'''
device = None
@ -79,4 +103,4 @@ def _create_iobinding(io_binding, inputs, model, device):
io_binding.bind_ortvalue_input(value_info.name, OrtValue(_ortvalue_from_torch_tensor(inputs[idx])))
for value_info in model.graph.output:
io_binding.bind_output(value_info.name, device.type, device_id=_utils.get_device_index(device))
io_binding.bind_output(value_info.name, device.type, device_id=get_device_index(device))

4
orttraining/orttraining/python/training/ortmodule.py → orttraining/orttraining/python/training/ortmodule/ortmodule.py
View file

@ -3,8 +3,8 @@
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from . import _ortmodule_io as _io
from ._ortmodule_graph_execution_manager_factory import GraphExecutionManagerFactory
from . import _io
from ._graph_execution_manager_factory import GraphExecutionManagerFactory
from onnxruntime.training import register_custom_ops_pytorch_exporter

13
orttraining/orttraining/python/training/runstateinfo.py
View file

@ -1,13 +0,0 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
class RunStateInfo(object):
def __init__(self, state, output_info):
"""
:param state: State of partial run that contains intermediate tensors needed to resume the run later.
:param output_info: Output info.
"""
self.state = state
self.output_info = output_info

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

@ -16,7 +16,7 @@ from collections import OrderedDict
from collections import namedtuple
from inspect import signature
from onnxruntime.training import _utils, ORTModule
from onnxruntime.training.ortmodule import ORTModule, _utils
import _test_helpers
# Import autocasting libs

2
orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py
View file

@ -17,7 +17,7 @@ import datetime
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule
def train(model, optimizer, scheduler, train_dataloader, epoch, device, args):
# ========================================

3
orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_pipeline_parallel.py
View file

@ -5,8 +5,7 @@ import deepspeed
from deepspeed.pipe import PipelineModule, LayerSpec
from deepspeed.utils import RepeatingLoader
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule, _utils
import argparse

3
orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py
View file

@ -10,14 +10,13 @@ $ deepspeed orttraining_test_ortmodule_deepspeed_zero_stage_1.py \
"""
import argparse
import logging
import os
import torch
import time
from torchvision import datasets, transforms
import torch.distributed as dist
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule
import deepspeed

3
orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py
View file

@ -9,8 +9,7 @@ from torchvision import datasets, transforms
import time
from torch.nn.parallel import DistributedDataParallel as DDP
import os
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule
import numpy as np
# Usage :

2
orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py
View file

@ -5,7 +5,7 @@ import time
from torchvision import datasets, transforms
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule
class NeuralNet(torch.nn.Module):

2
orttraining/orttraining/test/python/orttraining_test_ortmodule_torch_lightning_basic.py
View file

@ -10,7 +10,7 @@ from torch.utils.data import DataLoader
import pytorch_lightning as pl
import onnxruntime
from onnxruntime.training import ORTModule
from onnxruntime.training.ortmodule import ORTModule
class LitAutoEncoder(pl.LightningModule):