Options
All
  • Public
  • Public/Protected
  • All
Menu

tractjs

npm version Test Deploy to Github Pages

Run ONNX and TensorFlow inference in the browser. A thin wrapper on top of tract.

The Open Neural Network Exchange is a format which many popular libraries like PyTorch, TensorFlow and MXNet can export to which allows tractjs to run neural networks from (almost) any library.

Website | API Docs

Why tractjs instead of ONNX.js?

There is currently one other usable ONNX runner for the browser, ONNX.js. There are a couple of things tractjs does better:

  • tractjs supports more operators:
    • LSTMs (even bidirectional) are supported while ONNX.js does not support any recurrent networks.
    • Some ONNX-ML models like decision tree classifiers are also supported.
  • tractjs is more convenient to use. It can build to a single file tractjs.min.js which contains the inlined WASM and WebWorker. The WASM backend of ONNX.js can not as easily be used without a build system.

There are however also some downsides to tractjs. See the FAQ.

Getting started

Without a bundler

<html>
  <head>
    <meta charset="utf-8" />
    <script src="https://unpkg.com/tractjs/dist/tractjs.min.js"></script>
    <script>
      tractjs.load("path/to/your/model").then((model) => {
        model
          .predict([new tractjs.Tensor(new Float32Array([1, 2, 3, 4]), [2, 2])])
          .then((preds) => {
            console.log(preds);
          });
      });
    </script>
  </head>
</html>

With a bundler

npm install tractjs
import * as tractjs from "tractjs";

tractjs.load("path/to/your/model").then((model) => {
  model
    .predict([new tractjs.Tensor(new Float32Array([1, 2, 3, 4]), [2, 2])])
    .then((preds) => {
      console.log(preds);
    });
});

With Node.js

tractjs now runs in Node.js! Models are fetched from the file system.

const tractjs = require("tractjs");

tractjs.load("./path/to/your/model").then((model) => {
  model
    .predict([new tractjs.Tensor(new Float32Array([1, 2, 3, 4]), [2, 2])])
    .then((preds) => {
      console.log(preds);
    });
});

FAQ

Why does my model with dynamic input dimensions not work?

Currently, tract requires has some restrictions on dynamic dimensions. If your model has a dynamic dimension, there's multiple solutions:

  1. Declare a dynamic dimension via an input fact. Input facts are a way to provide additional information about input type and shape that can not be inferred via the model data:
const model = await tractjs.load("path/to/your/model", {
  inputFacts: {
    0: ["float32", [1, "s", 224, 224]],
  },
});
  1. Set fixed input dimensions via input facts. This is of course not ideal because subsequently the model can only be passed inputs with this exact shape:
const model = await tractjs.load("path/to/your/model", {
  inputFacts: {
    // be careful with image model input facts! here I use ONNX's NCHW format
    // if you are using TF you will probably need to use NHWC (`[1, 224, 224, 3]`).
    0: ["float32", [1, 3, 224, 224]],
  },
});
  1. Turn optimize off. This is the nuclear option. It will turn off all optimizations relying on information about input shape. This will make sure your model work (even with multiple dynamic dimensions) but significantly impact performance:
const model = await tractjs.load("path/to/your/model", {
  optimize: false,
});

What about size?

At the time of writing, tractjs is very large for web standards (6.2MB raw, 2.1MB gzipped). This is due to tract being quite large, and due to some overhead from inlining the WASM. But it's not as bad as it sounds. You can load tractjs lazily along your demo, where you will likely have to load significantly large weights too.

If you are working on a very size-sensitive application, get in touch and we can work on decreasing the size. There are some more optimizations to be done (e. g. an option not to inline WASM, and removing panics from the build). There is also ongoing work in tract to decrease size.

What about WebGL / WebNN support?

tractjs are bindings to the tract Rust library which was originally not intended to be run on the web. WebGL / WebNN support would be great, but would require lots of web-specific changes in tract so it is currently not under consideration.

License

Apache 2.0/MIT

All original work licensed under either of

Contribution

Contributions are very welcome! See CONTRIBUTING.md.

Index

Classes

Type aliases

Variables

Functions

Type aliases

DataType

DataType: "int8" | "uint8" | "int16" | "uint16" | "int32" | "uint32" | "float32" | "float64"

DimSymbol

DimSymbol: string | { id: string; intercept: number; slope: number }

Fact

Fact: [DataType, Shape]

Format

Format: "onnx" | "tensorflow"

Metadata

Metadata: {}

Type declaration

  • [key: string]: string

Options

Options: { format?: Format; inputFacts?: Record<number, Fact>; inputs?: Array<string>; optimize?: undefined | false | true; outputs?: Array<string> }

Model loading options.

Type declaration

  • Optional format?: Format

    The model format. Either "onnx" or "tensorflow". If undefined, will attempt to infer from URL file extension.

  • Optional inputFacts?: Record<number, Fact>

    Mapping of indices to facts to set for the input. Each fact is passed to set_input_fact.

  • Optional inputs?: Array<string>

    The node names of model inputs. Passed to set_input_names.

  • Optional optimize?: undefined | false | true

    Whether to optimize the model. Currently only works if the input shape is fully determined. If you need e. g. dynamic batch sizes set this to false. true by default.

  • Optional outputs?: Array<string>

    The node names of model outputs. Passed to set_output_names.

Shape

Shape: Array<number | DimSymbol>

TypedArray

TypedArray: Int8Array | Uint8Array | Int16Array | Uint16Array | Int32Array | Uint32Array | Uint8ClampedArray | Float32Array | Float64Array

Variables

Const load

load: typeof load

Loads a model. Alias for Model.load.

param

The URL to load the model from. Will be passed to fetch.

param

Additional options. See Options for details.

Functions

terminate

  • terminate(): void
  • Utility function for terminating the worker from outside of the module.

    Returns void

Legend

  • Constructor
  • Property
  • Method
  • Static method

Generated using TypeDoc