diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs index 1e2821ead4..f775f3ad49 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs @@ -52,6 +52,16 @@ namespace Microsoft.ML.OnnxRuntime NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtLoadCheckpoint(NativeOnnxValueHelper.GetPlatformSerializedString(checkpointPath), out handle)); } + /// + /// Saves the checkpoint + /// absolute path to the checkpoint file. + /// absolute path to the checkpoint file. + /// + public void SaveCheckpoint(string checkpointPath, bool includeOptimizerState = false) + { + NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtSaveCheckpoint(handle, NativeOnnxValueHelper.GetPlatformSerializedString(checkpointPath), includeOptimizerState)); + } + #region SafeHandle /// /// Overrides SafeHandle.ReleaseHandle() to properly dispose of diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/Training/NativeTrainingMethods.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/Training/NativeTrainingMethods.shared.cs index de55c5399c..5df0720022 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/Training/NativeTrainingMethods.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/NativeTrainingMethods.shared.cs @@ -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 /// /// Creates an instance of OrtSession with provided parameters /// - /// checkpoint string path - /// (Output) Loaded OrtCheckpointState instance + /// OrtCheckpointState instance to save + /// Checkpoint string path + /// Flag indicating whether to save the optimizer state. [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() diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/Training/OrtPropertyType.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/Training/OrtPropertyType.shared.cs new file mode 100644 index 0000000000..17505d6a27 --- /dev/null +++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/OrtPropertyType.shared.cs @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +namespace Microsoft.ML.OnnxRuntime +{ +#if __ENABLE_TRAINING_APIS__ + /// + /// Property types + /// + public enum OrtPropertyType + { + OrtIntProperty = 0, + OrtFloatProperty = 1, + OrtStringProperty = 2, + } +#endif +} diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs index 54a597a41e..9874863e91 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs @@ -319,16 +319,6 @@ namespace Microsoft.ML.OnnxRuntime } - /// - /// Saves a checkpoint to path. It can be loaded into - /// - /// Specify path for saving the checkpoint. - /// SFlag indicating whether to save optimizer state or not. - public void SaveCheckpoint(string path, bool saveOptimizerState = false) - { - NativeApiStatus.VerifySuccess(NativeTrainingMethods.OrtSaveCheckpoint(NativeOnnxValueHelper.GetPlatformSerializedString(path), _nativeHandle, saveOptimizerState)); - } - #endregion #region private methods diff --git a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/TrainingTest.cs b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/TrainingTest.cs index 3c1ba092c0..0d16203dbc 100644 --- a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/TrainingTest.cs +++ b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/TrainingTest.cs @@ -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); diff --git a/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java b/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java index 58f27e47cd..51641cc97d 100644 --- a/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java +++ b/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java @@ -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. - * - *

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. - * - *

\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. - * - *

