@ignitionai/backend-onnx

Single responsibility: turn a trained TF.js model into a portable ONNX model, and run inference against an ONNX model at runtime. This package is what makes IgnitionAI a serious tool rather than a browser-only toy.

Source: packages/backend-onnx/src/ on GitHub.

What “Train → Deploy” means in practice

The pipeline has four stages:

1. Train        → @ignitionai/backend-tfjs (browser, TF.js)
2. Export       → @ignitionai/backend-onnx (TF.js → SavedModel)
3. Convert      → tf2onnx (one-time Python step, SavedModel → .onnx)
4. Deploy       → anywhere that has an ONNX runtime

You stay in JavaScript for everything except stage 3, which is a one-time setup dance — install the Python tooling once, and after that every model you train can be converted with a one-liner.

Public API surface

Stage 1 — Train in the browser

Nothing new here — you’re using @ignitionai/backend-tfjs exactly as described in the Quickstart:

import { IgnitionEnvTFJS } from '@ignitionai/backend-tfjs'
import { CartPoleEnv } from '@ignitionai/environments'
 
const env = new IgnitionEnvTFJS(new CartPoleEnv())
env.train('dqn')
// ... wait for convergence ...
env.stop()

When training finishes, the env.agent property holds a DQNAgent with a trained TF.js model under .model. Stage 2 takes that model and prepares it for export.

Stage 2 — Export with saveForOnnxExport

import { saveForOnnxExport } from '@ignitionai/backend-onnx'
 
const { savedModelPath, conversionScript } = await saveForOnnxExport(
  env.agent.model,
  './export',
  undefined,   // onnxOutputPath (optional)
  15,          // opset version (default: 13, recommended: 15+)
)
 
console.log(`SavedModel written to: ${savedModelPath}`)
console.log(`Run this to convert: ${conversionScript}`)

saveForOnnxExport does two things:

1. Writes the TF.js model to a TensorFlow SavedModel directory — the format tf2onnx can consume.
2. Generates a convert.sh (or equivalent) script you run once to invoke tf2onnx.

You can also call generateConversionScript() directly if you want to customize the conversion flags (opset version, signature inputs, etc.).

Stage 3 — Convert with tf2onnx (one-time setup)

This is the only non-JS step. You install tf2onnx once via pip, then run the conversion script:

# one-time setup
pip install -U tf2onnx
 
# convert the SavedModel to ONNX
bash ./export/convert.sh
# produces ./export/model.onnx

The result is a platform-agnostic .onnx file — a few hundred kilobytes for a CartPole-class policy, a few megabytes for a bigger one.

Stage 4 — Deploy anywhere

Once you have the .onnx file, the world opens up. Non-exhaustive deploy targets:

The same .onnx file works everywhere. That’s the whole point of the ONNX format.

Bundle size — inference-only runtime

The @ignitionai/backend-onnx inference code (OnnxAgent + universal runtime) adds ~2KB gzipped to your bundle. The ONNX Runtime itself (onnxruntime-web or onnxruntime-node) is loaded as an optional peer dependency and is not bundled by default — you bring your own runtime matching your target platform.

For browser deployments, the ~800KB ONNX Runtime is a one-time cost — it serves every ONNX model you load, regardless of architecture.

Stage 5 — Run inference in JS via OnnxAgent

If you want to run the trained model back in the browser without the training dependencies, OnnxAgent is a drop-in AgentInterface that auto-detects the environment and loads the correct ONNX Runtime (Web in browser, Node in server):

import { OnnxAgent } from '@ignitionai/backend-onnx'
 
const agent = new OnnxAgent({
  modelPath: './model.onnx',
  actionSize: 2,
})
 
await agent.load()
 
const observation = [0.1, -0.02, 0.05, 0.01]
const action = await agent.getAction(observation)
console.log(`Agent picked action ${action}`)

Three lines. The model runs at full speed on WebGPU where available, and because ONNX Runtime is much lighter than TensorFlow.js, the bundle size for an inference-only page is a fraction of a training-capable page.

OnnxAgent satisfies AgentInterface, so you can also feed it to an InferenceEnv:

import { IgnitionEnv } from '@ignitionai/core'
import { OnnxAgent } from '@ignitionai/backend-onnx'
import { CartPoleEnv } from '@ignitionai/environments'
 
const env = new IgnitionEnv(new CartPoleEnv())
const agent = new OnnxAgent({ modelPath: './model.onnx', actionSize: 2 })
await agent.load()
 
// Manually drive the inference loop
for (let step = 0; step < 500; step++) {
  await env.inferStep()
}

Loading models from HuggingFace Hub

The loadOnnxModelFromHub helper fetches an .onnx file directly from a HF repo — handy for sharing pre-trained weights:

import { loadOnnxModelFromHub, OnnxAgent } from '@ignitionai/backend-onnx'
 
const modelBlob = await loadOnnxModelFromHub({
  repoId: 'ignitionai/cartpole-dqn-v1',
  filename: 'model.onnx',
})
 
const agent = new OnnxAgent({ modelBlob, actionSize: 2 })
await agent.load()

Inspecting an ONNX model

If you want to debug a converted model (wrong input shape, wrong output count, etc.), use inspectSession:

import { createOnnxSession, inspectSession } from '@ignitionai/backend-onnx'
 
const session = await createOnnxSession('./model.onnx')
console.log(inspectSession(session))
// → { inputs: [...], outputs: [...], opsetVersion: ... }

This is the first thing to try when getAction() throws an “input shape mismatch” error — it tells you exactly what the converted model expects.

Previous: ← @ignitionai/backend-tfjs · Next: @ignitionai/storage →