PyTorch 2.x models can run client-side with ONNX Runtime Web, WebGPU, and WASM fallback. Export, ship, verify, and tune them. Read now.
Why run a PyTorch model in the browser
Client-side inference keeps weights and inputs on the user’s device. That cuts round-trips to a backend, removes a class of network failure from the critical path, and lets demos and tools work offline once assets are cached. For many interactive apps—previews, form assistants, local classifiers—latency matters more than peak server throughput, so shipping a model with the page can be the right default.
PyTorch 2.x training and research still usually live on a GPU server. Browser runtimes are different: you export a fixed graph, load it through a web-native engine, and pick a compute path that the device actually supports. ONNX Runtime Web is the common bridge: export from PyTorch, run with WebGPU when available, and fall back to WASM when the GPU path is missing or unstable.
Export a graph the browser can load
Start from a model that is already frozen for inference—no training loops, no dynamic control flow you cannot express in ONNX. Trace or export to ONNX with fixed input shapes where you can; dynamic axes are fine when the product needs them, but they widen the surface you must test. Strip training-only ops, fuse what the exporter allows, and keep operator coverage in mind: an op that works in PyTorch can be unsupported or slow under ONNX Runtime Web.
Ship the smallest graph that still meets product quality. Quantize when accuracy holds under your own checks. Prefer one clear input/output contract (tensor names, dtypes, layout) and document it next to the export script so the web loader does not guess. Version the artifact with the page so a stale CDN cache cannot mix an old model with new preprocessing code.
Ship runtime, assets, and a fallback path
Load ONNX Runtime Web as you would any other critical dependency: pin a known build, host it yourself or from a CDN you control, and fail loudly if the script does not load. Fetch model weights with caching headers that match how often you re-export. On first visit, budget for a one-time download; on return visits, serve from the HTTP cache or a service worker so inference does not wait on the network again.
- Try WebGPU first when the browser exposes it and your session can create a device; it is the high-throughput path for many models.
- If WebGPU is absent, blocked, or fails at init, fall back to WASM so the feature still runs, even if slower.
- Surface a simple status (“GPU”, “CPU/WASM”, or “unavailable”) so users and support know which path is active.
Keep preprocessing on the same side as inference: normalize, resize, and tokenize in JavaScript or WASM with the same rules you used at export time. Mismatched mean/std, channel order, or token IDs produce silent quality bugs that look like a bad model.
Verify correctness, then tune
Before polish, prove parity. Run the same fixed inputs through the original PyTorch path and the browser path; compare outputs with a numeric tolerance suited to the task (classification scores vs. generative tokens need different bars). Automate a short smoke set in CI that loads the page assets in a headless browser, exercises WebGPU when the environment allows, and always covers WASM. Treat init failures, missing ops, and OOM-style aborts as first-class test cases, not edge notes.
Tune only after parity is green. Profile load time separately from first-inference time. Shrink the ONNX file, lazy-load the runtime until the UI needs it, and warm the session on idle if users reliably click the feature soon after. Cap concurrent inferences so one tab cannot thrash memory. Prefer clearer errors and a working WASM path over a fragile WebGPU-only experience. The goal is a pipeline you can re-export, ship, verify, and adjust without rewriting the product each time the weights change.