Struct SupervisedTraining 

pub struct SupervisedTraining<M: LearnerModel> { /* private fields */ }

Structure to configure and launch supervised learning trainings.

Implementations

impl<M: LearnerModel> SupervisedTraining<M>

pub fn new( directory: impl AsRef<Path>, dataloader_train: Arc<dyn DataLoader<<M as TrainStep>::Input>>, dataloader_valid: Arc<dyn DataLoader<<M as InferenceStep>::Input>>, ) -> Self

Creates a new runner for a supervised training.

Arguments

• directory - The directory to save the checkpoints.
• dataloader_train - The dataloader for the training split.
• dataloader_valid - The dataloader for the validation split.

impl<M: LearnerModel> SupervisedTraining<M>

pub fn with_training_strategy( self, training_strategy: TrainingStrategy<M>, ) -> Self

Replace the default training strategy (SingleDeviceTrainingStrategy) with the provided one.

Arguments

• training_strategy - The training strategy.

pub fn with_metric_logger<ML>(self, logger: ML) -> Self

where ML: MetricLogger + 'static,

Replace the default metric loggers with the provided ones.

Arguments

• logger - The training logger.

pub fn with_progress_logger<PL>(self, logger: PL) -> Self

where PL: TrainingProgressLogger + 'static,

Register a progress logger to track and store training progress.

Example

// `MyTrainingProgressLogger` is a user-defined type that implements
// `burn_train::logger::TrainingProgressLogger`.
let learner = SupervisedTraining::new(...)
    .with_progress_logger(MyTrainingProgressLogger);

pub fn with_checkpointing_strategy<CS: CheckpointingStrategy + 'static>( self, strategy: CS, ) -> Self

Update the checkpointing_strategy.

pub fn renderer(self, renderer: Box<dyn MetricsRenderer + 'static>) -> Self

Replace the default CLI renderer with a custom one.

Arguments

• renderer - The custom renderer.

pub fn metrics<Me: MetricRegistration<M>>(self, metrics: Me) -> Self

Register all metrics as numeric for the training and validation set.

pub fn metrics_text<Me: TextMetricRegistration<M>>(self, metrics: Me) -> Self

Register all metrics as text for the training and validation set.

pub fn metric_train<Me: Metric + 'static>(self, metric: Me) -> Self

where <M as TrainStep>::Output: Adaptor<Me::Input>,

Register a training metric.

pub fn metric_valid<Me: Metric + 'static>(self, metric: Me) -> Self

where <M as InferenceStep>::Output: Adaptor<Me::Input>,

Register a validation metric.

pub fn grads_accumulation(self, accumulation: usize) -> Self

Enable gradients accumulation.

Notes

When you enable gradients accumulation, the gradients object used by the optimizer will be the sum of all gradients generated by each backward pass. It might be a good idea to reduce the learning to compensate.

The effect is similar to increasing the batch size and the learning rate by the accumulation amount.

Panics

Panics if accumulation is zero.

pub fn gradient_checkpointing(self) -> Self

Enables autodiff checkpointing.

Notes

Gradient checkpointing recomputes activations during backpropagation for operations marked as memory-bound, while compute-bound operations still cache their output. This reduces peak memory usage at the cost of additional computation for memory-bound ops.

pub fn metric_train_numeric<Me>(self, metric: Me) -> Self

where Me: Metric + Numeric + 'static, <M as TrainStep>::Output: Adaptor<Me::Input>,

Register a numeric training metric.

pub fn metric_valid_numeric<Me: Metric + Numeric + 'static>( self, metric: Me, ) -> Self

where <M as InferenceStep>::Output: Adaptor<Me::Input>,

Register a numeric validation metric.

pub fn num_epochs(self, num_epochs: usize) -> Self

The number of epochs the training should last.

pub fn checkpoint(self, checkpoint: usize) -> Self

The epoch from which the training must resume.

pub fn interrupter(&self) -> Interrupter

Provides a handle that can be used to interrupt training.

pub fn with_interrupter(self, interrupter: Interrupter) -> Self

Override the handle for stopping training with an externally provided handle

pub fn early_stopping<Strategy>(self, strategy: Strategy) -> Self

where Strategy: EarlyStoppingStrategy + Clone + Send + Sync + 'static,

Register an early stopping strategy to stop the training when the conditions are meet.

pub fn with_application_logger( self, logger: Option<Box<dyn ApplicationLoggerInstaller>>, ) -> Self

By default, Rust logs are captured and written into experiment.log. If disabled, standard Rust log handling will apply.

pub fn with_default_checkpointers(self) -> Self

Register a checkpointer that will save the optimizer, the model and the learning rate scheduler to separate burnpack files.

pub fn with_custom_checkpointers<CM, CO, CL>( self, module_checkpointer: CM, optimizer_checkpointer: CO, lr_checkpointer: CL, ) -> Self

where CM: Checkpointer<ModuleRecord> + 'static, CO: Checkpointer<OptimizerRecord> + 'static, CL: Checkpointer<LrSchedulerRecord> + 'static,

Register your own checkpointers that will save the optimizer, the model and the learning rate scheduler to separate burnpack files.

pub fn summary(self) -> Self

Enable the training summary report.

The summary will be displayed after .fit(), when the renderer is dropped.

pub fn label(self, label: &str) -> Self

Set a label for this training, making it easier to differentiate multiple runs.

impl<M: LearnerModel> SupervisedTraining<M>

pub fn launch(self, learner: Learner<M>) -> LearningResult<M>

Launch this training with the given Learner.