onnxruntime/js/web/lib/onnxjs/session-handler-inference.ts
Caroline Zhu e3b043ba17
[js/web/training] runTrainStep implementation (#18006)
### Description
* based on design document & following InferenceSession's run
implementation, implemented TrainingSession.runTrainStep

### Motivation and Context
* Adding web bindings for training

#### Related work
* #16521 allowed for training artifacts to be built
* #17333 added interfaces for training
* #17474 allowed for training package to be built + added training
backend to web package
* #17891 implementation for createTrainingSession on the TypeScript side
**[SHOULD BE MERGED IN BEFORE THIS PR]**

---------

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
Co-authored-by: Ashwini Khade <askhade@microsoft.com>
2023-11-02 08:32:50 -07:00

45 lines
1.5 KiB
TypeScript

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
import {InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common';
import {Session} from './session';
import {Tensor as OnnxjsTensor} from './tensor';
export class OnnxjsSessionHandler implements InferenceSessionHandler {
constructor(private session: Session) {
this.inputNames = this.session.inputNames;
this.outputNames = this.session.outputNames;
}
async dispose(): Promise<void> {}
inputNames: readonly string[];
outputNames: readonly string[];
async run(
feeds: SessionHandler.FeedsType, _fetches: SessionHandler.FetchesType,
_options: InferenceSession.RunOptions): Promise<SessionHandler.ReturnType> {
const inputMap = new Map<string, OnnxjsTensor>();
for (const name in feeds) {
if (Object.hasOwnProperty.call(feeds, name)) {
const feed = feeds[name];
inputMap.set(
name,
new OnnxjsTensor(
feed.dims, feed.type as OnnxjsTensor.DataType, undefined, undefined,
feed.data as OnnxjsTensor.NumberType));
}
}
const outputMap = await this.session.run(inputMap);
const output: SessionHandler.ReturnType = {};
outputMap.forEach((tensor, name) => {
output[name] = new Tensor(tensor.type, tensor.data, tensor.dims);
});
return output;
}
startProfiling(): void {
this.session.startProfiling();
}
endProfiling(): void {
this.session.endProfiling();
}
}