Struct ModelGen 

pub struct ModelGen { /* private fields */ }

Builder for generating Burn model code from ONNX files.

ModelGen converts ONNX models into Burn-compatible Rust source code and model weights. It can be used from both build scripts and CLI applications.

Conversion Process

1. Parses ONNX model file(s)
2. Converts ONNX operations to Burn nodes using the node registry
3. Generates Rust source code with type-safe tensor operations
4. Saves model weights in BurnPack (.bpk) format

Examples

Using in a build script (build.rs)

use burn_onnx::ModelGen;

ModelGen::new()
    .input("path/to/model.onnx")
    .out_dir("model/")
    .run_from_script();

This generates code in $OUT_DIR/model/model.rs which can be included in your crate:

include!(concat!(env!("OUT_DIR"), "/model/model.rs"));

Using from CLI

use burn_onnx::ModelGen;

ModelGen::new()
    .input("path/to/model.onnx")
    .out_dir("src/model/")
    .run_from_cli();

Development mode for debugging

use burn_onnx::ModelGen;

ModelGen::new()
    .input("path/to/model.onnx")
    .out_dir("model/")
    .development(true)  // Generates .onnx.txt and .graph.txt debug files
    .run_from_cli();

Implementations

impl ModelGen

pub fn new() -> Self

Creates a new ModelGen builder with default settings.

Default configuration:

• Development mode: off
• Load strategy: LoadStrategy::File

Examples

use burn_onnx::ModelGen;

ModelGen::new()
    .input("model.onnx")
    .out_dir("./out")
    .run_from_cli();

pub fn out_dir(&mut self, out_dir: &str) -> &mut Self

Sets the output directory for generated files.

When used with run_from_script, this path is appended to $OUT_DIR. When used with run_from_cli, this is the absolute or relative path where files will be written.

Arguments

• out_dir - Directory path where generated .rs and record files will be saved

Examples

use burn_onnx::ModelGen;

ModelGen::new()
    .out_dir("model/")  // In build.rs: $OUT_DIR/model/
    .input("model.onnx")
    .run_from_script();

pub fn input(&mut self, input: &str) -> &mut Self

Adds an ONNX model file to convert.

Multiple input files can be added by calling this method multiple times. Each input file will generate a separate .rs file with the same base name.

Arguments

• input - Path to the ONNX model file (.onnx)

Examples

use burn_onnx::ModelGen;

ModelGen::new()
    .input("encoder.onnx")
    .input("decoder.onnx")  // Generate multiple models
    .out_dir("models/")
    .run_from_cli();

pub fn development(&mut self, development: bool) -> &mut Self

Enables development mode for debugging.

When enabled, generates additional debug files alongside the Rust source:

• <model>.onnx.txt - Debug representation of the parsed ONNX graph
• <model>.graph.txt - Debug representation of the converted Burn graph

Arguments

• development - If true, generate debug files

Examples

use burn_onnx::ModelGen;

ModelGen::new()
    .input("model.onnx")
    .out_dir("debug/")
    .development(true)  // Generates model.onnx.txt and model.graph.txt
    .run_from_cli();

pub fn load_strategy(&mut self, strategy: LoadStrategy) -> &mut Self

Sets the weight loading strategy for the generated model.

See LoadStrategy for available options.

Examples

use burn_onnx::{ModelGen, LoadStrategy};

// WASM or embedded: load weights from bytes at runtime
ModelGen::new()
    .input("model.onnx")
    .out_dir("model/")
    .load_strategy(LoadStrategy::Bytes)
    .run_from_script();

pub fn simplify(&mut self, simplify: bool) -> &mut Self

Enable or disable graph simplification passes (default: true).

When enabled, optimization passes like dead node elimination, common subexpression elimination (CSE), and pattern-based simplifications are applied to the ONNX IR before code generation.

pub fn partition(&mut self, partition: bool) -> &mut Self

Enable or disable submodule partitioning for large models (default: true).

When enabled, models with more than 200 nodes are automatically split into smaller submodule structs to keep generated code compilable. Each submodule gets its own forward() method, and the top-level Model delegates to them.

pub fn run_from_script(&self)

Runs code generation from a build script context.

Use this method when calling from build.rs. The output directory will be $OUT_DIR/<out_dir>, allowing the generated code to be included with:

include!(concat!(env!("OUT_DIR"), "/<out_dir>/<model>.rs"));

Panics

Panics if OUT_DIR environment variable is not set (should be set by Cargo).

Examples

In build.rs:

use burn_onnx::ModelGen;

ModelGen::new()
    .input("path/to/model.onnx")
    .out_dir("model/")
    .run_from_script();

pub fn run_from_cli(&self)

Runs code generation from a CLI or application context.

Use this method when calling from a CLI tool or regular application. The output directory is used as-is (relative or absolute path).

Panics

Panics if out_dir was not set via out_dir.

Examples

use burn_onnx::ModelGen;

ModelGen::new()
    .input("model.onnx")
    .out_dir("./generated/")
    .run_from_cli();

Trait Implementations

impl Debug for ModelGen

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for ModelGen

fn default() -> Self

Returns the “default value” for a type.