Model

The first step is to create a project and add the different Burn dependencies. Start by creating a new project with Cargo:

cargo new guide

As mentioned previously, this will initialize your guide project directory with a Cargo.toml and a src/main.rs file.

In the Cargo.toml file, add the burn dependency with train, vision and wgpu features. The default features automatically enable tui (for the dashboard), fusion and autotune for wgpu, among other defaults. Then run cargo build to build the project and import all the dependencies.

[package]
name = "guide"
version = "0.1.0"
edition = "2024"

[dependencies]
burn = { version = "~0.22", features = ["train", "vision", "wgpu"] }

Our goal will be to create a basic convolutional neural network used for image classification. We will keep the model simple by using two convolution layers followed by two linear layers, some pooling and ReLU activations. We will also use dropout to improve training performance.

Let us start by defining our model struct in a new file src/model.rs.

use burn::{
    nn::{
        Dropout, DropoutConfig, Linear, LinearConfig, Relu,
        conv::{Conv2d, Conv2dConfig},
        pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig},
    },
    prelude::*,
};

#[derive(Module, Debug)]
pub struct Model {
    conv1: Conv2d,
    conv2: Conv2d,
    pool: AdaptiveAvgPool2d,
    dropout: Dropout,
    linear1: Linear,
    linear2: Linear,
    activation: Relu,
}

There are two major things going on in this code sample.

1. You can create a deep learning module with the #[derive(Module)] attribute on top of a struct. This will generate the necessary code so that the struct implements the Module trait. This trait will make your module both trainable and (de)serializable while adding related functionalities. Like other attributes often used in Rust, such as Clone, PartialEq or Debug, each field within the struct must also implement the Module trait.

   🦀 Trait

   Traits are a powerful and flexible Rust language feature. They provide a way to define shared behavior for a particular type, which can be shared with other types.

   A type's behavior consists of the methods called on that type. Since all Modules should implement the same functionality, it is defined as a trait. Implementing a trait on a particular type usually requires the user to implement the defined behaviors of the trait for their types, though that is not the case here as explained above with the derive attribute. Check out the explainer below to learn why.

   For more details on traits, take a look at the associated chapter in the Rust Book.

   
   
   🦀 Derive Macro

   The derive attribute allows traits to be implemented easily by generating code that will implement a trait with its own default implementation on the type that was annotated with the derive syntax.

   This is accomplished through a feature of Rust called procedural macros, which allow us to run code at compile time that operates over Rust syntax, both consuming and producing Rust syntax. Using the attribute #[my_macro], you can effectively extend the provided code. You will see that the derive macro is very frequently employed to recursively implement traits, where the implementation consists of the composition of all fields.

   In this example, we want to derive the Module and Debug traits.

   #[derive(Module, Debug)]
   pub struct MyCustomModule {
       linear1: Linear,
       linear2: Linear,
       activation: Relu,
   }

   The basic Debug implementation is provided by the compiler to format a value using the {:?} formatter. For ease of use, the Module trait implementation is automatically handled by Burn so you don't have to do anything special. It essentially acts as parameter container.

   For more details on derivable traits, take a look at the Rust appendix, reference or example.

   
   

2. The model and its layers don't have a backend type parameter. Tensors carry a runtime Device, and Burn routes their operations through its Tensor → Bridge → Dispatch → Backend execution stack. This keeps models portable across backends without exposing backend generics in the user API. The device used to initialize the parameters determines where the model starts executing.

Note that each time you create a new file in the src directory you also need to explicitly add this module to the main.rs file. For instance after creating the model.rs, you need to add the following at the top of the main file:

mod model;

fn main() {
}

Next, we need to instantiate the model for training.

#[derive(Config, Debug)]
pub struct ModelConfig {
    num_classes: usize,
    hidden_size: usize,
    #[config(default = "0.5")]
    dropout: f64,
}

impl ModelConfig {
    /// Returns the initialized model.
    pub fn init(&self, device: &Device) -> Model {
        Model {
            conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device),
            conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device),
            pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(),
            activation: Relu::new(),
            linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device),
            linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device),
            dropout: DropoutConfig::new(self.dropout).init(),
        }
    }
}

At a glance, you can view the model configuration by printing the model instance:

mod model;

use crate::model::ModelConfig;
use burn::tensor::Device;

fn main() {
    let device = Device::wgpu(Default::default());
    let model = ModelConfig::new(10, 512).init(&device);

    println!("{model}");
}

Output:

Model {
  conv1: Conv2d {ch_in: 1, ch_out: 8, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 80}
  conv2: Conv2d {ch_in: 8, ch_out: 16, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 1168}
  pool: AdaptiveAvgPool2d {output_size: [8, 8]}
  dropout: Dropout {prob: 0.5}
  linear1: Linear {d_input: 1024, d_output: 512, bias: true, params: 524800}
  linear2: Linear {d_input: 512, d_output: 10, bias: true, params: 5130}
  activation: Relu
  params: 531178
}

pub fn init(&self, device: &Device) -> Model {
    Model {
        // ...
    }
}

When creating a custom neural network module, it is often a good idea to create a config alongside the model struct. This allows you to define default values for your network, thanks to the Config attribute. The benefit of this attribute is that it makes the configuration serializable, enabling you to painlessly save your model hyperparameters, enhancing your experimentation process. Note that a constructor will automatically be generated for your configuration, which will take in as input values the parameters which do not have default values: let config = ModelConfig::new(num_classes, hidden_size);. The default values can be overridden easily with builder-like methods: (e.g config.with_dropout(0.2);)

The first implementation block is related to the initialization method. As we can see, all fields are set using the configuration of the corresponding neural network's underlying module. In this specific case, we have chosen to expand the tensor channels from 1 to 8 with the first layer, then from 8 to 16 with the second layer, using a kernel size of 3 on all dimensions. We also use the adaptive average pooling module to reduce the dimensionality of the images to an 8 by 8 matrix, which we will flatten in the forward pass to have a 1024 (16 × 8 × 8) resulting tensor.

Now let's see how the forward pass is defined.

impl Model {
    /// # Shapes
    ///   - Images [batch_size, height, width]
    ///   - Output [batch_size, class_prob]
    pub fn forward(&self, images: Tensor<3>) -> Tensor<2> {
        let [batch_size, height, width] = images.dims();

        // Create a channel.
        let x = images.reshape([batch_size, 1, height, width]);

        let x = self.conv1.forward(x); // [batch_size, 8, _, _]
        let x = self.dropout.forward(x);
        let x = self.conv2.forward(x); // [batch_size, 16, _, _]
        let x = self.dropout.forward(x);
        let x = self.activation.forward(x);

        let x = self.pool.forward(x); // [batch_size, 16, 8, 8]
        let x = x.reshape([batch_size, 16 * 8 * 8]);
        let x = self.linear1.forward(x);
        let x = self.dropout.forward(x);
        let x = self.activation.forward(x);

        self.linear2.forward(x) // [batch_size, num_classes]
    }
}

For former PyTorch users, this might feel very intuitive, as each module is directly incorporated into the code using an eager API. Note that no abstraction is imposed for the forward method. You are free to define multiple forward functions with the names of your liking. Most of the neural network modules already built with Burn use the forward nomenclature, simply because it is standard in the field.

The Tensor struct takes its dimensionality and, optionally, its kind as generic arguments:

Tensor<3> // Float tensor (default)
Tensor<3, Float> // Float tensor (explicit)
Tensor<3, Int> // Int tensor
Tensor<3, Bool> // Bool tensor

The concrete element type, such as f16 or f32, is a runtime property that is specified in tensor creation operations. The default delement types can be configured per-device.