burn – tracel-ai

These crates are not maintained by Tracel, but they are part of the same Rust AI story. Anything
that helps you load data, build environments, or ship models belongs here. Built something that
fits? Open a PR to add it!

Contrary to the aforementioned backends, Autodiff is actually a backend _decorator_. This means that
it cannot exist by itself; it must encapsulate another backend.

The simple act of wrapping a base backend with Autodiff transparently equips it with
autodifferentiation support, making it possible to call backward on your model.

use burn::backend::{Autodiff, Wgpu};
use burn::tensor::{Distribution, Tensor};

fn main() {
    type Backend = Autodiff<Wgpu>;

    let device = Default::default();

    let x: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default, &device);
    let y: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default, &device).require_grad();

    let tmp = x.clone() + y.clone();
    let tmp = tmp.matmul(x);
    let tmp = tmp.exp();

    let grads = tmp.backward();
    let y_grad = y.grad(&grads).unwrap();
    println!("{y_grad}");
}

Of note, it is impossible to make the mistake of calling backward on a model that runs on a backend
that does not support autodiff (for inference), as this method is only offered by an Autodiff
backend.

See the Autodiff Backend README for more details.

This backend decorator enhances a backend with kernel fusion, provided that the inner backend
supports it. Note that you can compose this backend with other backend decorators such as Autodiff.
All first-party accelerated backends (like WGPU and CUDA) use Fusion by default (burn/fusion
feature flag), so you typically don't need to apply it manually.

#[cfg(not(feature = "fusion"))]
pub type Cuda<F = f32, I = i32> = CubeBackend<CudaRuntime, F, I, u8>;

#[cfg(feature = "fusion")]
pub type Cuda<F = f32, I = i32> = burn_fusion::Fusion<CubeBackend<CudaRuntime, F, I, u8>>;

Of note, we plan to implement automatic gradient checkpointing based on compute bound and memory
bound operations, which will work gracefully with the fusion backend to make your code run even
faster during training, see this issue.

See the Fusion Backend README for more details.

That backend has two parts, one client and one server. The client sends tensor operations over the
network to a remote compute backend. You can use any first-party backend as server in a single line
of code:

fn main_server() {
    // Start a server on port 3000.
    burn::server::start::<burn::backend::Cuda>(Default::default(), 3000);
}

fn main_client() {
    // Create a client that communicate with the server on port 3000.
    use burn::backend::{Autodiff, RemoteBackend};

    type Backend = Autodiff<RemoteDevice>;

    let device = RemoteDevice::new("ws://localhost:3000");
    let tensor_gpu =
        Tensor::<Backend, 2>::random([3, 3], Distribution::Default, &device);
}

As you can see in the previous video (click on the picture!), a new terminal UI dashboard based on
the Ratatui crate allows users to follow their training
with ease without having to connect to any external application.

You can visualize your training and validation metrics updating in real-time and analyze the
lifelong progression or recent history of any registered metrics using only the arrow keys. Break
from the training loop without crashing, allowing potential checkpoints to be fully written or
important pieces of code to complete without interruption 🛡

Burn supports importing ONNX (Open Neural Network Exchange) models through the
burn-onnx crate, allowing you to easily port models from
TensorFlow or PyTorch to Burn. The ONNX model is converted into Rust code that uses Burn's native
APIs, enabling the imported model to run on any Burn backend (CPU, GPU, WebAssembly) and benefit
from all of Burn's optimizations like automatic kernel fusion.

Our ONNX support is further described in
this section of the Burn Book 🔥.

> Note: This crate is in active development and currently supports a
> limited set of ONNX operators.

You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models.
This makes it easy to reuse existing models while benefiting from Burn's performance and deployment
features.

Learn more in the Saving & Loading Models
section of the Burn Book.

Several of our backends can run in WebAssembly environments: Flex for CPU execution, and WGPU for
GPU acceleration via WebGPU. This means that you can run inference directly within a browser. We
provide several examples of this:

• MNIST where you can draw digits and a small convnet tries to

find which one it is! 2️⃣ 7️⃣ 😰

• Image Classification

where you can upload images and classify them! 🌄

Burn's core components support no_std. This
means it can run in bare metal environment such as embedded devices without an operating system.

> As of now, only the Flex backend can be used in a _no_std_ environment.

To begin working effectively with Burn, it is crucial to understand its key components and
philosophy. This is why we highly recommend new users to read the first sections of
The Burn Book 🔥. It provides detailed examples and explanations
covering every facet of the framework, including building blocks like tensors, modules, and
optimizers, all the way to advanced usage, like coding your own GPU kernels.

