mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
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 <bmeswani@microsoft.com> Co-authored-by: carzh <carolinezhu@microsoft.com> Co-authored-by: Ashwini Khade <askhade@microsoft.com>
This commit is contained in:
Caroline
GitHub
parent
cbd97515cd
commit
228db24317
4 changed files with 299 additions and 11 deletions
|
|
@ -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()
|
||||
|
|
|
|||
32
js/web/lib/wasm/binding/ort-wasm.d.ts
vendored
32
js/web/lib/wasm/binding/ort-wasm.d.ts
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 <iostream>
|
||||
#include <vector>
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue