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

C API version 0.001 (#11758)

Browse source 
* C API version 0.001

* fix linker issues

* fixes for save checkpoint api

* plus fixes based on tests

* plus test_runner and other changes

* Plus cosmetic updates

* remove unnecessary headers

* plus some updates

* plus more changes

Co-authored-by: Ashwini Khade <askhade@microsoft.com@orttrainingdev10.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
This commit is contained in:
Ashwini Khade 2022-06-15 11:13:35 -07:00 committed by GitHub
parent fb88efbe18
commit f63e28c92f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 667 additions and 138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
cmake/onnxruntime_session.cmake
View file

@ -7,6 +7,15 @@ file(GLOB onnxruntime_session_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/core/session/*.cc"
)
if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
file(GLOB_RECURSE on_device_training_api_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/training_api/*.cc"
)
list(APPEND onnxruntime_session_srcs ${on_device_training_api_srcs})
endif()
if (onnxruntime_MINIMAL_BUILD)
set(onnxruntime_session_src_exclude
"${ONNXRUNTIME_ROOT}/core/session/provider_bridge_ort.cc"
@ -48,7 +57,7 @@ if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
endif()
if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
onnxruntime_add_include_to_target(onnxruntime_session Python::Module)
onnxruntime_add_include_to_target(onnxruntime_session Python::Module)
endif()
if (NOT onnxruntime_BUILD_SHARED_LIB)

8
cmake/onnxruntime_training.cmake
View file

@ -17,14 +17,6 @@ file(GLOB_RECURSE onnxruntime_training_srcs
"${ORTTRAINING_SOURCE_DIR}/core/agent/*.cc"
)
if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
file(GLOB_RECURSE onnxruntime_training_api_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/training_api/*.h"
"${ORTTRAINING_SOURCE_DIR}/training_api/*.cc"
)
list(APPEND onnxruntime_training_srcs ${onnxruntime_training_api_srcs})
endif()
# This needs to be built in framework.cmake
file(GLOB_RECURSE onnxruntime_training_framework_excluded_srcs CONFIGURE_DEPENDS

28
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -270,6 +270,11 @@ ORT_RUNTIME_CLASS(CUDAProviderOptionsV2);
ORT_RUNTIME_CLASS(Op);
ORT_RUNTIME_CLASS(OpAttr);
#ifdef ENABLE_TRAINING_ON_DEVICE
ORT_RUNTIME_CLASS(TrainingSession);
ORT_RUNTIME_CLASS(CheckpointState);
#endif
#ifdef _WIN32
typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr;
#else
@ -3343,13 +3348,13 @@ struct OrtApi {
_In_reads_(input_len) const OrtValue* const* initializers, size_t initializers_num);
/** \brief: Create attribute of onnxruntime operator
*
*
* \param[in] name of the attribute
* \param[in] data of the attribute
* \param[in] data length
* \param[in] data type
* \param[out] attribute that has been created, which must be released by OrtApi::ReleaseOpAttr
*
*
* \since Version 1.12.
*/
ORT_API2_STATUS(CreateOpAttr,
@ -3362,14 +3367,14 @@ struct OrtApi {
/* \brief: Release op attribute
*
* \param[in] attribute created by OrtApi::CreateOpAttr
*
*
* \since Version 1.12.
*/
ORT_CLASS_RELEASE(OpAttr);
/** \brief: Create onnxruntime native operator
*
* \param[in] kernel info
*
* \param[in] kernel info
* \param[in] operator name
* \param[in] operator domain
* \param[in] operator opset
@ -3379,7 +3384,7 @@ struct OrtApi {
* \param[in] attributes used to initialize the operator
* \param[in] number of the attributes
* \param[out] operator that has been created
*
*
* \since Version 1.12.
*/
ORT_API2_STATUS(CreateOp,
@ -3396,14 +3401,14 @@ struct OrtApi {
/** \brief: Invoke the operator created by OrtApi::CreateOp
* The inputs must follow the order as specified in onnx specification
*
*
* \param[in] kernel context
* \param[in] operator that has been created
* \param[in] inputs
* \param[in] number of inputs
* \param[in] outputs
* \param[in] number of outputs
*
*
* \since Version 1.12.
*/
ORT_API2_STATUS(InvokeOp,
@ -3417,10 +3422,15 @@ struct OrtApi {
/* \brief: Release an onnxruntime operator
*
* \param[in] operator created by OrtApi::CreateOp
*
*
* \since Version 1.12.
*/
ORT_CLASS_RELEASE(Op);
#ifdef ENABLE_TRAINING_ON_DEVICE
// defines c apis for on device training scenarios
#include "../../../orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h"
#endif
};
/*

16
onnxruntime/core/session/onnxruntime_c_api.cc
View file

@ -2527,6 +2527,22 @@ static constexpr OrtApi ort_api_1_to_12 = {
&OrtApis::CreateOp,
&OrtApis::InvokeOp,
&OrtApis::ReleaseOp,
#ifdef ENABLE_TRAINING_ON_DEVICE
// Experimental for on-device training. Always keep at the bottom.
&OrtApis::LoadCheckpoint,
&OrtApis::SaveCheckpoint,
&OrtApis::CreateTrainingSession,
&OrtApis::InitializeTrainingSession,
&OrtApis::TrainingSessionGetTrainModeOutputCount,
&OrtApis::TrainingSessionGetEvalModeOutputCount,
&OrtApis::ResetGrad,
&OrtApis::TrainStep,
&OrtApis::EvalStep,
&OrtApis::OptimizerStep,
&OrtApis::ReleaseTrainingSession,
&OrtApis::ReleaseCheckpointState,
#endif
};
// Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other)

33
onnxruntime/core/session/ort_apis.h
View file

@ -375,4 +375,37 @@ ORT_API_STATUS_IMPL(InvokeOp,
ORT_API(void, ReleaseOp, _Frees_ptr_opt_ OrtOp* op);
#ifdef ENABLE_TRAINING_ON_DEVICE
ORT_API_STATUS_IMPL(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options,
_Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out);
ORT_API_STATUS_IMPL(InitializeTrainingSession, _Inout_ OrtTrainingSession* session,
_In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path,
_In_ const ORTCHAR_T* optimizer_model_path);
ORT_API(void, ReleaseTrainingSession, _Frees_ptr_opt_ OrtTrainingSession* session);
ORT_API_STATUS_IMPL(TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
ORT_API_STATUS_IMPL(TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
ORT_API_STATUS_IMPL(ResetGrad, _Inout_ OrtTrainingSession* session);
ORT_API_STATUS_IMPL(TrainStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(input_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
ORT_API_STATUS_IMPL(EvalStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(input_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
ORT_API_STATUS_IMPL(OptimizerStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options);
ORT_API_STATUS_IMPL(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state);
ORT_API_STATUS_IMPL(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* session,
bool save_optimizer_state);
ORT_API(void, ReleaseCheckpointState, _Frees_ptr_opt_ OrtCheckpointState* session);
#endif
} // namespace OrtApis

48
orttraining/orttraining/test/training_api/common/synthetic_data_loader.cc
View file

@ -53,7 +53,20 @@ void SyntheticSampleBatch::AddFloatInput(const std::vector<int64_t>& shape) {
RandomFloats(data_vector_.back()->GetData<float>());
}
bool SyntheticDataLoader::GetNextSampleBatch(std::vector<Ort::Value>& batches) {
#define ORT_RETURN_ON_ERROR(expr) \
do { \
OrtStatus* onnx_status = (expr); \
if (onnx_status != NULL) { \
auto code = ort_api->GetErrorCode(onnx_status); \
const char* msg = ort_api->GetErrorMessage(onnx_status); \
ort_api->ReleaseStatus(onnx_status); \
printf("Run failed with error code :%d\n", code); \
printf("Error message :%s\n", msg); \
return false; \
} \
} while (0);
bool SyntheticDataLoader::GetNextSampleBatch(std::vector<OrtValue*>& batches) {
if (sample_batch_iter_index_ >= num_of_sample_batches) {
return false;
}
@ -62,31 +75,44 @@ bool SyntheticDataLoader::GetNextSampleBatch(std::vector<Ort::Value>& batches) {
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
auto& sample = sample_batch_collections_[sample_batch_iter_index_];
const auto* ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
for (size_t i = 0; i < sample->NumOfInput(); ++i) {
auto input_ptr = sample->GetInputAtIndex(i);
auto shape_vector = input_ptr->ShapeVector();
// Be noted: the created Ort::Value won't clean the raw data after its lifetime ended.
// Be noted: the created OrtValue won't clean the raw data after its lifetime ended.
auto ptr_flt = dynamic_cast<TypedSynctheticInput<float>*>(input_ptr);
if (ptr_flt) {
batches.push_back(Ort::Value::CreateTensor<float>(
memory_info, input_ptr->GetData<float>().data(),
input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size()));
OrtValue* value = NULL;
ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info,
input_ptr->GetData<float>().data(), (input_ptr->NumOfBytesPerSample() * sizeof(float)),
shape_vector.data(), shape_vector.size(),
ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
&value));
batches.emplace_back(value);
continue;
}
auto ptr_int = dynamic_cast<TypedSynctheticInput<int64_t>*>(input_ptr);
if (ptr_int) {
batches.push_back(Ort::Value::CreateTensor<int64_t>(
memory_info, input_ptr->GetData<int64_t>().data(),
input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size()));
OrtValue* value = NULL;
ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info,
input_ptr->GetData<int64_t>().data(), (input_ptr->NumOfBytesPerSample() * sizeof(int64_t)),
shape_vector.data(), shape_vector.size(),
ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,
&value));
batches.emplace_back(value);
continue;
}
auto ptr_int32 = dynamic_cast<TypedSynctheticInput<int32_t>*>(input_ptr);
if (ptr_int32) {
batches.push_back(Ort::Value::CreateTensor<int32_t>(
memory_info, input_ptr->GetData<int32_t>().data(),
input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size()));
OrtValue* value = nullptr;
ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info,
input_ptr->GetData<int32_t>().data(), (input_ptr->NumOfBytesPerSample() * sizeof(int32_t)),
shape_vector.data(), shape_vector.size(),
ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32,
&value));
batches.emplace_back(value);
continue;
}

6
orttraining/orttraining/test/training_api/common/synthetic_data_loader.h
View file

@ -95,7 +95,7 @@ struct SyntheticDataLoader {
num_of_sample_batches += 1;
}
bool GetNextSampleBatch(std::vector<Ort::Value>& batches);
bool GetNextSampleBatch(std::vector<OrtValue*>& batches);
size_t NumOfSampleBatches() {
return num_of_sample_batches;
@ -106,9 +106,9 @@ struct SyntheticDataLoader {
}
private:
// Be noted: all raw data MUST remain during the training, because all Ort::Value created as session inputs
// Be noted: all raw data MUST remain during the training, because all OrtValue created as session inputs
// did not explicitly copy the data in.
// And also, the created Ort::Value also won't clean the raw data pointer. The raw data should be removed when
// And also, the created OrtValue also won't clean the raw data pointer. The raw data should be removed when
// the life time of this struct ends.
std::vector<std::unique_ptr<SyntheticSampleBatch>> sample_batch_collections_;
int64_t sample_batch_count_;

11
orttraining/orttraining/test/training_api/core/checkpoint_test.cc
View file

@ -181,15 +181,14 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CPU) {
sample->AddFloatInput(fc2_bias_shape);
data_loader.AddSyntheticSampleBatch(std::move(sample));
std::vector<Ort::Value> all_weights_values;
std::vector<OrtValue*> all_weights_values;
data_loader.GetNextSampleBatch(all_weights_values);
ASSERT_EQ(all_weights_values.size(), 4);
Ort::Value* data_ptr = all_weights_values.data();
NameMLValMap name_to_ort_value{
{"fc1.weight", **reinterpret_cast<::OrtValue**>(data_ptr)},
{"fc1.bias", **reinterpret_cast<::OrtValue**>(data_ptr + 1)},
{"fc2.weight", **reinterpret_cast<::OrtValue**>(data_ptr + 2)},
{"fc2.bias", **reinterpret_cast<::OrtValue**>(data_ptr + 3)},
{"fc1.weight", *all_weights_values[0]},
{"fc1.bias", *all_weights_values[1]},
{"fc2.weight", *all_weights_values[2]},
{"fc2.bias", *all_weights_values[3]},
};
// Module/Optimizer creation and trainable parameter name definitions.

138
orttraining/orttraining/test/training_api/trainer/trainer.cc
View file

@ -1,8 +1,8 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <onnxruntime_cxx_api.h>
#include "orttraining/training_api/include/interfaces.h"
#include <onnxruntime_c_api.h>
#include "orttraining/training_api/include/utils.h"
#include "cxxopts.hpp"
#include "../common/synthetic_data_loader.h"
@ -16,6 +16,8 @@
using namespace onnxruntime::training::api;
using namespace std;
const OrtApi* g_ort_api = nullptr;
struct TestRunnerParameters {
// Models configs.
std::string model_training_graph_path;
@ -44,6 +46,19 @@ void EnforceCheck(bool run_ret, std::string err_msg) {
}
}
#define ORT_RETURN_ON_ERROR(expr) \
do { \
OrtStatus* onnx_status = (expr); \
if (onnx_status != NULL) { \
auto code = g_ort_api->GetErrorCode(onnx_status); \
const char* msg = g_ort_api->GetErrorMessage(onnx_status); \
g_ort_api->ReleaseStatus(onnx_status); \
printf("Run failed with error code :%d\n", code); \
printf("Error message :%s\n", msg); \
return -1; \
} \
} while (0);
bool ParseArguments(int argc, char* argv[], TestRunnerParameters& params) {
cxxopts::Options options("Training API Test", "Main Program to test training C++ APIs.");
// clang-format off
@ -152,38 +167,47 @@ void InitSyntheticDataLoader(
}
}
void RunTraining(const TestRunnerParameters& params) {
const auto& api = Ort::GetApi();
int RunTraining(const TestRunnerParameters& params) {
g_ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
CheckpointState state;
// TODO: update using public API's calling pattern, e.g. api.LoadCheckpoint().
EnforceCheck(LoadCheckpoint(params.checkpoint_to_load_path, state).IsOK(), "Failed to load checkpoint");
// Create Env
OrtEnv* env;
// TODO: enable global threadpool
OrtThreadingOptions* threading_options = nullptr;
ORT_RETURN_ON_ERROR(g_ort_api->CreateThreadingOptions(&threading_options));
ORT_RETURN_ON_ERROR(g_ort_api->CreateEnvWithGlobalThreadPools(
ORT_LOGGING_LEVEL_VERBOSE, "log", threading_options, &env));
g_ort_api->ReleaseThreadingOptions(threading_options);
OrtSessionOptions* session_options = nullptr;
EnforceCheck(api.CreateSessionOptions(&session_options) == nullptr, "Failed to create session options.");
// Load Checkpoint State
OrtCheckpointState* checkpoint_state;
ORT_RETURN_ON_ERROR(g_ort_api->LoadCheckpoint(params.checkpoint_to_load_path.c_str(), &checkpoint_state));
// Create TrainingSession
OrtSessionOptions* soptions;
ORT_RETURN_ON_ERROR(g_ort_api->CreateSessionOptions(&soptions));
#ifdef USE_CUDA
OrtCUDAProviderOptionsV2* cuda_options = nullptr;
EnforceCheck(api.CreateCUDAProviderOptions(&cuda_options) == nullptr, "Failed to create cuda provider options");
EnforceCheck(api.SessionOptionsAppendExecutionProvider_CUDA_V2(session_options, cuda_options) == nullptr,
"Failed to append cuda ep.");
ORT_RETURN_ON_ERROR(g_ort_api->CreateCUDAProviderOptions(&cuda_options));
ORT_RETURN_ON_ERROR(g_ort_api->SessionOptionsAppendExecutionProvider_CUDA_V2(soptions, cuda_options));
#endif
OrtEnv* env = nullptr;
EnforceCheck(api.CreateEnv(ORT_LOGGING_LEVEL_WARNING, "e2e_test_runner", &env) == nullptr, "Failed to create env");
// TODO: update using public API's calling pattern, e.g. api.CreateModule().
Ort::OrtModule module(env, session_options,
params.model_training_graph_path,
state.module_checkpoint_state.named_parameters,
params.model_evaluation_graph_path);
OrtTrainingSession* session;
ORT_RETURN_ON_ERROR(g_ort_api->CreateTrainingSession(env, soptions, checkpoint_state, &session));
// Initialize Training Session
bool do_eval = params.model_evaluation_graph_path.has_value();
ORT_RETURN_ON_ERROR(g_ort_api->InitializeTrainingSession(session, params.model_training_graph_path.c_str(),
do_eval ? params.model_evaluation_graph_path.value().c_str() : nullptr,
params.optimizer_training_graph_path.size() > 0 ? params.optimizer_training_graph_path.c_str() : nullptr));
// TODO: update using public API's calling pattern, e.g. api.CreateOptimizer().
Ort::OrtOptimizer optimizer(env, session_options,
params.optimizer_training_graph_path,
module.NamedParameters());
size_t train_mode_output_count, eval_mode_output_count = 0;
ORT_RETURN_ON_ERROR(g_ort_api->TrainingSessionGetTrainModeOutputCount(session, &train_mode_output_count));
if (do_eval) {
ORT_RETURN_ON_ERROR(g_ort_api->TrainingSessionGetEvalModeOutputCount(session, &eval_mode_output_count));
}
int64_t sample_batch_count_per_epoch = 4;
if (sample_batch_count_per_epoch < params.train_batch_size || sample_batch_count_per_epoch % params.train_batch_size != 0) {
@ -194,12 +218,12 @@ void RunTraining(const TestRunnerParameters& params) {
onnxruntime::training::test::training_api::SyntheticDataLoader data_loader;
InitSyntheticDataLoader(data_loader, params, num_of_batches_per_epoch);
int64_t total_step_count = params.num_train_epochs * num_of_batches_per_epoch;
int64_t warmup_step_count = total_step_count / 3;
// TODO: update using public API's calling pattern, e.g. api.CreateLinearLRScheduler().
Ort::OrtLinearLRScheduler scheduler = Ort::OrtLinearLRScheduler(optimizer, warmup_step_count, total_step_count);
// TODO: Add C API for LRScheduler
//int64_t total_step_count = params.num_train_epochs * num_of_batches_per_epoch;
//int64_t warmup_step_count = total_step_count / 3;
//Ort::OrtLinearLRScheduler scheduler = Ort::OrtLinearLRScheduler(optimizer, warmup_step_count, total_step_count);
std::cout << "Initialization completed. Now starting training loop." << std::endl;
const int64_t stabilized_perf_start_step = 0;
double stabilized_total_end_to_end_time{0};
auto end_to_end_start = std::chrono::high_resolution_clock::now();
@ -210,7 +234,7 @@ void RunTraining(const TestRunnerParameters& params) {
end_to_end_start = std::chrono::high_resolution_clock::now();
}
std::vector<Ort::Value> inputs;
std::vector<OrtValue*> inputs;
data_loader.GetNextSampleBatch(inputs);
#if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE)
@ -220,14 +244,15 @@ void RunTraining(const TestRunnerParameters& params) {
train_step_range.Begin();
#endif
std::vector<Ort::Value> fetches;
EnforceCheck(module.TrainStep(inputs, fetches), "Failed during module.TrainStep.");
std::vector<OrtValue*> fetches(train_mode_output_count);
ORT_RETURN_ON_ERROR(g_ort_api->TrainStep(session, nullptr,
inputs.size(), (const OrtValue* const*)inputs.data(),
train_mode_output_count, fetches.data()));
#if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE)
train_step_range.End();
#endif
float loss = *(fetches[0].GetTensorMutableData<float>());
float loss = onnxruntime::training::api::utils::GetValue<float>(*fetches[0]);
std::cout << "Batch # : " << batch_idx << " Loss: " << loss << std::endl;
if ((batch_idx + 1) % params.gradient_accumulation_steps == 0) {
@ -238,14 +263,14 @@ void RunTraining(const TestRunnerParameters& params) {
onnxruntime::profile::Color::Blue);
opt_step_range.Begin();
#endif
EnforceCheck(optimizer.Step(), "Failed during optimizer.Step().");
ORT_RETURN_ON_ERROR(g_ort_api->OptimizerStep(session, nullptr));
#if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE)
opt_step_range.End();
#endif
// Update learning rate.
EnforceCheck(scheduler.Step(), "Failed during shceduler.Step()");
//EnforceCheck(scheduler.Step(), "Failed during shceduler.Step()");
#if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE)
onnxruntime::profile::NvtxRangeCreator resetgrad_range(
@ -254,7 +279,7 @@ void RunTraining(const TestRunnerParameters& params) {
resetgrad_range.Begin();
#endif
EnforceCheck(module.ResetGrad(), "Failed during module.ResetGrad().");
ORT_RETURN_ON_ERROR(g_ort_api->ResetGrad(session));
#if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE)
resetgrad_range.End();
@ -262,41 +287,55 @@ void RunTraining(const TestRunnerParameters& params) {
}
if (do_eval && (batch_idx + 1) % params.eval_interval == 0) {
std::vector<Ort::Value> eval_results;
EnforceCheck(module.EvalStep(inputs, eval_results), "Failed during Module.EvalStep().");
std::vector<OrtValue*> eval_results(eval_mode_output_count);
ORT_RETURN_ON_ERROR(g_ort_api->EvalStep(session, nullptr,
inputs.size(), (const OrtValue* const*)inputs.data(),
train_mode_output_count, eval_results.data()));
}
if ((batch_idx + 1) % params.checkpoint_interval == 0) {
// Save trained weights
CheckpointState state_to_save;
EnforceCheck(module.GetStateDict(state_to_save.module_checkpoint_state), "Failed to load module states.");
EnforceCheck(optimizer.GetStateDict(state_to_save.optimizer_checkpoint_state), "Failed to load optimizer states.");
state_to_save.property_bag.AddProperty<int64_t>(std::string("epoch"), epoch);
std::string ckpt_file = params.output_dir + "/ckpt_" + params.model_name + std::to_string(batch_idx);
ORT_RETURN_ON_ERROR(g_ort_api->SaveCheckpoint(ckpt_file.c_str(), session, true));
// TODO: update using public API's calling pattern, e.g. api.SaveCheckpoint().
EnforceCheck(SaveCheckpoint(state_to_save, ckpt_file).IsOK(), "Failed to save checkpoint.");
// TODO: enable adding more properties to checkpoint
// state_to_save.property_bag.AddProperty<int64_t>(std::string("epoch"), epoch);
}
batch_idx++;
// release input ortvalues
for (size_t i = 0; i < inputs.size(); i++) {
g_ort_api->ReleaseValue(inputs[i]);
}
// TODO(askhade): release output values. Needs changes from Aishwarya's PR.
}
data_loader.ResetIterateIndex();
}
// Save trained weights
std::string ckpt_file = params.output_dir + "/ckpt_" + params.model_name;
ORT_RETURN_ON_ERROR(g_ort_api->SaveCheckpoint(ckpt_file.c_str(), session, true));
auto end = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> duration_seconds = end - end_to_end_start;
stabilized_total_end_to_end_time = duration_seconds.count();
std::cout << "Training completed - end to end latency: " << stabilized_total_end_to_end_time << "(s)" << std::endl;
api.ReleaseEnv(env);
// Delete all the ptrs
g_ort_api->ReleaseTrainingSession(session);
#ifdef USE_CUDA
// Finally, don't forget to release the provider options
api.ReleaseCUDAProviderOptions(cuda_options);
g_ort_api->ReleaseCUDAProviderOptions(cuda_options);
#endif
g_ort_api->ReleaseSessionOptions(soptions);
g_ort_api->ReleaseCheckpointState(checkpoint_state);
g_ort_api->ReleaseEnv(env);
api.ReleaseSessionOptions(session_options);
return 0;
}
int main(int argc, char* argv[]) {
@ -309,6 +348,5 @@ int main(int argc, char* argv[]) {
EnforceCheck(ParseArguments(argc, argv, params), "Parse arguments failed.");
// Start training session
RunTraining(params);
return 0;
return RunTraining(params);
}

21
orttraining/orttraining/training_api/include/checkpoint.h
View file

@ -49,6 +49,17 @@ struct CheckpointState {
PropertyBag property_bag;
};
/**
* @brief Save training states as ORT checkpoint.
*
* @param state parameter/optimizer and other user defined training states.
* @param checkpoint_path folder where checkpoint is saved.
* @return Status
* TODO: change state to const ref
*/
Status SaveCheckpoint(CheckpointState& state,
const PathString& checkpoint_path);
/**
* @brief Save ONNX initializers as ORT checkpoint.
*
@ -61,16 +72,6 @@ Status SaveCheckpoint(const std::vector<ONNX_NAMESPACE::TensorProto>& trainable_
const std::vector<ONNX_NAMESPACE::TensorProto>& non_trainable_tensor_protos,
const PathString& checkpoint_path);
/**
* @brief Save training states as ORT checkpoint.
*
* @param state parameter/optimizer and other user defined training states.
* @param checkpoint_path folder where checkpoint is saved.
* @return Status
*/
Status SaveCheckpoint(CheckpointState& state,
const PathString& checkpoint_path);
/**
* @brief Load training states from ORT checkpoint.
*

10
orttraining/orttraining/training_api/include/module.h
View file

@ -58,7 +58,7 @@ struct Module {
// Initialize a module from an ORT inference session with loaded
// training ONNX model and load parameters
Module(const std::string& train_model_path_or_bytes,
std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::optional<std::string>& eval_model_path_or_bytes = std::nullopt);
@ -84,10 +84,15 @@ struct Module {
// Return the states of the module as a map.
Status GetStateDict(ModuleCheckpointState& module_checkpoint_states);
// Returns the output count for training graph
size_t GetTrainModeOutputCount() const noexcept;
// Returns the output count for eval graph
size_t GetEvalModeOutputCount() const noexcept;
private:
std::unique_ptr<onnxruntime::InferenceSession> train_sess_{nullptr};
std::unique_ptr<onnxruntime::InferenceSession> eval_sess_{nullptr};
std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters_;
std::vector<std::string> train_input_names_;
std::vector<std::string> train_output_names_;
std::vector<std::string> eval_input_names_;
@ -95,6 +100,7 @@ struct Module {
std::vector<OrtValue> weights_;
std::vector<OrtValue> gradients_;
bool accumulate_gradient_ = true;
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters_;
};
} // namespace api

41
orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h Normal file
View file

@ -0,0 +1,41 @@
// This file contains c apis for on device training
// This file should never be included standalone
// It is included from within core/session/onnxruntime_c_api.h when
// on device training is enabled
// These apis can be moved to core/session/onnxruntime_c_api.h once they stabilize
// DO NOT UNCOMMENT
//#include "core/session/onnxruntime_c_api.h"
ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state);
ORT_API2_STATUS(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* session,
bool save_optimizer_state);
ORT_API2_STATUS(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options,
_Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out);
ORT_API2_STATUS(InitializeTrainingSession, _Inout_ OrtTrainingSession* session,
_In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path,
_In_ const ORTCHAR_T* optimizer_model_path);
ORT_API2_STATUS(TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
ORT_API2_STATUS(TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
ORT_API2_STATUS(ResetGrad, _Inout_ OrtTrainingSession* session);
ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
ORT_API2_STATUS(EvalStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
ORT_API2_STATUS(OptimizerStep, _Inout_ OrtTrainingSession* sess,
_In_opt_ const OrtRunOptions* run_options);
ORT_CLASS_RELEASE(TrainingSession);
ORT_CLASS_RELEASE(CheckpointState);

2
orttraining/orttraining/training_api/include/optimizer.h
View file

@ -86,7 +86,7 @@ struct Optimizer {
// TODO: load this info from checkpoint
OptimizerType optimizer_type_ = OptimizerType::AdamW;
std::unique_ptr<onnxruntime::InferenceSession> optim_sess_;
std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters_;
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters_;
GroupOptimizerState optimizer_state_;
std::vector<std::string> input_names_;
std::vector<std::string> output_names_;

55
orttraining/orttraining/training_api/include/training_session.h Normal file
View file

@ -0,0 +1,55 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "module.h"
#include "optimizer.h"
#include "checkpoint.h"
namespace onnxruntime {
namespace training {
namespace api {
using namespace common;
// Wrapper on top of module and optimizer classes and is the only class exposed via capis
class TrainingSession {
public:
TrainingSession(const Environment& session_env,
const SessionOptions& session_options,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& parameters);
Status Initialize(const std::string& train_model_uri,
const std::optional<std::string>& eval_model_uri,
const std::optional<std::string>& optim_model_uri);
size_t GetTrainModeOutputCount() const noexcept;
size_t GetEvalModeOutputCount() const noexcept;
Status TrainStep(const RunOptions& run_options,
const std::vector<OrtValue>& inputs,
std::vector<OrtValue>& fetches);
Status EvalStep(const RunOptions& run_options,
const std::vector<OrtValue>& inputs,
std::vector<OrtValue>& fetches);
Status ResetGrad();
Status OptimizerStep(const RunOptions& run_options);
Status CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state);
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TrainingSession);
const Environment& environment_;
SessionOptions session_options_;
const std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters_;
std::unique_ptr<Module> module_;
std::unique_ptr<Optimizer> optimizer_;
};
} // namespace api
} // namespace training
} // namespace onnxruntime

2
orttraining/orttraining/training_api/include/utils.h
View file

@ -29,7 +29,7 @@ Status OrtValueLike(const SessionState& sess_state, const OrtValue& input_val, O
// Create OrtValue from a single value of type T
template <typename T>
void WarpInOrtValue(T value,
void WrapInOrtValue(T value,
OrtValue* p_ortvalue,
AllocatorPtr alloc = nullptr) {
static CPUExecutionProviderInfo info;

84
orttraining/orttraining/training_api/module.cc
View file

@ -52,24 +52,27 @@ Status Parameter::ResetGrad() {
}
Module::Module(const std::string& train_model_path_or_bytes,
std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::optional<std::string>& eval_model_path_or_bytes) {
const std::optional<std::string>& eval_model_path_or_bytes) : named_parameters_{named_parameters} {
// Create session for training model
train_sess_ = std::make_unique<onnxruntime::InferenceSession>(session_options, env);
ORT_THROW_IF_ERROR(train_sess_->Load(train_model_path_or_bytes));
ORT_THROW_IF_ERROR(train_sess_->Initialize());
// Extract model input and output names
utils::GetGraphInputOutputNames(train_sess_, train_input_names_, train_output_names_);
auto& train_sess_state = train_sess_->GetSessionState();
std::vector<std::string> param_input_names, grad_input_names, user_input_names, reset_grad_name;
// Reorder the extracted input names in the following order:
// user inputs, weights, gradients, reset_grad
std::vector<std::string> user_input_names, param_input_names, grad_input_names, reset_grad_name;
std::string param_name;
std::unordered_map<std::string, size_t> param_name_to_grad_input_index_map;
for (const auto& input_name : train_input_names_) {
auto it = named_parameters.find(input_name);
if (it != named_parameters.end()) {
auto it = named_parameters_.find(input_name);
if (it != named_parameters_.end()) {
param_input_names.emplace_back(input_name);
} else if (input_name == ACCUMULATE_GRAD_CONTROL_INPUT_NAME) {
reset_grad_name.emplace_back(input_name);
@ -89,12 +92,12 @@ Module::Module(const std::string& train_model_path_or_bytes,
train_input_names_.insert(train_input_names_.end(), reset_grad_name.begin(), reset_grad_name.end());
// Loop each parameter, allocate it's memory based on user specified device.
auto& train_sess_state = train_sess_->GetSessionState();
for (auto& param_name : param_input_names) {
ORT_ENFORCE(named_parameters.find(param_name) != named_parameters.end());
OrtValue& source_ortvalue = named_parameters[param_name]->Data();
ORT_ENFORCE(source_ortvalue.IsTensor());
const Tensor& source_tensor = source_ortvalue.Get<Tensor>();
auto params_iter = named_parameters_.find(param_name);
ORT_ENFORCE(params_iter != named_parameters_.end());
// Retrieve the target device for "param_name"
std::vector<SessionState::NodeInfo> node_info_vec;
ORT_THROW_IF_ERROR(train_sess_state.GetInputNodeInfo(param_name, node_info_vec));
const auto& node_info = node_info_vec.front();
@ -103,39 +106,44 @@ Module::Module(const std::string& train_model_path_or_bytes,
ORT_ENFORCE(target_device == *(it->device), "Inconsistent device requirements found for input: ", param_name);
}
// Create parameter value copy with corresponding device user sets the session on.
// We did not re-use the data even CPU tensor is needed.
// TODO(pengwa): consider whether we should alloc contiguous buffer for parameters or gradients.
OrtValue target_ortvalue;
auto allocator = train_sess_state.GetAllocator(target_device);
ORT_ENFORCE(allocator != nullptr);
// Copy ortvalue buffer from CPU to target_device for this "param_name" (based on graph partitioning)
// Only copies data if target device is not the same as the current device the buffer is placed on
Tensor::InitOrtValue(source_tensor.DataType(),
source_tensor.Shape(),
allocator, target_ortvalue);
Tensor* target_tensor_ptr = target_ortvalue.GetMutable<Tensor>();
ORT_THROW_IF_ERROR(train_sess_state.GetDataTransferMgr().CopyTensor(source_tensor, *target_tensor_ptr));
OrtValue& param_data = params_iter->second->Data();
ORT_ENFORCE(param_data.IsTensor());
const Tensor& param_data_tensor = param_data.Get<Tensor>();
// If the source device type is already same as target device skip copy
if (param_data_tensor.Location().device.Type() != target_device.Type()) {
// TODO: move this outside of the for loop?
auto target_allocator = train_sess_state.GetAllocator(target_device);
ORT_ENFORCE(target_allocator != nullptr);
auto param_share_ptr =
std::make_shared<Parameter>(param_name, target_ortvalue, named_parameters[param_name]->RequiresGrad());
named_parameters_.insert({param_name, param_share_ptr});
weights_.push_back(param_share_ptr->Data());
// Create a new tensor on the target_device and switch the source_ortvalue to point to this new tensor
auto target_tensor = std::make_unique<Tensor>(param_data_tensor.DataType(), param_data_tensor.Shape(), target_allocator);
ORT_THROW_IF_ERROR(train_sess_state.GetDataTransferMgr().CopyTensor(param_data_tensor, *target_tensor.get()));
auto ml_tensor_type = DataTypeImpl::GetType<Tensor>();
// TODO test the original buffer is released.
param_data.Init(target_tensor.release(), ml_tensor_type, ml_tensor_type->GetDeleteFunc());
}
weights_.push_back(param_data);
// Create gradient buffer when parameter requires gradient.
if (param_share_ptr->RequiresGrad()) {
if (params_iter->second->RequiresGrad()) {
// Create gradient accumulation buffer.
auto it = param_name_to_grad_input_index_map.find(param_name);
ORT_ENFORCE(it != param_name_to_grad_input_index_map.end(), "Gradient buffer input not providered for param: ",
param_name);
const size_t grad_input_index = it->second;
auto& param_grad_buffer_name = grad_input_names[grad_input_index];
auto& param_grad_name = grad_input_names[grad_input_index];
// TODO: don't pre-allocate the gradient buffer.
// Gradient usually stays on the same device of its parameter.
OrtValue param_grad_buffer_ortvalue;
ORT_THROW_IF_ERROR(utils::OrtValueLike(train_sess_state, target_ortvalue, param_grad_buffer_ortvalue));
ORT_THROW_IF_ERROR(param_share_ptr->SetGrad(param_grad_buffer_name, param_grad_buffer_ortvalue));
gradients_[grad_input_index] = param_share_ptr->Gradient();
OrtValue param_grad;
ORT_THROW_IF_ERROR(utils::OrtValueLike(train_sess_state, param_data, param_grad));
ORT_THROW_IF_ERROR(params_iter->second->SetGrad(param_grad_name, param_grad));
gradients_[grad_input_index] = params_iter->second->Gradient();
}
}
@ -146,9 +154,9 @@ Module::Module(const std::string& train_model_path_or_bytes,
utils::GetGraphInputOutputNames(eval_sess_, eval_input_names_, eval_output_names_);
// Eval model validation
// We are making certain assumptions: Like the order in which parameters
// occur will be same between train and eval graphs,
// and all the weights present in both graphs match.
// We are making certain assumptions: Like the order in which parameters occur will be same between train and eval
// graphs, and all the weights present in both graphs match.
// TODO: Add the checks instead of making assumptions??
std::vector<std::string> eval_user_input_names, param_input_names;
for (const auto& input_name : eval_input_names_) {
if (named_parameters_.find(input_name) != named_parameters_.end()) {
@ -168,6 +176,14 @@ Module::Module(const std::string& train_model_path_or_bytes,
}
}
size_t Module::GetTrainModeOutputCount() const noexcept {
return train_output_names_.size();
}
size_t Module::GetEvalModeOutputCount() const noexcept {
return eval_output_names_.size();
}
std::vector<std::shared_ptr<Parameter>> Module::Parameters() const {
std::vector<std::shared_ptr<Parameter>> params;
for (auto& it : named_parameters_) {
@ -187,7 +203,7 @@ Status Module::TrainStep(const std::vector<OrtValue>& inputs, std::vector<OrtVal
feeds.insert(feeds.end(), gradients_.begin(), gradients_.end());
// TODO: consider maintaining this as ortvalue instead of bool
OrtValue do_update_input;
utils::WarpInOrtValue<bool>(accumulate_gradient_, &do_update_input);
utils::WrapInOrtValue<bool>(accumulate_gradient_, &do_update_input);
feeds.push_back(do_update_input);
// TODO: need to filter out the grads from the output ortvalues

218
orttraining/orttraining/training_api/onnxruntime_training_c_api.cc Normal file
View file

@ -0,0 +1,218 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/error_code_helper.h"
#include "core/framework/ort_value.h"
#include "core/session/ort_apis.h"
#include "core/session/ort_env.h"
#include "orttraining/training_api/include/checkpoint.h"
#include "orttraining/training_api/include/training_session.h"
#include "core/session/abi_session_options_impl.h"
ORT_API_STATUS_IMPL(OrtApis::CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options,
_Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out) {
API_IMPL_BEGIN
std::unique_ptr<onnxruntime::training::api::TrainingSession> train_sess;
auto chkpt_state = reinterpret_cast<onnxruntime::training::api::CheckpointState*>(checkpoint_state);
OrtStatus* status = nullptr;
*out = nullptr;
ORT_TRY {
train_sess = std::make_unique<onnxruntime::training::api::TrainingSession>(
env->GetEnvironment(),
options == nullptr ? onnxruntime::SessionOptions() : options->value,
chkpt_state->module_checkpoint_state.named_parameters);
*out = reinterpret_cast<OrtTrainingSession*>(train_sess.release());
}
ORT_CATCH(const std::exception& e) {
ORT_HANDLE_EXCEPTION([&]() {
status = OrtApis::CreateStatus(ORT_FAIL, e.what());
});
}
return status;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::InitializeTrainingSession, _Inout_ OrtTrainingSession* session,
_In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path,
_In_ const ORTCHAR_T* optimizer_model_path) {
API_IMPL_BEGIN
auto train_sess = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(session);
ORT_API_RETURN_IF_STATUS_NOT_OK(train_sess->Initialize(train_model_path,
eval_model_path ? std::optional<std::string>{eval_model_path}
: std::nullopt,
optimizer_model_path ? std::optional<std::string>{optimizer_model_path}
: std::nullopt));
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out) {
API_IMPL_BEGIN
auto session = reinterpret_cast<const onnxruntime::training::api::TrainingSession*>(sess);
*out = session->GetTrainModeOutputCount();
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out) {
API_IMPL_BEGIN
auto session = reinterpret_cast<const onnxruntime::training::api::TrainingSession*>(sess);
*out = session->GetEvalModeOutputCount();
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::ResetGrad, _Inout_ OrtTrainingSession* session) {
API_IMPL_BEGIN
auto train_session = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(session);
ORT_API_RETURN_IF_STATUS_NOT_OK(train_session->ResetGrad());
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs) {
API_IMPL_BEGIN
auto session = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(sess);
constexpr int queue_id = 0;
std::vector<OrtValue> feeds(inputs_len);
for (size_t i = 0; i != inputs_len; ++i) {
auto& ort_value = feeds[i] = *reinterpret_cast<const ::OrtValue*>(inputs[i]);
if (ort_value.Fence()) {
ort_value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id);
}
}
// Create output feed
std::vector<OrtValue> fetches(outputs_len);
for (size_t i = 0; i != outputs_len; ++i) {
if (outputs[i] != nullptr) {
::OrtValue& value = *(outputs[i]);
if (value.Fence())
value.Fence()->BeforeUsingAsOutput(onnxruntime::kCpuExecutionProvider, queue_id);
fetches[i] = value;
}
}
Status status;
if (run_options == nullptr) {
OrtRunOptions op;
status = session->TrainStep(op, feeds, fetches);
} else {
status = session->TrainStep(*run_options, feeds, fetches);
}
if (!status.IsOK())
return onnxruntime::ToOrtStatus(status);
for (size_t i = 0; i != outputs_len; ++i) {
::OrtValue& value = fetches[i];
if (value.Fence())
value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id);
if (outputs[i] == nullptr) {
outputs[i] = new OrtValue(value);
}
}
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::EvalStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs) {
API_IMPL_BEGIN
auto session = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(sess);
constexpr int queue_id = 0;
std::vector<OrtValue> feeds(inputs_len);
for (size_t i = 0; i != inputs_len; ++i) {
auto& ort_value = feeds[i] = *reinterpret_cast<const ::OrtValue*>(inputs[i]);
if (ort_value.Fence()) ort_value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id);
}
// Create output feed
std::vector<OrtValue> fetches(outputs_len);
for (size_t i = 0; i != outputs_len; ++i) {
if (outputs[i] != nullptr) {
::OrtValue& value = *(outputs[i]);
if (value.Fence())
value.Fence()->BeforeUsingAsOutput(onnxruntime::kCpuExecutionProvider, queue_id);
fetches[i] = value;
}
}
Status status;
if (run_options == nullptr) {
OrtRunOptions op;
status = session->EvalStep(op, feeds, fetches);
} else {
status = session->EvalStep(*run_options, feeds, fetches);
}
if (!status.IsOK())
return onnxruntime::ToOrtStatus(status);
for (size_t i = 0; i != outputs_len; ++i) {
::OrtValue& value = fetches[i];
if (value.Fence())
value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id);
if (outputs[i] == nullptr) {
outputs[i] = new OrtValue(value);
}
}
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::OptimizerStep, _Inout_ OrtTrainingSession* sess,
_In_opt_ const OrtRunOptions* run_options) {
API_IMPL_BEGIN
auto session = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(sess);
if (run_options == nullptr) {
OrtRunOptions op;
ORT_API_RETURN_IF_STATUS_NOT_OK(session->OptimizerStep(op));
} else {
ORT_API_RETURN_IF_STATUS_NOT_OK(session->OptimizerStep(*run_options));
}
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state) {
API_IMPL_BEGIN
*checkpoint_state = nullptr;
auto chkpt_state = std::make_unique<onnxruntime::training::api::CheckpointState>();
ORT_API_RETURN_IF_STATUS_NOT_OK(onnxruntime::training::api::LoadCheckpoint(checkpoint_path, *chkpt_state));
*checkpoint_state = reinterpret_cast<OrtCheckpointState*>(chkpt_state.release());
return nullptr;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* sess,
bool save_optimizer_state) {
API_IMPL_BEGIN
auto session = reinterpret_cast<onnxruntime::training::api::TrainingSession*>(sess);
onnxruntime::training::api::CheckpointState chkpt_state;
ORT_API_RETURN_IF_STATUS_NOT_OK(session->CreateCheckpointState(chkpt_state, save_optimizer_state));
ORT_API_RETURN_IF_STATUS_NOT_OK(onnxruntime::training::api::SaveCheckpoint(chkpt_state, checkpoint_path));
return nullptr;
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseTrainingSession, _Frees_ptr_opt_ OrtTrainingSession* session) {
delete reinterpret_cast<onnxruntime::training::api::TrainingSession*>(session);
}
ORT_API(void, OrtApis::ReleaseCheckpointState, _Frees_ptr_opt_ OrtCheckpointState* checkpoint_state) {
delete reinterpret_cast<onnxruntime::training::api::CheckpointState*>(checkpoint_state);
}

4
orttraining/orttraining/training_api/optimizer.cc
View file

@ -106,8 +106,8 @@ Optimizer::Optimizer(const std::string& optim_path_or_bytes,
Status Optimizer::Step() {
OrtValue learning_rate_input, step_input;
utils::WarpInOrtValue<float>(optimizer_state_.learning_rate, &learning_rate_input);
utils::WarpInOrtValue<int64_t>(optimizer_state_.step, &step_input);
utils::WrapInOrtValue<float>(optimizer_state_.learning_rate, &learning_rate_input);
utils::WrapInOrtValue<int64_t>(optimizer_state_.step, &step_input);
std::vector<OrtValue> feeds({learning_rate_input, step_input});
feeds.insert(feeds.end(), inputs_.begin(), inputs_.end());

69
orttraining/orttraining/training_api/training_session.cc Normal file
View file

@ -0,0 +1,69 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_api/include/training_session.h"
namespace onnxruntime {
namespace training {
namespace api {
TrainingSession::TrainingSession(const Environment& session_env,
const SessionOptions& session_options,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& parameters)
: environment_(session_env),
session_options_{session_options},
named_parameters_{parameters} {}
Status TrainingSession::Initialize(const std::string& train_model_uri, const std::optional<std::string>& eval_model_uri,
const std::optional<std::string>& optim_model_uri) {
module_ = std::move(std::make_unique<Module>(train_model_uri, named_parameters_, session_options_,
environment_, eval_model_uri));
if (optim_model_uri.has_value()) {
optimizer_ = std::move(std::make_unique<Optimizer>(optim_model_uri.value(), named_parameters_,
session_options_, environment_));
}
return Status::OK();
}
size_t TrainingSession::GetTrainModeOutputCount() const noexcept {
return module_->GetTrainModeOutputCount();
}
size_t TrainingSession::GetEvalModeOutputCount() const noexcept {
return module_->GetEvalModeOutputCount();
}
Status TrainingSession::TrainStep(const RunOptions&,
const std::vector<OrtValue>& inputs,
std::vector<OrtValue>& fetches) {
return module_->TrainStep(inputs, fetches);
}
Status TrainingSession::EvalStep(const RunOptions&,
const std::vector<OrtValue>& inputs,
std::vector<OrtValue>& fetches) {
return module_->EvalStep(inputs, fetches);
}
Status TrainingSession::ResetGrad() {
return module_->ResetGrad();
}
Status TrainingSession::OptimizerStep(const RunOptions&) {
return optimizer_->Step();
}
Status TrainingSession::CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state) {
ORT_RETURN_IF_ERROR(module_->GetStateDict(chkpt_state.module_checkpoint_state));
if (save_optimizer_state) {
ORT_RETURN_IF_ERROR(optimizer_->GetStateDict(chkpt_state.optimizer_checkpoint_state));
}
return Status::OK();
}
} // namespace api
} // namespace training
} // namespace onnxruntime