saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/web/training] lazyResetGrad implementation (#18711)

Browse source 
### Description
* implemented lazyResetGrad function

### Motivation and Context
* we are in the process of adding language bindings to enable training
on web
* lazyresetgrad ensures that the gradients are calculated correctly
after the first runTrainStep call

---------

Co-authored-by: Ashwini Khade <askhade@microsoft.com>
This commit is contained in:
Caroline Zhu 2023-12-11 17:36:54 -08:00 committed by GitHub
parent ccf3b2054b
commit eb03032925
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 27 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
js/common/lib/backend.ts
View file

@ -48,6 +48,7 @@ export interface TrainingSessionHandler extends SessionHandler {
readonly evalInputNames: readonly string[];
readonly evalOutputNames: readonly string[];
lazyResetGrad(): Promise<void>;
runTrainStep(
feeds: SessionHandler.FeedsType, fetches: SessionHandler.FetchesType,
options: InferenceSession.RunOptions): Promise<SessionHandler.ReturnType>;

4
js/common/lib/training-session-impl.ts
View file

@ -192,6 +192,10 @@ export class TrainingSession implements TrainingSessionInterface {
return returnValue;
}
async lazyResetGrad(): Promise<void> {
await this.handler.lazyResetGrad();
}
runTrainStep(feeds: FeedsType, options?: RunOptions): Promise<ReturnType>;
runTrainStep(feeds: FeedsType, fetches: FetchesType, options?: RunOptions): Promise<ReturnType>;
async runTrainStep(feeds: FeedsType, arg1?: FetchesType|RunOptions, arg2?: RunOptions): Promise<ReturnType> {

6
js/common/lib/training-session.ts
View file

@ -22,6 +22,12 @@ export declare namespace TrainingSession {
export interface TrainingSession {
// #region run()
/**
* Lazily resets the gradients of all trainable parameters to zero. Should happen after the invocation of
* runOptimizerStep.
*/
lazyResetGrad(): Promise<void>;
/**
* Run TrainStep asynchronously with the given feeds and options.
*

6
js/web/lib/wasm/session-handler-training.ts
View file

@ -6,7 +6,7 @@ import {env, InferenceSession, OnnxValue, SessionHandler, Tensor, TrainingSessio
import {SerializableModeldata, TensorMetadata} from './proxy-messages';
import {decodeTensorMetadata, encodeTensorMetadata} from './session-handler-inference';
import {createSessionAllocate, initRuntime, isOrtEnvInitialized} from './wasm-core-impl';
import {createCheckpointHandle, createTrainingSessionHandle, getContiguousParameters, getModelInputOutputNames, getParametersSize, loadParametersBuffer, releaseTrainingSessionAndCheckpoint, runEvalStep, runOptimizerStep, runTrainStep} from './wasm-training-core-impl';
import {createCheckpointHandle, createTrainingSessionHandle, getContiguousParameters, getModelInputOutputNames, getParametersSize, lazyResetGrad, loadParametersBuffer, releaseTrainingSessionAndCheckpoint, runEvalStep, runOptimizerStep, runTrainStep} from './wasm-training-core-impl';
export class OnnxruntimeWebAssemblyTrainingSessionHandler implements TrainingSessionHandler {
private sessionId: number;
@ -105,6 +105,10 @@ export class OnnxruntimeWebAssemblyTrainingSessionHandler implements TrainingSes
return resultMap;
}
async lazyResetGrad(): Promise<void> {
await lazyResetGrad(this.sessionId);
}
async runTrainStep(
feeds: SessionHandler.FeedsType, fetches: SessionHandler.FetchesType,
options: InferenceSession.RunOptions): Promise<SessionHandler.ReturnType> {

11
js/web/lib/wasm/wasm-training-core-impl.ts
View file

@ -253,6 +253,17 @@ const moveOutputToTensorMetadataArr =
return output;
};
export const lazyResetGrad = async(trainingSessionId: number): Promise<void> => {
const wasm = getInstance();
if (wasm._OrtTrainingLazyResetGrad) {
const errorCode = wasm._OrtTrainingLazyResetGrad(trainingSessionId);
ifErrCodeCheckLastError(errorCode, 'Can\'t call lazyResetGrad.');
} else {
throw new Error(NO_TRAIN_FUNCS_MSG);
}
};
export const runTrainStep = async(
trainingSessionId: number, inputIndices: number[], inputTensors: TensorMetadata[], outputIndices: number[],
outputTensors: Array<TensorMetadata|null>, options: InferenceSession.RunOptions): Promise<TensorMetadata[]> => {