burn

A flexible, backend-agnostic deep-learning framework - train and infer in Rust with pluggable compute backends.

Recipe

use burn::backend::NdArray;
use burn::module::Module;
use burn::nn::Linear;
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
#[derive(Module, Debug)]
struct Model<B: burn::backend::Backend> {
    linear: Linear<B>,
}
 
fn forward<B: burn::backend::Backend>(m: &Model<B>, x: Tensor<B, 2>) -> Tensor<B, 2> {
    m.linear.forward(x)
}

When to reach for this:

Research prototypes that may move from CPU to GPU later
Teaching neural nets with Rust type safety
Training small models entirely in Rust
Teams wanting one API across embedded and server targets

Working Example

use burn::backend::NdArray;
use burn::module::Module;
use burn::nn::{Linear, LinearConfig};
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
#[derive(Module, Debug)]
struct Mlp<B: burn::backend::Backend> {
    fc: Linear<B>,
}
 
fn main() {
    let device = <B as burn::backend::Backend>::Device::default();
    let mlp = Mlp::<B> {
        fc: LinearConfig::new(4, 2).init(&device),
    };
    let input = Tensor::<B, 2>::from_data([[1.0, 2.0, 3.0, 4.0]], &device);
    let out = mlp.fc.forward(input);
    println!("{:?}", out.into_data());
}

What this demonstrates:

Module derive for parameter registration
LinearConfig builder pattern for layer setup
Rank-2 tensor typing for batch features

Deep Dive

How It Works

Backends implement tensor ops; user code is generic over Backend.
Learner API wires datasets, optimizers, metrics, and checkpointing.
Import ONNX for inference-only paths; export when interoperating with Python stacks.
WGPU backend targets cross-vendor GPU without CUDA installs.

Backend Picks

Backend	Target
NdArray	CPU dev, CI
WGPU	Cross-platform GPU
CUDA	NVIDIA servers

Gotchas

Compile times with generic backends - monomorphization grows binaries. Fix: type alias one backend per crate binary.
ONNX op coverage gaps - exotic layers fail import. Fix: simplify graph in PyTorch export or fall back to ort.
WGPU driver issues on headless servers - no adapter found. Fix: use CUDA or CPU NdArray in datacenter images.
Learning rate too high on NdArray - NaNs without GPU-style mixed precision guardrails. Fix: start conservative, clip gradients.
Checkpoint format changes - breaks resume training across burn versions. Fix: version checkpoints in object storage metadata.

Alternatives

Alternative	Use When	Don't Use When
candle	HF transformer weights day one	You need generic backend swapping
ort	Production ONNX inference only	Training inside Rust
tch	Must reuse arbitrary PyTorch code	Pure Rust dependency policy
JAX/Python	Large-scale TPU training	Edge deploy without Python

FAQs

Is burn production-ready?

Inference paths mature faster than large-scale training. Validate accuracy and perf on your model before cutting Python.

How do I switch backends?

Change the type alias type B = ... and rebuild with appropriate feature flags - no source rewrite for module code.

Does burn support CNNs and transformers?

Yes via burn-import and custom modules - transformer support is active but check examples for your architecture.

How do I load datasets?

Implement Dataset trait or adapt CSV/Parquet readers from the data section into tensor batches.

Can I deploy to WASM?

Experimental paths exist; expect size and SIMD limits - profile before betting on browser training.

How does training loop look?

Use LearnerBuilder with optimizer config, metric trackers, and learner.fit(dataloader) - see burn book for full snippet.

ONNX vs native burn inference?

ONNX via ort is great for frozen graphs; native burn modules allow Rust-only hot paths.

How do I compare to PyTorch?

Run golden tensor comparisons on a fixed batch after import/export - do not trust loss curves alone.

GPU memory management?

WGPU/CUDA backends cache allocations; drop unused tensors and limit batch size when OOM.

Relation to candle?

candle is HF-focused; burn is general - see candle for transformer zoo shortcuts.

candle - HF transformers
ONNX Runtime (ort) - imported graphs
GPU & Acceleration - WGPU/CUDA setup
ML in Rust Basics - orientation

Stack versions: This page was written for Rust 1.97.0 (edition 2024), Tokio 1.x, Axum 0.8, serde 1.0, sqlx 0.8, clap 4, and Polars 0.46+.

burn

