onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { InferenceSession, TrainingSessionHandler } from 'onnxruntime-common';
5

6
import { OnnxruntimeWebAssemblyBackend } from './backend-wasm';
7
import { OnnxruntimeWebAssemblyTrainingSessionHandler } from './wasm/session-handler-training';
8

9
class OnnxruntimeTrainingWebAssemblyBackend extends OnnxruntimeWebAssemblyBackend {
10
  async createTrainingSessionHandler(
11
    checkpointStateUriOrBuffer: string | Uint8Array,
12
    trainModelUriOrBuffer: string | Uint8Array,
13
    evalModelUriOrBuffer: string | Uint8Array,
14
    optimizerModelUriOrBuffer: string | Uint8Array,
15
    options: InferenceSession.SessionOptions,
16
  ): Promise<TrainingSessionHandler> {
17
    const handler = new OnnxruntimeWebAssemblyTrainingSessionHandler();
18
    await handler.createTrainingSession(
19
      checkpointStateUriOrBuffer,
20
      trainModelUriOrBuffer,
21
      evalModelUriOrBuffer,
22
      optimizerModelUriOrBuffer,
23
      options,
24
    );
25
    return Promise.resolve(handler);
26
  }
27
}
28

29
export const wasmBackend = new OnnxruntimeTrainingWebAssemblyBackend();
30