<div align="center">
<img src="./assets/logo-burn-full.png" width="200px"/>
[](https://crates.io/crates/burn)
[](https://github.com/burn-rs/burn/actions/workflows/test-burn.yml)
[](https://docs.rs/burn)
[](https://releases.rs/docs/released/1.65.0)
[](https://github.com/burn-rs/burn/blob/master/LICENSE)
> This library aims to be a complete deep learning framework with extreme flexibility written in Rust.
> The goal would be to satisfy researchers as well as practitioners making it easier to experiment, train and deploy your models.
<div align="left">
__Sections__
* [Features](#features)
* [Get Started](#get-started)
* [Examples](#examples)
* [MNIST](#mnist)
* [Components](#components)
* [Backend](#backend)
* [Tensor](#tensor)
* [Module](#module)
* [Forward](#forward)
* [Config](#config)
* [Learner](#learner)
* [License](#license)
## Features
* Flexible and intuitive custom neural network module 🤖
* Stateless and thread safe forward pass 🚀
* Fast training with full support for `metric`, `logging` and `checkpointing` 🌟
* [Burn-Tensor](https://github.com/burn-rs/burn/tree/doc/readme/burn-tensor): Tensor library with autodiff, CPU and GPU support 🔥
* [Burn-Dataset](https://github.com/burn-rs/burn/tree/doc/readme/burn-dataset): Dataset library with multiple utilities and sources 📚
## Get Started
The best way to get started with burn is the look at the [examples](#examples).
Also, this may be a good idea to checkout the main [components](#components) to get a quick overview of how to use burn.
### Examples
For now there is only one example, but more to come 💪..
#### MNIST
The [MNIST](https://github.com/burn-rs/burn/blob/main/examples/mnist) example is not just of small script that shows you how to train a basic model, but it's a quick one showing you how to:
* Define your own custom [module](#module) (MLP).
* Create the data pipeline from a raw dataset to a batched multi-threaded fast DataLoader.
* Configure a [learner](#learner) to display and log metrics as well as to keep training checkpoints.
The example can be run like so:
```bash
git clone https://github.com/burn-rs/burn.git
cd burn
# Use the --release flag to really speed up training.
echo "Using ndarray backend"
cargo run --example mnist --release --features ndarray # CPU NdArray Backend - f32 - single thread
cargo run --example mnist --release --features ndarray-blas-openblas # CPU NdArray Backend - f32 - blas with openblas
cargo run --example mnist --release --features ndarray-blas-netlib # CPU NdArray Backend - f32 - blas with netlib
echo "Using tch backend"
export TORCH_CUDA_VERSION=cu113 # Set the cuda version
cargo run --example mnist --release --features tch-gpu # GPU Tch Backend - f16
cargo run --example mnist --release --features tch-cpu # CPU Tch Backend - f32
```
### Components
Knowing the main components will be of great help when starting playing with `burn`.
#### Backend
Almost everything is based on the `Backend` trait, which allows to run tensor operations with different implementations without having to change your code.
A backend does not necessary have autodiff capabilities, therefore you can use `ADBackend` when you require it.
#### Tensor
The `Tensor` struct is at the core of the `burn` framework.
It takes two generic parameters, the `Backend` and the number of dimensions `D`,
```rust
use burn::tensor::{Tensor, Shape, Data};
use burn::tensor::backend::{Backend, NdArrayBackend, TchBackend};
fn my_func<B: Backend>() {
let _my_tensor = Tensor::<B, 2>::ones(Shape::new([3, 3]));
}
fn main() {
my_func<NdArrayBackend<f32>>();
my_func<TchBackend<f32>>();
}
```
#### Module
The `Module` derive let your create your own neural network module similar to PyTorch.
```rust
use burn::nn;
use burn::module::{Param, Module};
use burn::tensor::backend::Backend;
#[derive(Module, Debug)]
struct MyModule<B: Backend> {
my_param: Param<nn::Linear<B>>,
repeat: usize,
}
```
Note that only the fields wrapped inside `Param` are updated during training, and the other ones should implement `Clone`.
#### Forward
The `Forward` trait can also be implemented by your module.
```rust
use burn::module::Forward;
use burn::tensor::Tensor;
impl<B: Backend> Forward<Tensor<B, 2>, Tensor<B, 2>> for MyModule<B> {
fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
let mut x = input;
for _ in 0..self.repeat {
x = self.my_param.forward(x);
}
x
}
}
```
Note that you can implement multiple time the `Forward` trait with different inputs and outputs.
#### Config
The `Config` derive lets you define serializable and deserializable configurations or hyper-parameters for your [modules](#module) or any components.
```rust
use burn::config::Config;
#[derive(Config)]
struct MyConfig {
#[config(default = 1.0e-6)]
pub epsilon: usize,
pub dim: usize,
}
```
The derive also adds useful methods to your config.
```rust
fn my_func() {
let config = MyConfig::new(100);
println!("{}", config.epsilon); // 1.0.e-6
println!("{}", config.dim); // 100
let config = MyConfig::new(100).with_epsilon(1.0e-8);
println!("{}", config.epsilon); // 1.0.e-8
}
```
#### Learner
The `Learner` is the main `struct` that let you train a neural network with support for `logging`, `metric`, `checkpointing` and more.
In order to create a learner, you must use the `LearnerBuilder`.
```rust
use burn::train::LearnerBuilder;
let learner = LearnerBuilder::new("/tmp/artifact_dir")
.metric_train_plot(AccuracyMetric::new())
.metric_valid_plot(AccuracyMetric::new())
.metric_train(LossMetric::new())
.metric_valid(LossMetric::new())
.with_file_checkpointer::<f32>(2)
.num_epochs(config.num_epochs)
.build(model, optim);
```
See this [example](https://github.com/burn-rs/burn/blob/main/examples/mnist) for a real usage.
## License
Burn is distributed under the terms of both the MIT license and the Apache License (Version 2.0).
See [LICENSE-APACHE](./LICENSE-APACHE) and [LICENSE-MIT](./LICENSE-MIT) for details.
Opening a pull request is assumed to signal agreement with these licensing terms.