Running TensorFlow.js with the Wasm Backend

This page answers one task: a web app runs a TensorFlow.js model — image classification, pose detection, a small text model — and needs it to run fast and consistently on devices where the GPU backends are slow, unavailable or unreliable, by using the WebAssembly backend.

Prerequisites

  • [ ] A TensorFlow.js app (@tensorflow/tfjs or @tensorflow/tfjs-core with a converted model).
  • [ ] npm install @tensorflow/tfjs-backend-wasm.
  • [ ] Ability to serve static files, including .wasm, from your origin or a CDN.

What the Wasm backend is

TensorFlow.js separates the model API from backends that execute operations: WebGL and WebGPU run kernels on the GPU, the CPU backend runs them in plain JavaScript, and the Wasm backend runs them in WebAssembly using XNNPACK, a library of highly optimised neural-network kernels. Compared with the plain CPU backend, it is typically an order of magnitude faster. Compared with the GPU backends, the picture depends on the model: for small and medium models — many mobile vision models and classifiers — the Wasm backend often matches or beats WebGL, because GPU backends pay overheads for uploading data, compiling shaders and synchronising results. For large models with big matrix multiplications, GPUs win.

The Wasm backend also behaves more predictably: results are consistent across devices (no GPU-driver float differences), the first inference is not delayed by shader compilation, and it works where WebGL is blocklisted or software-emulated. That makes it a strong default for many production apps, with GPU backends as an upgrade where they measurably help.

TensorFlow.js backends compared The WebGL and WebGPU backends run on the GPU and excel on large models but have shader-compilation warm-up and driver variability. The Wasm backend runs XNNPACK kernels on the CPU with SIMD and threads, starts fast and is consistent, and wins on small and medium models. The plain CPU backend is a slow fallback. WebGL / WebGPU GPU kernels, big-model winner shader compile on first run driver-dependent results large models Wasm (XNNPACK) SIMD + threads on the CPU fast first inference consistent across devices small/medium models CPU (plain JS) no special support needed 10× slower than Wasm last-resort fallback fallback only

Step 1 — register and select the backend

import * as tf from "@tensorflow/tfjs-core";
import "@tensorflow/tfjs-backend-wasm";
import { setWasmPaths } from "@tensorflow/tfjs-backend-wasm";

setWasmPaths("/tfjs-wasm/");                     // where tfjs-backend-wasm*.wasm files are served
await tf.setBackend("wasm");
await tf.ready();
console.log(tf.getBackend());                    // "wasm"

The backend ships three binaries: tfjs-backend-wasm.wasm (baseline), tfjs-backend-wasm-simd.wasm (SIMD) and tfjs-backend-wasm-threaded-simd.wasm (SIMD and threads). At startup it detects what the browser supports and loads the best one. Copy all three into your static output — bundlers do not discover them automatically — and point setWasmPaths at their location.

If you import the full @tensorflow/tfjs bundle, it registers the WebGL and CPU backends too; importing tfjs-core, the converter and only the backends you use keeps the JavaScript bundle considerably smaller.

Step 2 — enable threads with cross-origin isolation

The threaded build uses a pool of web workers sharing memory, which requires SharedArrayBuffer and therefore cross-origin isolation. Serve the page with:

Cross-Origin-Opener-Policy: same-origin
Cross-Origin-Embedder-Policy: require-corp

and serve model files and Wasm binaries same-origin or with CORP/CORS headers. Without isolation the backend falls back to the single-threaded SIMD build, which is still fast. Control the thread count with setThreadsCount(n) before initialising; on phones, two to four threads usually perform best. The isolation details are in detecting cross-origin isolation at runtime.

Step 3 — load and run the model

import * as tfconv from "@tensorflow/tfjs-converter";

const model = await tfconv.loadGraphModel("/models/mobilenet_v3/model.json");

function classify(image) {
  return tf.tidy(() => {
    const input = tf.browser.fromPixels(image).resizeBilinear([224, 224]).toFloat().div(255).expandDims(0);
    return model.predict(input);
  });
}

const logits = classify(imgElement);
const probs = await logits.data();                 // read results asynchronously
logits.dispose();

tf.tidy disposes intermediate tensors automatically; tensors that escape it — like logits — must be disposed by hand, or the backend’s memory grows with every call. Read results with data() rather than dataSync() to avoid blocking. Run a warm-up inference at startup so the first real request is not slower; with the Wasm backend this mostly allocates buffers and touches the kernels once.

Choosing a backend at runtime At startup the app tries WebGPU if available, then benchmarks the candidate backends on a warm-up inference of the actual model, and keeps the fastest. The Wasm backend is the reliable fallback when GPU backends are missing or slow. startup load model once candidates webgpu, webgl, wasm warm-up + time 3 runs each pick fastest tf.setBackend(best) fallback wasm if GPU fails

Step 4 — choose the backend by measuring

Do not guess which backend is faster for your model on your users’ devices; measure at startup with the actual model:

