Using Burn means having your models optimized on any backend. When possible, we provide a way toautomatically and dynamically create custom kernels that minimize data relocation between differentmemory spaces, extremely useful when moving memory is the bottleneck.

As an example, you could write your own GELU activation function with the high level tensor api (seeRust code snippet below).

fngelu_custom<B:Backend,constD:usize>(x:Tensor<B,D>) ->Tensor<B,D>{let x = x.clone()*((x /SQRT_2).erf() +1);    x /2}

Then, at runtime, a custom low-level kernel will be automatically created for your specificimplementation and will rival a handcrafted GPU implementation. The kernel consists of about 60lines of WGSL WebGPU Shading Language,an extremely verbose lower level shader language you probably don't want to program your deeplearning models in!

Forfirst-party backends, an asynchronous execution styleis used, which allows to perform various optimizations, such as the previously mentioned automatickernel fusion.

Asynchronous execution also ensures that the normal execution of the framework does not block themodel computations, which implies that the framework overhead won't impact the speed of executionsignificantly. Conversely, the intense computations in the model do not interfere with theresponsiveness of the framework. For more information about our asynchronous backends, seethis blog post.

Burn emphasizes thread safety by leveraging theownership system of Rust.With Burn, each module is the owner of its weights. It is therefore possible to send a module toanother thread for computing the gradients, then send the gradients to the main thread that canaggregate them, andvoilà, you get multi-device training.

This is a very different approach from what PyTorch does, where backpropagation actually mutates thegrad attribute of each tensor parameter. This is not a thread-safe operation and thereforerequires lower level synchronization primitives, seedistributed training for reference. Note thatthis is still very fast, but not compatible across different backends and quite hard to implement.

One of the main roles of a deep learning framework is to reduce the amount of memory necessary torun models. The naive way of handling memory is that each tensor has its own memory space, which isallocated when the tensor is created then deallocated as the tensor gets out of scope. However,allocating and deallocating data is very costly, so a memory pool is often required to achieve goodthroughput. Burn offers an infrastructure that allows for easily creating and selecting memorymanagement strategies for backends. For more details on memory management in Burn, seethis blog post.

Another very important memory optimization of Burn is that we keep track of when a tensor can bemutated in-place just by using the ownership system well. Even though it is a rather small memoryoptimization on its own, it adds up considerably when training or running inference with largermodels and contributes to reduce the memory usage even more. For more information, seethis blog post about tensor handling.

A good deep learning framework should ensure that models run smoothly on all hardware. However, notall hardware share the same behavior in terms of execution speed. For instance, a matrixmultiplication kernel can be launched with many different parameters, which are highly sensitive tothe size of the matrices and the hardware. Using the wrong configuration could reduce the speed ofexecution by a large factor (10 times or even more in extreme cases), so choosing the right kernelsbecomes a priority.

With our home-made backends, we run benchmarks automatically and choose the best configuration forthe current hardware and matrix sizes with a reasonable caching strategy.

This adds a small overhead by increasing the warmup execution time, but stabilizes quickly after afew forward and backward passes, saving lots of time in the long run. Note that this feature isn'tmandatory, and can be disabled when cold starts are a priority over optimized throughput.

It is no secret that deep learning is mostly relying on matrix multiplication as its core operation,since this is how fully-connected neural networks are modeled.

More and more, hardware manufacturers optimize their chips specifically for matrix multiplicationworkloads. For instance, Nvidia has itsTensor Cores and today most cellphones have AI specializedchips. As of this moment, we support Tensor Cores with our LibTorch, Candle, CUDA, Metal and WGPU/SPIR-Vbackends, but not other accelerators yet. We hopethis issue gets resolved at some point to bringsupport to our WGPU backend.

Burn aims to be the most flexible deep learning framework. While it's crucial to maintaincompatibility with a wide variety of backends, Burn also provides the ability to extend thefunctionalities of a backend implementation to suit your personal modeling requirements.

This versatility is advantageous in numerous ways, such as supporting custom operations like flashattention or manually writing your own kernel for a specific backend to enhance performance. Seethis section in the Burn Book 🔥for more details.

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

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

