diff --git a/js/common/lib/backend-impl.ts b/js/common/lib/backend-impl.ts index 75feba1d0a..e129c6971a 100644 --- a/js/common/lib/backend-impl.ts +++ b/js/common/lib/backend-impl.ts @@ -26,7 +26,7 @@ const backendsSortedByPriority: string[] = []; * @ignore */ export const registerBackend = (name: string, backend: Backend, priority: number): void => { - if (backend && typeof backend.init === 'function' && typeof backend.createSessionHandler === 'function') { + if (backend && typeof backend.init === 'function' && typeof backend.createInferenceSessionHandler === 'function') { const currentBackend = backends.get(name); if (currentBackend === undefined) { backends.set(name, {backend, priority}); diff --git a/js/common/lib/backend.ts b/js/common/lib/backend.ts index 804f33f00d..dd04ef3f15 100644 --- a/js/common/lib/backend.ts +++ b/js/common/lib/backend.ts @@ -3,6 +3,7 @@ import {InferenceSession} from './inference-session.js'; import {OnnxValue} from './onnx-value.js'; +import {TrainingSession} from './training-session.js'; /** * @ignore @@ -14,16 +15,23 @@ export declare namespace SessionHandler { } /** - * Represent a handler instance of an inference session. + * Represents shared SessionHandler functionality * * @ignore */ -export interface SessionHandler { +interface SessionHandler { dispose(): Promise; readonly inputNames: readonly string[]; readonly outputNames: readonly string[]; +} +/** + * Represent a handler instance of an inference session. + * + * @ignore + */ +export interface InferenceSessionHandler extends SessionHandler { startProfiling(): void; endProfiling(): void; @@ -31,6 +39,20 @@ export interface SessionHandler { options: InferenceSession.RunOptions): Promise; } +/** + * Represent a handler instance of a training inference session. + * + * @ignore + */ +export interface TrainingSessionHandler extends SessionHandler { + runTrainStep( + feeds: SessionHandler.FeedsType, fetches: SessionHandler.FetchesType, + options: InferenceSession.RunOptions): Promise; + + loadParametersBuffer(array: Uint8Array, trainableOnly: boolean): Promise; + getContiguousParameters(trainableOnly: boolean): Promise; +} + /** * Represent a backend that provides implementation of model inferencing. * @@ -42,8 +64,13 @@ export interface Backend { */ init(): Promise; - createSessionHandler(uriOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): - Promise; + createInferenceSessionHandler(uriOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): + Promise; + + createTrainingSessionHandler? + (checkpointStateUriOrBuffer: TrainingSession.URIorBuffer, trainModelUriOrBuffer: TrainingSession.URIorBuffer, + evalModelUriOrBuffer: TrainingSession.URIorBuffer, optimizerModelUriOrBuffer: TrainingSession.URIorBuffer, + options: InferenceSession.SessionOptions): Promise; } export {registerBackend} from './backend-impl.js'; diff --git a/js/common/lib/env.ts b/js/common/lib/env.ts index 525272294c..c78ae0fc83 100644 --- a/js/common/lib/env.ts +++ b/js/common/lib/env.ts @@ -9,6 +9,7 @@ export declare namespace Env { 'ort-wasm.wasm'?: string; 'ort-wasm-threaded.wasm'?: string; 'ort-wasm-simd.wasm'?: string; + 'ort-training-wasm-simd.wasm'?: string; 'ort-wasm-simd-threaded.wasm'?: string; /* eslint-enable @typescript-eslint/naming-convention */ }; diff --git a/js/common/lib/index.ts b/js/common/lib/index.ts index 85df1747f8..9cbfcc4e8b 100644 --- a/js/common/lib/index.ts +++ b/js/common/lib/index.ts @@ -22,3 +22,4 @@ export * from './env.js'; export * from './inference-session.js'; export * from './tensor.js'; export * from './onnx-value.js'; +export * from './training-session.js'; diff --git a/js/common/lib/inference-session-impl.ts b/js/common/lib/inference-session-impl.ts index 06949b4a26..9bc2088f20 100644 --- a/js/common/lib/inference-session-impl.ts +++ b/js/common/lib/inference-session-impl.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. import {resolveBackend} from './backend-impl.js'; -import {SessionHandler} from './backend.js'; +import {InferenceSessionHandler} from './backend.js'; import {InferenceSession as InferenceSessionInterface} from './inference-session.js'; import {OnnxValue} from './onnx-value.js'; import {Tensor} from './tensor.js'; @@ -14,7 +14,7 @@ type FetchesType = InferenceSessionInterface.FetchesType; type ReturnType = InferenceSessionInterface.ReturnType; export class InferenceSession implements InferenceSessionInterface { - private constructor(handler: SessionHandler) { + private constructor(handler: InferenceSessionHandler) { this.handler = handler; } run(feeds: FeedsType, options?: RunOptions): Promise; @@ -195,7 +195,7 @@ export class InferenceSession implements InferenceSessionInterface { const eps = options.executionProviders || []; const backendHints = eps.map(i => typeof i === 'string' ? i : i.name); const backend = await resolveBackend(backendHints); - const handler = await backend.createSessionHandler(filePathOrUint8Array, options); + const handler = await backend.createInferenceSessionHandler(filePathOrUint8Array, options); return new InferenceSession(handler); } @@ -213,5 +213,5 @@ export class InferenceSession implements InferenceSessionInterface { return this.handler.outputNames; } - private handler: SessionHandler; + private handler: InferenceSessionHandler; } diff --git a/js/common/lib/training-session-impl.ts b/js/common/lib/training-session-impl.ts new file mode 100644 index 0000000000..f06d06bda0 --- /dev/null +++ b/js/common/lib/training-session-impl.ts @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import {TrainingSessionHandler} from './backend.js'; +import {InferenceSession as InferenceSession} from './inference-session.js'; +import {TrainingSession as TrainingSessionInterface, TrainingSessionCreateOptions} from './training-session.js'; + +type SessionOptions = InferenceSession.SessionOptions; + +export class TrainingSession implements TrainingSessionInterface { + private constructor(handler: TrainingSessionHandler) { + this.handler = handler; + } + private handler: TrainingSessionHandler; + + get inputNames(): readonly string[] { + return this.handler.inputNames; + } + get outputNames(): readonly string[] { + return this.handler.outputNames; + } + + static async create(_trainingOptions: TrainingSessionCreateOptions, _sessionOptions?: SessionOptions): + Promise { + throw new Error('Method not implemented'); + } + + async loadParametersBuffer(_array: Uint8Array, _trainableOnly: boolean): Promise { + throw new Error('Method not implemented.'); + } + + async getContiguousParameters(_trainableOnly: boolean): Promise { + throw new Error('Method not implemented.'); + } + + runTrainStep(feeds: InferenceSession.OnnxValueMapType, options?: InferenceSession.RunOptions|undefined): + Promise; + runTrainStep( + feeds: InferenceSession.OnnxValueMapType, fetches: InferenceSession.FetchesType, + options?: InferenceSession.RunOptions|undefined): Promise; + async runTrainStep(_feeds: unknown, _fetches?: unknown, _options?: unknown): + Promise { + throw new Error('Method not implemented.'); + } + + async release(): Promise { + return this.handler.dispose(); + } +} diff --git a/js/common/lib/training-session.ts b/js/common/lib/training-session.ts new file mode 100644 index 0000000000..0967d79b33 --- /dev/null +++ b/js/common/lib/training-session.ts @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import {InferenceSession} from './inference-session.js'; +import {TrainingSession as TrainingSessionImpl} from './training-session-impl.js'; + +/* eslint-disable @typescript-eslint/no-redeclare */ + +export declare namespace TrainingSession { + /** + * Either URI file path (string) or Uint8Array containing model or checkpoint information. + */ + type URIorBuffer = string|Uint8Array; +} + +/** + * Represent a runtime instance of an ONNX training session, + * which contains a model that can be trained, and, optionally, + * an eval and optimizer model. + */ +export interface TrainingSession { + // #region run() + + /** + * Run TrainStep asynchronously with the given feeds and options. + * + * @param feeds - Representation of the model input. See type description of `InferenceSession.InputType` for + detail. + * @param options - Optional. A set of options that controls the behavior of model training. + * @returns A promise that resolves to a map, which uses output names as keys and OnnxValue as corresponding values. + */ + runTrainStep(feeds: InferenceSession.FeedsType, options?: InferenceSession.RunOptions): + Promise; + + /** + * Run a single train step with the given inputs and options. + * + * @param feeds - Representation of the model input. + * @param fetches - Representation of the model output. + * detail. + * @param options - Optional. A set of options that controls the behavior of model inference. + * @returns A promise that resolves to a map, which uses output names as keys and OnnxValue as corresponding + values. + */ + runTrainStep( + feeds: InferenceSession.FeedsType, fetches: InferenceSession.FetchesType, + options?: InferenceSession.RunOptions): Promise; + + // #endregion + + // #region copy parameters + /** + * Copies from a buffer containing parameters to the TrainingSession parameters. + * + * @param buffer - buffer containing parameters + * @param trainableOnly - True if trainable parameters only to be modified, false otherwise. + */ + loadParametersBuffer(array: Uint8Array, trainableOnly: boolean): Promise; + + /** + * Copies from the TrainingSession parameters to a buffer. + * + * @param trainableOnly - True if trainable parameters only to be copied, false othrwise. + * @returns A promise that resolves to a buffer of the requested parameters. + */ + getContiguousParameters(trainableOnly: boolean): Promise; + // #endregion + + // #region release() + + /** + * Release the inference session and the underlying resources. + */ + release(): Promise; + // #endregion + + // #region metadata + + /** + * Get input names of the loaded model. + */ + readonly inputNames: readonly string[]; + + /** + * Get output names of the loaded model. + */ + readonly outputNames: readonly string[]; + // #endregion +} + +/** + * Represents the optional parameters that can be passed into the TrainingSessionFactory. + */ +export interface TrainingSessionCreateOptions { + /** + * URI or buffer for a .ckpt file that contains the checkpoint for the training model. + */ + checkpointState: TrainingSession.URIorBuffer; + /** + * URI or buffer for the .onnx training file. + */ + trainModel: TrainingSession.URIorBuffer; + /** + * Optional. URI or buffer for the .onnx optimizer model file. + */ + optimizerModel?: TrainingSession.URIorBuffer; + /** + * Optional. URI or buffer for the .onnx eval model file. + */ + evalModel?: TrainingSession.URIorBuffer; +} + +/** + * Defines method overload possibilities for creating a TrainingSession. + */ +export interface TrainingSessionFactory { + // #region create() + + /** + * Creates a new TrainingSession and asynchronously loads any models passed in through trainingOptions + * + * @param trainingOptions specify models and checkpoints to load into the Training Session + * @param sessionOptions specify configuration for training session behavior + * + * @returns Promise that resolves to a TrainingSession object + */ + create(trainingOptions: TrainingSessionCreateOptions, sessionOptions?: InferenceSession.SessionOptions): + Promise; + + // #endregion +} + +// eslint-disable-next-line @typescript-eslint/naming-convention +export const TrainingSession: TrainingSessionFactory = TrainingSessionImpl; diff --git a/js/node/lib/backend.ts b/js/node/lib/backend.ts index d3680f9d44..5f5ad49a2d 100644 --- a/js/node/lib/backend.ts +++ b/js/node/lib/backend.ts @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import {Backend, InferenceSession, SessionHandler} from 'onnxruntime-common'; +import {Backend, InferenceSession, InferenceSessionHandler, SessionHandler} from 'onnxruntime-common'; import {Binding, binding} from './binding'; -class OnnxruntimeSessionHandler implements SessionHandler { +class OnnxruntimeSessionHandler implements InferenceSessionHandler { #inferenceSession: Binding.InferenceSession; constructor(pathOrBuffer: string|Uint8Array, options: InferenceSession.SessionOptions) { @@ -53,8 +53,8 @@ class OnnxruntimeBackend implements Backend { return Promise.resolve(); } - async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): - Promise { + async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): + Promise { return new Promise((resolve, reject) => { process.nextTick(() => { try { diff --git a/js/react_native/lib/backend.ts b/js/react_native/lib/backend.ts index 058531f415..3b1852699a 100644 --- a/js/react_native/lib/backend.ts +++ b/js/react_native/lib/backend.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import {Backend, InferenceSession, SessionHandler, Tensor,} from 'onnxruntime-common'; +import {Backend, InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common'; import {Platform} from 'react-native'; import {binding, Binding, JSIBlob, jsiHelper} from './binding'; @@ -43,7 +43,7 @@ const normalizePath = (path: string): string => { return path; }; -class OnnxruntimeSessionHandler implements SessionHandler { +class OnnxruntimeSessionHandler implements InferenceSessionHandler { #inferenceSession: Binding.InferenceSession; #key: string; @@ -165,8 +165,8 @@ class OnnxruntimeBackend implements Backend { return Promise.resolve(); } - async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): - Promise { + async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): + Promise { const handler = new OnnxruntimeSessionHandler(pathOrBuffer); await handler.loadModel(options || {}); return handler; diff --git a/js/web/lib/backend-onnxjs.ts b/js/web/lib/backend-onnxjs.ts index 18a068e0ce..5ea7de809a 100644 --- a/js/web/lib/backend-onnxjs.ts +++ b/js/web/lib/backend-onnxjs.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. /* eslint-disable import/no-internal-modules */ -import {Backend, InferenceSession, SessionHandler} from 'onnxruntime-common'; +import {Backend, InferenceSession, InferenceSessionHandler} from 'onnxruntime-common'; import {Session} from './onnxjs/session'; import {OnnxjsSessionHandler} from './onnxjs/session-handler'; @@ -11,8 +11,8 @@ class OnnxjsBackend implements Backend { // eslint-disable-next-line @typescript-eslint/no-empty-function async init(): Promise {} - async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): - Promise { + async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): + Promise { // NOTE: Session.Config(from onnx.js) is not compatible with InferenceSession.SessionOptions(from // onnxruntime-common). // In future we should remove Session.Config and use InferenceSession.SessionOptions. diff --git a/js/web/lib/backend-wasm.ts b/js/web/lib/backend-wasm.ts index ceb20044d9..04108c2ad0 100644 --- a/js/web/lib/backend-wasm.ts +++ b/js/web/lib/backend-wasm.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import {Backend, env, InferenceSession, SessionHandler} from 'onnxruntime-common'; +import {Backend, env, InferenceSession, InferenceSessionHandler} from 'onnxruntime-common'; import {cpus} from 'os'; import {initializeWebAssemblyInstance} from './wasm/proxy-wrapper'; @@ -40,10 +40,12 @@ class OnnxruntimeWebAssemblyBackend implements Backend { // init wasm await initializeWebAssemblyInstance(); } - createSessionHandler(path: string, options?: InferenceSession.SessionOptions): Promise; - createSessionHandler(buffer: Uint8Array, options?: InferenceSession.SessionOptions): Promise; - async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): - Promise { + createInferenceSessionHandler(path: string, options?: InferenceSession.SessionOptions): + Promise; + createInferenceSessionHandler(buffer: Uint8Array, options?: InferenceSession.SessionOptions): + Promise; + async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): + Promise { const handler = new OnnxruntimeWebAssemblySessionHandler(); await handler.loadModel(pathOrBuffer, options); return Promise.resolve(handler); diff --git a/js/web/lib/onnxjs/session-handler.ts b/js/web/lib/onnxjs/session-handler.ts index 0b06a7a747..47e50aeab6 100644 --- a/js/web/lib/onnxjs/session-handler.ts +++ b/js/web/lib/onnxjs/session-handler.ts @@ -1,12 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import {InferenceSession, SessionHandler, Tensor} from 'onnxruntime-common'; +import {InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common'; import {Session} from './session'; import {Tensor as OnnxjsTensor} from './tensor'; -export class OnnxjsSessionHandler implements SessionHandler { +export class OnnxjsSessionHandler implements InferenceSessionHandler { constructor(private session: Session) { this.inputNames = this.session.inputNames; this.outputNames = this.session.outputNames; diff --git a/js/web/lib/wasm/session-handler.ts b/js/web/lib/wasm/session-handler.ts index 4e00878d00..7bc467449c 100644 --- a/js/web/lib/wasm/session-handler.ts +++ b/js/web/lib/wasm/session-handler.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. import {readFile} from 'fs'; -import {env, InferenceSession, SessionHandler, Tensor} from 'onnxruntime-common'; +import {env, InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common'; import {promisify} from 'util'; import {SerializableModeldata, TensorMetadata} from './proxy-messages'; @@ -40,7 +40,7 @@ const decodeTensorMetadata = (tensor: TensorMetadata): Tensor => { } }; -export class OnnxruntimeWebAssemblySessionHandler implements SessionHandler { +export class OnnxruntimeWebAssemblySessionHandler implements InferenceSessionHandler { private sessionId: number; inputNames: string[];