Disclamer 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.

Sections

• Features
• Get Started
  • Examples
  • Components
    • Backend
    • Tensor
    • Autodiff
    • Module
    • Config
    • Learner
• no_std support
• License

Features

• Flexible and intuitive custom neural network module 🔥
• Training with full support for metric, logging and checkpointing 📈
• Tensor crate with backends as pluging 🔧
  • Tch backend with CPU/GPU support 🚀
  • NdArray backend with no_std support, running on any platform 👌
  • Autodiff backend making any backend differentiable 🌟
• Dataset crate with multiple utilities and sources 📚

Get Started

The best way to get started with burn is to clone the repo and play with the examples. This may also be a good idea to take a look the main components of burn to get a quick overview of the fundamental building blocks.

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.

Components

Understanding the key components and philosophy of burn can greatly help when beginning to work with the framework.

Backend

Nearly everything in burn is based on the Backend trait, which enables you to run tensor operations using different implementations without having to modify your code. While a backend may not necessarily have autodiff capabilities, the ADBackend trait specifies when autodiff is needed. This trait not only abstracts operations but also tensor, device and element types, providing each backend the flexibility they need. It's worth noting that the trait assumes eager mode since burn fully supports dynamic graphs. However, we may create another API to assist with integrating graph-based backends, without requiring any changes to the user's code.

Tensor

At the core of burn lies the Tensor struct, which encompasses multiple types of tensors, including Float, Int, and Bool. The element types of these tensors are specified by the backend and are usually designated as a generic argument (e.g., NdArrayBackend<f32>). Although the same struct is used for all tensors, the available methods differ depending on the tensor kind. You can specify the desired tensor kind by setting the third generic argument, which defaults to Float. The first generic argument specifies the backend, while the second specifies the number of dimensions.

```rust use burn::tensor::backend::Backend; use burn::tensor::{Tensor, Int};

fn function(tensorfloat: Tensor) { let _tensorbool = tensorfloat.clone().equalelem(2.0); // Tensor let tensorint = tensor_float.argmax(1) // Tensor } ```

As demonstrated in the previous example, nearly all operations require owned tensors as parameters, which means that calling Clone explicitly is necessary when reusing the same tensor multiple times. However, there's no need to worry since the tensor's data won't be copied, it will be flagged as readonly when multiple tensors use the same allocated memory. This enables backends to reuse tensor data when possible, similar to a copy-on-write pattern, while remaining completely transparent to the user.

Autodiff

The 'Backend' trait is highly flexible, enabling backpropagation to be implemented using a simple backend decorator, which makes any backend differentiable.

```rust use burn::tensor::backend::{ADBackend, Backend}; use burn::tensor::{Distribution, Tensor}; use burnautodiff::ADBackendDecorator; use burnndarray::NdArrayBackend;

fn linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor { x.matmul(weight) + bias }

fn main() { type Backend = NdArrayBackend;

let weight = Tensor::random([3, 3], Distribution::Standard);
let bias = Tensor::zeros([1, 3]);
let x = Tensor::random([3, 3], Distribution::Standard);

let y = linear::<Backend>(x.clone(), weight.clone(), bias.clone());
// y.backward() // Method backward doesn't exist

let y = linear::<ADBackendDecorator<Backend>>(
    Tensor::from_inner(x),
    Tensor::from_inner(weight).require_grad(),
    Tensor::from_inner(bias).require_grad(),
);
let grads = y.backward(); // Method exists

}

```

Module

The Module derive allows you to create your own neural network modules, similar to PyTorch. Note that the Module derive generates all the necessary methods to make your type essentially a parameter container. It makes no assumptions about how the forward function is declared.

```rust use burn::nn; use burn::module::{Param, Module}; use burn::tensor::backend::Backend;

[derive(Module, Debug)]

pub struct PositionWiseFeedForward { linearinner: Param>, linearouter: Param>, dropout: Dropout, gelu: GELU, }

impl PositionWiseFeedForward { pub fn forward(&self, input: Tensor) -> Tensor { let x = self.linear_inner.forward(input); let x = self.gelu.forward(x); let x = self.dropout.forward(x);

    self.linear_outer.forward(x)
}

let model = ...;
let optim = ...;

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(10)
    .build(model, optim);

let _model_trained = learner.fit(dataloader_train, dataloader_valid);