use burn::backend::{Autodiff,Wgpu};use burn::tensor::{Distribution,Tensor};fnmain(){typeBackend =Autodiff<Wgpu>;let x:Tensor<Backend,2> =Tensor::random([32,32],Distribution::Default);let y:Tensor<Backend,2> =Tensor::random([32,32],Distribution::Default).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 backendthat does not support autodiff (for inference), as this method is only offered by an Autodiffbackend.

See theAutodiff Backend README for more details.

This backend decorator enhances a backend with kernel fusion, provided that the inner backendsupports it. Note that you can compose this backend with other backend decorators such as Autodiff.For now, only the WGPU and CUDA backends have support for fused kernels.

use burn::backend::{Autodiff,Fusion,Wgpu};use burn::tensor::{Distribution,Tensor};fnmain(){typeBackend =Autodiff<Fusion<Wgpu>>;let x:Tensor<Backend,2> =Tensor::random([32,32],Distribution::Default);let y:Tensor<Backend,2> =Tensor::random([32,32],Distribution::Default).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, we plan to implement automatic gradient checkpointing based on compute bound and memorybound operations, which will work gracefully with the fusion backend to make your code run evenfaster during training, seethis issue.

See theFusion Backend README for more details.

That backend simplifies hardware operability, if for instance you want to execute some operations on the CPU and other operations on the GPU.

use burn::tensor::{Distribution,Tensor};use burn::backend::{NdArray,Router,Wgpu, ndarray::NdArrayDevice, router::duo::MultiDevice, wgpu::WgpuDevice,};fnmain(){typeBackend =Router<(Wgpu,NdArray)>;let device_0 =MultiDevice::B1(WgpuDevice::DiscreteGpu(0));let device_1 =MultiDevice::B2(NdArrayDevice::Cpu);let tensor_gpu =Tensor::<Backend,2>::random([3,3], burn::tensor::Distribution::Default,&device_0);let tensor_cpu =Tensor::<Backend,2>::random([3,3], burn::tensor::Distribution::Default,&device_1);}

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:

fnmain_server(){// Start a server on port 3000.    burn::server::start::<burn::backend::Cuda>(Default::default(),3000);}fnmain_client(){// Create a client that communicate with the server on port 3000.use burn::backend::{Autodiff,RemoteBackend};typeBackend =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 ontheRatatui crate allows users to follow their trainingwith ease without having to connect to any external application.

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

ONNX (Open Neural Network Exchange) is an open-standard format that exports both the architectureand the weights of a deep learning model.

Burn supports the importation of models that follow the ONNX standard so you can easily port a modelyou have written in another framework like TensorFlow or PyTorch to Burn to benefit from all theadvantages our framework offers.

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

Note: This crate is in active development and currently supports alimited 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:

• Import pre-trained PyTorch models into Burn
• Load models from Safetensors format

Several of our backends can compile to Web Assembly: Candle and NdArray for CPU, and WGPU for GPU.This means that you can run inference directly within a browser. We provide several examples ofthis:

• MNIST where you can draw digits and a small convnet tries tofind which one it is! 2️⃣ 7️⃣ 😰
• Image Classification where you can upload images andclassify them! 🌄

Burn's core components supportno_std. Thismeans it can run in bare metal environment such as embedded devices without an operating system.

As of now, only the NdArray backend can be used in ano_std environment.

To begin working effectively with Burn, it is crucial to understand its key components andphilosophy. This is why we highly recommend new users to read the first sections ofThe Burn Book 🔥. It provides detailed examples and explanationscovering every facet of the framework, including building blocks like tensors, modules, andoptimizers, 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 datewith 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)]pubstructPositionWiseFeedForward<B:Backend>{linear_inner: nn::Linear<B>,linear_outer: nn::Linear<B>,dropout: nn::Dropout,gelu: nn::Gelu,}impl<B:Backend>PositionWiseFeedForward<B>{pubfnforward<constD: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 ofexamples in the repository that shows how to usethe framework in different scenarios.

Followingthe book:

• Basic Workflow : Creates a custom CNNModule to train on the MNIST datasetand use for inference.
• Custom Training Loop : Implements a basic training loop insteadof using theLearner.
• Custom WGPU Kernel : Learn how to create your own customoperation with the WGPU backend.

Additional examples:

• Custom CSV Dataset : Implements a dataset to parse CSV data for aregression task.
• Regression : Trains a simple MLP on the California Housing datasetto predict the median house value for a district.
• Custom Image Dataset : Trains a simple CNN on custom imagedataset following a simple folder structure.
• Custom Renderer : Implements a custom renderer to display theLearner progress.
• Image Classification Web : Image classification web browserdemo using Burn, WGPU and WebAssembly.
• MNIST Inference on Web : An interactive MNIST inference demo inthe browser. The demo is availableonline.
• MNIST Training : Demonstrates how to train a customModule (MLP) with theLearner configured to log metrics and keep training checkpoints.
• Named Tensor : Performs operations with the experimentalNamedTensorfeature.
• ONNX Import Inference : Imports an ONNX model pre-trained on MNIST toperform inference on a sample image with Burn.
• PyTorch Import Inference : Imports a PyTorch model pre-trained onMNIST to perform inference on a sample image with Burn.
• Text Classification : Trains a text classification transformermodel on the AG News or DbPedia dataset. The trained model can then be used to classify a textsample.
• Text Generation : Trains a text generation transformer model on theDbPedia dataset.
• Wasserstein GAN MNIST : Trains a WGAN model to generate new handwritten digitsbased on MNIST.

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

We keep an updated and curated list of models and examples built with Burn, see thetracel-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 amodel using Burn and want to share it? You can also open a Pull Request and add your model under thecommunity section!

Deep Learning is a special form of software where you need very high level abstractions as well asextremely fast execution time. Rust is the perfect candidate for that use case since it provideszero-cost abstractions to easily create neural network modules, and fine-grained control over memoryto optimize every detail.

It's important that a framework be easy to use at a high level so that its users can focus oninnovating in the AI field. However, since running models relies so heavily on computations,performance can't be neglected.

To this day, the mainstream solution to this problem has been to offer APIs in Python, but rely onbindings to low-level languages such as C/C++. This reduces portability, increases complexity andcreates frictions between researchers and engineers. We feel like Rust's approach to abstractionsmakes it versatile enough to tackle this two languages dichotomy.

Rust also comes with the Cargo package manager, which makes it incredibly easy to build, test, anddeploy from any environment, which is usually a pain in Python.

Although Rust has the reputation of being a difficult language at first, we strongly believe itleads to more reliable, bug-free solutions built faster (after some practice 😅)!

In the event that you are trying to load a model record saved in a version older than0.14.0, makesure to use a compatible version (0.14,0.15 or0.16) with therecord-backward-compatfeature flag.

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

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

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

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