jax-js-nonconsuming (Fork)
    Preparing search index...

    Class ONNXModel

    Loads an ONNX model (.onnx file) and provides a jax-js function that evaluates it.

    The returned function takes input tensors and returns output tensors. Input tensors are NOT consumed — dispose them yourself when done. Initializers (model weights, data) are cached and reused across calls.

    import { ONNXModel } from "@hamk-uas/jax-js-nonconsuming-onnx";
    import { numpy as np } from "@hamk-uas/jax-js-nonconsuming";

    const modelBytes = await fetch("./model.onnx").then((r) => r.bytes());
    const model = new ONNXModel(modelBytes);

    try {
      using input = np.ones([1, 3, 224, 224]);
      const { output } = model.run({ input });
      output.dispose();
    } finally {
      model.dispose();
    }

    Constructors

    • Load a new model from binary contents of an .onnx file.

      Parameters

      • modelBytes: Uint8Array<ArrayBuffer>

      Returns ONNXModel

    Properties

    model: ModelProto

    The parsed model as a Protobuf object.

    run: (
        inputs: Record<string, np.Array>,
        options?: ONNXRunOptions,
    ) => Record<string, np.Array>

    Run a forward pass on the model. This function is bound to this, so you don't need to create a separate closure to pass it to transformations such as jit() and grad().

    Methods

    • Dispose of this model and free model weights.

      After disposing, run() should not be called anymore, it will not be able to find the missing variables.

      Returns void

    Constructors
    Properties
    Methods