onnxruntime

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

4
import { InferenceSession, InferenceSessionHandler, SessionHandler, Tensor } from 'onnxruntime-common';
5

6
import { Session } from './session';
7
import { Tensor as OnnxjsTensor } from './tensor';
8

9
export class OnnxjsSessionHandler implements InferenceSessionHandler {
10
  constructor(private session: Session) {
11
    this.inputNames = this.session.inputNames;
12
    this.outputNames = this.session.outputNames;
13
  }
14

15
  async dispose(): Promise<void> {}
16
  inputNames: readonly string[];
17
  outputNames: readonly string[];
18
  async run(
19
    feeds: SessionHandler.FeedsType,
20
    _fetches: SessionHandler.FetchesType,
21
    _options: InferenceSession.RunOptions,
22
  ): Promise<SessionHandler.ReturnType> {
23
    const inputMap = new Map<string, OnnxjsTensor>();
24
    for (const name in feeds) {
25
      if (Object.hasOwnProperty.call(feeds, name)) {
26
        const feed = feeds[name];
27
        inputMap.set(
28
          name,
29
          new OnnxjsTensor(
30
            feed.dims,
31
            feed.type as OnnxjsTensor.DataType,
32
            undefined,
33
            undefined,
34
            feed.data as OnnxjsTensor.NumberType,
35
          ),
36
        );
37
      }
38
    }
39
    const outputMap = await this.session.run(inputMap);
40
    const output: SessionHandler.ReturnType = {};
41
    outputMap.forEach((tensor, name) => {
42
      output[name] = new Tensor(tensor.type, tensor.data, tensor.dims);
43
    });
44
    return output;
45
  }
46
  startProfiling(): void {
47
    this.session.startProfiling();
48
  }
49
  endProfiling(): void {
50
    this.session.endProfiling();
51
  }
52
}
53