A flexible, backend-agnostic deep-learning framework - train and infer in Rust with pluggable compute backends.

Recipe

use burn::backend::NdArray;
use burn::module::Module;
use burn::nn::Linear;
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
#[derive(Module, Debug)]
struct Model<B: burn::backend::Backend> {
    linear: Linear<B>,
}
 
fn forward<B: burn::backend::Backend>(m: &Model<B>, x: Tensor<B, 2>) -> Tensor<B, 2> {
    m.linear.forward(x)
}

When to reach for this:

Research prototypes that may move from CPU to GPU later
Teaching neural nets with Rust type safety
Training small models entirely in Rust
Teams wanting one API across embedded and server targets

Working Example

use burn::backend::NdArray;
use burn::module::Module;
use burn::nn::{Linear, LinearConfig};
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
#[derive(Module, Debug)]
struct Mlp<B: burn::backend::Backend> {
    fc: Linear<B>,
}
 
fn main() {
    let device = <B as burn::backend::Backend>::Device::default();
    let mlp = Mlp::<B> {
        fc: LinearConfig::new(4, 2).init(&device),
    };
    let input = Tensor::<B, 2>::from_data([[1.0, 2.0, 3.0, 4.0]], &device);
    let out = mlp.fc.forward(input);
    println!("{:?}", out.into_data());
}

What this demonstrates:

Module derive for parameter registration
LinearConfig builder pattern for layer setup
Rank-2 tensor typing for batch features

Deep Dive

How It Works

Backends implement tensor ops; user code is generic over Backend.
Learner API wires datasets, optimizers, metrics, and checkpointing.
Import ONNX for inference-only paths; export when interoperating with Python stacks.
WGPU backend targets cross-vendor GPU without CUDA installs.

Backend Picks

Backend	Target
NdArray	CPU dev, CI
WGPU	Cross-platform GPU
CUDA	NVIDIA servers

Gotchas

Compile times with generic backends - monomorphization grows binaries. Fix: type alias one backend per crate binary.
ONNX op coverage gaps - exotic layers fail import. Fix: simplify graph in PyTorch export or fall back to ort.
WGPU driver issues on headless servers - no adapter found. Fix: use CUDA or CPU NdArray in datacenter images.
Learning rate too high on NdArray - NaNs without GPU-style mixed precision guardrails. Fix: start conservative, clip gradients.
Checkpoint format changes - breaks resume training across burn versions. Fix: version checkpoints in object storage metadata.

Alternatives

Alternative	Use When	Don't Use When
candle	HF transformer weights day one	You need generic backend swapping
ort	Production ONNX inference only	Training inside Rust
tch	Must reuse arbitrary PyTorch code	Pure Rust dependency policy
JAX/Python	Large-scale TPU training	Edge deploy without Python

FAQs

Is burn production-ready?

Inference paths mature faster than large-scale training. Validate accuracy and perf on your model before cutting Python.

How do I switch backends?

Change the type alias type B = ... and rebuild with appropriate feature flags - no source rewrite for module code.

Does burn support CNNs and transformers?

Yes via burn-import and custom modules - transformer support is active but check examples for your architecture.

How do I load datasets?

Implement Dataset trait or adapt CSV/Parquet readers from the data section into tensor batches.

Can I deploy to WASM?

Experimental paths exist; expect size and SIMD limits - profile before betting on browser training.

How does training loop look?

Use LearnerBuilder with optimizer config, metric trackers, and learner.fit(dataloader) - see burn book for full snippet.

ONNX vs native burn inference?

ONNX via ort is great for frozen graphs; native burn modules allow Rust-only hot paths.

How do I compare to PyTorch?

Run golden tensor comparisons on a fixed batch after import/export - do not trust loss curves alone.

GPU memory management?

WGPU/CUDA backends cache allocations; drop unused tensors and limit batch size when OOM.

Relation to candle?

candle is HF-focused; burn is general - see candle for transformer zoo shortcuts.

candle - HF transformers
ONNX Runtime (ort) - imported graphs
GPU & Acceleration - WGPU/CUDA setup
ML in Rust Basics - orientation

Stack versions: This page was written for Rust 1.97.0 (edition 2024), Tokio 1.x, Axum 0.8, serde 1.0, sqlx 0.8, clap 4, and Polars 0.46+.

burn

Search Documentation

burn