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

C, C++, Python, C# API update for on device training (#15518)

Browse source
This commit is contained in:
Baiju Meswani 2023-04-21 11:36:01 -07:00 committed by GitHub
parent a6d6e45be2
commit b5a1941835
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
39 changed files with 1380 additions and 536 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs
View file

@ -52,6 +52,16 @@ namespace Microsoft.ML.OnnxRuntime
NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtLoadCheckpoint(NativeOnnxValueHelper.GetPlatformSerializedString(checkpointPath), out handle));
}
/// <summary>
/// Saves the checkpoint
/// <param name="checkpointPath"> absolute path to the checkpoint file.</param>
/// <param name="includeOptimizerState"> absolute path to the checkpoint file.</param>
/// </summary>
public void SaveCheckpoint(string checkpointPath, bool includeOptimizerState = false)
{
NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtSaveCheckpoint(handle, NativeOnnxValueHelper.GetPlatformSerializedString(checkpointPath), includeOptimizerState));
}
#region SafeHandle
/// <summary>
/// Overrides SafeHandle.ReleaseHandle() to properly dispose of

112
csharp/src/Microsoft.ML.OnnxRuntime/Training/NativeTrainingMethods.shared.cs
View file

@ -33,6 +33,13 @@ namespace Microsoft.ML.OnnxRuntime
public IntPtr ReleaseTrainingSession;
public IntPtr ReleaseCheckpointState;
public IntPtr ExportModelForInferencing;
public IntPtr SetSeed;
public IntPtr TrainingSessionGetTrainingModelInputCount;
public IntPtr TrainingSessionGetEvalModelInputCount;
public IntPtr TrainingSessionGetTrainingModelInputName;
public IntPtr TrainingSessionGetEvalModelInputName;
public IntPtr AddProperty;
public IntPtr GetProperty;
}
internal static class NativeTrainingMethods
@ -77,6 +84,14 @@ namespace Microsoft.ML.OnnxRuntime
OrtSchedulerStep = (DOrtSchedulerStep)Marshal.GetDelegateForFunctionPointer(trainingApi_.SchedulerStep, typeof(DOrtSchedulerStep));
OrtReleaseTrainingSession = (DOrtReleaseTrainingSession)Marshal.GetDelegateForFunctionPointer(trainingApi_.ReleaseTrainingSession, typeof(DOrtReleaseTrainingSession));
OrtReleaseCheckpointState = (DOrtReleaseCheckpointState)Marshal.GetDelegateForFunctionPointer(trainingApi_.ReleaseCheckpointState, typeof(DOrtReleaseCheckpointState));
OrtExportModelForInferencing = (DOrtExportModelForInferencing)Marshal.GetDelegateForFunctionPointer(trainingApi_.ExportModelForInferencing, typeof(DOrtExportModelForInferencing));
OrtSetSeed = (DOrtSetSeed)Marshal.GetDelegateForFunctionPointer(trainingApi_.SetSeed, typeof(DOrtSetSeed));
OrtGetTrainingModelInputCount = (DOrtGetTrainingModelInputCount)Marshal.GetDelegateForFunctionPointer(trainingApi_.TrainingSessionGetTrainingModelInputCount, typeof(DOrtGetTrainingModelInputCount));
OrtGetEvalModelInputCount = (DOrtGetEvalModelInputCount)Marshal.GetDelegateForFunctionPointer(trainingApi_.TrainingSessionGetEvalModelInputCount, typeof(DOrtGetEvalModelInputCount));
OrtGetTrainingModelInputName = (DOrtGetTrainingModelInputName)Marshal.GetDelegateForFunctionPointer(trainingApi_.TrainingSessionGetTrainingModelInputName, typeof(DOrtGetTrainingModelInputName));
OrtGetEvalModelInputName = (DOrtGetEvalModelInputName)Marshal.GetDelegateForFunctionPointer(trainingApi_.TrainingSessionGetEvalModelInputName, typeof(DOrtGetEvalModelInputName));
OrtAddProperty = (DOrtAddProperty)Marshal.GetDelegateForFunctionPointer(trainingApi_.AddProperty, typeof(DOrtAddProperty));
OrtGetProperty = (DOrtGetProperty)Marshal.GetDelegateForFunctionPointer(trainingApi_.GetProperty, typeof(DOrtGetProperty));
}
}
@ -98,13 +113,14 @@ namespace Microsoft.ML.OnnxRuntime
/// <summary>
/// Creates an instance of OrtSession with provided parameters
/// </summary>
/// <param name="checkpointPath">checkpoint string path</param>
/// <param name="checkpointState">(Output) Loaded OrtCheckpointState instance</param>
/// <param name="checkpointState">OrtCheckpointState instance to save</param>
/// <param name="checkpointPath">Checkpoint string path</param>
/// <param name="includeOptimizerState">Flag indicating whether to save the optimizer state.</param>
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /* OrtStatus* */DOrtSaveCheckpoint(
IntPtr /*(OrtCheckpointState*)*/ checkpointState,
byte[] checkpointPath,
IntPtr /*(OrtTrainingSession*)*/ session,
bool saveOptimizerState);
bool includeOptimizerState);
public static DOrtSaveCheckpoint OrtSaveCheckpoint;
@ -170,7 +186,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtLazyResetGrad OrtLazyResetGrad;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtTrainStep(
public delegate IntPtr /*(OrtStatus*)*/ DOrtTrainStep(
IntPtr /*(OrtTrainingSession*)*/ session,
IntPtr /*(OrtSessionRunOptions*)*/ runOptions, // can be null to use the default options
UIntPtr inputCount,
@ -182,7 +198,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtTrainStep OrtTrainStep;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtEvalStep(
public delegate IntPtr /*(OrtStatus*)*/ DOrtEvalStep(
IntPtr /*(OrtTrainingSession*)*/ session,
IntPtr /*(OrtSessionRunOptions*)*/ runOptions, // can be null to use the default options
UIntPtr inputCount,
@ -194,7 +210,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtEvalStep OrtEvalStep;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtOptimizerStep(
public delegate IntPtr /*(OrtStatus*)*/ DOrtOptimizerStep(
IntPtr /*(OrtTrainingSession*)*/ session,
IntPtr /*(OrtSessionRunOptions*)*/ runOptions // can be null to use the default options
);
@ -202,7 +218,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtOptimizerStep OrtOptimizerStep;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtSetLearningRate(
public delegate IntPtr /*(OrtStatus*)*/ DOrtSetLearningRate(
IntPtr /*(OrtTrainingSession*)*/ session,
float learningRate
);
@ -210,7 +226,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtSetLearningRate OrtSetLearningRate;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtGetLearningRate(
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetLearningRate(
IntPtr /*(OrtTrainingSession*)*/ session,
out float learningRate
);
@ -218,7 +234,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtGetLearningRate OrtGetLearningRate;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtRegisterLinearLRScheduler(
public delegate IntPtr /*(OrtStatus*)*/ DOrtRegisterLinearLRScheduler(
IntPtr /*(OrtTrainingSession*)*/ session,
long warmupStepCount,
long totalStepCount,
@ -227,7 +243,7 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtRegisterLinearLRScheduler OrtRegisterLinearLRScheduler;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(ONNStatus*)*/ DOrtSchedulerStep(
public delegate IntPtr /*(OrtStatus*)*/ DOrtSchedulerStep(
IntPtr /*(OrtTrainingSession*)*/ session
);
public static DOrtSchedulerStep OrtSchedulerStep;
@ -240,6 +256,80 @@ namespace Microsoft.ML.OnnxRuntime
public delegate void DOrtReleaseCheckpointState(IntPtr /*(OrtCheckpointState*)*/checkpointState);
public static DOrtReleaseCheckpointState OrtReleaseCheckpointState;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtExportModelForInferencing(
IntPtr /*(OrtTrainingSession*)*/ session,
byte[] inferenceModelPath,
UIntPtr graphOutputCount,
IntPtr[] /*(const char* const*)*/ graphOutputNames
);
public static DOrtExportModelForInferencing OrtExportModelForInferencing;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtSetSeed(
long seed
);
public static DOrtSetSeed OrtSetSeed;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetTrainingModelInputCount(
IntPtr /*(OrtTrainingSession*)*/ session,
out UIntPtr inputCount
);
public static DOrtGetTrainingModelInputCount OrtGetTrainingModelInputCount;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetEvalModelInputCount(
IntPtr /*(OrtTrainingSession*)*/ session,
out UIntPtr inputCount
);
public static DOrtGetEvalModelInputCount OrtGetEvalModelInputCount;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetTrainingModelInputName(
IntPtr /*(OrtTrainingSession*)*/ session,
UIntPtr index,
IntPtr /*(OrtAllocator*)*/ allocator,
out IntPtr /*(char**)*/name
);
public static DOrtGetTrainingModelInputName OrtGetTrainingModelInputName;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetEvalModelInputName(
IntPtr /*(OrtTrainingSession*)*/ session,
UIntPtr index,
IntPtr /*(OrtAllocator*)*/ allocator,
out IntPtr /*(char**)*/name
);
public static DOrtGetEvalModelInputName OrtGetEvalModelInputName;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtAddProperty(
IntPtr /*(OrtCheckpointState*)*/ checkpointState,
IntPtr /*(const char*)*/ propertyName,
OrtPropertyType propertyType,
IntPtr /*(const void*)*/ propertyValue
);
public static DOrtAddProperty OrtAddProperty;
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /*(OrtStatus*)*/ DOrtGetProperty(
IntPtr /*(OrtCheckpointState*)*/ checkpointState,
IntPtr /*(const char*)*/ propertyName,
IntPtr /*(OrtAllocator*)*/ allocator,
out OrtPropertyType propertyType,
out IntPtr /*(const void**)*/ propertyValue
);
public static DOrtGetProperty OrtGetProperty;
#endregion TrainingSession API
public static bool TrainingEnabled()

17
csharp/src/Microsoft.ML.OnnxRuntime/Training/OrtPropertyType.shared.cs Normal file
View file

@ -0,0 +1,17 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
namespace Microsoft.ML.OnnxRuntime
{
#if __ENABLE_TRAINING_APIS__
/// <summary>
/// Property types
/// </summary>
public enum OrtPropertyType
{
OrtIntProperty = 0,
OrtFloatProperty = 1,
OrtStringProperty = 2,
}
#endif
}

10
csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs
View file

@ -319,16 +319,6 @@ namespace Microsoft.ML.OnnxRuntime
}
/// <summary>
/// Saves a checkpoint to path. It can be loaded into <see cref="CheckpointState"/>
/// </summary>
/// <param name="path">Specify path for saving the checkpoint.</param>
/// <param name="saveOptimizerState">SFlag indicating whether to save optimizer state or not.</param>
public void SaveCheckpoint(string path, bool saveOptimizerState = false)
{
NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtSaveCheckpoint(NativeOnnxValueHelper.GetPlatformSerializedString(path), _nativeHandle, saveOptimizerState));
}
#endregion
#region private methods

2
csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/TrainingTest.cs
View file

@ -176,7 +176,7 @@ namespace Microsoft.ML.OnnxRuntime.Tests
// Save checkpoint
string savedCheckpointPath = Path.Combine(Directory.GetCurrentDirectory(), "saved_checkpoint.ckpt");
trainingSession.SaveCheckpoint(savedCheckpointPath, false);
state.SaveCheckpoint(savedCheckpointPath, true);
// Load checkpoint and run train step
var loadedState = new CheckpointState(savedCheckpointPath);

64
java/src/main/java/ai/onnxruntime/OrtTrainingSession.java
View file

@ -250,35 +250,9 @@ public final class OrtTrainingSession implements AutoCloseable {
*/
public void saveCheckpoint(Path outputPath, boolean saveOptimizer) throws OrtException {
checkClosed();
String outputStr = outputPath.toString();
saveCheckpoint(
OnnxRuntime.ortApiHandle,
OnnxRuntime.ortTrainingApiHandle,
nativeHandle,
outputStr,
saveOptimizer);
checkpoint.saveCheckpoint(outputPath, saveOptimizer);
}
/*
* \brief Save the training session states to a checkpoint directory on disk.
*
* <p>This function retrieves the training session states from the training session and serializes
* them to a checkpoint directory on disk. This checkpoint can later be loaded by invoking
* LoadCheckpoint to continue the training with the same states.
*
* <p>\param[in] checkpoint_path Path to the checkpoint directory \param[in] session The training
* session from where the checkpoint states are to be retrieved. \param[in] save_optimizer_state
* Boolean flag indicating whether or not to save the optimizer states to the checkpoint.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>ORT_API2_STATUS(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _In_ const
* OrtTrainingSession* session, bool save_optimizer_state);
*/
private native void saveCheckpoint(
long apiHandle, long trainingHandle, long nativeHandle, String path, boolean saveOptimizer)
throws OrtException;
/*
* \brief Retrieves the number of user outputs in the training model.
*
@ -935,6 +909,24 @@ public final class OrtTrainingSession implements AutoCloseable {
}
}
/**
* Saves the checkpoint out to disk.
*
* @param outputPath The path to save.
* @param saveOptimizer Save the optimizer state as well?
* @throws OrtException If the checkpoint failed to save.
*/
public void saveCheckpoint(Path outputPath, boolean saveOptimizer) throws OrtException {
Objects.requireNonNull(outputPath, "checkpoint path must not be null");
String outputStr = outputPath.toString();
saveCheckpoint(
OnnxRuntime.ortApiHandle,
OnnxRuntime.ortTrainingApiHandle,
nativeHandle,
outputStr,
saveOptimizer);
}
@Override
public void close() {
close(OnnxRuntime.ortTrainingApiHandle, nativeHandle);
@ -959,6 +951,24 @@ public final class OrtTrainingSession implements AutoCloseable {
private static native long loadCheckpoint(long apiHandle, long trainingApiHandle, String path)
throws OrtException;
/* \brief Save the given state to a checkpoint directory on disk.
*
* This function serializes the provided checkpoint state to a directory on disk.
* This checkpoint can later be loaded by invoking LoadCheckpoint to continue the training with the same state.
*
* \param[in] checkpoint_state The checkpoint state to save.
* \param[in] checkpoint_path Path to the checkpoint directory.
* \param[in] include_optimizer_state Flag to indicate whether to save the optimizer state or not.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
* ORT_API2_STATUS(SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* checkpoint_path,
* const bool include_optimizer_state);
*/
private native void saveCheckpoint(
long apiHandle, long trainingHandle, long nativeHandle, String path, boolean saveOptimizer)
throws OrtException;
private native void close(long trainingApiHandle, long nativeHandle);
}
}

36
java/src/main/native/ai_onnxruntime_OrtTrainingSession.c
View file

@ -113,42 +113,6 @@ JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_closeSession
trainApi->ReleaseTrainingSession((OrtTrainingSession*)nativeHandle);
}
/*
* Class: ai_onnxruntime_OrtTrainingSession
* Method: saveCheckpoint
* Signature: (JJJLjava/lang/String;Z)V
*/
JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_saveCheckpoint
(JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainingApiHandle, jlong nativeHandle, jstring outputPath, jboolean overwrite) {
(void) jobj; // Required JNI parameters not needed by functions which don't need to access their host object.
const OrtApi* api = (const OrtApi*) apiHandle;
const OrtTrainingApi* trainApi = (const OrtTrainingApi*) trainingApiHandle;
const OrtTrainingSession* trainSession = (const OrtTrainingSession*) nativeHandle;
#ifdef _WIN32
// The output of GetStringChars is not null-terminated, so we copy it and add a terminator
const jchar* cPath = (*jniEnv)->GetStringChars(jniEnv, outputPath, NULL);
size_t stringLength = (*jniEnv)->GetStringLength(jniEnv, outputPath);
wchar_t* newString = (wchar_t*)calloc(stringLength + 1, sizeof(wchar_t));
if (newString == NULL) {
(*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath);
throwOrtException(jniEnv, 1, "Not enough memory");
} else {
wcsncpy_s(newString, stringLength + 1, (const wchar_t*)cPath, stringLength);
checkOrtStatus(jniEnv, api,
trainApi->SaveCheckpoint(newString, trainSession, overwrite));
free(newString);
(*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath);
}
#else
// GetStringUTFChars is null terminated, so can be used directly
const char* cPath = (*jniEnv)->GetStringUTFChars(jniEnv, outputPath, NULL);
checkOrtStatus(jniEnv, api, trainApi->SaveCheckpoint(cPath, trainSession, overwrite));
(*jniEnv)->ReleaseStringUTFChars(jniEnv, outputPath, cPath);
#endif
}
/*
* Class: ai_onnxruntime_OrtTrainingSession
* Method: getTrainInputNames

36
java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c
View file

@ -45,6 +45,42 @@ JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OrtTrainingSession_00024OrtCheckpoin
return (jlong) checkpoint;
}
/*
* Class: ai_onnxruntime_OrtTrainingSession
* Method: saveCheckpoint
* Signature: (JJJLjava/lang/String;Z)V
*/
JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_00024OrtCheckpointState_saveCheckpoint
(JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainingApiHandle, jlong nativeHandle, jstring outputPath, jboolean saveOptimizer) {
(void) jobj; // Required JNI parameters not needed by functions which don't need to access their host object.
const OrtApi* api = (const OrtApi*) apiHandle;
const OrtTrainingApi* trainApi = (const OrtTrainingApi*) trainingApiHandle;
OrtCheckpointState* checkpointState = (OrtCheckpointState*) nativeHandle;
#ifdef _WIN32
// The output of GetStringChars is not null-terminated, so we copy it and add a terminator
const jchar* cPath = (*jniEnv)->GetStringChars(jniEnv, outputPath, NULL);
size_t stringLength = (*jniEnv)->GetStringLength(jniEnv, outputPath);
wchar_t* newString = (wchar_t*)calloc(stringLength + 1, sizeof(wchar_t));
if (newString == NULL) {
(*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath);
throwOrtException(jniEnv, 1, "Not enough memory");
} else {
wcsncpy_s(newString, stringLength + 1, (const wchar_t*)cPath, stringLength);
checkOrtStatus(jniEnv, api,
trainApi->SaveCheckpoint(checkpointState, newString, saveOptimizer));
free(newString);
(*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath);
}
#else
// GetStringUTFChars is null terminated, so can be used directly
const char* cPath = (*jniEnv)->GetStringUTFChars(jniEnv, outputPath, NULL);
checkOrtStatus(jniEnv, api, trainApi->SaveCheckpoint(checkpointState, cPath, saveOptimizer));
(*jniEnv)->ReleaseStringUTFChars(jniEnv, outputPath, cPath);
#endif
}
/*
* Class: ai_onnxruntime_OrtTrainingSession_OrtCheckpointState
* Method: close

159
orttraining/orttraining/python/orttraining_pybind_state.cc
View file

@ -166,14 +166,13 @@ struct TrainingConfigurationResult {
#ifdef ENABLE_TRAINING_APIS
// Thin wrapper over internal C++ Optimizer
struct PyOptimizer {
PyOptimizer(const std::string optimizer_model_uri,
onnxruntime::training::api::Module* model, std::vector<std::shared_ptr<IExecutionProvider>> provider)
PyOptimizer(const std::string optimizer_model_uri, onnxruntime::training::api::CheckpointState* state,
std::vector<std::shared_ptr<IExecutionProvider>> providers)
: optimizer_() {
auto env = GetTrainingEnv().GetORTEnv();
// XXX: We hope that env will be around when optimizer needs it.
optimizer_ = std::make_shared<onnxruntime::training::api::Optimizer>(optimizer_model_uri,
model->NamedParameters(), onnxruntime::SessionOptions(),
*env, provider);
optimizer_ = std::make_shared<onnxruntime::training::api::Optimizer>(
optimizer_model_uri, state, onnxruntime::SessionOptions(), *env, providers);
}
std::shared_ptr<onnxruntime::training::api::Optimizer> optimizer_;
@ -882,7 +881,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
py::class_<onnxruntime::training::api::Module> training_module(m, "Module", R"pbdoc(Training Module.)pbdoc");
training_module
.def(py::init([](const std::string& model_uri,
onnxruntime::training::api::CheckpointState& state,
onnxruntime::training::api::CheckpointState* state,
std::optional<std::string> eval_model_uri,
OrtDevice device) {
onnxruntime::SessionOptions session_option;
@ -890,9 +889,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
auto env = GetTrainingEnv().GetORTEnv();
return std::make_unique<onnxruntime::training::api::Module>(
model_uri,
state.module_checkpoint_state.named_parameters, session_option,
*env, provider, eval_model_uri);
model_uri, state, session_option, *env, provider, eval_model_uri);
}))
.def("train_step",
[](onnxruntime::training::api::Module* model,
@ -920,43 +917,74 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
[](onnxruntime::training::api::Module* model, bool trainable_only) -> size_t {
return model->GetParametersSize(trainable_only);
})
.def("save_checkpoint",
[](onnxruntime::training::api::Module* model, const std::string& checkpoint_path) -> void {
onnxruntime::training::api::CheckpointState state;
ORT_THROW_IF_ERROR(model->GetStateDict(state.module_checkpoint_state));
ORT_THROW_IF_ERROR(onnxruntime::training::api::SaveCheckpoint(state,
ToPathString(checkpoint_path)));
})
.def("export_model_for_inferencing",
[](onnxruntime::training::api::Module* model, const std::string& inference_model_path,
const std::vector<std::string>& graph_output_names) -> void {
ORT_ENFORCE(model, "Received a nullptr for expected pointer to class training::api::Module");
ORT_THROW_IF_ERROR(model->ExportModelForInferencing(inference_model_path,
graph_output_names));
})
.def("input_names",
[](onnxruntime::training::api::Module* model, const bool is_training) {
auto count_method = [&model, is_training]() -> size_t {
return is_training ? model->GetTrainingModelInputCount() : model->GetEvalModelInputCount();
};
auto name_method = [&model, is_training](const size_t index) -> std::string {
return is_training ? model->GetTrainingModelInputName(index) : model->GetEvalModelInputName(index);
};
std::vector<std::string> names;
for (size_t index = 0; index < count_method(); ++index) {
names.push_back(name_method(index));
}
return names;
})
.def("output_names",
[](onnxruntime::training::api::Module* model, const bool is_training) {
auto count_method = [&model, is_training]() -> size_t {
return is_training ? model->GetTrainingModelOutputCount() : model->GetEvalModelOutputCount();
};
auto name_method = [&model, is_training](const size_t index) -> std::string {
return is_training ? model->GetTrainingModelOutputName(index) : model->GetEvalModelOutputName(index);
};
std::vector<std::string> names;
for (size_t index = 0; index < count_method(); ++index) {
names.push_back(name_method(index));
}
return names;
});
py::class_<onnxruntime::training::api::CheckpointState>
checkpoint_state(m, "CheckpointState", R"pbdoc(CheckpointState.)pbdoc");
checkpoint_state.def(py::init([](
const std::string& ckpt_uri) {
onnxruntime::training::api::CheckpointState state;
ORT_THROW_IF_ERROR(onnxruntime::training::api::LoadCheckpoint(ToPathString(ckpt_uri), state));
return state;
}));
checkpoint_state
.def(py::init())
.def("add_property", [](onnxruntime::training::api::CheckpointState* state,
const std::string& property_name,
const std::variant<int64_t, float, std::string>& property_value) {
state->property_bag.AddProperty(property_name, property_value);
})
.def("get_property", [](onnxruntime::training::api::CheckpointState* state, const std::string& property_name) {
return state->property_bag.GetProperty<onnxruntime::training::api::PropertyDataType>(property_name);
})
.def("has_property", [](onnxruntime::training::api::CheckpointState* state, const std::string& property_name) {
return state->property_bag.HasProperty(property_name);
});
py::class_<PyOptimizer>
training_optimizer(m, "Optimizer", R"pbdoc(Training Optimizer.)pbdoc");
training_optimizer.def(py::init([](
const std::string optimizer_model_uri,
onnxruntime::training::api::Module* model,
OrtDevice device) {
onnxruntime::SessionOptions session_option;
std::vector<std::shared_ptr<IExecutionProvider>> provider = GetExecutionProvidersForTrainingApis(device);
training_optimizer
.def(py::init([](const std::string optimizer_model_uri,
onnxruntime::training::api::CheckpointState* state,
OrtDevice device) {
std::vector<std::shared_ptr<IExecutionProvider>> providers = GetExecutionProvidersForTrainingApis(device);
return std::make_unique<PyOptimizer>(
optimizer_model_uri,
model, provider);
}))
return std::make_unique<PyOptimizer>(optimizer_model_uri, state, providers);
}))
.def("optimizer_step", [](PyOptimizer* optimizer) -> void {
ORT_THROW_IF_ERROR(optimizer->optimizer_->Step());
})
@ -980,32 +1008,51 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
.def("scheduler_step", [](onnxruntime::training::api::LinearLRScheduler* scheduler) -> void {
ORT_THROW_IF_ERROR(scheduler->Step());
});
m.def(
"save_checkpoint",
[](const std::vector<py::bytes>& trainable_tensor_protos_pybytes,
const std::vector<py::bytes>& non_trainable_tensor_protos_pybytes,
const std::string& checkpoint_path) {
std::vector<TensorProto> trainable_tensor_protos(trainable_tensor_protos_pybytes.size());
std::vector<TensorProto> non_trainable_tensor_protos(non_trainable_tensor_protos_pybytes.size());
auto parse_pybytes_to_tensor_proto =
[](const std::vector<py::bytes>& tensor_protos_pybytes, std::vector<TensorProto>& tensor_protos) {
for (size_t i = 0; i < tensor_protos_pybytes.size(); ++i) {
std::istringstream tensor_proto_istream(tensor_protos_pybytes[i]);
ORT_ENFORCE(tensor_proto_istream.good(), "Broken tensor proto istream to read.");
google::protobuf::io::IstreamInputStream zero_copy_input(&tensor_proto_istream);
const bool result =
tensor_protos[i].ParseFromZeroCopyStream(&zero_copy_input) && tensor_proto_istream.eof();
ORT_ENFORCE(result, "Parse tensor proto failed.");
}
};
parse_pybytes_to_tensor_proto(trainable_tensor_protos_pybytes, trainable_tensor_protos);
parse_pybytes_to_tensor_proto(non_trainable_tensor_protos_pybytes, non_trainable_tensor_protos);
ORT_THROW_IF_ERROR(onnxruntime::training::api::SaveCheckpoint(trainable_tensor_protos,
non_trainable_tensor_protos,
ToPathString(checkpoint_path)));
});
m.def("save_checkpoint",
[](const std::vector<py::bytes>& trainable_tensor_protos_pybytes,
const std::vector<py::bytes>& non_trainable_tensor_protos_pybytes,
const std::string& checkpoint_path) {
std::vector<TensorProto> trainable_tensor_protos(trainable_tensor_protos_pybytes.size());
std::vector<TensorProto> non_trainable_tensor_protos(non_trainable_tensor_protos_pybytes.size());
auto parse_pybytes_to_tensor_proto =
[](const std::vector<py::bytes>& tensor_protos_pybytes, std::vector<TensorProto>& tensor_protos) {
for (size_t i = 0; i < tensor_protos_pybytes.size(); ++i) {
std::istringstream tensor_proto_istream(tensor_protos_pybytes[i]);
ORT_ENFORCE(tensor_proto_istream.good(), "Broken tensor proto istream to read.");
google::protobuf::io::IstreamInputStream zero_copy_input(&tensor_proto_istream);
const bool result =
tensor_protos[i].ParseFromZeroCopyStream(&zero_copy_input) && tensor_proto_istream.eof();
ORT_ENFORCE(result, "Parse tensor proto failed.");
}
};
parse_pybytes_to_tensor_proto(trainable_tensor_protos_pybytes, trainable_tensor_protos);
parse_pybytes_to_tensor_proto(non_trainable_tensor_protos_pybytes, non_trainable_tensor_protos);
ORT_THROW_IF_ERROR(onnxruntime::training::api::SaveCheckpoint(trainable_tensor_protos,
non_trainable_tensor_protos,
ToPathString(checkpoint_path)));
[](onnxruntime::training::api::CheckpointState* checkpoint_state,
const std::string& checkpoint_path, const bool include_optimizer_state) -> void {
ORT_THROW_IF_ERROR(
onnxruntime::training::api::SaveCheckpoint(*checkpoint_state, ToPathString(checkpoint_path),
include_optimizer_state));
});
m.def("load_checkpoint",
[](const std::string& checkpoint_path) -> onnxruntime::training::api::CheckpointState {
onnxruntime::training::api::CheckpointState state;
ORT_THROW_IF_ERROR(
onnxruntime::training::api::LoadCheckpoint(ToPathString(checkpoint_path), state));
return state;
});
m.def("get_model_after_loading_checkpoint",
[](const std::string& checkpoint_path, const py::bytes& serialized_model) {
ONNX_NAMESPACE::ModelProto model_proto;

19
orttraining/orttraining/python/training/api/README.md
View file

@ -6,7 +6,7 @@ This is a simple guide on how to use onnxruntime training APIs.
The ort training APIs need the following files for performing training
1. The training onnx model.
2. The eval onnx model (optional).
3. The optimizer onnx model.
3. The optimizer onnx model (optional).
4. The checkpoint file.
To generate these files, refer to this [onnxblock's README](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md)
@ -18,18 +18,17 @@ Once the onnx models are generated, you can use the training APIs to run your tr
```py
from onnxruntime.training.api import Module, Optimizer, CheckpointState
# Create Checkpoint State.
state = CheckpointState("checkpoint.ckpt")
state = CheckpointState.load_checkpoint("checkpoint.ckpt")
# Create Module and Optimizer.
model = Module("training_model.onnx", state, "eval_model.onnx")
optimizer = Optimizer("optimizer.onnx", model)
# Data should be a list of numpy arrays.
forward_inputs = ...
# Set model in training mode and run a Train step.
model.train()
model(forward_inputs)
training_model_outputs = model(<inputs to your training model>)
# Optimizer step
optimizer.step()
@ -37,13 +36,13 @@ optimizer.step()
# Set Model in eval mode and run an Eval step.
model.eval()
loss = model(forward_inputs)
eval_model_outputs = model(<inputs to your eval model>)
# Assuming that the loss is the first element of the output in our case.
print("Loss : ", loss[0])
# Assuming that the loss is the first element of the output in the training model.
print("Loss : ", training_model_outputs[0])
# Saving checkpoint.
model.save_checkpoint("checkpoint_export.ckpt")
CheckpointState.save_checkpoint(state, "checkpoint_export.ckpt")
```

18
orttraining/orttraining/python/training/api/__init__.py
View file

@ -1,4 +1,14 @@
from .checkpoint_state import CheckpointState # noqa: F401
from .lr_scheduler import LinearLRScheduler # noqa: F401
from .module import Module # noqa: F401
from .optimizer import Optimizer # noqa: F401
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from onnxruntime.training.api.checkpoint_state import CheckpointState
from onnxruntime.training.api.lr_scheduler import LinearLRScheduler
from onnxruntime.training.api.module import Module
from onnxruntime.training.api.optimizer import Optimizer
__all__ = [
"CheckpointState",
"LinearLRScheduler",
"Module",
"Optimizer",
]

74
orttraining/orttraining/python/training/api/checkpoint_state.py
View file

@ -1,19 +1,77 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# checkpoint_state.py
from __future__ import annotations
import os
from onnxruntime.capi import _pybind_state as C
class CheckpointState:
"""
Class for Loading CheckpointState.
This class is a wrapper of CheckpointState Class.
"""Class that holds the state of the training session state
Args:
state (CheckpointState): The C.Checkpoint state object that holds the underlying session state.
"""
def __init__(self, ckpt_uri) -> None:
def __init__(self, state: C.CheckpointState):
if not isinstance(state, C.CheckpointState):
raise TypeError(f"Invalid argument for CheckpointState received {type(state)}")
self._state = state
@classmethod
def load_checkpoint(cls, checkpoint_uri: str | os.PathLike) -> CheckpointState:
"""Loads the checkpoint state from the checkpoint file
Args:
checkpoint_uri: The path to the checkpoint file.
Returns:
CheckpointState: The checkpoint state object.
"""
Initializes CheckpointState object with the given checkpoint uri.
The returned object will be used to initialize the Module.
return cls(C.load_checkpoint(os.fspath(checkpoint_uri)))
@classmethod
def save_checkpoint(
cls, state: CheckpointState, checkpoint_uri: str | os.PathLike, include_optimizer_state: bool = False
) -> None:
"""Saves the checkpoint state to the checkpoint file
Args:
state: The checkpoint state object.
checkpoint_uri: The path to the checkpoint file.
include_optimizer_state: If True, the optimizer state is also saved to the checkpoint file.
"""
self._state = C.CheckpointState(ckpt_uri)
C.save_checkpoint(state._state, os.fspath(checkpoint_uri), include_optimizer_state)
def __getitem__(self, name: str) -> int | float | str:
"""Gets the property associated with the given name
Args:
name: The name of the property
Returns:
The value of the property
"""
return self._state.get_property(name)
def __setitem__(self, name: str, value: int | float | str) -> None:
"""Sets the property value for the given name
Args:
name: The name of the property
value: The value of the property
"""
self._state.add_property(name, value)
def __contains__(self, name: str) -> bool:
"""Checks if the property exists in the state
Args:
name: The name of the property
Returns:
True if the property exists, False otherwise
"""
return self._state.has_property(name)

24
orttraining/orttraining/python/training/api/lr_scheduler.py
View file

@ -1,33 +1,31 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# lr_scheduler.py
from onnxruntime.capi import _pybind_state as C
from onnxruntime.training.api.optimizer import Optimizer
class LinearLRScheduler:
"""
Linearly updates the learning rate in the optimizer
"""Linearly updates the learning rate in the optimizer
The linear learning rate scheduler decays the learning rate by linearly updated
multiplicative factor from the initial learning rate set on the training session to 0. The decay
is performed after the initial warm up phase where the learning rate is linearly incremented
from to the initial learning rate provided.
from 0 to the initial learning rate provided.
Args:
optimizer (:obj:`training_api.Optimizer`): User's onnxruntime training Optimizer
warmup_step_count (int): The number of steps in the warm up phase.
total_step_count (int): The total number of training steps.
initial_lr (float): The initial learning rate.
optimizer: User's onnxruntime training Optimizer
warmup_step_count: The number of steps in the warm up phase.
total_step_count: The total number of training steps.
initial_lr: The initial learning rate.
"""
def __init__(self, optimizer, warmup_step_count, total_step_count, initial_lr) -> None:
def __init__(self, optimizer: Optimizer, warmup_step_count: int, total_step_count: int, initial_lr: float):
self._scheduler = C.LinearLRScheduler(optimizer._optimizer, warmup_step_count, total_step_count, initial_lr)
def step(self):
"""
The step method of the LinearLRScheduler class is used to update the learning rate of the optimizer according
to the scheduler's strategy.
def step(self) -> None:
"""Updates the learning rate of the optimizer linearly.
This method should be called at each step of training to ensure that the learning rate is properly adjusted.
"""
self._scheduler.scheduler_step()

142
orttraining/orttraining/python/training/api/module.py
View file

@ -1,30 +1,46 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# module.py
from typing import List
from __future__ import annotations
import os
import numpy as np
from onnxruntime.capi import _pybind_state as C
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue, get_ort_device_type
from onnxruntime.capi.onnxruntime_pybind11_state import OrtValueVector
from onnxruntime.training.api.checkpoint_state import CheckpointState
class Module:
"""
Class for running Training.
This class is a wrapper of Module Class.
"""Trainer class that provides training and evaluation methods for ONNX models.
Before instantiating the Module class, it is expected that the training artifacts have been
generated using the `onnxruntime.training.artifacts.generate_artifacts` utility.
The training artifacts include:
- The training model
- The evaluation model (optional)
- The optimizer model (optional)
- The checkpoint file
Args:
train_model_uri: The path to the training model.
state: The checkpoint state object.
eval_model_uri: The path to the evaluation model.
device: The device to run the model on. Default is "cpu".
"""
training: bool
def __init__(self, train_model_uri, state, eval_model_uri=None, device: str = "cpu") -> None:
"""
Initializes Model for Training.
__init__ will call an internatl function to create the model.
"""
# TODO : Add support for bytes on train_model_uri and eval_model_uri.
def __init__(
self,
train_model_uri: os.PathLike,
state: CheckpointState,
eval_model_uri: os.PathLike | None = None,
device: str = "cpu",
) -> None:
self.training = True
options = device.split(":")
self._device_type = options[0]
@ -35,20 +51,33 @@ class Module:
C.OrtDevice.default_memory(),
device_id,
)
self._model = C.Module(train_model_uri, state._state, eval_model_uri, self._device)
self._model = C.Module(
os.fspath(train_model_uri),
state._state,
os.fspath(eval_model_uri) if eval_model_uri is not None else None,
self._device,
)
self._state = state
def __call__(self, *user_inputs) -> tuple[np.ndarray] | np.ndarray:
"""Invokes either the training or the evaluation step of the model.
def __call__(self, user_inputs):
"""
This method enables calling Module as a function to run the model.
Args:
user_inputs : list of numpy objects.
user_inputs: The inputs to the model.
Returns:
fetches : list of numpy objects.
fetches : The outputs of the model.
"""
is_np_input = False
forward_inputs = OrtValueVector()
forward_inputs.reserve(len(user_inputs))
for element in user_inputs:
forward_inputs.push_back(OrtValue.ortvalue_from_numpy(element)._ortvalue)
for tensor in user_inputs:
if isinstance(tensor, np.ndarray):
is_np_input = True
forward_inputs.push_back(OrtValue.ortvalue_from_numpy(tensor)._ortvalue)
elif isinstance(tensor, OrtValue):
forward_inputs.push_back(tensor._ortvalue)
else:
raise ValueError(f"Expected input of type: numpy array or OrtValue, actual: {type(tensor)}")
fetches = OrtValueVector()
if self.training:
@ -56,29 +85,32 @@ class Module:
else:
self._model.eval_step(forward_inputs, fetches)
return [val.numpy() for val in fetches]
if len(fetches) == 1:
if is_np_input:
return fetches[0].numpy()
def train(self, mode: bool = True):
return fetches[0]
return tuple(val.numpy() for val in fetches) if is_np_input else tuple(fetches)
def train(self, mode: bool = True) -> Module:
"""Sets the Module in training mode.
This has any effect only on Module Class.
Args:
mode (bool): whether to set training mode (``True``) or evaluation
mode (``False``). Default: ``True``.
mode: whether to set training mode (True) or evaluation
mode (False). Default: True.
Returns:
Module: self
self
"""
self.training = mode
return self
def eval(self):
def eval(self) -> Module:
"""Sets the Module in evaluation mode.
This has any effect only on Module Class.
Returns:
Module: self
self
"""
return self.train(False)
@ -91,17 +123,14 @@ class Module:
"""
return self._model.lazy_reset_grad()
def save_checkpoint(self, ckpt_uri):
"""
Saves the checkpoint.
"""
# TODO : move this out of Module Class.
self._model.save_checkpoint(ckpt_uri)
# This function will change when the parameters will be exposed.
def get_contiguous_parameters(self, trainable_only: bool = False) -> OrtValue:
"""
Returns contiguous parameters object.
"""Creates a contiguous buffer of the training session parameters
Args:
trainable_only: If True, only trainable parameters are considered. Otherwise, all parameters are considered.
Returns:
The contiguous buffer of the training session parameters.
"""
parameters = OrtValue.ortvalue_from_shape_and_type(
[
@ -116,23 +145,44 @@ class Module:
return parameters
def get_parameters_size(self, trainable_only: bool = False) -> int:
"""
Returns the size of the parameters.
"""Returns the size of the parameters.
Args:
trainable_only: If True, only trainable parameters are considered. Otherwise, all parameters are considered.
Returns:
The number of primitive (example floating point) elements in the parameters.
"""
return self._model.get_parameters_size(trainable_only)
def copy_buffer_to_parameters(self, buffer) -> None:
"""
Copies buffer to parameters.
def copy_buffer_to_parameters(self, buffer: OrtValue) -> None:
"""Copies the OrtValue buffer to the training session parameters.
Args:
buffer: The OrtValue buffer to copy to the training session parameters.
"""
self._model.copy_buffer_to_parameters(buffer)
def export_model_for_inferencing(self, inference_model_uri: str, graph_output_names: List[str]) -> None:
def export_model_for_inferencing(
self, inference_model_uri: str | os.PathLike, graph_output_names: list[str]
) -> None:
"""Exports the model for inferencing.
Once training is complete, this function can be used to drop the training specific nodes in the onnx model.
In particular, this function does the following:
- Parse over the training graph and identify nodes that generate the given output names.
- Drop all subsequent nodes in the graph since they are not relevant to the inference graph.
Args:
inference_model_uri: The path to the inference model.
graph_output_names: The list of output names that are required for inferencing.
"""
self._model.export_model_for_inferencing(inference_model_uri, graph_output_names)
self._model.export_model_for_inferencing(os.fspath(inference_model_uri), graph_output_names)
def input_names(self) -> list[str]:
"""Returns the input names of the training or eval model."""
return self._model.input_names(self.training)
def output_names(self) -> list[str]:
"""Returns the output names of the training or eval model."""
return self._model.output_names(self.training)

44
orttraining/orttraining/python/training/api/optimizer.py
View file

@ -1,36 +1,48 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# optimizer.py
from __future__ import annotations
import os
from typing import TYPE_CHECKING
from onnxruntime.capi import _pybind_state as C
if TYPE_CHECKING:
from onnxruntime.training.api.module import Module
class Optimizer:
"""
Class for running Optimize Step in Training.
This class is a wrapper of Optimizer Class.
"""Class that provides methods to update the model parameters based on the computed gradients.
Args:
optimizer_uri: The path to the optimizer model.
model: The module to be trained.
"""
def __init__(self, train_optimizer_uri, model) -> None:
"""
Initializes Optimizer with the optimizer onnx and the parameters from the model.
"""
self._optimizer = C.Optimizer(train_optimizer_uri, model._model, model._device)
def __init__(self, optimizer_uri: str | os.PathLike, module: Module):
self._optimizer = C.Optimizer(os.fspath(optimizer_uri), module._state._state, module._device)
def step(self):
"""
Run Optimizer Step.
def step(self) -> None:
"""Updates the model parameters based on the computed gradients.
This method updates the model parameters by taking a step in the direction of the computed gradients.
The optimizer used depends on the optimizer model provided.
"""
self._optimizer.optimizer_step()
def set_learning_rate(self, learning_rate: float) -> None:
"""
Set Learning Rate.
"""Sets the learning rate for the optimizer.
Args:
learning_rate: The learning rate to be set.
"""
self._optimizer.set_learning_rate(learning_rate)
def get_learning_rate(self) -> float:
"""
Get Learning Rate.
"""Gets the current learning rate of the optimizer.
Returns:
float: The current learning rate.
"""
return self._optimizer.get_learning_rate()

2
orttraining/orttraining/python/training/artifacts.py
View file

@ -149,7 +149,7 @@ def generate_artifacts(
checkpoint_path = artifact_directory / f"{prefix}checkpoint"
if os.path.exists(checkpoint_path):
logging.info("Checkpoint path %s already exists. Overwriting.", checkpoint_path)
onnxblock.save_checkpoint(training_block.parameters(), str(checkpoint_path))
onnxblock.save_checkpoint(training_block.parameters(), checkpoint_path)
logging.info("Saved checkpoint to %s", checkpoint_path)
# If optimizer is not specified, skip creating the optimizer model

4
orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py
View file

@ -26,7 +26,7 @@ def save_checkpoint(
trainable_params, non_trainable_params = parameters
trainable_params = [param.SerializeToString() for param in trainable_params]
non_trainable_params = [param.SerializeToString() for param in non_trainable_params]
_internal_save_checkpoint(trainable_params, non_trainable_params, path_to_checkpoint)
_internal_save_checkpoint(trainable_params, non_trainable_params, os.fspath(path_to_checkpoint))
def load_checkpoint_to_model(path_to_checkpoint: Union[str, os.PathLike], model: onnx.ModelProto) -> None:
@ -37,4 +37,4 @@ def load_checkpoint_to_model(path_to_checkpoint: Union[str, os.PathLike], model:
model (onnx.ModelProto): The model to load the checkpoint to.
"""
model.ParseFromString(_internal_load_checkpoint_to_model(path_to_checkpoint, model.SerializeToString()))
model.ParseFromString(_internal_load_checkpoint_to_model(os.fspath(path_to_checkpoint), model.SerializeToString()))

5
orttraining/orttraining/test/python/orttraining_test_ort_apis.py
View file

@ -54,9 +54,10 @@ def run_onnxruntime_test_all_ctest(cwd, log, filter):
def run_training_api_tests(cwd, log):
"""Runs the onnxruntime_test_all executable with the TrainingApiTest* gtest filter."""
log.debug("Running: TrainingApi tests")
log.debug("Running: TrainingApi and TrainingCApi tests")
run_onnxruntime_test_all_ctest(cwd, log, "TrainingApiTest*")
run_onnxruntime_test_all_ctest(cwd, log, "TrainingCApiTest*")
def run_checkpoint_api_tests(cwd, log):
@ -75,6 +76,8 @@ def main():
run_onnxblock_tests(cwd, log)
run_training_apis_python_api_tests(cwd, log)
run_training_api_tests(cwd, log)
run_checkpoint_api_tests(cwd, log)

322
orttraining/orttraining/test/python/orttraining_test_python_bindings.py
View file

@ -1,16 +1,22 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from __future__ import annotations
import os
import tempfile
import numpy as np
import onnx
import pytest
import torch
from orttraining_test_onnxblock import _get_models
import onnxruntime.training.onnxblock as onnxblock
from onnxruntime.training import artifacts
from onnxruntime.training.api import CheckpointState, LinearLRScheduler, Module, Optimizer
class SimpleModelWithCrossEntropyLoss(onnxblock.TrainingModel):
class SimpleModelWithCrossEntropyLoss(onnxblock.TrainingBlock):
def __init__(self):
super().__init__()
self.loss = onnxblock.loss.CrossEntropyLoss()
@ -19,126 +25,107 @@ class SimpleModelWithCrossEntropyLoss(onnxblock.TrainingModel):
return self.loss(output_name)
def _create_training_models():
# Given
def _create_training_artifacts(artifact_directory: str | os.PathLike):
device = "cpu"
batch_size, input_size, hidden_size, output_size = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, batch_size, input_size, hidden_size, output_size)
# Build the onnx model with loss
simple_model = SimpleModelWithCrossEntropyLoss()
with onnxblock.onnx_model(onnx_model) as accessor:
_ = simple_model(onnx_model.graph.output[0].name)
eval_model = accessor.eval_model
requires_grad = [name for name, param in pt_model.named_parameters() if param.requires_grad]
frozen_params = [name for name, param in pt_model.named_parameters() if not param.requires_grad]
optimizer = onnxblock.optim.AdamW()
with onnxblock.onnx_model() as accessor:
_ = optimizer(simple_model.parameters())
optimizer_model = accessor.model
artifacts.generate_artifacts(
onnx_model,
optimizer=artifacts.OptimType.AdamW,
loss=artifacts.LossType.CrossEntropyLoss,
requires_grad=requires_grad,
frozen_params=frozen_params,
artifact_directory=artifact_directory,
)
return simple_model, onnx_model, optimizer_model, eval_model, pt_model
training_model_file = os.path.join(artifact_directory, "training_model.onnx")
eval_model_file = os.path.join(artifact_directory, "eval_model.onnx")
optimizer_model_file = os.path.join(artifact_directory, "optimizer_model.onnx")
checkpoint_file = os.path.join(artifact_directory, "checkpoint")
def _get_test_models_path(directory, simple_model, onnx_model, optimizer_model=None, eval_model=None):
trainable_params, non_trainable_params = simple_model.parameters()
paths = []
checkpoint_file_path = os.path.join(directory, "checkpoint")
onnxblock.save_checkpoint((trainable_params, non_trainable_params), checkpoint_file_path)
paths.append(checkpoint_file_path)
model_file_path = os.path.join(directory, "training_model.onnx")
onnx.save(onnx_model, model_file_path)
paths.append(model_file_path)
if optimizer_model:
optimizer_file_path = os.path.join(directory, "optimizer.onnx")
onnx.save(optimizer_model, optimizer_file_path)
paths.append(optimizer_file_path)
if eval_model:
eval_model_file_path = os.path.join(directory, "eval_model.onnx")
onnx.save(eval_model, eval_model_file_path)
paths.append(eval_model_file_path)
return tuple(paths)
return checkpoint_file, training_model_file, eval_model_file, optimizer_model_file, pt_model
def test_train_step():
# Initialize Models
simple_model, onnx_model, _, _, pt_model = _create_training_models()
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path = _get_test_models_path(temp_dir, simple_model, onnx_model)
(
checkpoint_file_path,
training_model_file_path,
_,
_,
pt_model,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
model = Module(model_file_path, state)
print(training_model_file_path)
model = Module(training_model_file_path, state)
model.train()
fetches = model(forward_inputs)
ort_loss = model(inputs, labels)
# Calculate loss using pytorch model to compare it with Module's output.
pt_outputs = pt_model(torch.from_numpy(inputs))
loss_fn = torch.nn.CrossEntropyLoss()
pt_loss = loss_fn(pt_outputs, torch.from_numpy(labels).long())
assert np.allclose(fetches[0], pt_loss.detach().numpy())
assert np.allclose(ort_loss, pt_loss.detach().numpy())
def test_eval_step():
# Initialize Models
simple_model, onnx_model, _, eval_model, _ = _create_training_models()
# Generating random data for testing.
# TODO : add utility function to convert numpy arrays to OrtValueVector.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, eval_model_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, eval_model=eval_model
)
(
checkpoint_file_path,
training_model_file_path,
eval_model_file_path,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
model = Module(model_file_path, state, eval_model_file_path)
model = Module(training_model_file_path, state, eval_model_file_path)
model.train()
model(forward_inputs)
model(inputs, labels)
model.eval()
fetches = model(forward_inputs)
fetches = model(inputs, labels)
assert fetches
def test_optimizer_step():
# Initialize Models
simple_model, onnx_model, optimizer_model, _, _ = _create_training_models()
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, optimizer_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, optimizer_model=optimizer_model
)
(
checkpoint_file_path,
training_model_file_path,
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
model = Module(model_file_path, state)
optimizer = Optimizer(optimizer_file_path, model)
model = Module(training_model_file_path, state)
optimizer = Optimizer(optimizer_model_file_path, model)
model.train()
old_flatten_params = model.get_contiguous_parameters()
model(forward_inputs)
model(inputs, labels)
optimizer.step()
new_params = model.get_contiguous_parameters()
@ -147,19 +134,19 @@ def test_optimizer_step():
def test_get_and_set_lr():
# Initialize Models
simple_model, onnx_model, optimizer_model, _, _ = _create_training_models()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, optimizer_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, optimizer_model=optimizer_model
)
(
checkpoint_file_path,
training_model_file_path,
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
model = Module(model_file_path, state)
optimizer = Optimizer(optimizer_file_path, model)
model = Module(training_model_file_path, state)
optimizer = Optimizer(optimizer_model_file_path, model)
# Test get and set learning rate.
lr = optimizer.get_learning_rate()
@ -173,24 +160,23 @@ def test_get_and_set_lr():
def test_scheduler_step():
# Initialize Models
simple_model, onnx_model, optimizer_model, _, _ = _create_training_models()
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, optimizer_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, optimizer_model=optimizer_model
)
(
checkpoint_file_path,
training_model_file_path,
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
model = Module(model_file_path, state)
optimizer = Optimizer(optimizer_file_path, model)
model = Module(training_model_file_path, state)
optimizer = Optimizer(optimizer_model_file_path, model)
scheduler = LinearLRScheduler(optimizer, 1, 2, 0.2)
# Test get and set learning rate.
@ -198,7 +184,7 @@ def test_scheduler_step():
assert np.allclose(lr, 0.0)
model.train()
model(forward_inputs)
model(inputs, labels)
optimizer.step()
scheduler.step()
@ -208,36 +194,37 @@ def test_scheduler_step():
def test_training_module_checkpoint():
# Initialize Models
simple_model, onnx_model, _, _, _ = _create_training_models()
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path = _get_test_models_path(temp_dir, simple_model, onnx_model)
(
checkpoint_file_path,
training_model_file_path,
_,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Training Module and Training Optimizer.
model = Module(model_file_path, state)
model = Module(training_model_file_path, state)
model.train()
model(forward_inputs)
model(inputs, labels)
checkpoint_save_path = os.path.join(temp_dir, "checkpoint_export.ckpt")
model.save_checkpoint(checkpoint_save_path)
CheckpointState.save_checkpoint(state, checkpoint_save_path)
old_flatten_params = model.get_contiguous_parameters()
# Assert the checkpoint was saved.
assert os.path.exists(checkpoint_save_path)
# Assert the checkpoint parameters remain after saving.
state = CheckpointState(checkpoint_save_path)
new_model = Module(model_file_path, state)
new_state = CheckpointState.load_checkpoint(checkpoint_save_path)
new_model = Module(training_model_file_path, new_state)
new_params = new_model.get_contiguous_parameters()
@ -245,31 +232,30 @@ def test_training_module_checkpoint():
def test_copy_buffer_to_parameters():
# Initialize Models
simple_model, onnx_model, optimizer_model, _, _ = _create_training_models()
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int32).numpy()
forward_inputs = [inputs, labels]
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, optimizer_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, optimizer_model=optimizer_model
)
state = CheckpointState(checkpoint_file_path)
(
checkpoint_file_path,
training_model_file_path,
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
model = Module(model_file_path, state)
optimizer = Optimizer(optimizer_file_path, model)
model = Module(training_model_file_path, state)
optimizer = Optimizer(optimizer_model_file_path, model)
# Keep a copy of the parameters.
old_output_params = model.get_contiguous_parameters()
# Run a Training Step.
model.train()
model(forward_inputs)
model(inputs, labels)
optimizer.step()
# Get the new parameters.
@ -288,20 +274,20 @@ def test_copy_buffer_to_parameters():
def test_export_model_for_inferencing():
# Initialize Models
simple_model, onnx_model, _, eval_model, _ = _create_training_models()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path, eval_model_file_path = _get_test_models_path(
temp_dir, simple_model, onnx_model, eval_model=eval_model
)
(
checkpoint_file_path,
training_model_file_path,
eval_model_file_path,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
model = Module(model_file_path, state, eval_model_file_path)
model = Module(training_model_file_path, state, eval_model_file_path)
# Export inference model
inference_model_file_path = os.path.join(temp_dir, "inference_model.onnx")
@ -310,17 +296,75 @@ def test_export_model_for_inferencing():
def test_cuda_execution_provider():
# Initialize Models
simple_model, onnx_model, _, _, pt_model = _create_training_models()
with tempfile.TemporaryDirectory() as temp_dir:
# Save models & checkpoint files to load them later.
checkpoint_file_path, model_file_path = _get_test_models_path(temp_dir, simple_model, onnx_model)
(
checkpoint_file_path,
training_model_file_path,
_,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState(checkpoint_file_path)
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
model = Module(model_file_path, state, device="cuda")
model = Module(training_model_file_path, state, device="cuda")
params = model.get_contiguous_parameters()
# Check if parameters are moved to cuda.
assert params.device_name() == "Cuda"
@pytest.mark.parametrize(
"property_value",
[-1, 0, 1, 1234567890, -1.0, -0.1, 0.1, 1.0, 1234.0, "hello", "world", "onnxruntime"],
)
def test_add_get_property(property_value):
with tempfile.TemporaryDirectory() as temp_dir:
(
checkpoint_file_path,
training_model_file_path,
_,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
_ = Module(training_model_file_path, state)
# Float values in python are double precision.
# Convert to float32 to match the type of the property.
if isinstance(property_value, float):
property_value = float(np.float32(property_value))
state["property"] = property_value
assert "property" in state
assert state["property"] == property_value
CheckpointState.save_checkpoint(state, checkpoint_file_path)
new_state = CheckpointState.load_checkpoint(checkpoint_file_path)
assert "property" in new_state
assert new_state["property"] == property_value
def test_get_input_output_names():
with tempfile.TemporaryDirectory() as temp_dir:
(
checkpoint_file_path,
training_model_file_path,
eval_model_file_path,
_,
_,
) = _create_training_artifacts(temp_dir)
# Create Checkpoint State.
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module.
model = Module(training_model_file_path, state, eval_model_file_path)
assert model.input_names() == ["input-0", "labels"]
assert model.output_names() == ["onnx::loss::128"]

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

@ -237,13 +237,16 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
named_parameters.insert({it->first, param});
}
auto state = CheckpointState();
state.module_checkpoint_state.named_parameters = named_parameters;
onnxruntime::SessionOptions session_option;
std::unique_ptr<Environment> env;
ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));
std::vector<std::shared_ptr<IExecutionProvider>> cuda_provider{onnxruntime::test::DefaultCudaExecutionProvider()};
auto model = std::make_unique<Module>(model_uri, named_parameters, session_option,
auto model = std::make_unique<Module>(model_uri, &state, session_option,
*env, cuda_provider);
auto optimizer = std::make_unique<Optimizer>(optim_uri, model->NamedParameters(), session_option,
auto optimizer = std::make_unique<Optimizer>(optim_uri, &state, session_option,
*env, cuda_provider);
/// Phase 2 - Run Optimizer.GetStateDict and call save checkpoint APIs.
@ -262,7 +265,7 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
// Call Save APIs.
PathString checkpoint_path{
ConcatPathComponent<PathChar>(tmp_dir.Path(), ORT_TSTR("e2e_ckpt_save_cpu"))};
ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path));
ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path, true));
// Check the ckpt files in the directory.
std::set<PathString> expected_file_names{
@ -365,7 +368,7 @@ TEST(CheckpointApiTest, SaveCustomPropertyAsCheckpoint_ThenLoad_CPU) {
// Call Save APIs.
PathString checkpoint_path{
ConcatPathComponent<PathChar>(tmp_dir.Path(), ORT_TSTR("e2e_ckpt_save_cpu"))};
ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path));
ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path, true));
// Check the ckpt files in the directory.
std::set<PathString> expected_file_names{

29
orttraining/orttraining/test/training_api/core/data_utils.h
View file

@ -3,14 +3,15 @@
#pragma once
#include <random>
#include <vector>
#include "core/framework/ort_value.h"
#include "core/framework/tensor.h"
#include "test/framework/test_utils.h"
#include "test/util/include/test_utils.h"
namespace onnxruntime {
namespace training {
namespace test {
namespace onnxruntime::training::test {
template <typename T>
void OrtValueToVec(const OrtValue& val, std::vector<T>& output) {
@ -39,6 +40,22 @@ void CudaOrtValueToCpuVec(const OrtValue& val, std::vector<T>& output,
output.assign(val_ptr, val_ptr + src_tensor.Shape().Size());
}
} // namespace test
} // namespace training
} // namespace onnxruntime
inline void GenerateRandomData(std::vector<float>& data) {
float scale = 1.f;
float mean = 0.f;
float seed = 123.f;
std::default_random_engine generator_float{gsl::narrow_cast<uint32_t>(seed)};
std::normal_distribution<float> distribution_float{mean, scale};
std::for_each(data.begin(), data.end(),
[&generator_float, &distribution_float](float& value) { value = distribution_float(generator_float); });
}
inline void GenerateRandomInput(gsl::span<const int64_t> dims, OrtValue& input) {
TensorShape shape(dims);
std::vector<float> data(shape.Size());
GenerateRandomData(data);
onnxruntime::test::CreateInputOrtValueOnCPU<float>(dims, data, &input);
}
} // namespace onnxruntime::training::test

37
orttraining/orttraining/test/training_api/core/training_api_tests.cc
View file

@ -2,14 +2,11 @@
// Licensed under the MIT License.
#include <thread>
#include <random>
#include "gtest/gtest.h"
#include "nlohmann/json.hpp"
#include "test/framework/test_utils.h"
#include "test/util/include/asserts.h"
#include "test/util/include/test_utils.h"
#include "core/framework/tensorprotoutils.h"
#include "orttraining/training_api/utils.h"
#include "orttraining/training_api/module.h"
@ -31,24 +28,6 @@ namespace {
#define MODEL_FOLDER ORT_TSTR("testdata/training_api/")
void GenerateRandomData(std::vector<float>& data) {
float scale = 1.f;
float mean = 0.f;
float seed = 123.f;
std::default_random_engine generator_float{gsl::narrow_cast<uint32_t>(seed)};
std::normal_distribution<float> distribution_float{mean, scale};
std::for_each(data.begin(), data.end(),
[&generator_float, &distribution_float](float& value) { value = distribution_float(generator_float); });
}
void GenerateRandomInput(gsl::span<const int64_t> dims, OrtValue& input) {
TensorShape shape(dims);
std::vector<float> data(shape.Size());
GenerateRandomData(data);
onnxruntime::test::CreateInputOrtValueOnCPU<float>(dims, data, &input);
}
void TestModuleExport(const std::vector<std::shared_ptr<IExecutionProvider>>& providers) {
auto training_model_uri = MODEL_FOLDER "training_model.onnx";
auto eval_model_uri = MODEL_FOLDER "eval_model.onnx";
@ -60,7 +39,7 @@ void TestModuleExport(const std::vector<std::shared_ptr<IExecutionProvider>>& pr
std::unique_ptr<Environment> env;
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
auto model = std::make_unique<onnxruntime::training::api::Module>(
ToUTF8String(training_model_uri), state.module_checkpoint_state.named_parameters, onnxruntime::SessionOptions(),
ToUTF8String(training_model_uri), &state, onnxruntime::SessionOptions(),
*env, providers, ToUTF8String(eval_model_uri));
auto test_dir = ORT_TSTR("export_model_for_inferencing_test_dir");
@ -138,10 +117,10 @@ void TestLRSchduler(const std::basic_string<ORTCHAR_T>& test_file_name, float in
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
const std::vector<std::shared_ptr<IExecutionProvider>> providers{onnxruntime::test::DefaultCudaExecutionProvider()};
auto model = std::make_unique<onnxruntime::training::api::Module>(
ToUTF8String(model_uri), state.module_checkpoint_state.named_parameters,
ToUTF8String(model_uri), &state,
session_option, *env, providers);
auto optim = std::make_shared<onnxruntime::training::api::Optimizer>(
ToUTF8String(optim_uri), model->NamedParameters(), session_option,
ToUTF8String(optim_uri), &state, session_option,
*env, providers);
OrtValue input, target;
@ -207,7 +186,7 @@ TEST(TrainingApiTest, ModuleParametersSize) {
std::unique_ptr<Environment> env;
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
auto model = std::make_unique<onnxruntime::training::api::Module>(ToUTF8String(model_uri),
state.module_checkpoint_state.named_parameters, session_option,
&state, session_option,
*env, std::vector<std::shared_ptr<IExecutionProvider>>());
size_t params_size = 0;
for (auto& param : model->Parameters()) {
@ -230,7 +209,7 @@ TEST(TrainingApiTest, ModuleCopyBufferToParameters) {
std::unique_ptr<Environment> env;
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
auto model = std::make_unique<onnxruntime::training::api::Module>(ToUTF8String(model_uri),
state.module_checkpoint_state.named_parameters, session_option,
&state, session_option,
*env, std::vector<std::shared_ptr<IExecutionProvider>>());
int64_t params_size = static_cast<int64_t>(model->GetParametersSize());
std::vector<float> expected_param_buffer(params_size);
@ -268,7 +247,7 @@ TEST(TrainingApiTest, ModuleTrainStep) {
std::unique_ptr<Environment> env;
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
auto model = std::make_unique<onnxruntime::training::api::Module>(ToUTF8String(model_uri),
state.module_checkpoint_state.named_parameters, session_option,
&state, session_option,
*env, std::vector<std::shared_ptr<IExecutionProvider>>());
ASSERT_EQ(model->GetTrainingModelOutputCount(), 1);
OrtValue input, target;
@ -340,10 +319,10 @@ TEST(TrainingApiTest, OptimStep) {
std::shared_ptr<IExecutionProvider> cpu_provider = onnxruntime::test::DefaultCpuExecutionProvider();
ASSERT_STATUS_OK(Environment::Create(nullptr, env));
auto model = std::make_unique<onnxruntime::training::api::Module>(
ToUTF8String(model_uri), state.module_checkpoint_state.named_parameters, session_option,
ToUTF8String(model_uri), &state, session_option,
*env, providers);
auto optim = std::make_unique<onnxruntime::training::api::Optimizer>(
ToUTF8String(optim_uri), model->NamedParameters(), session_option,
ToUTF8String(optim_uri), &state, session_option,
*env, providers);
OrtValue input, target;

135
orttraining/orttraining/test/training_api/core/training_capi_tests.cc Normal file
View file

@ -0,0 +1,135 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "gtest/gtest.h"
#include "onnxruntime_c_api.h"
#include "onnxruntime_training_c_api.h"
#include "onnxruntime_training_cxx_api.h"
#include "orttraining/training_api/checkpoint.h"
#include "orttraining/test/training_api/core/data_utils.h"
#include "test/util/include/temp_dir.h"
namespace onnxruntime::training::test {
#define MODEL_FOLDER ORT_TSTR("testdata/training_api/")
TEST(TrainingCApiTest, SaveCheckpoint) {
auto model_uri = MODEL_FOLDER "training_model.onnx";
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
Ort::TrainingSession training_session = Ort::TrainingSession(Ort::SessionOptions(), checkpoint_state, model_uri);
auto test_dir = ORT_TSTR("save_checkpoint_dir");
if (Env::Default().FolderExists(test_dir)) {
ORT_ENFORCE(Env::Default().DeleteFolder(test_dir).IsOK());
}
onnxruntime::test::TemporaryDirectory tmp_dir{test_dir};
PathString checkpoint_path{
ConcatPathComponent<PathChar>(tmp_dir.Path(), ORT_TSTR("new_checkpoint.ckpt"))};
Ort::CheckpointState::SaveCheckpoint(checkpoint_state, checkpoint_path);
Ort::CheckpointState new_checkpoint_state = Ort::CheckpointState::LoadCheckpoint(checkpoint_path);
Ort::TrainingSession new_training_session = Ort::TrainingSession(Ort::SessionOptions(), new_checkpoint_state, model_uri);
}
TEST(TrainingCApiTest, AddIntProperty) {
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
int64_t value = 365 * 24;
checkpoint_state.AddProperty("hours in a year", value);
auto property = checkpoint_state.GetProperty("hours in a year");
ASSERT_EQ(std::get<int64_t>(property), value);
}
TEST(TrainingCApiTest, AddFloatProperty) {
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
float value = 3.14f;
checkpoint_state.AddProperty("pi", value);
auto property = checkpoint_state.GetProperty("pi");
ASSERT_EQ(std::get<float>(property), value);
}
TEST(TrainingCApiTest, AddStringProperty) {
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
std::string value("onnxruntime");
checkpoint_state.AddProperty("framework", value);
auto property = checkpoint_state.GetProperty("framework");
ASSERT_EQ(std::get<std::string>(property), value);
}
TEST(TrainingCApiTest, InputNames) {
auto model_uri = MODEL_FOLDER "training_model.onnx";
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
Ort::TrainingSession training_session = Ort::TrainingSession(Ort::SessionOptions(), checkpoint_state, model_uri);
const auto input_names = training_session.InputNames(true);
ASSERT_EQ(input_names.size(), 2U);
ASSERT_EQ(input_names.front(), "input-0");
ASSERT_EQ(input_names.back(), "labels");
}
TEST(TrainingCApiTest, OutputNames) {
auto model_uri = MODEL_FOLDER "training_model.onnx";
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
Ort::TrainingSession training_session = Ort::TrainingSession(Ort::SessionOptions(), checkpoint_state, model_uri);
const auto output_names = training_session.OutputNames(true);
ASSERT_EQ(output_names.size(), 1U);
ASSERT_EQ(output_names.front(), "onnx::loss::21273");
}
TEST(TrainingCApiTest, ToBuffer) {
auto model_uri = MODEL_FOLDER "training_model.onnx";
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
Ort::TrainingSession training_session = Ort::TrainingSession(Ort::SessionOptions(), checkpoint_state, model_uri);
Ort::Value buffer = training_session.ToBuffer(true);
ASSERT_TRUE(buffer.IsTensor());
auto tensor_info = buffer.GetTensorTypeAndShapeInfo();
auto shape = tensor_info.GetShape();
ASSERT_EQ(shape.size(), 1U);
ASSERT_EQ(shape.front(), static_cast<int64_t>(397510));
buffer = training_session.ToBuffer(false);
ASSERT_TRUE(buffer.IsTensor());
tensor_info = buffer.GetTensorTypeAndShapeInfo();
shape = tensor_info.GetShape();
ASSERT_EQ(shape.size(), 1U);
ASSERT_EQ(shape.front(), static_cast<int64_t>(397510));
}
TEST(TrainingCApiTest, FromBuffer) {
auto model_uri = MODEL_FOLDER "training_model.onnx";
Ort::CheckpointState checkpoint_state = Ort::CheckpointState::LoadCheckpoint(MODEL_FOLDER "checkpoint.ckpt");
Ort::TrainingSession training_session = Ort::TrainingSession(Ort::SessionOptions(), checkpoint_state, model_uri);
OrtValue* buffer_impl = std::make_unique<OrtValue>().release();
GenerateRandomInput(std::array<int64_t, 1>{397510}, *buffer_impl);
Ort::Value buffer(buffer_impl);
training_session.FromBuffer(buffer);
}
} // namespace onnxruntime::training::test

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

@ -322,10 +322,10 @@ int RunTraining(const TestRunnerParameters& params) {
std::ostringstream oss;
oss << "ckpt_" << params.model_name << std::to_string(batch_idx);
PathString ckpt_file = ConcatPathComponent<PathChar>(params.output_dir, ToPathString(oss.str()));
Ort::CheckpointState::SaveCheckpoint(session, ckpt_file, true);
// TODO(baiju): enable adding more properties to checkpoint
// state_to_save.property_bag.AddProperty<int64_t>(std::string("epoch"), epoch);
checkpoint_state.AddProperty("epoch", epoch);
checkpoint_state.AddProperty("loss", *loss);
checkpoint_state.AddProperty("framework", "onnxruntime");
Ort::CheckpointState::SaveCheckpoint(checkpoint_state, ckpt_file);
}
batch_idx++;
}
@ -337,7 +337,7 @@ int RunTraining(const TestRunnerParameters& params) {
std::ostringstream oss;
oss << "ckpt_" << params.model_name;
PathString ckpt_file = ConcatPathComponent<PathChar>(params.output_dir, ToPathString(oss.str()));
Ort::CheckpointState::SaveCheckpoint(session, ckpt_file, true);
Ort::CheckpointState::SaveCheckpoint(checkpoint_state, ckpt_file);
auto end = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> duration_seconds = end - end_to_end_start;

11
orttraining/orttraining/training_api/checkpoint.cc
View file

@ -306,7 +306,7 @@ Status OrtSaveOptimizerStatesInternal(OptimizerCheckpointState& optimizer_state,
}
Status OrtSaveInternal(
CheckpointState& state, const PathString& checkpoint_path) {
CheckpointState& state, const PathString& checkpoint_path, const bool include_optimizer_state) {
LOGS_DEFAULT(INFO) << "Saving model checkpoint files to " << ToUTF8String(checkpoint_path);
LOGS_DEFAULT_IF(Env::Default().FolderExists(checkpoint_path), WARNING)
<< "Checkpoint directory exists - data may be overwritten.";
@ -316,7 +316,9 @@ Status OrtSaveInternal(
ORT_RETURN_IF_ERROR(OrtSaveModuleStatesInternal(state.module_checkpoint_state, checkpoint_path));
// Write optimizer state tensors files.
ORT_RETURN_IF_ERROR(OrtSaveOptimizerStatesInternal(state.optimizer_checkpoint_state, checkpoint_path));
if (include_optimizer_state) {
ORT_RETURN_IF_ERROR(OrtSaveOptimizerStatesInternal(state.optimizer_checkpoint_state, checkpoint_path));
}
// Write properties file
const PropertyBag& property_bag = state.property_bag;
@ -560,8 +562,9 @@ Status SaveCheckpoint(const std::vector<ONNX_NAMESPACE::TensorProto>& trainable_
return OrtSaveInternal(trainable_tensor_protos, non_trainable_tensor_protos, checkpoint_path);
}
Status SaveCheckpoint(CheckpointState& states, const PathString& checkpoint_path) {
return OrtSaveInternal(states, checkpoint_path);
Status SaveCheckpoint(CheckpointState& states, const PathString& checkpoint_path,
const bool include_optimizer_state) {
return OrtSaveInternal(states, checkpoint_path, include_optimizer_state);
}
Status LoadCheckpoint(const PathString& checkpoint_path, CheckpointState& checkpoint_states) {

4
orttraining/orttraining/training_api/checkpoint.h
View file

@ -57,8 +57,8 @@ struct CheckpointState {
* @return Status
* TODO: change state to const ref
*/
Status SaveCheckpoint(CheckpointState& state,
const PathString& checkpoint_path);
Status SaveCheckpoint(CheckpointState& state, const PathString& checkpoint_path,
const bool include_optimizer_state);
/**
* @brief Save ONNX initializers as ORT checkpoint.

12
orttraining/orttraining/training_api/checkpoint_property.h
View file

@ -64,6 +64,10 @@ struct PropertyBag {
return named_properties_.size();
}
bool HasProperty(const std::string& property_name) const {
return named_properties_.count(property_name);
}
private:
const InlinedVector<int32_t> supported_data_types{
ONNX_NAMESPACE::TensorProto::FLOAT,
@ -77,6 +81,14 @@ struct PropertyBag {
InlinedHashMap<std::string, PropertyDataType> named_properties_;
};
template <>
inline PropertyDataType PropertyBag::GetProperty<PropertyDataType>(const std::string& name) const {
auto it = named_properties_.find(name);
ORT_ENFORCE(it != named_properties_.end(), "No property named ", name);
return it->second;
}
} // namespace api
} // namespace training
} // namespace onnxruntime

69
orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h
View file

@ -10,16 +10,25 @@
ORT_RUNTIME_CLASS(TrainingSession); /// Type that enables performing training for the given user models.
ORT_RUNTIME_CLASS(CheckpointState); /// Type that holds the training states for the training session.
typedef enum OrtPropertyType {
OrtIntProperty = 0,
OrtFloatProperty = 1,
OrtStringProperty = 2,
} OrtPropertyType;
struct OrtTrainingApi {
/** \brief Load a checkpoint state from directory on disk into checkpoint_state.
*
* This function will parse a checkpoint directory, pull relevant files and load the training
* states into the checkpoint_state. This checkpoint state can then be used to create the
* state into the checkpoint_state. This checkpoint state can then be used to create the
* training session by invoking CreateTrainingSession. By doing so, the training session will resume
* training from the given checkpoint.
* training from the given checkpoint state.
* Note that the training session created with a checkpoint state uses this state to store the entire
* training state (including model parameters, its gradients, the optimizer states and the properties).
* As a result, it is required that the checkpoint state outlive the lifetime of the training session.
*
* \param[in] checkpoint_path Path to the checkpoint directory
* \param[out] checkpoint_state Checkpoint states that contains the states of the training session.
* \param[out] checkpoint_state Checkpoint state that contains the states of the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
@ -27,21 +36,20 @@ struct OrtTrainingApi {
ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path,
_Outptr_ OrtCheckpointState** checkpoint_state);
/** \brief Save the training session states to a checkpoint directory on disk.
/** \brief Save the given state to a checkpoint directory on disk.
*
* This function retrieves the training session states from the training session and serializes them
* to a checkpoint directory on disk. This checkpoint can later be loaded by invoking LoadCheckpoint
* to continue the training with the same states.
* This function serializes the provided checkpoint state to a directory on disk.
* This checkpoint can later be loaded by invoking LoadCheckpoint to continue the training with the same state.
*
* \param[in] checkpoint_path Path to the checkpoint directory
* \param[in] session The training session from where the checkpoint states are to be retrieved.
* \param[in] save_optimizer_state Boolean flag indicating whether or not to save the optimizer states to the checkpoint.
* \param[in] checkpoint_state The checkpoint state to save.
* \param[in] checkpoint_path Path to the checkpoint directory.
* \param[in] include_optimizer_state Flag to indicate whether to save the optimizer state or not.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _In_ const OrtTrainingSession* session,
bool save_optimizer_state);
ORT_API2_STATUS(SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* checkpoint_path,
const bool include_optimizer_state);
/** \brief Create a training session that can be used to begin or resume training.
*
@ -371,6 +379,43 @@ struct OrtTrainingApi {
*/
ORT_API2_STATUS(TrainingSessionGetEvalModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
_In_ OrtAllocator* allocator, _Outptr_ char** output);
/** \brief Adds the given property to the checkpoint state.
*
* Runtime properties such as epoch, training step, best score, and others can be added to the checkpoint
* state by the user if they desire by calling this function with the appropriate property name and
* value. The given property name must be unique to be able to successfully add the property.
*
* \param[in] checkpoint_state The checkpoint state which should hold the property.
* \param[in] property_name Unique name of the property being added.
* \param[in] property_type Type of the property associated with the given name.
* \param[in] property_value Property value associated with the given name.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(AddProperty, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _In_ enum OrtPropertyType property_type,
_In_ void* property_value);
/** \brief Gets the property value associated with the given name from the checkpoint state.
*
* Gets the property value from an existing entry in the checkpoint state. The property must
* exist in the checkpoint state to be able to retrieve it successfully.
* Property values are allocated on the heap. The user must free up the memory as needed by their application.
*
* \param[in] checkpoint_state The checkpoint state that is currently holding the property.
* \param[in] property_name Unique name of the property being retrieved.
* \param[in] allocator Allocator used to allocate the memory for the property_value.
* \param[out] property_type Type of the property associated with the given name.
* \param[out] property_value Property value associated with the given name.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetProperty, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _Inout_ OrtAllocator* allocator,
_Out_ enum OrtPropertyType* property_type, _Outptr_ void** property_value);
};
typedef struct OrtTrainingApi OrtTrainingApi;

90
orttraining/orttraining/training_api/include/onnxruntime_training_cxx_api.h
View file

@ -4,6 +4,7 @@
#pragma once
#include "onnxruntime_training_c_api.h"
#include <optional>
#include <variant>
namespace Ort::detail {
@ -36,21 +37,26 @@ ORT_DEFINE_TRAINING_RELEASE(TrainingSession);
} // namespace detail
// TODO(bmeswani): remove forward declaration when the SaveCheckpoint no longer depends on TrainingSession
class TrainingSession;
using Property = std::variant<int64_t, float, std::string>;
/** \brief Class that holds the state of the training session state
*
* Wraps OrtCheckpointState
*
*/
class CheckpointState : public detail::Base<OrtCheckpointState> {
private:
CheckpointState(OrtCheckpointState* checkpoint_state) { p_ = checkpoint_state; }
public:
// Construct the checkpoint state by loading the checkpoint by calling LoadCheckpoint
CheckpointState() = delete;
/** \brief Loads the checkpoint at provided path and returns the checkpoint state
*
* Wraps OrtTrainingApi::LoadCheckpoint
*
* \param[in] path_to_checkpoint Path to the checkpoint file to load
* \return CheckpointState object which holds the state of the training session parameters.
*/
static CheckpointState LoadCheckpoint(const std::basic_string<ORTCHAR_T>& path_to_checkpoint);
@ -58,12 +64,34 @@ class CheckpointState : public detail::Base<OrtCheckpointState> {
*
* Wraps OrtTrainingApi::SaveCheckpoint
*
* \param[in] checkpoint_state Training session checkpoint state to save to the checkpoint file
* \param[in] path_to_checkpoint Path to the checkpoint file to load
*/
static void SaveCheckpoint(const TrainingSession& session, const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
bool include_optimizer_states);
static void SaveCheckpoint(const CheckpointState& checkpoint_state,
const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
const bool include_optimizer_state = false);
/** \brief Adds the given property to the state.
*
* Wraps OrtTrainingApi::AddProperty
*
* \param[in] property_name Name of the property to add to the state.
* \param[in] property_value Value of the property to add to the state.
*/
void AddProperty(const std::string& property_name, const Property& property_value);
/** \brief Gets the property associated with the given name from the state.
*
* Wraps OrtTrainingApi::GetProperty
*
* \param[in] property_name Name of the property to get from the state.
* \return Property value associated with the property name.
*/
Property GetProperty(const std::string& property_name);
};
/** \brief Manage the training loop using this class
/** \brief Trainer class that provides training, evaluation and optimizer methods for
* executing ONNX models.
*
* Wraps OrtTrainingSession
*
@ -148,11 +176,61 @@ class TrainingSession : public detail::Base<OrtTrainingSession> {
*
* Wraps OrtTrainingApi::ExportModelForInferencing
*
* \param[in] inference_model_path Path to a location where the inference ready onnx model should be saved to.
* \param[in] graph_output_names Vector of output names that the inference model should have.
*/
void ExportModelForInferencing(const std::basic_string<ORTCHAR_T>& inference_model_path,
const std::vector<std::string>& graph_output_names);
/** \brief Gets the graph input names.
*
* Wraps OrtTrainingApi::TrainingSessionGetTrainingModelInputName,
* OrtTrainingApi::TrainingSessionGetEvalModelInputName,
* OrtTrainingApi::TrainingSessionGetTrainingModelInputCount
* OrtTrainingApi::TrainingSessionGetEvalModelInputCount
*
* \param[in] training Whether the training model input names are requested or eval model input names.
* \return Graph input names for either the training model or the eval model.
*
*/
std::vector<std::string> InputNames(const bool training);
/** \brief Gets the graph output names.
*
* Wraps OrtTrainingApi::TrainingSessionGetTrainingModelOutputName,
* OrtTrainingApi::TrainingSessionGetEvalModelOutputName,
* OrtTrainingApi::TrainingSessionGetTrainingModelOutputCount
* OrtTrainingApi::TrainingSessionGetEvalModelOutputCount
*
* \param[in] training Whether the training model output names are requested or eval model output names.
* \return Graph output names for either the training model or the eval model.
*/
std::vector<std::string> OutputNames(const bool training);
/** \brief Copies the training session model parameters to a contiguous buffer
*
* Wraps OrtTrainingApi::CopyParametersToBuffer
*
* \param[in] only_trainable Whether to only copy trainable parameters or to copy all parameters.
* \return Contiguous buffer to the model parameters.
*/
Value ToBuffer(const bool only_trainable);
/** \brief Loads the training session model parameters from a contiguous buffer
*
* Wraps OrtTrainingApi::CopyBufferToParameters
*
* \param[in] buffer Contiguous buffer to load the parameters from.
*/
void FromBuffer(Value& buffer);
};
/** \brief Sets the given seed for random number generation.
*
* Wraps OrtTrainingApi::SetSeed
*
* \param[in] seed Manual seed to use for random number generation.
*/
void SetSeed(const int64_t seed);
} // namespace Ort

151
orttraining/orttraining/training_api/include/onnxruntime_training_cxx_inline.h
View file

@ -82,16 +82,105 @@ inline void TrainingSession::OptimizerStep() {
ThrowOnError(GetTrainingApi().OptimizerStep(p_, run_options));
}
inline std::vector<std::string> TrainingSession::InputNames(const bool training) {
auto& input_count_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelInputCount
: GetTrainingApi().TrainingSessionGetEvalModelInputCount;
auto& input_name_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelInputName
: GetTrainingApi().TrainingSessionGetEvalModelInputName;
size_t input_count = 0;
ThrowOnError(input_count_function(p_, &input_count));
std::vector<std::string> input_names(input_count);
AllocatorWithDefaultOptions allocator;
for (size_t index = 0; index < input_count; ++index) {
char* input_name;
ThrowOnError(input_name_function(p_, index, allocator, &input_name));
input_names[index] = std::string(input_name);
allocator.Free(input_name);
}
return input_names;
}
inline std::vector<std::string> TrainingSession::OutputNames(const bool training) {
auto& output_count_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelOutputCount
: GetTrainingApi().TrainingSessionGetEvalModelOutputCount;
auto& output_name_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelOutputName
: GetTrainingApi().TrainingSessionGetEvalModelOutputName;
size_t output_count = 0;
ThrowOnError(output_count_function(p_, &output_count));
std::vector<std::string> output_names(output_count);
AllocatorWithDefaultOptions allocator;
for (size_t index = 0; index < output_count; ++index) {
char* output_name;
ThrowOnError(output_name_function(p_, index, allocator, &output_name));
output_names[index] = std::string(output_name);
allocator.Free(output_name);
}
return output_names;
}
inline Value TrainingSession::ToBuffer(const bool only_trainable) {
size_t buffer_size = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &buffer_size, only_trainable));
std::array<int64_t, 1> buffer_shape{static_cast<int64_t>(buffer_size)};
AllocatorWithDefaultOptions allocator;
Value buffer = Value::CreateTensor(allocator, buffer_shape.data(), 1U,
ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);
ThrowOnError(GetTrainingApi().CopyParametersToBuffer(p_, buffer, only_trainable));
return buffer;
}
inline void TrainingSession::FromBuffer(Value& buffer) {
if (!buffer.IsTensor()) {
ThrowStatus(Status("Incorrect buffer received. Expected a tensor buffer.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
auto tensor_info = buffer.GetTensorTypeAndShapeInfo();
auto buffer_shape = tensor_info.GetShape();
if (buffer_shape.size() != 1U) {
ThrowStatus(Status("Incorrect buffer received. Expected a contiguous tensor buffer.",
OrtErrorCode::ORT_INVALID_ARGUMENT));
}
auto buffer_size = buffer_shape.front();
size_t session_buffer_size_trainable_only = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &session_buffer_size_trainable_only, true));
if (buffer_size == static_cast<int64_t>(session_buffer_size_trainable_only)) {
ThrowOnError(GetTrainingApi().CopyBufferToParameters(p_, buffer, true));
return;
}
size_t session_buffer_size = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &session_buffer_size, false));
if (buffer_size != static_cast<int64_t>(session_buffer_size)) {
ThrowStatus(Status("Incorrect buffer size received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
ThrowOnError(GetTrainingApi().CopyBufferToParameters(p_, buffer, false));
}
inline CheckpointState CheckpointState::LoadCheckpoint(const std::basic_string<ORTCHAR_T>& path_to_checkpoint) {
OrtCheckpointState* checkpoint_state;
ThrowOnError(GetTrainingApi().LoadCheckpoint(path_to_checkpoint.c_str(), &checkpoint_state));
return CheckpointState(checkpoint_state);
}
inline void CheckpointState::SaveCheckpoint(const TrainingSession& session,
inline void CheckpointState::SaveCheckpoint(const CheckpointState& checkpoint_states,
const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
bool include_optimizer_states) {
ThrowOnError(GetTrainingApi().SaveCheckpoint(path_to_checkpoint.c_str(), session, include_optimizer_states));
const bool include_optimizer_state) {
ThrowOnError(GetTrainingApi().SaveCheckpoint(checkpoint_states, path_to_checkpoint.c_str(),
include_optimizer_state));
}
inline void TrainingSession::ExportModelForInferencing(const std::basic_string<ORTCHAR_T>& inference_model_path,
@ -109,4 +198,60 @@ inline void SetSeed(const int64_t seed) {
ThrowOnError(GetTrainingApi().SetSeed(seed));
}
inline void CheckpointState::AddProperty(const std::string& property_name, const Property& property_value) {
if (std::holds_alternative<int64_t>(property_value)) {
int64_t value = std::get<int64_t>(property_value);
void* value_p = &value;
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtIntProperty, value_p));
} else if (std::holds_alternative<float>(property_value)) {
float value = std::get<float>(property_value);
void* value_p = &value;
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtFloatProperty, value_p));
} else if (std::holds_alternative<std::string>(property_value)) {
std::string value = std::get<std::string>(property_value);
auto buffer = std::make_unique<char[]>(value.length() + 1).release();
memcpy(buffer, value.c_str(), value.length());
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtStringProperty, buffer));
} else {
ThrowStatus(Status("Unknown property type received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
}
inline Property CheckpointState::GetProperty(const std::string& property_name) {
void* property_value = nullptr;
OrtPropertyType property_type;
AllocatorWithDefaultOptions allocator;
ThrowOnError(GetTrainingApi().GetProperty(p_, property_name.c_str(), allocator, &property_type, &property_value));
Property property;
switch (property_type) {
case OrtPropertyType::OrtIntProperty: {
auto value_p = reinterpret_cast<int64_t*>(property_value);
property = *value_p;
allocator.Free(property_value);
break;
}
case OrtPropertyType::OrtFloatProperty: {
auto value_p = reinterpret_cast<float*>(property_value);
property = *value_p;
allocator.Free(property_value);
break;
}
case OrtPropertyType::OrtStringProperty: {
auto value_p = reinterpret_cast<char*>(property_value);
property = std::string(value_p);
allocator.Free(property_value);
break;
}
default: {
ThrowStatus(Status("Unknown property type received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
break;
}
}
return property;
}
} // namespace Ort

4
orttraining/orttraining/training_api/lr_scheduler.h
View file

@ -30,11 +30,11 @@ struct LRSchedulerBase {
protected:
int64_t GetStepInternal() {
return optim_->optimizer_state_.step;
return optim_->optimizer_state_->step;
}
float GetInitialLRInternal() {
return optim_->optimizer_state_.initial_lr;
return optim_->optimizer_state_->initial_lr;
}
private:

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

@ -1,6 +1,8 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_api/module.h"
#include "core/common/safeint.h"
#include "core/common/string_utils.h"
#include "core/framework/execution_provider.h"
@ -9,7 +11,7 @@
#include "core/session/onnxruntime_session_options_config_keys.h"
#include "core/graph/graph_utils.h"
#include "orttraining/training_api/module.h"
#include "orttraining/training_api/checkpoint.h"
#include "orttraining/training_api/utils.h"
using namespace onnxruntime;
@ -148,12 +150,12 @@ Status Parameter::ResetGrad() {
}
Module::Module(const std::string& train_model_path_or_bytes,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
CheckpointState* state,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
const std::optional<std::string>& eval_model_path_or_bytes)
: named_parameters_{named_parameters} {
: state_{state} {
// Enforce weight prepacking is disabled
// If user explicitly enabled weight prepacking then return error.
// Default value is enabled. Therefore, explicitly disable it if the value is not set by user.
@ -161,7 +163,8 @@ Module::Module(const std::string& train_model_path_or_bytes,
if (session_options.config_options.TryGetConfigEntry(kOrtSessionOptionsConfigDisablePrepacking, disable_prepacking)) {
ORT_ENFORCE(disable_prepacking == "1", "Prepacking is not supported for training scenarios.");
} else {
const_cast<SessionOptions&>(session_options).config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "1";
const_cast<SessionOptions&>(session_options)
.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "1";
}
train_sess_ = std::make_unique<onnxruntime::InferenceSession>(session_options, env);
@ -171,6 +174,9 @@ Module::Module(const std::string& train_model_path_or_bytes,
}
ORT_THROW_IF_ERROR(train_sess_->Initialize());
// Make sure that the checkpoint state can copy tensors
state_->module_checkpoint_state.train_session_data_transfer_mgr = &train_sess_->GetDataTransferManager();
// Extract model input and output names
std::vector<std::string> train_input_names, train_output_names;
utils::GetGraphInputOutputNames(train_sess_, train_input_names, train_output_names);
@ -181,8 +187,8 @@ Module::Module(const std::string& train_model_path_or_bytes,
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 = state_->module_checkpoint_state.named_parameters.find(input_name);
if (it != state_->module_checkpoint_state.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);
@ -211,8 +217,8 @@ Module::Module(const std::string& train_model_path_or_bytes,
// 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) {
auto params_iter = named_parameters_.find(param_name);
ORT_ENFORCE(params_iter != named_parameters_.end());
auto params_iter = state_->module_checkpoint_state.named_parameters.find(param_name);
ORT_ENFORCE(params_iter != state_->module_checkpoint_state.named_parameters.end());
// Retrieve the target device for "param_name"
InlinedVector<SessionState::NodeInfo> node_info_vec;
@ -280,7 +286,8 @@ Module::Module(const std::string& train_model_path_or_bytes,
// TODO: Add the checks instead of making assumptions??
std::vector<std::string> eval_user_input_names, eval_param_input_names;
for (const auto& input_name : eval_input_names_) {
if (named_parameters_.find(input_name) != named_parameters_.end()) {
if (state_->module_checkpoint_state.named_parameters.find(input_name) !=
state_->module_checkpoint_state.named_parameters.end()) {
// it is a parameter
eval_param_input_names.emplace_back(input_name);
continue;
@ -316,13 +323,14 @@ std::string Module::GetTrainingModelOutputName(size_t index) const {
}
std::string Module::GetEvalModelOutputName(size_t index) const {
ORT_ENFORCE(index < eval_output_names_.size(), "Eval output name index out of range. Expected in range [0-", eval_output_names_.size(), "). Actual: ", index);
ORT_ENFORCE(index < eval_output_names_.size(), "Eval output name index out of range. Expected in range [0-",
eval_output_names_.size(), "). Actual: ", index);
return eval_output_names_.at(index);
}
size_t Module::GetParametersSize(const bool trainable_only) const {
SafeInt<size_t> parameters_size = 0;
for (const auto& it : named_parameters_) {
for (const auto& it : state_->module_checkpoint_state.named_parameters) {
if (trainable_only && !it.second->RequiresGrad()) {
continue;
}
@ -333,12 +341,16 @@ size_t Module::GetParametersSize(const bool trainable_only) const {
std::vector<std::shared_ptr<Parameter>> Module::Parameters() const {
std::vector<std::shared_ptr<Parameter>> params;
for (auto& it : named_parameters_) {
for (auto& it : state_->module_checkpoint_state.named_parameters) {
params.push_back(it.second);
}
return params;
}
std::unordered_map<std::string, std::shared_ptr<Parameter>> Module::NamedParameters() const {
return state_->module_checkpoint_state.named_parameters;
}
Status Module::CopyParametersToBuffer(OrtValue& parameters_buffer, const bool trainable_only) {
ORT_ENFORCE(parameters_buffer.IsAllocated(), "Parameters buffer should be pre-allocated.");
ORT_ENFORCE(parameters_buffer.IsTensor(), "Parameters buffer should be of tensor type.");
@ -353,7 +365,7 @@ Status Module::CopyParametersToBuffer(OrtValue& parameters_buffer, const bool tr
size_t offset = 0;
for (const auto& param_name : weight_names_) {
auto& param = named_parameters_.at(param_name);
auto& param = state_->module_checkpoint_state.named_parameters.at(param_name);
if (trainable_only && !param->RequiresGrad()) {
continue;
}
@ -396,7 +408,7 @@ Status Module::CopyBufferToParameters(OrtValue& parameters_buffer, const bool tr
size_t offset = 0;
for (const auto& param_name : weight_names_) {
auto& param = named_parameters_.at(param_name);
auto& param = state_->module_checkpoint_state.named_parameters.at(param_name);
if (trainable_only && !param->RequiresGrad()) {
continue;
}
@ -459,17 +471,6 @@ Status Module::EvalStep(const std::vector<OrtValue>& inputs, std::vector<OrtValu
return Status::OK();
}
Status Module::GetStateDict(ModuleCheckpointState& module_checkpoint_state) {
module_checkpoint_state.named_parameters = NamedParameters();
// Pass the training session data transfer manager for data copying when saving.
// An alternative is, we can do copy at this stage.
ORT_RETURN_IF_NOT(train_sess_, "training session not initialized");
const DataTransferManager& sess_data_transfer_manager = train_sess_->GetDataTransferManager();
module_checkpoint_state.train_session_data_transfer_mgr = &sess_data_transfer_manager;
return Status::OK();
}
Status Module::ExportModelForInferencing(const std::string& inference_model_path,
gsl::span<const std::string> graph_output_names) const {
ORT_RETURN_IF(!eval_sess_ || eval_model_path_.empty(),
@ -488,7 +489,7 @@ Status Module::ExportModelForInferencing(const std::string& inference_model_path
// The cloned model's inputs are transformed such that the model has only user defined inputs. All parameters
// are moved to be constant initializers for the model.
ORT_RETURN_IF_ERROR(TransformModelInputsForInference(inference_model->MainGraph(), named_parameters_,
ORT_RETURN_IF_ERROR(TransformModelInputsForInference(inference_model->MainGraph(), state_->module_checkpoint_state.named_parameters,
eval_sess_->GetDataTransferManager()));
// Save the model at desired location.

13
orttraining/orttraining/training_api/module.h
View file

@ -52,12 +52,14 @@ struct ModuleCheckpointState {
const DataTransferManager* train_session_data_transfer_mgr;
};
struct CheckpointState;
struct Module {
public:
// 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,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
CheckpointState* state,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
@ -66,9 +68,7 @@ struct Module {
// Return the trainable/nontrainable parameters
std::vector<std::shared_ptr<Parameter>> Parameters() const;
std::unordered_map<std::string, std::shared_ptr<Parameter>> NamedParameters() const {
return named_parameters_;
}
std::unordered_map<std::string, std::shared_ptr<Parameter>> NamedParameters() const;
// Reset and release the gradient buffer of all trainable params lazily.
Status LazyResetGrad();
@ -81,9 +81,6 @@ struct Module {
// and take in a separate inference graph, while sharing the parameters
Status EvalStep(const std::vector<OrtValue>& inputs, std::vector<OrtValue>& outputs);
// Return the states of the module as a map.
Status GetStateDict(ModuleCheckpointState& module_checkpoint_states);
// Returns the output count for training graph
size_t GetTrainingModelOutputCount() const noexcept;
@ -133,7 +130,7 @@ struct Module {
std::vector<OrtValue> weights_;
std::vector<OrtValue> gradients_;
bool accumulate_gradient_ = false;
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters_;
CheckpointState* state_; // Non owning pointer to the state.
std::string eval_model_path_;
size_t train_user_input_count_ = 0U;
size_t eval_user_input_count_ = 0U;

97
orttraining/orttraining/training_api/onnxruntime_training_c_api.cc
View file

@ -43,7 +43,7 @@ ORT_API_STATUS_IMPL(OrtTrainingApis::CreateTrainingSession, _In_ const OrtEnv* e
env->GetEnvironment(),
options == nullptr ? onnxruntime::SessionOptions() : options->value,
options == nullptr ? ProvidersType() : CreateProviders(options->provider_factories),
chkpt_state->module_checkpoint_state.named_parameters,
chkpt_state,
onnxruntime::training::api::ModelIdentifiers(
onnxruntime::ToUTF8String(train_model_path),
eval_model_path ? std::optional<std::string>(onnxruntime::ToUTF8String(eval_model_path))
@ -270,13 +270,12 @@ ORT_API_STATUS_IMPL(OrtTrainingApis::LoadCheckpoint, _In_ const ORTCHAR_T* check
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtTrainingApis::SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path,
_In_ const OrtTrainingSession* sess, bool save_optimizer_state) {
ORT_API_STATUS_IMPL(OrtTrainingApis::SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state,
_In_ const ORTCHAR_T* checkpoint_path, const bool include_optimizer_state) {
API_IMPL_BEGIN
auto session = reinterpret_cast<const 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));
auto chkpt_state = reinterpret_cast<onnxruntime::training::api::CheckpointState*>(checkpoint_state);
ORT_API_RETURN_IF_STATUS_NOT_OK(
onnxruntime::training::api::SaveCheckpoint(*chkpt_state, checkpoint_path, include_optimizer_state));
return nullptr;
API_IMPL_END
@ -403,6 +402,88 @@ ORT_API_STATUS_IMPL(OrtTrainingApis::TrainingSessionGetEvalModelInputName, _In_
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtTrainingApis::AddProperty, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _In_ enum OrtPropertyType property_type,
_In_ void* property_value) {
API_IMPL_BEGIN
OrtStatus* status = nullptr;
auto chkpt_state = reinterpret_cast<onnxruntime::training::api::CheckpointState*>(checkpoint_state);
switch (property_type) {
case OrtPropertyType::OrtIntProperty: {
int64_t* value = reinterpret_cast<int64_t*>(property_value);
chkpt_state->property_bag.AddProperty(property_name, *value);
break;
}
case OrtPropertyType::OrtFloatProperty: {
float* value = reinterpret_cast<float*>(property_value);
chkpt_state->property_bag.AddProperty(property_name, *value);
break;
}
case OrtPropertyType::OrtStringProperty: {
char* value = reinterpret_cast<char*>(property_value);
chkpt_state->property_bag.AddProperty(property_name, value);
break;
}
default: {
std::ostringstream stream;
stream << "Given property type: " << property_type << " is not supported.";
status = OrtApis::CreateStatus(ORT_FAIL, stream.str().c_str());
break;
}
}
return status;
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtTrainingApis::GetProperty, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _Inout_ OrtAllocator* allocator,
_Out_ enum OrtPropertyType* property_type, _Outptr_ void** property_value) {
API_IMPL_BEGIN
OrtStatus* status = nullptr;
auto chkpt_state = reinterpret_cast<const onnxruntime::training::api::CheckpointState*>(checkpoint_state);
const auto value = chkpt_state->property_bag.GetProperty<
onnxruntime::training::api::PropertyDataType>(property_name);
if (std::holds_alternative<int64_t>(value)) {
int64_t* value_p = reinterpret_cast<int64_t*>(allocator->Alloc(allocator, sizeof(int64_t)));
if (!value_p) {
return OrtApis::CreateStatus(ORT_FAIL, "Int property value buffer allocation failed.");
}
*value_p = std::get<int64_t>(value);
*(reinterpret_cast<int64_t**>(property_value)) = value_p;
*property_type = OrtPropertyType::OrtIntProperty;
} else if (std::holds_alternative<float>(value)) {
float* value_p = reinterpret_cast<float*>(allocator->Alloc(allocator, sizeof(float)));
if (!value_p) {
return OrtApis::CreateStatus(ORT_FAIL, "Float property value buffer allocation failed.");
}
*value_p = std::get<float>(value);
*(reinterpret_cast<float**>(property_value)) = value_p;
*property_type = OrtPropertyType::OrtFloatProperty;
} else if (std::holds_alternative<std::string>(value)) {
auto property_value_str = std::get<std::string>(value);
// property_value_str.length() + 1 for null termination of c strings
auto buffer = reinterpret_cast<char*>(allocator->Alloc(allocator, property_value_str.length() + 1));
memcpy(buffer, property_value_str.c_str(), property_value_str.length());
buffer[property_value_str.length()] = '\0';
*(reinterpret_cast<char**>(property_value)) = buffer;
*property_type = OrtPropertyType::OrtStringProperty;
} else {
std::ostringstream stream;
stream << "Unknown type for property: " << property_name;
status = OrtApis::CreateStatus(ORT_FAIL, stream.str().c_str());
}
return status;
API_IMPL_END
}
static constexpr OrtTrainingApi ort_training_api = {
// NOTE: The C# bindings depend on the API order within this struct. Since Training APIs are not officially
// released, it is OK to change the order here, however a corresponding matching change should also be done in the
@ -433,6 +514,8 @@ static constexpr OrtTrainingApi ort_training_api = {
&OrtTrainingApis::TrainingSessionGetEvalModelInputCount,
&OrtTrainingApis::TrainingSessionGetTrainingModelInputName,
&OrtTrainingApis::TrainingSessionGetEvalModelInputName,
&OrtTrainingApis::AddProperty,
&OrtTrainingApis::GetProperty,
};
ORT_API(const OrtTrainingApi*, OrtTrainingApis::GetTrainingApi, uint32_t) {

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

@ -1,14 +1,15 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_api/optimizer.h"
#include "core/framework/execution_provider.h"
#include "core/framework/TensorSeq.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include "core/session/inference_session.h"
#include "core/session/environment.h"
#include "orttraining/training_api/checkpoint.h"
#include "orttraining/training_api/utils.h"
#include "orttraining/training_api/optimizer.h"
namespace onnxruntime {
namespace training {
@ -71,9 +72,9 @@ Status GraphInputsAreExpected(gsl::span<std::string> actual_graph_inputs,
} // namespace
Status Optimizer::GenerateMomentumNamedStates() {
auto& param_named_optimizer_states = optimizer_state_.param_named_optimizer_states;
auto& param_named_optimizer_states = optimizer_state_->param_named_optimizer_states;
auto& optim_sess_state = optim_sess_->GetSessionState();
for (auto& pair : named_parameters_) {
for (auto& pair : state_->module_checkpoint_state.named_parameters) {
if (pair.second->RequiresGrad()) {
param_named_optimizer_states.insert({pair.first, ParameterOptimizerState()});
ParameterOptimizerState& cur_param_optimizer_states = param_named_optimizer_states[pair.first];
@ -91,12 +92,12 @@ Status Optimizer::GenerateMomentumNamedStates() {
// Constructs the ortvalue inputs to be fed to the graph at each step
Status Optimizer::ConstructInputs() {
if (optimizer_type_ == OptimizerType::AdamW) {
auto& param_named_optimizer_states = optimizer_state_.param_named_optimizer_states;
auto& param_named_optimizer_states = optimizer_state_->param_named_optimizer_states;
std::vector<Tensor> params, grads, first_order_moments, second_order_moments;
// Collect all the non user defined inputs from the named_parameters_.
for (auto& [parameter_name, parameter] : named_parameters_) {
// Collect all the non user defined inputs from the state_->module_checkpoint_state.named_parameters.
for (auto& [parameter_name, parameter] : state_->module_checkpoint_state.named_parameters) {
if (parameter->RequiresGrad()) {
// Collect parameters and prepare for tensorseq creation
auto* param_tensor = parameter->Data().GetMutable<Tensor>();
@ -152,11 +153,16 @@ Status Optimizer::ConstructInputs() {
}
Optimizer::Optimizer(const std::string& optim_path_or_bytes,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
CheckpointState* state,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers)
: optim_sess_(std::make_unique<InferenceSession>(session_options, env)), named_parameters_(named_parameters) {
: optim_sess_(std::make_unique<InferenceSession>(session_options, env)), state_(state) {
if (state_->optimizer_checkpoint_state.group_named_optimizer_states.empty()) {
state_->optimizer_checkpoint_state.group_named_optimizer_states.insert(
{GROUP_ZERO_NAME, std::make_shared<GroupOptimizerState>()});
}
optimizer_state_ = state_->optimizer_checkpoint_state.group_named_optimizer_states.at(GROUP_ZERO_NAME);
for (const auto& execution_provider : providers) {
ORT_THROW_IF_ERROR(optim_sess_->RegisterExecutionProvider(execution_provider));
}
@ -164,6 +170,9 @@ Optimizer::Optimizer(const std::string& optim_path_or_bytes,
ORT_THROW_IF_ERROR(optim_sess_->Load(optim_path_or_bytes));
ORT_THROW_IF_ERROR(optim_sess_->Initialize());
// Make sure that the checkpoint state can copy tensors
state_->optimizer_checkpoint_state.optimizer_session_data_transfer_mgr = &optim_sess_->GetDataTransferManager();
utils::GetGraphInputOutputNames(optim_sess_, input_names_, output_names_);
if (optimizer_type_ == OptimizerType::AdamW) {
@ -178,11 +187,11 @@ Optimizer::Optimizer(const std::string& optim_path_or_bytes,
Status Optimizer::Step() {
OrtValue learning_rate_input, step_input;
utils::WrapInOrtValue<float>(optimizer_state_.learning_rate, &learning_rate_input);
utils::WrapInOrtValue<float>(optimizer_state_->learning_rate, &learning_rate_input);
// Use step count + 1 before running optimizer step.
// This is necessary since bias correction uses the step
// as a power. Using power of 0 is wrong.
utils::WrapInOrtValue<int64_t>(optimizer_state_.step + 1, &step_input);
utils::WrapInOrtValue<int64_t>(optimizer_state_->step + 1, &step_input);
std::vector<OrtValue> feeds({learning_rate_input, step_input});
feeds.insert(feeds.end(), inputs_.begin(), inputs_.end());
@ -192,7 +201,7 @@ Status Optimizer::Step() {
// extract step output and update
if (utils::GetValue<int64_t>(outputs[0]) == 1LL) {
optimizer_state_.step++;
optimizer_state_->step++;
}
return Status::OK();
@ -202,7 +211,7 @@ Status Optimizer::GetStateDict(OptimizerCheckpointState& optimizer_checkpoint_st
auto& grouped_optimizer_states = optimizer_checkpoint_state.group_named_optimizer_states;
// To support multiple groups, Optimizer constructor need accept informations for groupping.
grouped_optimizer_states.insert({GROUP_ZERO_NAME, std::make_shared<GroupOptimizerState>(optimizer_state_)});
grouped_optimizer_states.insert({GROUP_ZERO_NAME, std::make_shared<GroupOptimizerState>(*optimizer_state_)});
// Pass the optimizer session data transfer manager for data copying when saving.
// An alternative is, we can do copy at this stage.
@ -217,8 +226,8 @@ Status Optimizer::LoadStateDict(const OptimizerCheckpointState& optimizer_checkp
optimizer_checkpoint_states.group_named_optimizer_states.find(GROUP_ZERO_NAME);
ORT_ENFORCE(group_optimizer_state_it != optimizer_checkpoint_states.group_named_optimizer_states.cend(),
"Group 0 not found in the optimizer checkpoint states.");
optimizer_state_.initial_lr = group_optimizer_state_it->second->initial_lr;
optimizer_state_.step = group_optimizer_state_it->second->step;
optimizer_state_->initial_lr = group_optimizer_state_it->second->initial_lr;
optimizer_state_->step = group_optimizer_state_it->second->step;
// TODO(pengwa): restore the momentums state from checkpoint.
return Status::OK();

18
orttraining/orttraining/training_api/optimizer.h
View file

@ -49,6 +49,8 @@ enum class OptimizerType {
// Lamb,
};
struct CheckpointState;
struct Optimizer {
friend struct LRSchedulerBase;
@ -57,7 +59,7 @@ struct Optimizer {
// training ONNX model For each parameter, initialize the OptimizerState based
// on the graph input's ValueInfoProto if the parameter doesn't have it already.
Optimizer(const std::string& optim_path_or_bytes,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
CheckpointState* state,
const onnxruntime::SessionOptions& session_options,
const Environment& env,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers);
@ -69,23 +71,23 @@ struct Optimizer {
Status LoadStateDict(const OptimizerCheckpointState& optimizer_checkpoint_states);
Status SetLearningRate(float lr) {
optimizer_state_.learning_rate = lr;
optimizer_state_->learning_rate = lr;
return Status::OK();
}
float GetLearningRate() const noexcept {
return optimizer_state_.learning_rate;
return optimizer_state_->learning_rate;
}
Status SetInitialLearningRate(float initial_lr) {
optimizer_state_.initial_lr = initial_lr;
optimizer_state_.learning_rate = initial_lr;
optimizer_state_->initial_lr = initial_lr;
optimizer_state_->learning_rate = initial_lr;
return Status::OK();
}
private:
int64_t GetStep() const {
return optimizer_state_.step;
return optimizer_state_->step;
}
// Generates optimizer momentum states for applicable optimizer types
@ -97,8 +99,8 @@ struct Optimizer {
// TODO: load this info from checkpoint
OptimizerType optimizer_type_ = OptimizerType::AdamW;
std::unique_ptr<onnxruntime::InferenceSession> optim_sess_;
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters_;
GroupOptimizerState optimizer_state_;
CheckpointState* state_; // Non owning pointer to the state.
std::shared_ptr<GroupOptimizerState> optimizer_state_;
std::vector<std::string> input_names_;
std::vector<std::string> output_names_;
std::vector<OrtValue> inputs_;

12
orttraining/orttraining/training_api/ort_training_apis.h
View file

@ -44,8 +44,8 @@ ORT_API_STATUS_IMPL(SchedulerStep, _Inout_ OrtTrainingSession* sess);
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, _In_ const OrtTrainingSession* session,
bool save_optimizer_state);
ORT_API_STATUS_IMPL(SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* checkpoint_path,
const bool include_optimizer_state);
ORT_API_STATUS_IMPL(GetParametersSize, _Inout_ OrtTrainingSession* sess,
_Out_ size_t* out, bool trainable_only);
@ -76,4 +76,12 @@ ORT_API_STATUS_IMPL(TrainingSessionGetTrainingModelInputName, _In_ const OrtTrai
ORT_API_STATUS_IMPL(TrainingSessionGetEvalModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
_In_ OrtAllocator* allocator, _Outptr_ char** output);
ORT_API_STATUS_IMPL(AddProperty, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _In_ enum OrtPropertyType property_type,
_In_ void* property_value);
ORT_API_STATUS_IMPL(GetProperty, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _Inout_ OrtAllocator* allocator,
_Out_ enum OrtPropertyType* property_type, _Outptr_ void** property_value);
} // namespace OrtTrainingApis

17
orttraining/orttraining/training_api/training_session.cc
View file

@ -10,14 +10,14 @@ namespace api {
TrainingSession::TrainingSession(const Environment& session_env,
const SessionOptions& session_options,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& parameters,
CheckpointState* state,
const ModelIdentifiers& model_identifiers)
: named_parameters_{parameters},
module_{std::make_unique<Module>(model_identifiers.train_model, named_parameters_,
: state_{state},
module_{std::make_unique<Module>(model_identifiers.train_model, state_,
session_options, session_env, providers, model_identifiers.eval_model)},
optimizer_{model_identifiers.optim_model.has_value()
? std::make_unique<Optimizer>(
model_identifiers.optim_model.value(), named_parameters_,
model_identifiers.optim_model.value(), state_,
session_options, session_env, providers)
: std::unique_ptr<Optimizer>()} {}
@ -86,15 +86,6 @@ Status TrainingSession::OptimizerStep(const RunOptions&) {
return optimizer_->Step();
}
Status TrainingSession::CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state) const {
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();
}
Status TrainingSession::SetLearningRate(float learning_rate) noexcept {
ORT_RETURN_IF_NOT(optimizer_, "No optimizer session initialized.");
ORT_RETURN_IF_ERROR(optimizer_->SetLearningRate(learning_rate));

6
orttraining/orttraining/training_api/training_session.h
View file

@ -28,7 +28,7 @@ class TrainingSession {
TrainingSession(const Environment& session_env,
const SessionOptions& session_options,
const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& parameters,
CheckpointState* state,
const ModelIdentifiers& model_identifiers);
Status RegisterScheduler(const std::function<
@ -69,8 +69,6 @@ class TrainingSession {
Status SchedulerStep() noexcept;
Status CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state) const;
size_t GetParametersSize(const bool trainable_only = true) const;
Status CopyParametersToBuffer(OrtValue& parameters_buffer, const bool trainable_only = true);
@ -83,7 +81,7 @@ class TrainingSession {
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TrainingSession);
const std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters_;
CheckpointState* state_; // Non owning pointer to the checkpoint state. It must outlive the training session.
std::unique_ptr<Module> module_;
std::shared_ptr<Optimizer> optimizer_;
std::unique_ptr<LRSchedulerBase> scheduler_;