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.
- Customizable, intuitive and user-friendly neural network module 🔥
- Comprehensive training tools, including
metrics
,logging
, andcheckpointing
📈 - 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 📦
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 🔥.
We keep an updated and curated list of models and examples built with Burn, see the burn-rs/models repository for more details.
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.
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 |
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.
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.
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.
Thanks to all current sponsors 🙏.
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.