![](./assets/logo-burn-full.png)
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.
Sections
- Flexible and intuitive custom neural network module ๐ค
- Stateless and thread safe forward pass ๐
- Fast training with full support for
metric
,logging
andcheckpointing
๐ - Burn-Tensor: Tensor library with autodiff, CPU and GPU support ๐ฅ
- Burn-Dataset: Dataset library with multiple utilities and sources ๐
The best way to get started with burn is the look at the examples. Also, this may be a good idea to checkout the main components to get a quick overview of how to use burn.
For now there is only one example, but more to come ๐ช..
The 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 (MLP).
- Create the data pipeline from a raw dataset to a batched multi-threaded fast DataLoader.
- Configure a learner to display and log metrics as well as to keep training checkpoints.
The example can be run like so:
$ git clone https://github.com/burn-rs/burn.git
$ cd burn
$ export TORCH_CUDA_VERSION=cu113 # Set the cuda version
$ # Use the --release flag to really speed up training.
$ cargo run --example mnist --release # CPU NdArray Backend
$ cargo run --example mnist_cuda_gpu --release # GPU Tch Backend
Knowing the main components will be of great help when starting playing with burn
.
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.
The Tensor
struct is at the core of the burn
framework.
It takes two generic parameters, the Backend
and the number of dimensions D
,
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>>();
}
The Module
derive let your create your own neural network module similar to PyTorch.
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
.
The Forward
trait can also be implemented by your module.
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.
The Config
derive lets you define serializable and deserializable configurations or hyper-parameters for your modules or any components.
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.
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
}
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
.
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 for a real usage.
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.