From 228db2431785afd0244e156210bfc6d0af24c1da Mon Sep 17 00:00:00 2001 From: Caroline Date: Mon, 28 Aug 2023 11:05:02 -0700 Subject: [PATCH] Add training API functions to WASM API (#16521) ### Description * Created `wasm/training_api` source and header files & modified WebAssembly CMake to include training flags * The `wasm/training_api` files use an `OrtTrainingManager` handle which is a struct of an OrtCheckpointState and an OrtTrainingSession, rather than creating a CheckpointState handle & a separate TrainingSession handle. * This is so that the TypeScript side only has to manage one handle that will be passed between TrainingSession & CheckpointState representations, rather than the TypeScript side managing separate CheckpointStateHandle and TrainingSessionHandle. ### Motivation and Context WASM API needs to be updated with ORT training API function calls so that ORT training web bindings can be added for on-device training. --------- Co-authored-by: Baiju Meswani Co-authored-by: carzh Co-authored-by: Ashwini Khade --- cmake/onnxruntime_webassembly.cmake | 30 +++-- js/web/lib/wasm/binding/ort-wasm.d.ts | 32 ++++++ onnxruntime/wasm/api.cc | 93 +++++++++++++++- onnxruntime/wasm/api.h | 155 ++++++++++++++++++++++++++ 4 files changed, 299 insertions(+), 11 deletions(-) diff --git a/cmake/onnxruntime_webassembly.cmake b/cmake/onnxruntime_webassembly.cmake index 4243031045..d7712a7b70 100644 --- a/cmake/onnxruntime_webassembly.cmake +++ b/cmake/onnxruntime_webassembly.cmake @@ -277,19 +277,29 @@ else() "SHELL:-s EXPORT_NAME=ortWasmThreaded" "SHELL:-s DEFAULT_PTHREAD_STACK_SIZE=131072" ) - if (onnxruntime_ENABLE_WEBASSEMBLY_SIMD) - set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME "ort-wasm-simd-threaded") - else() - set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME "ort-wasm-threaded") - endif() else() target_link_options(onnxruntime_webassembly PRIVATE "SHELL:-s EXPORT_NAME=ortWasm" ) - if (onnxruntime_ENABLE_WEBASSEMBLY_SIMD) - set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME "ort-wasm-simd") - else() - set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME "ort-wasm") - endif() endif() + + set(target_name ort) + + if (onnxruntime_ENABLE_TRAINING_APIS) + list(APPEND target_name "training") + endif() + + list(APPEND target_name "wasm") + + if (onnxruntime_ENABLE_WEBASSEMBLY_SIMD) + list(APPEND target_name "simd") + endif() + + if (onnxruntime_ENABLE_WEBASSEMBLY_THREADS) + list(APPEND target_name "threaded") + endif() + + list(JOIN target_name "-" target_name) + + set_target_properties(onnxruntime_webassembly PROPERTIES OUTPUT_NAME ${target_name}) endif() diff --git a/js/web/lib/wasm/binding/ort-wasm.d.ts b/js/web/lib/wasm/binding/ort-wasm.d.ts index 06fcbf6344..7f0430b7b2 100644 --- a/js/web/lib/wasm/binding/ort-wasm.d.ts +++ b/js/web/lib/wasm/binding/ort-wasm.d.ts @@ -64,6 +64,38 @@ export interface OrtWasmModule extends EmscriptenModule { _OrtEndProfiling(sessionHandle: number): number; // #endregion + // #region ORT Training APIs + _OrtTrainingLoadCheckpoint?(dataOffset: number, dataLength: number): number; + + _OrtTrainingReleaseCheckpoint?(checkpointHandle: number): void; + + _OrtTrainingCreateSession? + (sessionOptionsHandle: number, checkpointHandle: number, trainOffset: number, trainLength: number, + evalOffset: number, evalLength: number, optimizerOffset: number, optimizerLength: number): number; + + _OrtTrainingLazyResetGrad?(trainingHandle: number): number; + + _OrtTrainingRunTrainStep? + (trainingHandle: number, inputsOffset: number, inputCount: number, outputsOffset: number, outputCount: number, + runOptionsHandle: number): number; + + _OrtTrainingOptimizerStep?(trainingHandle: number, runOptionsHandle: number): number; + + _OrtTrainingEvalStep? + (trainingHandle: number, inputsOffset: number, inputCount: number, outputsOffset: number, outputCount: number, + runOptionsHandle: number): number; + + _OrtTrainingGetParametersSize?(trainingHandle: number, paramSizeT: number, trainableOnly: boolean): number; + + _OrtTrainingCopyParametersToBuffer? + (trainingHandle: number, parametersBuffer: number, parameterCount: number, trainableOnly: boolean): number; + + _OrtTrainingCopyParametersFromBuffer? + (trainingHandle: number, parametersBuffer: number, parameterCount: number, trainableOnly: boolean): number; + + _OrtTrainingReleaseSession?(trainingHandle: number): void; + // #endregion + // #region config mainScriptUrlOrBlob?: string|Blob; // #endregion diff --git a/onnxruntime/wasm/api.cc b/onnxruntime/wasm/api.cc index 496c9c401f..aabefeaa7a 100644 --- a/onnxruntime/wasm/api.cc +++ b/onnxruntime/wasm/api.cc @@ -1,9 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "api.h" +#ifdef ENABLE_TRAINING_APIS +#include "onnxruntime_training_cxx_api.h" +#endif #include "core/session/onnxruntime_cxx_api.h" +#include "api.h" #include #include @@ -384,3 +387,91 @@ char* OrtEndProfiling(ort_session_handle_t session) { ? file_name : nullptr; } + +// Training API Section + +#ifdef ENABLE_TRAINING_APIS +#define CHECK_TRAINING_STATUS(ORT_API_NAME, ...) \ + CheckStatus(Ort::GetTrainingApi().ORT_API_NAME(__VA_ARGS__)) + +ort_training_checkpoint_handle_t EMSCRIPTEN_KEEPALIVE OrtTrainingLoadCheckpoint(void* checkpoint_data_buffer, size_t checkpoint_size) { + OrtCheckpointState* checkpoint_state = nullptr; + return (CHECK_TRAINING_STATUS(LoadCheckpointFromBuffer, checkpoint_data_buffer, checkpoint_size, &checkpoint_state) == ORT_OK) + ? checkpoint_state + : nullptr; +} + +void EMSCRIPTEN_KEEPALIVE OrtTrainingReleaseCheckpoint(ort_training_checkpoint_handle_t training_checkpoint_state_handle) { + Ort::GetTrainingApi().ReleaseCheckpointState(training_checkpoint_state_handle); +} + +ort_training_session_handle_t EMSCRIPTEN_KEEPALIVE OrtTrainingCreateSession(const ort_session_options_handle_t options, + ort_training_checkpoint_handle_t training_checkpoint_state_handle, + void* train_model, + size_t train_size, + void* eval_model, + size_t eval_size, + void* optimizer_model, + size_t optimizer_size) { + OrtTrainingSession* training_session = nullptr; + return (CHECK_TRAINING_STATUS(CreateTrainingSessionFromArray, g_env, options, + training_checkpoint_state_handle, train_model, train_size, + eval_model, eval_size, optimizer_model, optimizer_size, + &training_session) == ORT_OK) + ? training_session + : nullptr; +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingLazyResetGrad(ort_training_session_handle_t training_handle) { + return CHECK_TRAINING_STATUS(LazyResetGrad, training_handle); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingRunTrainStep(ort_training_session_handle_t training_handle, + ort_tensor_handle_t* inputs, + size_t input_count, + ort_tensor_handle_t* outputs, + size_t output_count, + ort_run_options_handle_t options) { + return CHECK_TRAINING_STATUS(TrainStep, training_handle, options, input_count, inputs, output_count, outputs); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingOptimizerStep(ort_training_session_handle_t training_handle, + const ort_run_options_handle_t run_options) { + return CHECK_TRAINING_STATUS(OptimizerStep, training_handle, run_options); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingEvalStep(ort_training_session_handle_t training_handle, + ort_tensor_handle_t* inputs, + size_t input_count, + ort_tensor_handle_t* outputs, + size_t output_count, + ort_run_options_handle_t options) { + return CHECK_TRAINING_STATUS(EvalStep, training_handle, + options, input_count, inputs, output_count, outputs); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingGetParametersSize(ort_training_session_handle_t training_handle, + size_t* param_size, + bool trainable_only) { + return CHECK_TRAINING_STATUS(GetParametersSize, training_handle, param_size, trainable_only); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingCopyParametersToBuffer(ort_training_session_handle_t training_handle, + ort_tensor_handle_t parameters_buffer, + size_t parameter_count, + bool trainable_only) { + return CHECK_TRAINING_STATUS(CopyParametersToBuffer, training_handle, parameters_buffer, trainable_only); +} + +int EMSCRIPTEN_KEEPALIVE OrtTrainingCopyParametersFromBuffer(ort_training_session_handle_t training_handle, + ort_tensor_handle_t parameters_buffer, + size_t parameter_count, + bool trainable_only) { + return CHECK_TRAINING_STATUS(CopyBufferToParameters, training_handle, parameters_buffer, trainable_only); +} + +void EMSCRIPTEN_KEEPALIVE OrtTrainingReleaseSession(ort_training_session_handle_t training_handle) { + Ort::GetTrainingApi().ReleaseTrainingSession(training_handle); +} + +#endif diff --git a/onnxruntime/wasm/api.h b/onnxruntime/wasm/api.h index 5494a9e1b4..b9103414aa 100644 --- a/onnxruntime/wasm/api.h +++ b/onnxruntime/wasm/api.h @@ -24,6 +24,14 @@ using ort_run_options_handle_t = OrtRunOptions*; struct OrtValue; using ort_tensor_handle_t = OrtValue*; +#ifdef ENABLE_TRAINING_APIS +struct OrtTrainingSession; +using ort_training_session_handle_t = OrtTrainingSession*; + +struct OrtCheckpointState; +using ort_training_checkpoint_handle_t = OrtCheckpointState*; +#endif + extern "C" { /** @@ -222,4 +230,151 @@ int EMSCRIPTEN_KEEPALIVE OrtRun(ort_session_handle_t session, * Caller must release the C style string after use by calling OrtFree(). */ char* EMSCRIPTEN_KEEPALIVE OrtEndProfiling(ort_session_handle_t session); + +// Training API Section + +#ifdef ENABLE_TRAINING_APIS +/** + * @brief Load the checkpoint for training. + * + * @param checkpoint_data_buffer pointer to a buffer containing the CheckpointState + * @param checkpoint_size size of the CheckpointState in bytes + * @return ort_training_checkpoint_handle_t + */ +ort_training_checkpoint_handle_t EMSCRIPTEN_KEEPALIVE OrtTrainingLoadCheckpoint(void* checkpoint_data_buffer, size_t checkpoint_size); + +/** + * @brief Release the specified ORT training checkpoint state. + * + * @param training_checkpoint_state_handle handle for the CheckpointState + */ +void EMSCRIPTEN_KEEPALIVE OrtTrainingReleaseCheckpoint(ort_training_checkpoint_handle_t training_checkpoint_state_handle); + +/** + * Creates an instance of a training session that can be used to begin or resume training from a given checkpoint state + * for the given onnx models. + * @param options Session options that the user can customize for this training session. + * @param training_checkpoint_state_handle Training states that the training session uses as a starting point for training. + * @param train_model pointer to a buffer containing the ONNX training model + * @param train_size size of the train_model buffer in bytes + * @param eval_model pointer to a buffer containing the ONNX evaluation model + * @param eval_size size of the eval_model buffer in bytes + * @param optimizer_model pointer to a buffer containing the ONNX optimizer model + * @param optimizer_size size of the optimizer_model buffer in bytes + * @return a handle of the ORT training session + * + */ +ort_training_session_handle_t EMSCRIPTEN_KEEPALIVE OrtTrainingCreateSession(ort_session_options_handle_t options, + ort_training_checkpoint_handle_t training_checkpoint_state_handle, + void* train_model, + size_t train_size, + void* eval_model, + size_t eval_size, + void* optimizer_model, + size_t optimizer_size); + +/** + * Resets the gradients of all trainable parameters to zero for the specified TrainingSession + * @param training_handle handle of the training session + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingLazyResetGrad(ort_training_session_handle_t training_handle); + +/** + * @brief Run a single training step. + * + * @param training_handle session handle of the specified session + * @param inputs user inputs to the training model + * @param input_count number of user inputs to the training model + * @param outputs [out] user outputs computed by train step + * @param output_count [out] number of user outputs expected from this train step + * @param run_options handle of the run options + * @return int ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingRunTrainStep(ort_training_session_handle_t training_handle, + ort_tensor_handle_t* inputs, size_t input_count, + ort_tensor_handle_t* outputs, + size_t output_count, + ort_run_options_handle_t run_options = nullptr); + +/** + * Performs weight updates for the trainable parameters in the given training session using the optimizer model. + * @param training_handle handle of the training session + * @param run_options optional parameter of run options for this training step + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingOptimizerStep(ort_training_session_handle_t training_handle, + ort_run_options_handle_t run_options = nullptr); + +/** + * Computs outputs for the eval model associated with the given training session. + * @param training_handle handle of the training session + * @param options run options for this eval step + * @param input_count number of user inputs to the eval model + * @param inputs the user inputs to the eval model + * @param output_count [out] number of user outputs expected from this eval step + * @param outputs [out] user outputs computed by the eval step + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingEvalStep(ort_training_session_handle_t training_handle, + ort_tensor_handle_t* inputs, + size_t input_count, + ort_tensor_handle_t* outputs, + size_t output_count, + ort_run_options_handle_t options = nullptr); + +/** + * Retrieves the size of all parameters for the training state. + * When the trainable_only argument is true, the size is calculated for trainable params only. + * + * @param training_handle handle of the training session + * @param param_size [out] size of all parameter elements + * @param trainable_only skips non-trainable parameters when true. + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingGetParametersSize(ort_training_session_handle_t training_handle, + size_t* param_size, + bool trainable_only); + +/** + * Copy all parameters to a contiguous buffer held by the argument parameters_buffer + * + * User is responsible for allocating and freeing resources used by the parameters_buffer. + * Parameter ordering is preserved. + * + * @param training_handle handle of the training session + * @param parameters_buffer [out] pre-allocated OrtValue buffer to copy onto. Must be same size as results of + * GetParametersSize api call + * @param parameter_count number of parameters expected in the parameters_buffer + * @param trainable_only whether to skip non-trainable parameters + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingCopyParametersToBuffer(ort_training_session_handle_t training_handle, + ort_tensor_handle_t parameters_buffer, + size_t parameter_count, + bool trainable_only); + +/** + * Copy parameters values from given contiguous buffer held by parameters_buffer to the training state. + * Parameter ordering is preserved. + * @param training_handle handle of the training session + * @param parameters_buffer OrtValue buffer to copy from. Must be same size as results of + * GetParametersSize api call + * @param parameter_count number of parameters expected in the parameters_buffer + * @param trainable_only whether to skip non-trainable parameters + * @returns ORT error code. If not zero, call OrtGetLastError() to get detailed error message. + */ +int EMSCRIPTEN_KEEPALIVE OrtTrainingCopyParametersFromBuffer(ort_training_session_handle_t training_handle, + ort_tensor_handle_t parameters_buffer, + size_t parameter_count, + bool trainable_only); + +/** + * @brief Release the specified ORT training session. + * + * @param training_session_handle handle of the training session + */ +void EMSCRIPTEN_KEEPALIVE OrtTrainingReleaseSession(ort_training_session_handle_t training_session_handle); + +#endif };