> The project is constantly evolving, and we try as much as possible to keep the book up to date
> with new additions. However, we might miss some details sometimes, so if you see something weird,
> let us know! We also gladly accept Pull Requests 😄

Let's start with a code snippet that shows how intuitive the framework is to use! In the following,
we declare a neural network module with some parameters along with its forward pass.

use burn::nn;
use burn::module::Module;
use burn::tensor::backend::Backend;

#[derive(Module, Debug)]
pub struct PositionWiseFeedForward<B: Backend> {
    linear_inner: nn::Linear<B>,
    linear_outer: nn::Linear<B>,
    dropout: nn::Dropout,
    gelu: nn::Gelu,
}

impl<B: Backend> PositionWiseFeedForward<B> {
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
        let x = self.linear_inner.forward(input);
        let x = self.gelu.forward(x);
        let x = self.dropout.forward(x);

        self.linear_outer.forward(x)
    }
}

We have a somewhat large amount of examples in the repository that shows how to use
the framework in different scenarios.

Following the book:

• Basic Workflow : Creates a custom CNN Module to train on the MNIST dataset

and use for inference.

• Custom Training Loop : Implements a basic training loop instead

of using the Learner.

• Custom WGPU Kernel : Learn how to create your own custom

operation with the WGPU backend.

Additional examples:

• Custom CSV Dataset : Implements a dataset to parse CSV data for a

regression task.

• Regression : Trains a simple MLP on the California Housing dataset

to predict the median house value for a district.

• Custom Image Dataset : Trains a simple CNN on custom image

dataset following a simple folder structure.

• Custom Renderer : Implements a custom renderer to display the

<code class="ra0-md-code">Learner</code> progress.

• Image Classification Web : Image classification web browser

demo using Burn, WGPU and WebAssembly.

• MNIST Inference on Web : An interactive MNIST inference demo in

the browser. The demo is available online.

• MNIST Training : Demonstrates how to train a custom Module (MLP) with the

Learner configured to log metrics and keep training checkpoints.

• PyTorch Import Inference : Imports a PyTorch model pre-trained

on MNIST to perform inference on a sample image with Burn.

• Text Classification : Trains a text classification transformer

model on the AG News or DbPedia dataset. The trained model can then be used to classify a text
sample.

• Text Generation : Trains a text generation transformer model on the

DbPedia dataset.

• Wasserstein GAN MNIST : Trains a WGAN model to generate new handwritten digits

based on MNIST.

For more practical insights, you can clone the repository and run any of them directly on your
computer!

We keep an updated and curated list of models and examples built with Burn, see the
tracel-ai/models repository for more details.

Don't see the model you want? Don't hesitate to open an issue, and we may prioritize it. Built a
model using Burn and want to share it? You can also open a Pull Request and add your model under the
community section!

Deep Learning is a special form of software where you need very high level abstractions as well as
extremely fast execution time. Rust is the perfect candidate for that use case since it provides
zero-cost abstractions to easily create neural network modules, and fine-grained control over memory
to optimize every detail. To this day, the mainstream solution has been to offer APIs in Python but
rely on bindings to low-level languages such as C/C++. This reduces portability, increases
complexity and creates friction between researchers and engineers. Rust's approach to abstractions
is versatile enough to tackle this two-language dichotomy, and Cargo makes it easy to build, test
and deploy from any environment, which is usually a pain in Python.

Rust's AI ecosystem is young, but it is real and growing quickly. Foundational pieces are already
here: Burn and CubeCL for training and compute,
candle for inference, Hugging Face's tokenizers and
safetensors, and polars and ndarray for data. Betting on Rust today means betting on a stack
that is growing, and one where contributors still shape the direction. The pieces that don't exist
yet are opportunities rather than dead-ends (see Contributing).

Rust is also what makes one-stack-everywhere possible: a single self-contained binary with no Python
runtime to ship, running from servers down to no_std embedded targets.

In the event that you are trying to load a model record saved in a version older than 0.14.0, make
sure to use a compatible version (0.14, 0.15 or 0.16) with the record-backward-compat
feature flag.

features = [..., "record-backward-compat"]

Otherwise, the record won't be deserialized correctly and you will get an error message. This error
will also point you to the backward compatible feature flag.

The backward compatibility was maintained for deserialization when loading records. Therefore, as
soon as you have saved the record again it will be saved according to the new structure and you can
upgrade back to the current version

Please note that binary formats are not backward compatible. Thus, you will need to load your record
in a previous version and save it in any of the other self-describing record format (e.g., using the
NamedMpkFileRecorder) before using a compatible version (as described) with the
record-backward-compat feature flag.