async function pickBackend(model, sample) {
  const results = [];
  for (const name of ["webgpu", "webgl", "wasm"]) {
    if (!(await tf.setBackend(name).catch(() => false))) continue;
    await tf.ready();
    tf.tidy(() => model.predict(sample));                     // warm-up (shader compile, allocation)
    const t0 = performance.now();
    for (let i = 0; i < 3; i++) await tf.tidy(() => model.predict(sample)).data();
    results.push([name, (performance.now() - t0) / 3]);
  }
  results.sort((a, b) => a[1] - b[1]);
  await tf.setBackend(results[0][0]);
  return results;
}

Cache the choice per device in localStorage with the model version, so the benchmark runs once. Report the chosen backend and timings in telemetry; the distribution across real devices is usually surprising.

Step 5 — run inference in a worker

Even with the Wasm backend’s own thread pool, the main thread coordinates each inference. For video-rate workloads, run TensorFlow.js in a dedicated worker: transfer frames as ImageBitmaps, run the model there, and post back small results. The Wasm backend works in workers without changes, and the page stays responsive. The worker pattern is described in keeping the UI responsive during long Wasm tasks.

Loading models quickly

Model weights often dominate startup. A converted graph model is a model.json describing the graph plus one or more binary weight shards; the browser downloads them in parallel, then the backend uploads or copies the weights into its memory. Serve shards with long-lived cache headers and content-hashed names, and precompress them, though already-quantised weights compress less. Cache the model with TensorFlow.js’s model.save("indexeddb://…") after the first download and load from IndexedDB on later visits, checking a version string so updates replace the cached copy. For an app whose model is not needed on the first screen, start downloading it after the page becomes interactive, and show a clear “preparing model” state until the warm-up inference completes. With the Wasm backend, warm-up is short — there are no shaders to compile — so most of the perceived startup is the download, which caching removes entirely after the first visit.

Making models fit the CPU

The Wasm backend rewards models designed for CPUs. Prefer architectures built for mobile inference — MobileNet, EfficientNet-Lite, small transformer variants — over large server models. Quantised weights shrink downloads; TensorFlow.js can load uint8 or float16 quantised weights and dequantise them on load, cutting model size by two to four times with little accuracy loss, as described in quantizing models for Wasm inference. Reduce input resolution where accuracy allows: halving the side of an image input quarters the computation. Batch work when latency permits — classifying eight thumbnails in one call uses the kernels more efficiently than eight calls. And check that every operation in the model is supported by the Wasm backend; unsupported operations either fail or fall back, and the backend’s documentation lists the kernels it implements. When a model needs operations the Wasm backend lacks, ONNX Runtime Web with its own Wasm backend is an alternative worth comparing.

Expected output

MobileNet v3 classification runs in about 12 ms per image with the threaded SIMD Wasm backend on a laptop and about 45 ms on a mid-range phone; the startup benchmark picks Wasm on most phones and WebGPU on desktops with discrete GPUs; and memory stays flat over thousands of inferences.

Gotchas

  • .wasm files not served. The backend fails to initialise. Copy all three binaries and call setWasmPaths.
  • No threads. The page is not cross-origin isolated. Add COOP/COEP or accept single-threaded SIMD.
  • Growing memory. Tensors not disposed. Use tf.tidy and dispose escaping tensors.
  • dataSync() on the main thread. Blocks the UI. Use data().
  • Benchmarking without warm-up. The first GPU run includes shader compilation. Warm up every backend before timing.
  • Re-downloading the model each visit. Save it to IndexedDB and load from there with a version check.
  • Assuming the GPU is faster. For small models it often is not. Measure per device.

Performance note

MobileNet v3 (224×224) on a mid-range Android phone: plain CPU backend 520 ms, Wasm baseline 98 ms, Wasm SIMD 58 ms, Wasm SIMD with threads 41 ms, WebGL 63 ms after a 900 ms first-run shader compile.

MobileNet v3 inference on a mid-range phone Milliseconds per inference for MobileNet v3 at 224 by 224 on a mid-range Android phone with the plain CPU backend, the Wasm backend without SIMD, with SIMD, with SIMD and threads, and the WebGL backend after warm-up. ms per inference CPU (plain JS) 520 ms Wasm baseline 98 ms Wasm SIMD 58 ms Wasm SIMD + threads 41 ms WebGL (after warm-up) 63 ms

Frequently Asked Questions

Does the Wasm backend work in Node? Yes — the same package runs in Node, and @tensorflow/tfjs-node offers native TensorFlow for heavier server workloads.

Can I mix backends for different models? The backend is global per TensorFlow.js instance; run models with different backends in different workers.

How large are the Wasm binaries? Each is a few hundred kilobytes compressed; only one is downloaded per device.

Are results identical across backends? Close but not bit-identical; small floating-point differences between kernels are normal.

What if setBackend("wasm") returns false? The binaries failed to load — check the network panel for 404s or wrong MIME types — and fall back to another backend rather than failing the feature.

Does the Wasm backend support training? It supports the operations needed for many training tasks, but in-browser training is usually better on GPU backends.

Can I run several models at once? Yes, on one backend instance. They share the thread pool, so run them sequentially or in separate workers if latency matters.

← Back to Machine Learning Inference in the Browser