LiteRT for Web with LiteRT.js

LiteRT.js is Google's high performance WebAI runtime, targeting production Web applications. It is a continuation of the LiteRT stack, ensuring multi-framework support and unifying our core runtime across all platforms.

LiteRT.js supports the following core features:

  1. In-browser hardware-accelerated inference: Run models with exceptional CPU performance, accelerated by XNNPack mapped to lightweight WebAssembly (Wasm). For GPU and dedicated hardware scaling (such as NPUs), LiteRT.js natively surfaces both the WebGPU API and the emerging WebNN API empowering fine grained platform-specific optimization.
  2. Multi-framework compatibility: Streamline development semantics by compiling from your preferred ML Framework natively: PyTorch, JAX or TensorFlow.
  3. Iterate on existing pipelines: Out-the-box integration with existing TensorFlow.js architectures by parsing natively supported TensorFlow.js Tensors as direct boundary inputs and outputs.

Installation

Install the @litertjs/core package from npm:

npm install @litertjs/core

The Wasm files are located in node_modules/@litertjs/core/wasm/. For convenience, copy and serve the entire wasm/ folder. Then, import the package and load the Wasm files:

import {loadLiteRt} from '@litertjs/core';

// Load the LiteRT.js Wasm files from a CDN.
await loadLiteRt('https://cdn.jsdelivr.net/npm/@litertjs/core/wasm/')
// Alternatively, host them from your server.
// They are located in node_modules/@litertjs/core/wasm/
await loadLiteRt(`your/path/to/wasm/`);

Model conversion

LiteRT.js uses the same .tflite format as the rest of the LiteRT ecosystem, and it supports existing models on Kaggle and Huggingface. If you have a new PyTorch model, you'll need to convert it.

Convert a PyTorch Model to LiteRT

To convert a PyTorch model to LiteRT, use the litert-torch converter.

import litert_torch

# Load your torch model. We're using resnet for this example.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1)

sample_inputs = (torch.randn(1, 3, 224, 224),)

# Convert the model to LiteRT.
edge_model = litert_torch.convert(resnet18.eval(), sample_inputs)

# Export the model.
edge_model.export('resnet.tflite')

Run the Converted Model

After converting the model to a .tflite file, you can run it in the browser.

import {loadAndCompile} from '@litertjs/core';

// Load the model hosted from your server. This makes an http(s) request.
const model = await loadAndCompile('/path/to/model.tflite', {
    accelerator: 'webgpu',
    // Can select from 'webnn', 'webgpu', & 'wasm'.
    // Unsupported ops on webgpu & webnn automatically fallback to CPU.
});
// The model can also be loaded from a Uint8Array if you want to fetch it yourself.

// Create image input data
const image = new Float32Array(224 * 224 * 3).fill(0);
const inputTensor = new Tensor(image, /* shape */ [1, 3, 224, 224]);

// Run the model
const outputs = await model.run