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.
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.jswhich 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.
<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>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);
});
});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);
});
});Currently, tract requires has some restrictions on dynamic dimensions. If your model has a dynamic dimension, there's multiple solutions:
- 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]],
},
});- 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]],
},
});- Turn
optimizeoff. 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,
});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.
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.
All original work licensed under either of
- Apache License, Version 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
- MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT) at your option.
Contributions are very welcome! See CONTRIBUTING.md.