Skip to content

Repository files navigation

Burn ONNX

Current Crates.io Version Documentation Test Status license Ask DeepWiki

Import ONNX models into the Burn deep learning framework.

Repository | Burn Repository

Overview

burn-onnx converts ONNX models to native Burn Rust code, allowing you to run models from PyTorch, TensorFlow, and other frameworks on any Burn backend - from WebAssembly to CUDA.

Key features:

  • Generates readable, modifiable Rust source code from ONNX models
  • Produces burnpack weight files for efficient loading
  • Works with any Burn backend (CPU, GPU, WebGPU, embedded)
  • Supports both std and no_std environments
  • Full opset compliance: all supported operators work across ONNX opset versions 1 through 24
  • Extensible: hook operators outside the supported set (custom/vendor domains) to your own Rust kernels, or override the generated code for a built-in operator
  • Graph simplification (enabled by default): attention coalescing, constant folding, constant shape propagation, idempotent-op elimination, identity-element elimination, CSE, dead code elimination, and permute-reshape detection

Quick Start

Add to your Cargo.toml:

[build-dependencies]
burn-onnx = "0.21"

In your build.rs:

use burn_onnx::ModelGen;

fn main() {
    ModelGen::new()
        .input("src/model/my_model.onnx")
        .out_dir("model/")
        .run_from_script();
}

Include the generated code in src/model/mod.rs:

pub mod my_model {
    include!(concat!(env!("OUT_DIR"), "/model/my_model.rs"));
}

Then use the model:

use crate::model::my_model::Model;

let model: Model = Model::default();
let output = model.forward(input_tensor);

For detailed usage instructions, see the ONNX Import Guide in the Burn Book.

Custom Operators

Operators outside the supported set do not have to block an import. That covers vendor domains such as com.microsoft, ops from a framework's custom export, and op types not implemented yet. Register a hook and supply the Rust yourself:

// build.rs
ModelGen::new()
    .input("src/model/my_model.onnx")
    .out_dir("model/")
    .register_custom_op(FftReal)      // handles my_domain::FftReal
    .register_op_override(MyMatMul)   // replaces the generated code for every MatMul
    .run_from_script();
  • CustomOp supplies type inference and code generation for one ONNX (op_type, domain). It can read the node's attributes and its constant inputs.
  • OpOverride replaces the generated code for a built-in operator, to route it to a fused, quantized, or hardware-specific kernel of your own. Type inference still comes from the built-in.

Everything a hook needs is re-exported from burn_onnx::ext, so implementing one does not mean depending on onnx-ir or matching proc-macro2/quote versions by hand.

Not sure which operators a given model needs? Run the import with no hooks registered: it fails with a list of every unsupported operator, its domain, and how many nodes use it.

Runnable example: custom-op-hooks. Full reference: "Custom Operators and Overrides" in the Development Guide.

Examples

Example Description
onnx-inference Basic ONNX model inference
image-classification-web WebAssembly/WebGPU image classifier
custom-op-hooks Hooks for custom ops and built-in op overrides

Model Validation

We validate burn-onnx against 26 real-world models spanning image classification, object detection, depth estimation, NLP, speech, and generative AI. Each model check verifies the full pipeline: ONNX import, Rust codegen, weight loading, and numerical accuracy against ONNX Runtime reference outputs.

Supported Operators

See the Supported ONNX Operators table for the complete list of supported operators. Anything outside that table can still be imported by registering a hook — see Custom Operators above.

Contributing

We welcome contributions! Please read the Contributing Guidelines before opening a PR, and the Development Guide for architecture and implementation details.

For questions and discussions, join us on Discord.

License

Licensed under either of Apache License, Version 2.0 or MIT license at your option.

About

Convert ONNX models into native, backend-agnostic Burn code for inference and fine-tuning.

Resources

Contributing

Stars

120 stars

Watchers

4 watching

Forks

Releases

Packages

Used by

Contributors

Languages