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

refactor frontend (#3235)

Browse source 
* refactor frontend

* remove training python files from inferencing build

* update according to reviewer's comments

* merge pybind_state.cc

* refactor pybind_state.cc

* code clean up

* missed a forward declaration in ort_pybind_state.cc

* passed pytest

* move training_session.py into a subfolder per reviewer's comment

* add copyright

Co-authored-by: liqun <liqun@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
liqunfu 2020-03-19 20:59:41 -07:00 committed by GitHub
parent d9f628cb1d
commit d521efd904
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 386 additions and 294 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
cmake/onnxruntime_python.cmake
View file

@ -36,6 +36,13 @@ set(onnxruntime_pybind_srcs_pattern
"${ONNXRUNTIME_ROOT}/python/*.h"
)
if (onnxruntime_ENABLE_TRAINING)
list(APPEND onnxruntime_pybind_srcs_pattern
"${ORTTRAINING_ROOT}/orttraining/python/*.cc"
"${ORTTRAINING_ROOT}/orttraining/python/*.h"
)
endif()
file(GLOB onnxruntime_pybind_srcs CONFIGURE_DEPENDS
${onnxruntime_pybind_srcs_pattern}
)
@ -142,9 +149,28 @@ endif()
file(GLOB onnxruntime_backend_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/python/backend/*.py"
)
file(GLOB onnxruntime_python_srcs CONFIGURE_DEPENDS
if (onnxruntime_ENABLE_TRAINING)
file(GLOB onnxruntime_python_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/python/*.py"
)
"${ORTTRAINING_SOURCE_DIR}/python/*.py"
)
else()
file(GLOB onnxruntime_python_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/python/*.py"
)
endif()
if (onnxruntime_ENABLE_TRAINING)
file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/*.py"
)
else()
file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/python/training/*.py"
)
endif()
file(GLOB onnxruntime_python_test_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/test/python/*.py"
)
@ -168,6 +194,7 @@ add_custom_command(
TARGET onnxruntime_pybind11_state POST_BUILD
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/backend
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi/training
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/datasets
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/tools
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/tools/featurizer_ops
@ -192,6 +219,9 @@ add_custom_command(
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_srcs}
$<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_capi_training_srcs}
$<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi/training/
COMMAND ${CMAKE_COMMAND} -E copy
$<TARGET_FILE:onnxruntime_pybind11_state>
$<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi/

6
onnxruntime/__init__.py
View file

@ -13,10 +13,6 @@ from onnxruntime.capi._pybind_state import get_all_providers, get_available_prov
from onnxruntime.capi.session import InferenceSession, IOBinding
from onnxruntime.capi import onnxruntime_validation
try:
from onnxruntime.capi._pybind_state import TrainingParameters
from onnxruntime.capi.session import TrainingSession
except ImportError:
print("ONNX Runtime Training identifiers are unavailable.")
from onnxruntime.capi.training import *
onnxruntime_validation.check_distro_info()

5
onnxruntime/python/onnxruntime_pybind_exceptions.h
View file

@ -1,3 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <pybind11/pybind11.h>
#include <stdexcept>
#include "core/common/status.h"
@ -40,7 +43,7 @@ struct EPFail : std::runtime_error {
explicit EPFail(const std::string& what) : std::runtime_error(what) {}
};
void RegisterExceptions(pybind11::module& m) {
inline void RegisterExceptions(pybind11::module& m) {
pybind11::register_exception<Fail>(m, "Fail");
pybind11::register_exception<InvalidArgument>(m, "InvalidArgument");
pybind11::register_exception<NoSuchFile>(m, "NoSuchFile");

256
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -6,6 +6,7 @@
#include "python/onnxruntime_pybind_exceptions.h"
#include "python/onnxruntime_pybind_mlvalue.h"
#include "python/onnxruntime_pybind_state_common.h"
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#define PY_ARRAY_UNIQUE_SYMBOL onnxruntime_python_ARRAY_API
@ -90,13 +91,6 @@
#include "core/providers/cpu/cpu_execution_provider.h"
#include "core/providers/cpu/cpu_provider_factory.h"
#ifdef ENABLE_TRAINING
#include "core/framework/random_seed.h"
#include "orttraining/core/session/training_session.h"
#include "orttraining/core/graph/optimizer_config.h"
#include "orttraining/core/framework/mpi_setup.h"
#endif // ENABLE_TRAINING
#ifdef USE_CUDA
#include "core/providers/cuda/cuda_provider_factory.h"
OrtDevice::DeviceId cuda_device_id = 0;
@ -149,16 +143,6 @@ namespace py = pybind11;
using namespace onnxruntime;
using namespace onnxruntime::logging;
static AllocatorPtr& GetAllocator() {
static AllocatorPtr alloc = std::make_shared<TAllocator>();
return alloc;
}
static const SessionOptions& GetDefaultCPUSessionOptions() {
static SessionOptions so;
return so;
}
template <typename T>
void AddNonTensor(OrtValue& val, std::vector<py::object>& pyobjs) {
pyobjs.push_back(py::cast(val.Get<T>()));
@ -272,27 +256,6 @@ void AddTensorAsPyObj(OrtValue& val, std::vector<py::object>& pyobjs) {
pyobjs.push_back(obj);
}
class SessionObjectInitializer {
public:
typedef const SessionOptions& Arg1;
typedef logging::LoggingManager* Arg2;
operator Arg1() {
return GetDefaultCPUSessionOptions();
}
operator Arg2() {
static std::string default_logger_id{"Default"};
static LoggingManager default_logging_manager{std::unique_ptr<ISink>{new CErrSink{}},
Severity::kWARNING, false, LoggingManager::InstanceType::Default,
&default_logger_id};
return &default_logging_manager;
}
static SessionObjectInitializer Get() {
return SessionObjectInitializer();
}
};
inline void RegisterExecutionProvider(InferenceSession* sess, onnxruntime::IExecutionProviderFactory& f) {
auto p = f.CreateProvider();
OrtPybindThrowIfError(sess->RegisterExecutionProvider(std::move(p)));
@ -932,222 +895,9 @@ static struct {
#endif
#ifdef ENABLE_TRAINING
namespace {
struct TrainingParameters {
std::string loss_output_name;
std::unordered_set<std::string> weights_to_train;
std::unordered_set<std::string> weights_not_to_train;
onnxruntime::training::TrainingSession::ImmutableWeights immutable_weights;
// optimizer
std::string training_optimizer_name;
std::string loss_scale_input_name;
std::string scaled_loss_output_name;
std::string lr_params_feed_name = "Learning_Rate";
std::unordered_map<std::string, std::unordered_map<std::string, float>> optimizer_attributes_map;
bool use_fp16_moments = false;
bool use_mixed_precision = false;
bool allreduce_post_accumulation = false;
float loss_scale = 0.0f;
int world_rank = 0;
int world_size = 1;
int local_rank = 0;
int local_size = 1;
int gradient_accumulation_steps = 1;
int data_parallel_size = 1;
int horizontal_parallel_size = 1;
bool partition_optimizer = false;
int seed = -1;
};
// TODO: this method does not handle parallel optimization.
void ConfigureSessionForTraining(
training::TrainingSession* sess, TrainingParameters& parameters) {
//TODO tix, refactor the mpi related code to populate all fields correctly by default.
ORT_ENFORCE(parameters.horizontal_parallel_size <= parameters.world_size);
ORT_ENFORCE(parameters.data_parallel_size <= parameters.world_size);
if (parameters.world_size % parameters.horizontal_parallel_size != 0) {
throw std::runtime_error("Cannot split horizontal parallel group because world_size is not divisible");
}
auto data_group_size = parameters.world_size / parameters.horizontal_parallel_size;
if (data_group_size != parameters.data_parallel_size) {
std::cout << "WARNING: data_parallel_size is not correct, tuned automatically to "
<< data_group_size << std::endl;
parameters.data_parallel_size = data_group_size;
}
#ifdef USE_HOROVOD
// this condition block is temporary.
// For now, nccl allreduce kernel only implements for allreduce_post_accumulation
// hovorod allreduce kernel only implements for not allreduce_post_accumulation.
bool use_nccl = parameters.allreduce_post_accumulation;
if (!use_nccl && parameters.world_size > 1) {
auto mpi_context = training::setup_horovod();
std::cout << "mpi_context.world_rank: " << mpi_context.world_rank << std::endl;
std::cout << "mpi_context.local_rank: " << mpi_context.local_rank << std::endl;
std::cout << "mpi_context.world_size: " << mpi_context.world_size << std::endl;
std::cout << "mpi_context.local_size: " << mpi_context.local_size << std::endl;
parameters.local_size = mpi_context.local_size;
parameters.local_rank = mpi_context.local_rank;
}
void addObjectMethodsForTraining(py::module& m);
#endif
training::TrainingSession::TrainingConfiguration config{};
config.weight_names_to_train = parameters.weights_to_train;
config.weight_names_to_not_train = parameters.weights_not_to_train;
config.immutable_weights = parameters.immutable_weights;
config.set_gradients_as_graph_outputs = true;
config.gradient_accumulation_steps = parameters.gradient_accumulation_steps;
config.distributed_config.world_rank = parameters.world_rank;
config.distributed_config.world_size = parameters.world_size;
config.distributed_config.local_rank = parameters.local_rank;
config.distributed_config.local_size = parameters.local_size;
config.distributed_config.data_parallel_size = parameters.data_parallel_size;
config.distributed_config.horizontal_parallel_size = parameters.horizontal_parallel_size;
if (parameters.use_mixed_precision) {
training::TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mp{};
mp.add_loss_scaling = false;
mp.use_fp16_initializers = true;
config.mixed_precision_config = mp;
}
config.loss_name =
parameters.use_mixed_precision ? parameters.scaled_loss_output_name : parameters.loss_output_name;
if (!parameters.training_optimizer_name.empty()) {
training::TrainingSession::TrainingConfiguration::OptimizerConfiguration opt{};
opt.name = parameters.training_optimizer_name;
opt.learning_rate_input_name = parameters.lr_params_feed_name;
opt.weight_attributes_generator = [&parameters](const std::string& weight_name) {
const auto it = parameters.optimizer_attributes_map.find(weight_name);
ORT_ENFORCE(
it != parameters.optimizer_attributes_map.end(),
"Failed to find attribute map for weight ", weight_name);
return it->second;
};
opt.use_fp16_moments = parameters.use_fp16_moments;
opt.do_all_reduce_in_fp16 = true;
// TODO: this mapping is temporary.
// For now, nccl allreduce kernel only implements for allreduce_post_accumulation
// hovorod allreduce kernel only implements for not allreduce_post_accumulation.
// eventually we will have one all reduce kernel and let opt to have
// an allreduce_post_accumulation option and remove the use_nccl option.
opt.use_nccl = parameters.allreduce_post_accumulation;
opt.partition_optimizer = parameters.partition_optimizer;
config.optimizer_config = opt;
}
if (parameters.seed > 0) {
utils::SetStaticRandomSeed(static_cast<uint32_t>(parameters.seed));
std::cout << "Random seed is set to " << parameters.seed << std::endl;
}
training::TrainingSession::TrainingConfigurationResult config_result{};
OrtPybindThrowIfError(sess->ConfigureForTraining(config, config_result));
}
void addObjectMethodsForTraining(py::module& m) {
py::class_<TrainingParameters> parameters(m, "TrainingParameters", R"pbdoc(Configuration information for training.)pbdoc");
parameters.def(py::init())
.def_readwrite("loss_output_name", &TrainingParameters::loss_output_name)
.def_readwrite("immutable_weights", &TrainingParameters::immutable_weights)
.def_readwrite("weights_not_to_train", &TrainingParameters::weights_not_to_train)
.def_readwrite("weights_to_train", &TrainingParameters::weights_to_train)
.def_readwrite("loss_scale_input_name", &TrainingParameters::loss_scale_input_name)
.def_readwrite("scaled_loss_output_name", &TrainingParameters::scaled_loss_output_name)
.def_readwrite("training_optimizer_name", &TrainingParameters::training_optimizer_name)
.def_readwrite("lr_params_feed_name", &TrainingParameters::lr_params_feed_name)
.def_readwrite("optimizer_attributes_map", &TrainingParameters::optimizer_attributes_map)
.def_readwrite("use_fp16_moments", &TrainingParameters::use_fp16_moments)
.def_readwrite("use_mixed_precision", &TrainingParameters::use_mixed_precision)
.def_readwrite("allreduce_post_accumulation", &TrainingParameters::allreduce_post_accumulation)
.def_readwrite("loss_scale", &TrainingParameters::loss_scale)
.def_readwrite("world_rank", &TrainingParameters::world_rank)
.def_readwrite("world_size", &TrainingParameters::world_size)
.def_readwrite("gradient_accumulation_steps", &TrainingParameters::gradient_accumulation_steps)
.def_readwrite("partition_optimizer", &TrainingParameters::partition_optimizer);
py::class_<onnxruntime::training::TrainingSession, InferenceSession> training_session(m, "TrainingSession");
training_session.def(py::init<SessionOptions, SessionObjectInitializer>())
.def(py::init<SessionObjectInitializer, SessionObjectInitializer>())
.def("finalize", [](py::object) {
#ifdef USE_HOROVOD
training::shutdown_horovod();
#endif
})
.def("load_model", [](onnxruntime::training::TrainingSession* sess, const std::string& path, TrainingParameters& parameters) {
OrtPybindThrowIfError(sess->Load(path));
ConfigureSessionForTraining(sess, parameters);
std::vector<std::string> provider_types = {};
InitializeSession(sess, provider_types);
})
.def("read_bytes", [](onnxruntime::training::TrainingSession* sess, const py::bytes& serialized_model, TrainingParameters& parameters) {
std::istringstream buffer(serialized_model);
OrtPybindThrowIfError(sess->Load(buffer));
ConfigureSessionForTraining(sess, parameters);
std::vector<std::string> provider_types = {};
InitializeSession(sess, provider_types);
})
.def("get_state", [](onnxruntime::training::TrainingSession* sess) {
NameMLValMap state_tensors;
ORT_THROW_IF_ERROR(sess->GetStateTensors(state_tensors));
auto& data_transfer_manager = sess->GetDataTransferManager();
//convert to numpy array
std::map<std::string, py::object> rmap;
for (auto& kv : state_tensors) {
if (kv.second.IsTensor()) {
py::object obj;
const Tensor& rtensor = kv.second.Get<Tensor>();
GetPyObjFromTensor(rtensor, obj, &data_transfer_manager);
rmap.insert({kv.first, obj});
} else {
throw std::runtime_error("Non tensor type in session state tensors is not expected.");
}
}
return rmap;
})
.def("load_state", [](onnxruntime::training::TrainingSession* sess, std::unordered_map<std::string, py::object>& state, bool strict) {
NameMLValMap state_tensors;
for (auto initializer : state) {
OrtValue ml_value;
auto px = sess->GetModelInputs();
if (!px.first.IsOK() || !px.second) {
throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null");
}
CreateGenericMLValue(px.second, GetAllocator(), initializer.first, initializer.second, &ml_value);
if (PyErr_Occurred()) {
PyObject *ptype, *pvalue, *ptraceback;
PyErr_Fetch(&ptype, &pvalue, &ptraceback);
PyObject* pStr = PyObject_Str(ptype);
std::string sType = py::reinterpret_borrow<py::str>(pStr);
Py_XDECREF(pStr);
pStr = PyObject_Str(pvalue);
sType += ": ";
sType += py::reinterpret_borrow<py::str>(pStr);
Py_XDECREF(pStr);
throw std::runtime_error(sType);
}
state_tensors.insert(std::make_pair(initializer.first, ml_value));
}
ORT_THROW_IF_ERROR(sess->SetStateTensors(state_tensors, strict));
});
}
} // namespace
#endif // ENABLE_TRAINING
PYBIND11_MODULE(onnxruntime_pybind11_state, m) {
m.doc() = "pybind11 stateful interface to ONNX runtime";
RegisterExceptions(m);
@ -1228,4 +978,4 @@ PYBIND11_MODULE(onnxruntime_pybind11_state, m) {
}
} // namespace python
} // namespace onnxruntime
} // namespace onnxruntime

47
onnxruntime/python/onnxruntime_pybind_state_common.h Normal file
View file

@ -0,0 +1,47 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/common/logging/logging.h"
#include "core/common/logging/sinks/cerr_sink.h"
#include "core/framework/allocator.h"
#include "core/framework/session_options.h"
namespace onnxruntime {
namespace python {
using namespace onnxruntime;
using namespace onnxruntime::logging;
inline const SessionOptions& GetDefaultCPUSessionOptions() {
static SessionOptions so;
return so;
}
inline AllocatorPtr& GetAllocator() {
static AllocatorPtr alloc = std::make_shared<TAllocator>();
return alloc;
}
class SessionObjectInitializer {
public:
typedef const SessionOptions& Arg1;
typedef logging::LoggingManager* Arg2;
operator Arg1() {
return GetDefaultCPUSessionOptions();
}
operator Arg2() {
static std::string default_logger_id{"Default"};
static LoggingManager default_logging_manager{std::unique_ptr<ISink>{new CErrSink{}},
Severity::kWARNING, false, LoggingManager::InstanceType::Default,
&default_logger_id};
return &default_logging_manager;
}
static SessionObjectInitializer Get() {
return SessionObjectInitializer();
}
};
}
}

34
onnxruntime/python/session.py
View file

@ -207,36 +207,4 @@ class IOBinding:
self._iobinding.clear_binding_inputs()
def clear_binding_outputs(self):
self._iobinding.clear_binding_outputs()
class TrainingSession(InferenceSession):
def __init__(self, path_or_bytes, parameters, sess_options=None):
if sess_options:
self._sess = C.TrainingSession(
sess_options, C.get_session_initializer())
else:
self._sess = C.TrainingSession(
C.get_session_initializer(), C.get_session_initializer())
if isinstance(path_or_bytes, str):
self._sess.load_model(path_or_bytes, parameters)
elif isinstance(path_or_bytes, bytes):
self._sess.read_bytes(path_or_bytes, parameters)
else:
raise TypeError("Unable to load from type '{0}'".format(type(path_or_bytes)))
self._inputs_meta = self._sess.inputs_meta
self._outputs_meta = self._sess.outputs_meta
Session.__init__(self, self._sess)
def __del__(self):
if self._sess:
self._sess.finalize()
def get_state(self):
return self._sess.get_state()
def load_state(self, dict, strict=False):
self._sess.load_state(dict, strict)
self._iobinding.clear_binding_outputs()

7
onnxruntime/python/training/__init__.py Normal file
View file

@ -0,0 +1,7 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
# This is a placeholder file. It is deployed with inference only wheel.
print("An inference only version of ONNX Runtime is installed. Training functionalities are unavailable.")

0
orttraining/orttraining/python/dummy.txt
View file

0
onnxruntime/python/ort_trainer.py → orttraining/orttraining/python/ort_trainer.py
View file

242
orttraining/orttraining/python/orttraining_pybind_state.cc Normal file
View file

@ -0,0 +1,242 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "python/onnxruntime_pybind_exceptions.h"
#include "python/onnxruntime_pybind_state_common.h"
// pybind11/stl.h is needed to support std::unordered_set, etc.
#include <pybind11/stl.h>
#include "core/framework/random_seed.h"
#include "core/framework/session_options.h"
#include "orttraining/core/session/training_session.h"
#include "orttraining/core/graph/optimizer_config.h"
#include "orttraining/core/framework/mpi_setup.h"
namespace onnxruntime {
namespace python {
namespace py = pybind11;
using namespace onnxruntime;
using namespace onnxruntime::logging;
// BEGIN: forward declaration for stuff in onnxruntime_pybind_state
void InitializeSession(InferenceSession* sess, const std::vector<std::string>& provider_types);
void GetPyObjFromTensor(const Tensor& rtensor, py::object& obj, const DataTransferManager* data_transfer_manager = nullptr);
void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, AllocatorPtr alloc, const std::string& name_input,
py::object& value, OrtValue* p_mlvalue);
// END: forward declaration
struct TrainingParameters {
std::string loss_output_name;
std::unordered_set<std::string> weights_to_train;
std::unordered_set<std::string> weights_not_to_train;
onnxruntime::training::TrainingSession::ImmutableWeights immutable_weights;
// optimizer
std::string training_optimizer_name;
std::string loss_scale_input_name;
std::string scaled_loss_output_name;
std::string lr_params_feed_name = "Learning_Rate";
std::unordered_map<std::string, std::unordered_map<std::string, float>> optimizer_attributes_map;
bool use_fp16_moments = false;
bool use_mixed_precision = false;
bool allreduce_post_accumulation = false;
float loss_scale = 0.0f;
int world_rank = 0;
int world_size = 1;
int local_rank = 0;
int local_size = 1;
int gradient_accumulation_steps = 1;
int data_parallel_size = 1;
int horizontal_parallel_size = 1;
bool partition_optimizer = false;
int seed = -1;
};
// TODO: this method does not handle parallel optimization.
void ConfigureSessionForTraining(
training::TrainingSession* sess, TrainingParameters& parameters) {
//TODO tix, refactor the mpi related code to populate all fields correctly by default.
ORT_ENFORCE(parameters.horizontal_parallel_size <= parameters.world_size);
ORT_ENFORCE(parameters.data_parallel_size <= parameters.world_size);
if (parameters.world_size % parameters.horizontal_parallel_size != 0) {
throw std::runtime_error("Cannot split horizontal parallel group because world_size is not divisible");
}
auto data_group_size = parameters.world_size / parameters.horizontal_parallel_size;
if (data_group_size != parameters.data_parallel_size) {
std::cout << "WARNING: data_parallel_size is not correct, tuned automatically to "
<< data_group_size << std::endl;
parameters.data_parallel_size = data_group_size;
}
#ifdef USE_HOROVOD
// this condition block is temporary.
// For now, nccl allreduce kernel only implements for allreduce_post_accumulation
// hovorod allreduce kernel only implements for not allreduce_post_accumulation.
bool use_nccl = parameters.allreduce_post_accumulation;
if (!use_nccl && parameters.world_size > 1) {
auto mpi_context = training::setup_horovod();
std::cout << "mpi_context.world_rank: " << mpi_context.world_rank << std::endl;
std::cout << "mpi_context.local_rank: " << mpi_context.local_rank << std::endl;
std::cout << "mpi_context.world_size: " << mpi_context.world_size << std::endl;
std::cout << "mpi_context.local_size: " << mpi_context.local_size << std::endl;
parameters.local_size = mpi_context.local_size;
parameters.local_rank = mpi_context.local_rank;
}
#endif
training::TrainingSession::TrainingConfiguration config{};
config.weight_names_to_train = parameters.weights_to_train;
config.weight_names_to_not_train = parameters.weights_not_to_train;
config.immutable_weights = parameters.immutable_weights;
config.set_gradients_as_graph_outputs = true;
config.gradient_accumulation_steps = parameters.gradient_accumulation_steps;
config.distributed_config.world_rank = parameters.world_rank;
config.distributed_config.world_size = parameters.world_size;
config.distributed_config.local_rank = parameters.local_rank;
config.distributed_config.local_size = parameters.local_size;
config.distributed_config.data_parallel_size = parameters.data_parallel_size;
config.distributed_config.horizontal_parallel_size = parameters.horizontal_parallel_size;
if (parameters.use_mixed_precision) {
training::TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mp{};
mp.add_loss_scaling = false;
mp.use_fp16_initializers = true;
config.mixed_precision_config = mp;
}
config.loss_name =
parameters.use_mixed_precision ? parameters.scaled_loss_output_name : parameters.loss_output_name;
if (!parameters.training_optimizer_name.empty()) {
training::TrainingSession::TrainingConfiguration::OptimizerConfiguration opt{};
opt.name = parameters.training_optimizer_name;
opt.learning_rate_input_name = parameters.lr_params_feed_name;
opt.weight_attributes_generator = [&parameters](const std::string& weight_name) {
const auto it = parameters.optimizer_attributes_map.find(weight_name);
ORT_ENFORCE(
it != parameters.optimizer_attributes_map.end(),
"Failed to find attribute map for weight ", weight_name);
return it->second;
};
opt.use_fp16_moments = parameters.use_fp16_moments;
opt.do_all_reduce_in_fp16 = true;
// TODO: this mapping is temporary.
// For now, nccl allreduce kernel only implements for allreduce_post_accumulation
// hovorod allreduce kernel only implements for not allreduce_post_accumulation.
// eventually we will have one all reduce kernel and let opt to have
// an allreduce_post_accumulation option and remove the use_nccl option.
opt.use_nccl = parameters.allreduce_post_accumulation;
opt.partition_optimizer = parameters.partition_optimizer;
config.optimizer_config = opt;
}
if (parameters.seed > 0) {
utils::SetStaticRandomSeed(static_cast<uint32_t>(parameters.seed));
std::cout << "Random seed is set to " << parameters.seed << std::endl;
}
training::TrainingSession::TrainingConfigurationResult config_result{};
OrtPybindThrowIfError(sess->ConfigureForTraining(config, config_result));
}
void addObjectMethodsForTraining(py::module& m) {
py::class_<TrainingParameters> parameters(m, "TrainingParameters", R"pbdoc(Configuration information for training.)pbdoc");
parameters.def(py::init())
.def_readwrite("loss_output_name", &TrainingParameters::loss_output_name)
.def_readwrite("immutable_weights", &TrainingParameters::immutable_weights)
.def_readwrite("weights_not_to_train", &TrainingParameters::weights_not_to_train)
.def_readwrite("weights_to_train", &TrainingParameters::weights_to_train)
.def_readwrite("loss_scale_input_name", &TrainingParameters::loss_scale_input_name)
.def_readwrite("scaled_loss_output_name", &TrainingParameters::scaled_loss_output_name)
.def_readwrite("training_optimizer_name", &TrainingParameters::training_optimizer_name)
.def_readwrite("lr_params_feed_name", &TrainingParameters::lr_params_feed_name)
.def_readwrite("optimizer_attributes_map", &TrainingParameters::optimizer_attributes_map)
.def_readwrite("use_fp16_moments", &TrainingParameters::use_fp16_moments)
.def_readwrite("use_mixed_precision", &TrainingParameters::use_mixed_precision)
.def_readwrite("allreduce_post_accumulation", &TrainingParameters::allreduce_post_accumulation)
.def_readwrite("loss_scale", &TrainingParameters::loss_scale)
.def_readwrite("world_rank", &TrainingParameters::world_rank)
.def_readwrite("world_size", &TrainingParameters::world_size)
.def_readwrite("gradient_accumulation_steps", &TrainingParameters::gradient_accumulation_steps)
.def_readwrite("partition_optimizer", &TrainingParameters::partition_optimizer);
py::class_<onnxruntime::training::TrainingSession, InferenceSession> training_session(m, "TrainingSession");
training_session.def(py::init<SessionOptions, SessionObjectInitializer>())
.def(py::init<SessionObjectInitializer, SessionObjectInitializer>())
.def("finalize", [](py::object) {
#ifdef USE_HOROVOD
training::shutdown_horovod();
#endif
})
.def("load_model", [](onnxruntime::training::TrainingSession* sess, const std::string& path, TrainingParameters& parameters) {
OrtPybindThrowIfError(sess->Load(path));
ConfigureSessionForTraining(sess, parameters);
std::vector<std::string> provider_types = {};
InitializeSession(sess, provider_types);
})
.def("read_bytes", [](onnxruntime::training::TrainingSession* sess, const py::bytes& serialized_model, TrainingParameters& parameters) {
std::istringstream buffer(serialized_model);
OrtPybindThrowIfError(sess->Load(buffer));
ConfigureSessionForTraining(sess, parameters);
std::vector<std::string> provider_types = {};
InitializeSession(sess, provider_types);
})
.def("get_state", [](onnxruntime::training::TrainingSession* sess) {
NameMLValMap state_tensors;
ORT_THROW_IF_ERROR(sess->GetStateTensors(state_tensors));
auto& data_transfer_manager = sess->GetDataTransferManager();
//convert to numpy array
std::map<std::string, py::object> rmap;
for (auto& kv : state_tensors) {
if (kv.second.IsTensor()) {
py::object obj;
const Tensor& rtensor = kv.second.Get<Tensor>();
GetPyObjFromTensor(rtensor, obj, &data_transfer_manager);
rmap.insert({kv.first, obj});
} else {
throw std::runtime_error("Non tensor type in session state tensors is not expected.");
}
}
return rmap;
})
.def("load_state", [](onnxruntime::training::TrainingSession* sess, std::unordered_map<std::string, py::object>& state, bool strict) {
NameMLValMap state_tensors;
for (auto initializer : state) {
OrtValue ml_value;
auto px = sess->GetModelInputs();
if (!px.first.IsOK() || !px.second) {
throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null");
}
CreateGenericMLValue(px.second, GetAllocator(), initializer.first, initializer.second, &ml_value);
if (PyErr_Occurred()) {
PyObject *ptype, *pvalue, *ptraceback;
PyErr_Fetch(&ptype, &pvalue, &ptraceback);
PyObject* pStr = PyObject_Str(ptype);
std::string sType = py::reinterpret_borrow<py::str>(pStr);
Py_XDECREF(pStr);
pStr = PyObject_Str(pvalue);
sType += ": ";
sType += py::reinterpret_borrow<py::str>(pStr);
Py_XDECREF(pStr);
throw std::runtime_error(sType);
}
state_tensors.insert(std::make_pair(initializer.first, ml_value));
}
ORT_THROW_IF_ERROR(sess->SetStateTensors(state_tensors, strict));
});
}
} // namespace python
} // namespace onnxruntime

6
orttraining/orttraining/python/training/__init__.py Normal file
View file

@ -0,0 +1,6 @@
#-------------------------------------------------------------------------
# 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

42
orttraining/orttraining/python/training/training_session.py Normal file
View file

@ -0,0 +1,42 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
import sys
import os
from onnxruntime.capi import _pybind_state as C
from onnxruntime.capi.session import Session, InferenceSession, IOBinding
class TrainingSession(InferenceSession):
def __init__(self, path_or_bytes, parameters, sess_options=None):
if sess_options:
self._sess = C.TrainingSession(
sess_options, C.get_session_initializer())
else:
self._sess = C.TrainingSession(
C.get_session_initializer(), C.get_session_initializer())
if isinstance(path_or_bytes, str):
self._sess.load_model(path_or_bytes, parameters)
elif isinstance(path_or_bytes, bytes):
self._sess.read_bytes(path_or_bytes, parameters)
else:
raise TypeError("Unable to load from type '{0}'".format(type(path_or_bytes)))
self._inputs_meta = self._sess.inputs_meta
self._outputs_meta = self._sess.outputs_meta
Session.__init__(self, self._sess)
def __del__(self):
if self._sess:
self._sess.finalize()
def get_state(self):
return self._sess.get_state()
def load_state(self, dict, strict=False):
self._sess.load_state(dict, strict)

0
orttraining/pytorch_frontend/mnist_training.py → orttraining/pytorch_frontend_examples/mnist_training.py
View file

1
setup.py
View file

@ -214,6 +214,7 @@ setup(
packages=['onnxruntime',
'onnxruntime.backend',
'onnxruntime.capi',
'onnxruntime.capi.training',
'onnxruntime.datasets',
'onnxruntime.tools',
] + (['onnxruntime.nuphar'] if package_name == 'onnxruntime-nuphar' else []),