filopedraz/burn

Burn - A Flexible and Comprehensive Deep Learning Framework in Rust

This library strives to serve as a comprehensive deep learning framework, offering exceptional flexibility and written in Rust. Our objective is to cater to both researchers and practitioners by simplifying the process of experimenting, training, and deploying models.

Features

• Customizable, intuitive and user-friendly neural network module 🔥
• Comprehensive training tools, including metrics, logging, and checkpointing 📈
• Versatile Tensor crate equipped with pluggable backends 🔧
  • Torch backend, supporting both CPU and GPU 🚀
  • Ndarray backend with no_std compatibility, ensuring universal platform adaptability 👌
  • WebGPU backend, offering cross-platform, browser-inclusive, GPU-based computations 🌐
  • Candle backend (alpha) 🕯️
  • Autodiff backend that enables differentiability across all backends 🌟
• Dataset crate containing a diverse range of utilities and sources 📚
• Import crate that simplifies the integration of pretrained models 📦

Get Started

The Burn Book 🔥

To begin working effectively with burn, it is crucial to understand its key components and philosophy. For detailed examples and explanations covering every facet of the framework, please refer to The Burn Book 🔥.

Pre-trained Models

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

Examples

Here is a code snippet showing how intuitive the framework is to use, where we declare a position-wise feed-forward module 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: Linear<B>,
    linear_outer: Linear<B>,
    dropout: Dropout,
    gelu: 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)
    }
}

For more practical insights, you can clone the repository and experiment with the following examples:

• MNIST train a model on CPU/GPU using different backends.
• MNIST Inference Web run trained model in the browser for inference.
• Text Classification train a transformer encoder from scratch on GPU.
• Text Generation train an autoregressive transformer from scratch on GPU.

Supported Platforms

Burn-ndarray Backend

Option	CPU	GPU	Linux	MacOS	Windows	Android	iOS	WASM
Pure Rust	Yes	No	Yes	Yes	Yes	Yes	Yes	Yes
Accelerate	Yes	No	No	Yes	No	No	Yes	No
Netlib	Yes	No	Yes	Yes	Yes	No	No	No
Openblas	Yes	No	Yes	Yes	Yes	Yes	Yes	No

Burn-tch Backend

Option	CPU	GPU	Linux	MacOS	Windows	Android	iOS	WASM
CPU	Yes	No	Yes	Yes	Yes	Yes	Yes	No
CUDA	No	Yes	Yes	No	Yes	No	No	No
MPS	No	Yes	No	Yes	No	No	No	No
Vulkan	Yes	Yes	Yes	Yes	Yes	Yes	No	No

Burn-wgpu Backend

Option	CPU	GPU	Linux	MacOS	Windows	Android	iOS	WASM
Metal	No	Yes	No	Yes	No	No	Yes	No
Vulkan	Yes	Yes	Yes	Yes	Yes	Yes	Yes	No
OpenGL	No	Yes	Yes	Yes	Yes	Yes	Yes	No
WebGpu	No	Yes	No	No	No	No	No	Yes
Dx11/Dx12	No	Yes	No	No	Yes	No	No	No

Support for no_std

Burn, including its burn-ndarray backend, can work in a no_std environment, provided alloc is available for the inference mode. To accomplish this, simply turn off the default features in burn and burn-ndarray (which is the minimum requirement for running the inference mode). You can find a reference example in burn-no-std-tests.

The burn-core and burn-tensor crates also support no_std with alloc. These crates can be directly added as dependencies if necessary, as they are reexported by the burn crate.

Please be aware that when using the no_std mode, a random seed will be generated at build time if one hasn't been set using the Backend::seed method. Also, the spin::mutex::Mutex is used instead of std::sync::Mutex in this mode.

Contributing

Before contributing, please take a moment to review our code of conduct. It's also highly recommended to read our architecture document, which explains our architectural decisions. Please see more details in our contributing guide.

Disclaimer

Burn is currently in active development, and there will be breaking changes. While any resulting issues are likely to be easy to fix, there are no guarantees at this stage.

Sponsors

Thanks to all current sponsors 🙏.

License

Burn is distributed under the terms of both the MIT license and the Apache License (Version 2.0). See LICENSE-APACHE and LICENSE-MIT for details. Opening a pull request is assumed to signal agreement with these licensing terms.