\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); - */ - 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); } } diff --git a/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c b/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c index 591dae7888..938eb1c944 100644 --- a/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c +++ b/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c @@ -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 diff --git a/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c b/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c index 7678f28ffa..0e7ddeb356 100644 --- a/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c +++ b/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c @@ -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 diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index b9d5d78417..23b02372da 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -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> provider) + PyOptimizer(const std::string optimizer_model_uri, onnxruntime::training::api::CheckpointState* state, + std::vector> providers) : optimizer_() { auto env = GetTrainingEnv().GetORTEnv(); // XXX: We hope that env will be around when optimizer needs it. - optimizer_ = std::make_shared(optimizer_model_uri, - model->NamedParameters(), onnxruntime::SessionOptions(), - *env, provider); + optimizer_ = std::make_shared( + optimizer_model_uri, state, onnxruntime::SessionOptions(), *env, providers); } std::shared_ptr optimizer_; @@ -882,7 +881,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn py::class_ 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 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( - 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& 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 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 names; + for (size_t index = 0; index < count_method(); ++index) { + names.push_back(name_method(index)); + } + + return names; }); py::class_ 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& 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(property_name); + }) + .def("has_property", [](onnxruntime::training::api::CheckpointState* state, const std::string& property_name) { + return state->property_bag.HasProperty(property_name); + }); py::class_ 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> provider = GetExecutionProvidersForTrainingApis(device); + training_optimizer + .def(py::init([](const std::string optimizer_model_uri, + onnxruntime::training::api::CheckpointState* state, + OrtDevice device) { + std::vector> providers = GetExecutionProvidersForTrainingApis(device); - return std::make_unique( - optimizer_model_uri, - model, provider); - })) + return std::make_unique(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& trainable_tensor_protos_pybytes, + const std::vector& non_trainable_tensor_protos_pybytes, + const std::string& checkpoint_path) { + std::vector trainable_tensor_protos(trainable_tensor_protos_pybytes.size()); + std::vector non_trainable_tensor_protos(non_trainable_tensor_protos_pybytes.size()); + + auto parse_pybytes_to_tensor_proto = + [](const std::vector& tensor_protos_pybytes, std::vector& 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& trainable_tensor_protos_pybytes, - const std::vector& non_trainable_tensor_protos_pybytes, - const std::string& checkpoint_path) { - std::vector trainable_tensor_protos(trainable_tensor_protos_pybytes.size()); - std::vector non_trainable_tensor_protos(non_trainable_tensor_protos_pybytes.size()); - - auto parse_pybytes_to_tensor_proto = - [](const std::vector& tensor_protos_pybytes, std::vector& 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; diff --git a/orttraining/orttraining/python/training/api/README.md b/orttraining/orttraining/python/training/api/README.md index 0738d718f6..ccd72c101f 100644 --- a/orttraining/orttraining/python/training/api/README.md +++ b/orttraining/orttraining/python/training/api/README.md @@ -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() # 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() -# 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") ``` diff --git a/orttraining/orttraining/python/training/api/__init__.py b/orttraining/orttraining/python/training/api/__init__.py index 4f7499ed30..bc5def0df2 100644 --- a/orttraining/orttraining/python/training/api/__init__.py +++ b/orttraining/orttraining/python/training/api/__init__.py @@ -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", +] diff --git a/orttraining/orttraining/python/training/api/checkpoint_state.py b/orttraining/orttraining/python/training/api/checkpoint_state.py index 99d7554dc5..3ee82a9360 100644 --- a/orttraining/orttraining/python/training/api/checkpoint_state.py +++ b/orttraining/orttraining/python/training/api/checkpoint_state.py @@ -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) diff --git a/orttraining/orttraining/python/training/api/lr_scheduler.py b/orttraining/orttraining/python/training/api/lr_scheduler.py index 5783ee316d..645d5490bd 100644 --- a/orttraining/orttraining/python/training/api/lr_scheduler.py +++ b/orttraining/orttraining/python/training/api/lr_scheduler.py @@ -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() diff --git a/orttraining/orttraining/python/training/api/module.py b/orttraining/orttraining/python/training/api/module.py index da3889ac0e..2c87924e49 100644 --- a/orttraining/orttraining/python/training/api/module.py +++ b/orttraining/orttraining/python/training/api/module.py @@ -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) diff --git a/orttraining/orttraining/python/training/api/optimizer.py b/orttraining/orttraining/python/training/api/optimizer.py index 3a89b27e9f..e62ee71441 100644 --- a/orttraining/orttraining/python/training/api/optimizer.py +++ b/orttraining/orttraining/python/training/api/optimizer.py @@ -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() diff --git a/orttraining/orttraining/python/training/artifacts.py b/orttraining/orttraining/python/training/artifacts.py index bcbbbb75f0..ea519da89a 100644 --- a/orttraining/orttraining/python/training/artifacts.py +++ b/orttraining/orttraining/python/training/artifacts.py @@ -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 diff --git a/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py b/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py index eca97b1be0..bc50d4afa2 100644 --- a/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py +++ b/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py @@ -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())) diff --git a/orttraining/orttraining/test/python/orttraining_test_ort_apis.py b/orttraining/orttraining/test/python/orttraining_test_ort_apis.py index c12d7f634f..ce96ee3da6 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ort_apis.py +++ b/orttraining/orttraining/test/python/orttraining_test_ort_apis.py @@ -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) diff --git a/orttraining/orttraining/test/python/orttraining_test_python_bindings.py b/orttraining/orttraining/test/python/orttraining_test_python_bindings.py index 14d7b54ff3..65cf259327 100644 --- a/orttraining/orttraining/test/python/orttraining_test_python_bindings.py +++ b/orttraining/orttraining/test/python/orttraining_test_python_bindings.py @@ -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"] diff --git a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc index 9786a1f8c5..4fca303b26 100644 --- a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc +++ b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc @@ -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 env; ORT_THROW_IF_ERROR(Environment::Create(nullptr, env)); std::vector> cuda_provider{onnxruntime::test::DefaultCudaExecutionProvider()}; - auto model = std::make_unique(model_uri, named_parameters, session_option, + auto model = std::make_unique(model_uri, &state, session_option, *env, cuda_provider); - auto optimizer = std::make_unique(optim_uri, model->NamedParameters(), session_option, + auto optimizer = std::make_unique(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(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 expected_file_names{ @@ -365,7 +368,7 @@ TEST(CheckpointApiTest, SaveCustomPropertyAsCheckpoint_ThenLoad_CPU) { // Call Save APIs. PathString checkpoint_path{ ConcatPathComponent(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 expected_file_names{ diff --git a/orttraining/orttraining/test/training_api/core/data_utils.h b/orttraining/orttraining/test/training_api/core/data_utils.h index 627bc5a6d4..815fbd1b8a 100644 --- a/orttraining/orttraining/test/training_api/core/data_utils.h +++ b/orttraining/orttraining/test/training_api/core/data_utils.h @@ -3,14 +3,15 @@ #pragma once +#include #include #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 void OrtValueToVec(const OrtValue& val, std::vector& output) { @@ -39,6 +40,22 @@ void CudaOrtValueToCpuVec(const OrtValue& val, std::vector& output, output.assign(val_ptr, val_ptr + src_tensor.Shape().Size()); } -} // namespace test -} // namespace training -} // namespace onnxruntime +inline void GenerateRandomData(std::vector& data) { + float scale = 1.f; + float mean = 0.f; + float seed = 123.f; + + std::default_random_engine generator_float{gsl::narrow_cast(seed)}; + std::normal_distribution 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 dims, OrtValue& input) { + TensorShape shape(dims); + std::vector data(shape.Size()); + GenerateRandomData(data); + onnxruntime::test::CreateInputOrtValueOnCPU(dims, data, &input); +} + +} // namespace onnxruntime::training::test diff --git a/orttraining/orttraining/test/training_api/core/training_api_tests.cc b/orttraining/orttraining/test/training_api/core/training_api_tests.cc index 6d99098464..1487c2b29d 100644 --- a/orttraining/orttraining/test/training_api/core/training_api_tests.cc +++ b/orttraining/orttraining/test/training_api/core/training_api_tests.cc @@ -2,14 +2,11 @@ // Licensed under the MIT License. #include -#include #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& data) { - float scale = 1.f; - float mean = 0.f; - float seed = 123.f; - - std::default_random_engine generator_float{gsl::narrow_cast(seed)}; - std::normal_distribution 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 dims, OrtValue& input) { - TensorShape shape(dims); - std::vector data(shape.Size()); - GenerateRandomData(data); - onnxruntime::test::CreateInputOrtValueOnCPU(dims, data, &input); -} - void TestModuleExport(const std::vector>& 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>& pr std::unique_ptr env; ASSERT_STATUS_OK(Environment::Create(nullptr, env)); auto model = std::make_unique( - 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& test_file_name, float in ASSERT_STATUS_OK(Environment::Create(nullptr, env)); const std::vector> providers{onnxruntime::test::DefaultCudaExecutionProvider()}; auto model = std::make_unique( - ToUTF8String(model_uri), state.module_checkpoint_state.named_parameters, + ToUTF8String(model_uri), &state, session_option, *env, providers); auto optim = std::make_shared( - 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 env; ASSERT_STATUS_OK(Environment::Create(nullptr, env)); auto model = std::make_unique(ToUTF8String(model_uri), - state.module_checkpoint_state.named_parameters, session_option, + &state, session_option, *env, std::vector>()); size_t params_size = 0; for (auto& param : model->Parameters()) { @@ -230,7 +209,7 @@ TEST(TrainingApiTest, ModuleCopyBufferToParameters) { std::unique_ptr env; ASSERT_STATUS_OK(Environment::Create(nullptr, env)); auto model = std::make_unique(ToUTF8String(model_uri), - state.module_checkpoint_state.named_parameters, session_option, + &state, session_option, *env, std::vector>()); int64_t params_size = static_cast(model->GetParametersSize()); std::vector expected_param_buffer(params_size); @@ -268,7 +247,7 @@ TEST(TrainingApiTest, ModuleTrainStep) { std::unique_ptr env; ASSERT_STATUS_OK(Environment::Create(nullptr, env)); auto model = std::make_unique(ToUTF8String(model_uri), - state.module_checkpoint_state.named_parameters, session_option, + &state, session_option, *env, std::vector>()); ASSERT_EQ(model->GetTrainingModelOutputCount(), 1); OrtValue input, target; @@ -340,10 +319,10 @@ TEST(TrainingApiTest, OptimStep) { std::shared_ptr cpu_provider = onnxruntime::test::DefaultCpuExecutionProvider(); ASSERT_STATUS_OK(Environment::Create(nullptr, env)); auto model = std::make_unique( - ToUTF8String(model_uri), state.module_checkpoint_state.named_parameters, session_option, + ToUTF8String(model_uri), &state, session_option, *env, providers); auto optim = std::make_unique( - ToUTF8String(optim_uri), model->NamedParameters(), session_option, + ToUTF8String(optim_uri), &state, session_option, *env, providers); OrtValue input, target; diff --git a/orttraining/orttraining/test/training_api/core/training_capi_tests.cc b/orttraining/orttraining/test/training_api/core/training_capi_tests.cc new file mode 100644 index 0000000000..d87d2ad8ae --- /dev/null +++ b/orttraining/orttraining/test/training_api/core/training_capi_tests.cc @@ -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(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(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(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(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(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(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().release(); + GenerateRandomInput(std::array{397510}, *buffer_impl); + + Ort::Value buffer(buffer_impl); + + training_session.FromBuffer(buffer); +} + +} // namespace onnxruntime::training::test diff --git a/orttraining/orttraining/test/training_api/trainer/trainer.cc b/orttraining/orttraining/test/training_api/trainer/trainer.cc index 8c1f2f45e5..e330a89524 100644 --- a/orttraining/orttraining/test/training_api/trainer/trainer.cc +++ b/orttraining/orttraining/test/training_api/trainer/trainer.cc @@ -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(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(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(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 duration_seconds = end - end_to_end_start; diff --git a/orttraining/orttraining/training_api/checkpoint.cc b/orttraining/orttraining/training_api/checkpoint.cc index e09577e992..73b24ae49f 100644 --- a/orttraining/orttraining/training_api/checkpoint.cc +++ b/orttraining/orttraining/training_api/checkpoint.cc @@ -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& 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) { diff --git a/orttraining/orttraining/training_api/checkpoint.h b/orttraining/orttraining/training_api/checkpoint.h index cbe9b379b9..5e3d469d9a 100644 --- a/orttraining/orttraining/training_api/checkpoint.h +++ b/orttraining/orttraining/training_api/checkpoint.h @@ -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. diff --git a/orttraining/orttraining/training_api/checkpoint_property.h b/orttraining/orttraining/training_api/checkpoint_property.h index 6e0ce3babb..2f4460f1fb 100644 --- a/orttraining/orttraining/training_api/checkpoint_property.h +++ b/orttraining/orttraining/training_api/checkpoint_property.h @@ -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 supported_data_types{ ONNX_NAMESPACE::TensorProto::FLOAT, @@ -77,6 +81,14 @@ struct PropertyBag { InlinedHashMap named_properties_; }; +template <> +inline PropertyDataType PropertyBag::GetProperty(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 diff --git a/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h index ed4a20f7cd..93e844c9e5 100644 --- a/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h +++ b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h @@ -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; diff --git a/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_api.h b/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_api.h index 58a9b5ade8..ffaf579e6a 100644 --- a/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_api.h +++ b/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_api.h @@ -4,6 +4,7 @@ #pragma once #include "onnxruntime_training_c_api.h" #include +#include 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; +/** \brief Class that holds the state of the training session state + * + * Wraps OrtCheckpointState + * + */ class CheckpointState : public detail::Base { 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& path_to_checkpoint); @@ -58,12 +64,34 @@ class CheckpointState : public detail::Base { * * 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& path_to_checkpoint, - bool include_optimizer_states); + static void SaveCheckpoint(const CheckpointState& checkpoint_state, + const std::basic_string& 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 { * * 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& inference_model_path, const std::vector& 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 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 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 diff --git a/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_inline.h b/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_inline.h index 75452eab46..f947293722 100644 --- a/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_inline.h +++ b/orttraining/orttraining/training_api/include/onnxruntime_training_cxx_inline.h @@ -82,16 +82,105 @@ inline void TrainingSession::OptimizerStep() { ThrowOnError(GetTrainingApi().OptimizerStep(p_, run_options)); } +inline std::vector 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 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 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 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 buffer_shape{static_cast(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(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(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& 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& 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& 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(property_value)) { + int64_t value = std::get(property_value); + void* value_p = &value; + ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtIntProperty, value_p)); + } else if (std::holds_alternative(property_value)) { + float value = std::get(property_value); + void* value_p = &value; + ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtFloatProperty, value_p)); + } else if (std::holds_alternative(property_value)) { + std::string value = std::get(property_value); + auto buffer = std::make_unique(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(property_value); + property = *value_p; + allocator.Free(property_value); + break; + } + case OrtPropertyType::OrtFloatProperty: { + auto value_p = reinterpret_cast(property_value); + property = *value_p; + allocator.Free(property_value); + break; + } + case OrtPropertyType::OrtStringProperty: { + auto value_p = reinterpret_cast(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 diff --git a/orttraining/orttraining/training_api/lr_scheduler.h b/orttraining/orttraining/training_api/lr_scheduler.h index 738cda30a4..38f4e6b36c 100644 --- a/orttraining/orttraining/training_api/lr_scheduler.h +++ b/orttraining/orttraining/training_api/lr_scheduler.h @@ -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: diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc index 67185b488e..0987e87dd8 100644 --- a/orttraining/orttraining/training_api/module.cc +++ b/orttraining/orttraining/training_api/module.cc @@ -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>& named_parameters, + CheckpointState* state, const onnxruntime::SessionOptions& session_options, const Environment& env, const std::vector>& providers, const std::optional& 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(session_options).config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "1"; + const_cast(session_options) + .config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "1"; } train_sess_ = std::make_unique(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 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 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 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 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 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> Module::Parameters() const { std::vector> 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> 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& inputs, std::vectorGetDataTransferManager(); - 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 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. diff --git a/orttraining/orttraining/training_api/module.h b/orttraining/orttraining/training_api/module.h index c84a99d8cb..71f7659e8e 100644 --- a/orttraining/orttraining/training_api/module.h +++ b/orttraining/orttraining/training_api/module.h @@ -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>& named_parameters, + CheckpointState* state, const onnxruntime::SessionOptions& session_options, const Environment& env, const std::vector>& providers, @@ -66,9 +68,7 @@ struct Module { // Return the trainable/nontrainable parameters std::vector> Parameters() const; - std::unordered_map> NamedParameters() const { - return named_parameters_; - } + std::unordered_map> 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& inputs, std::vector& 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 weights_; std::vector gradients_; bool accumulate_gradient_ = false; - const std::unordered_map>& 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; diff --git a/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc b/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc index f8552de920..286c38196e 100644 --- a/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc +++ b/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc @@ -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(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(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(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(checkpoint_state); + + switch (property_type) { + case OrtPropertyType::OrtIntProperty: { + int64_t* value = reinterpret_cast(property_value); + chkpt_state->property_bag.AddProperty(property_name, *value); + break; + } + case OrtPropertyType::OrtFloatProperty: { + float* value = reinterpret_cast(property_value); + chkpt_state->property_bag.AddProperty(property_name, *value); + break; + } + case OrtPropertyType::OrtStringProperty: { + char* value = reinterpret_cast(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(checkpoint_state); + const auto value = chkpt_state->property_bag.GetProperty< + onnxruntime::training::api::PropertyDataType>(property_name); + + if (std::holds_alternative(value)) { + int64_t* value_p = reinterpret_cast(allocator->Alloc(allocator, sizeof(int64_t))); + if (!value_p) { + return OrtApis::CreateStatus(ORT_FAIL, "Int property value buffer allocation failed."); + } + *value_p = std::get(value); + *(reinterpret_cast(property_value)) = value_p; + *property_type = OrtPropertyType::OrtIntProperty; + } else if (std::holds_alternative(value)) { + float* value_p = reinterpret_cast(allocator->Alloc(allocator, sizeof(float))); + if (!value_p) { + return OrtApis::CreateStatus(ORT_FAIL, "Float property value buffer allocation failed."); + } + *value_p = std::get(value); + *(reinterpret_cast(property_value)) = value_p; + *property_type = OrtPropertyType::OrtFloatProperty; + } else if (std::holds_alternative(value)) { + auto property_value_str = std::get(value); + // property_value_str.length() + 1 for null termination of c strings + auto buffer = reinterpret_cast(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(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) { diff --git a/orttraining/orttraining/training_api/optimizer.cc b/orttraining/orttraining/training_api/optimizer.cc index 2d9635ca55..2008e4c7c3 100644 --- a/orttraining/orttraining/training_api/optimizer.cc +++ b/orttraining/orttraining/training_api/optimizer.cc @@ -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 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 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(); @@ -152,11 +153,16 @@ Status Optimizer::ConstructInputs() { } Optimizer::Optimizer(const std::string& optim_path_or_bytes, - const std::unordered_map>& named_parameters, + CheckpointState* state, const onnxruntime::SessionOptions& session_options, const Environment& env, const std::vector>& providers) - : optim_sess_(std::make_unique(session_options, env)), named_parameters_(named_parameters) { + : optim_sess_(std::make_unique(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()}); + } + 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(optimizer_state_.learning_rate, &learning_rate_input); + utils::WrapInOrtValue(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(optimizer_state_.step + 1, &step_input); + utils::WrapInOrtValue(optimizer_state_->step + 1, &step_input); std::vector 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(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(optimizer_state_)}); + grouped_optimizer_states.insert({GROUP_ZERO_NAME, std::make_shared(*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(); diff --git a/orttraining/orttraining/training_api/optimizer.h b/orttraining/orttraining/training_api/optimizer.h index ac5d35215e..4c28a600ba 100644 --- a/orttraining/orttraining/training_api/optimizer.h +++ b/orttraining/orttraining/training_api/optimizer.h @@ -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>& named_parameters, + CheckpointState* state, const onnxruntime::SessionOptions& session_options, const Environment& env, const std::vector>& 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 optim_sess_; - const std::unordered_map>& named_parameters_; - GroupOptimizerState optimizer_state_; + CheckpointState* state_; // Non owning pointer to the state. + std::shared_ptr optimizer_state_; std::vector input_names_; std::vector output_names_; std::vector inputs_; diff --git a/orttraining/orttraining/training_api/ort_training_apis.h b/orttraining/orttraining/training_api/ort_training_apis.h index 2d55c40b25..3f8edede36 100644 --- a/orttraining/orttraining/training_api/ort_training_apis.h +++ b/orttraining/orttraining/training_api/ort_training_apis.h @@ -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 diff --git a/orttraining/orttraining/training_api/training_session.cc b/orttraining/orttraining/training_api/training_session.cc index f05fea1c5f..8332558f9a 100644 --- a/orttraining/orttraining/training_api/training_session.cc +++ b/orttraining/orttraining/training_api/training_session.cc @@ -10,14 +10,14 @@ namespace api { TrainingSession::TrainingSession(const Environment& session_env, const SessionOptions& session_options, const std::vector>& providers, - const std::unordered_map>& parameters, + CheckpointState* state, const ModelIdentifiers& model_identifiers) - : named_parameters_{parameters}, - module_{std::make_unique(model_identifiers.train_model, named_parameters_, + : state_{state}, + module_{std::make_unique(model_identifiers.train_model, state_, session_options, session_env, providers, model_identifiers.eval_model)}, optimizer_{model_identifiers.optim_model.has_value() ? std::make_unique( - model_identifiers.optim_model.value(), named_parameters_, + model_identifiers.optim_model.value(), state_, session_options, session_env, providers) : std::unique_ptr()} {} @@ -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)); diff --git a/orttraining/orttraining/training_api/training_session.h b/orttraining/orttraining/training_api/training_session.h index 6a06522e19..2e7f06055f 100644 --- a/orttraining/orttraining/training_api/training_session.h +++ b/orttraining/orttraining/training_api/training_session.h @@ -28,7 +28,7 @@ class TrainingSession { TrainingSession(const Environment& session_env, const SessionOptions& session_options, const std::vector>& providers, - const std::unordered_map>& 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> named_parameters_; + CheckpointState* state_; // Non owning pointer to the checkpoint state. It must outlive the training session. std::unique_ptr module_; std::shared_ptr optimizer_; std::unique_ptr scheduler_;