@hoff97/tensor-js/dist/lib/onnx/model.d.ts

PyTorch like deep learning inferrence library

import { onnx } from 'onnx-proto';
import { Variable } from '../autograd/variable';
import { Mode, Module } from '../model/module';
import Tensor from '../types';
import { OnnxNode } from './node';
export declare type NodeId = number;
interface IntermediaryRes {
    value: Tensor<any>;
    used: number;
}
export interface ModelArgs {
    /**
     * Precision with which float tensors should be loaded. If 16 is specified,
     * all 32 bit floats are casted to 16 bit floats.
     *
     * Defaults to 32.
     */
    precision?: 16 | 32;
    /**
     * Constants that should not be transferred to another device.
     *
     * Useful for operations that should only happen only on the CPU
     */
    noConvertConstants?: string[];
    /**
     * Nodes that should not be transferred to another device
     *
     * Useful for operations that should only happen only on the CPU
     */
    noConvertNodes?: NodeId[];
    /**
     * If the model should be loaded in training or inference mode.
     * In inference mode, no gradients can be supported and
     * the model expects tensors as input.
     *
     * In training mode, gradients can be computed and
     * the inputs should be variables (although with noGrad: true if no gradient
     * is needed for the respective input)
     */
    mode?: Mode;
}
export declare class OnnxModel extends Module {
    private version;
    private inputSet;
    private nodes;
    private nodeIds;
    private defaultReady;
    private intermediaries;
    private constants;
    private noConvertConstants;
    private noConvertNodes;
    private modelProto;
    private nodeIdCounter;
    private precision;
    inputs: onnx.IValueInfoProto[];
    outputs: string[];
    /**
     * Builds a new onnx model
     *
     * @param buffer Onnx model
     * @param args Optional arguments for the model
     */
    constructor(buffer: ArrayBuffer | Uint8Array, args?: ModelArgs);
    private initNodes;
    private initializer;
    /**
     * Do a forward pass for the specified inputs
     *
     * @param wait Number of milliseconds to wait between each layer. This
     *             is especially useful, if your model is complex and
     *             you dont want your model to block your whole application.
     * @param returnIntermediary return after the given intermediary result
     *                           has been computed.
     */
    forward(inputs: Tensor<any>[], wait?: number): Promise<Tensor<any>[]>;
    protected initializeForward(inputs: Tensor<any>[], intermediaryRes: {
        [name: string]: IntermediaryRes;
    }, nodes: {
        [id: number]: {
            variableInputs: number;
        };
    }, nodesReady: NodeId[]): void;
    protected getInputsToNode(node: OnnxNode, intermediaryRes: {
        [name: string]: IntermediaryRes;
    }): {
        inputs: Tensor<any>[];
        toDelete: string[];
    };
    protected propagateResults(node: OnnxNode, intermediaryRes: {
        [name: string]: IntermediaryRes;
    }, outputs: Tensor<any>[], nodes: {
        [id: number]: {
            variableInputs: number;
        };
    }, nodesReady: NodeId[]): void;
    /**
     * Transfer the model to the CPU
     */
    toCPU(): Promise<void>;
    /**
     * Transfer the model to WASM
     */
    toWASM(): Promise<void>;
    /**
     * Transfer the model to the GPU
     */
    toGPU(): Promise<void>;
    /**
     * Optimize the model.
     */
    optimize(): void;
    prune(intermediariesToDelete?: string[]): void;
    private pruneIntermediaries;
    private removeNode;
    private insertNode;
    getNodeWithOutput(output: string): number | undefined;
    getNodeWithInput(output: string): number | undefined;
    resolveConstant(name: string): Tensor<any> | undefined;
    getNodes(): {
        [id: number]: OnnxNode;
    };
    /**
     * Deletes the model
     *
     * This will release the memory/framebuffers (depending on the backend)
     */
    delete(): void;
    getSubModules(): Module[];
    getParameters(): Variable<any>[];
}
export {};