Burn Framework (Part 2): Models & Multi-Backend
Learn how to construct neural network architectures with Burn Module derive macro and leverage flexible computation backends.
One of the standout features of Burn ↗ is the clean decoupling of model architectural definitions from hardware computation backends.
Defining Neural Networks with the Module Macro#
In Burn, model components implement the Module trait. Thanks to the #[derive(Module)] proc-macro, defining network layers is clean and strongly typed.
src/model.rs
use burn::nn::{Linear, LinearConfig, Relu};
use burn::module::Module;
use burn::tensor::backend::Backend;
use burn::tensor::Tensor;
#[derive(Module, Debug)]
pub struct Model<B: Backend> {
linear1: Linear<B>,
linear2: Linear<B>,
activation: Relu,
}
impl<B: Backend> Model<B> {
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, device: &B::Device) -> Self {
let linear1 = LinearConfig::new(input_dim, hidden_dim).init(device);
let linear2 = LinearConfig::new(hidden_dim, output_dim).init(device);
Self {
linear1,
linear2,
activation: Relu::new(),
}
}
pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
let x = self.linear1.forward(input);
let x = self.activation.forward(x);
// [!code focus]
self.linear2.forward(x)
}
}The Power of Multi-Backend Abstraction#
You can swap backends seamlessly by altering the generic backend parameter B:
- WGPU Backend: High performance across Vulkan, Metal, and DirectX.
- LibTorch Backend: Interoperability with PyTorch C++ bindings.
- Candle Backend: Lightweight execution based on Hugging Face Candle.
- NdArray Backend: Pure CPU fallbacks.
src/main.rs
use burn::backend::wgpu::{Wgpu, WgpuDevice};
fn main() {
let device = WgpuDevice::default();
let model = Model::<Wgpu>::new(784, 128, 10, &device);
println!("Model initialized successfully!");
}