Learner

The burn-train crate encapsulates multiple utilities for training deep learning models. The goal of the crate is to provide users with a well-crafted and flexible training loop, so that projects do not have to write such components from the ground up. Most of the interactions with burn-train will be with the SupervisedTraining struct, briefly presented in the previous training section. This struct enables you to configure the training loop, offering support for registering metrics, enabling logging, checkpointing states, using multiple devices, and so on.

There are still some assumptions in the current provided APIs, which may make them inappropriate for your learning requirements. Indeed, they assume your model will learn from a training dataset and be validated against another dataset. This is the most common paradigm, allowing users to do both supervised and unsupervised learning as well as fine-tuning. However, for more complex requirements, creating a custom training loop might be what you need.

Usage

The SupervisedLearning struct must be created with the training and validation dataloaders. It provides you with numerous options when it comes to configurations.

When the training is configured to your liking, you can then move forward to running the training. The launch method requires a learner object providing: the model, the optimizer and the learning rate scheduler. Note that the latter can be a simple float if you want it to be constant during training. See the learning rate scheduler section for the available schedulers.

The launch method will start the training and return the trained model once finished.

Again, please refer to the training section for a relevant code snippet.

Gradient Accumulation

Gradient accumulation lets each optimizer update use gradients from several batches without increasing the dataloader's batch size. This is useful when a larger batch would exceed device memory.

Enable it with SupervisedTraining::grads_accumulation(n), where n must be greater than zero. On a single device, Burn computes gradients for each batch and sums them over n batches before updating the model parameters.

Burn applies any remaining accumulated gradients at the end of each epoch. Learning-rate schedulers advance once per optimizer update, so warmup and decay durations should use that unit.

Accumulated gradients are summed without normalization. Account for this when choosing loss scaling and learning rates.

Custom training strategies and loops are responsible for their own accumulation and scheduler timing.

Parameter Groups

It's common to use different learning rates or optimizer settings for different parts of a model. Burn's ParamGroup routes module parameters by path or ID. Optimizers and learning-rate schedulers use the same matching rules but can be configured independently.

use burn::{
    module::ParamGroup,
    optim::{
        AdamWConfig,
        lr_scheduler::{
            composed::ComposedLrSchedulerConfig,
            cosine::CosineAnnealingLrSchedulerConfig,
            linear::LinearLrSchedulerConfig,
            module_lr_scheduler::ModuleLrSchedulerConfig,
        },
    },
    train::Learner,
};

let lr_scheduler_base = ComposedLrSchedulerConfig::new()
    .cosine(CosineAnnealingLrSchedulerConfig::new(1.0, 2000))
    .linear(LinearLrSchedulerConfig::new(1e-8, 1.0, 2000))
    .linear(LinearLrSchedulerConfig::new(1e-2, 1e-6, 10000));
let lr_scheduler = ModuleLrSchedulerConfig::new(lr_scheduler_base.into())
    .with_group(
        ParamGroup::from_predicate("conv"),
        LinearLrSchedulerConfig::new(1e-6, 1e-3, 14000),
    )
    .init()
    .unwrap();

let optim = AdamWConfig::new()
    .with_cautious_weight_decay(true)
    .with_weight_decay(5e-5)
    .init();

let result = training.launch(Learner::new(
    model,
    optim,
    lr_scheduler,
));

The composed base schedule multiplies the values of its three component schedules at every step. ModuleLrSchedulerConfig assigns that base policy to parameters outside the conv group and the separate linear policy to parameters inside it. Configure the groups before calling init(); the individual scheduler configurations' build() methods are internal APIs.

For group-specific optimizers, matching precedence, gradient clipping, and optimizer state, see Optimizer.

Errors and interruptions

Training can stop early on an error or a user-triggered interruption. The LearningResult returned by launch reports the reason:

• error holds the TrainingError that stopped training.
• interrupted holds the Interruption when training was stopped by a user, without an error.

Both are None when training ran to completion or early stopping ended it. The epoch during which training stopped is not checkpointed.

Evaluation reports the same fields in the EvaluationResult returned by eval and eval_all.

Artifacts

When creating a SupervisedTraining instance, all the collected data will be saved under the directory provided as the argument to the new method. Here is an example of the data layout for a model checkpointed to the burnpack format, with the accuracy and loss metrics registered:

├── experiment.log
├── checkpoint
│   ├── model-1.bpk
│   ├── optim-1.bpk
│   └── scheduler-1.bpk
│   ├── model-2.bpk
│   ├── optim-2.bpk
│   └── scheduler-2.bpk
├── train
│   ├── epoch-1
│   │   ├── Accuracy.log
│   │   └── Loss.log
│   └── epoch-2
│       ├── Accuracy.log
│       └── Loss.log
└── valid
    ├── epoch-1
    │   ├── Accuracy.log
    │   └── Loss.log
    └── epoch-2
        ├── Accuracy.log
        └── Loss.log

You can choose to save or synchronize that local directory with a remote file system, if desired. The file checkpointer is capable of automatically deleting old checkpoints according to a specified configuration.