burn_train/learner/supervised/
paradigm.rs

1use crate::checkpoint::{
2    AsyncCheckpointer, Checkpointer, CheckpointingStrategy, ComposedCheckpointingStrategy,
3    FileCheckpointer, KeepLastNCheckpoints, MetricCheckpointingStrategy,
4};
5use crate::learner::EarlyStoppingStrategy;
6use crate::learner::base::Interrupter;
7use crate::logger::{FileMetricLogger, MetricLogger, TrainingProgressLogger};
8use crate::metric::processor::{
9    AsyncProcessorTraining, FullEventProcessorTraining, MetricsTraining,
10};
11use crate::metric::store::{Aggregate, Direction, EventStoreClient, LogEventStore, Split};
12use crate::metric::{Adaptor, LossMetric, Metric, Numeric};
13use crate::multi::MultiDeviceLearningStrategy;
14use crate::renderer::{MetricsRenderer, default_renderer};
15use crate::single::SingleDeviceTrainingStrategy;
16use crate::{
17    ApplicationLoggerInstaller, EarlyStoppingStrategyRef, ExecutionStrategy,
18    FileApplicationLoggerInstaller, InferenceModelInput, InferenceModelOutput, InferenceStep,
19    LearnerEvent, LearnerModel, LearnerSummaryConfig, LearningCheckpointer, LearningResult,
20    TrainStep, TrainingComponents, TrainingModelInput, TrainingModelOutput, TrainingStrategy,
21};
22use crate::{Learner, SupervisedLearningStrategy};
23use burn_core::data::dataloader::DataLoader;
24use burn_core::store::ModuleRecord;
25use burn_core::tensor::Device;
26use burn_optim::OptimizerRecord;
27use burn_optim::lr_scheduler::LrSchedulerRecord;
28use std::collections::BTreeSet;
29use std::path::{Path, PathBuf};
30use std::sync::Arc;
31
32/// A reference to the training split [DataLoader](DataLoader).
33pub type TrainLoader<M> = Arc<dyn DataLoader<TrainingModelInput<M>>>;
34/// A reference to the validation split [DataLoader](DataLoader).
35pub type ValidLoader<M> = Arc<dyn DataLoader<InferenceModelInput<M>>>;
36/// The event processor type for supervised learning.
37pub type SupervisedTrainingEventProcessor<M> = AsyncProcessorTraining<
38    LearnerEvent<TrainingModelOutput<M>>,
39    LearnerEvent<InferenceModelOutput<M>>,
40>;
41
42/// Structure to configure and launch supervised learning trainings.
43pub struct SupervisedTraining<M: LearnerModel> {
44    // Not that complex. Extracting into another type would only make it more confusing.
45    #[allow(clippy::type_complexity)]
46    checkpointers: Option<(
47        AsyncCheckpointer<ModuleRecord>,
48        AsyncCheckpointer<OptimizerRecord>,
49        AsyncCheckpointer<LrSchedulerRecord>,
50    )>,
51    num_epochs: usize,
52    checkpoint: Option<usize>,
53    directory: PathBuf,
54    grad_accumulation: Option<usize>,
55    grad_checkpointing: bool,
56    renderer: Option<Box<dyn MetricsRenderer + 'static>>,
57    metrics: MetricsTraining<TrainingModelOutput<M>, InferenceModelOutput<M>>,
58    event_store: LogEventStore,
59    interrupter: Interrupter,
60    tracing_logger: Option<Box<dyn ApplicationLoggerInstaller>>,
61    checkpointer_strategy: Box<dyn CheckpointingStrategy>,
62    early_stopping: Option<EarlyStoppingStrategyRef>,
63    training_strategy: Option<TrainingStrategy<M>>,
64    dataloader_train: TrainLoader<M>,
65    dataloader_valid: ValidLoader<M>,
66    // Use BTreeSet instead of HashSet for consistent (alphabetical) iteration order
67    summary_metrics: BTreeSet<String>,
68    summary: bool,
69    progress_logger: Option<Box<dyn TrainingProgressLogger>>,
70    label: Option<String>,
71}
72
73impl<M: LearnerModel> SupervisedTraining<M> {
74    /// Creates a new runner for a supervised training.
75    ///
76    /// # Arguments
77    ///
78    /// * `directory` - The directory to save the checkpoints.
79    /// * `dataloader_train` - The dataloader for the training split.
80    /// * `dataloader_valid` - The dataloader for the validation split.
81    pub fn new(
82        directory: impl AsRef<Path>,
83        dataloader_train: Arc<dyn DataLoader<<M as TrainStep>::Input>>,
84        dataloader_valid: Arc<dyn DataLoader<<M as InferenceStep>::Input>>,
85    ) -> Self {
86        let directory = directory.as_ref().to_path_buf();
87        let experiment_log_file = directory.join("experiment.log");
88        Self {
89            num_epochs: 1,
90            checkpoint: None,
91            checkpointers: None,
92            directory,
93            grad_accumulation: None,
94            grad_checkpointing: false,
95            metrics: MetricsTraining::default(),
96            event_store: LogEventStore::default(),
97            renderer: None,
98            interrupter: Interrupter::new(),
99            tracing_logger: Some(Box::new(FileApplicationLoggerInstaller::new(
100                experiment_log_file,
101            ))),
102            checkpointer_strategy: Box::new(
103                ComposedCheckpointingStrategy::builder()
104                    .add(KeepLastNCheckpoints::new(2))
105                    .add(MetricCheckpointingStrategy::new(
106                        &LossMetric::new(), // default to valid loss
107                        Aggregate::Mean,
108                        Direction::Lowest,
109                        Split::Valid,
110                    ))
111                    .build(),
112            ),
113            early_stopping: None,
114            training_strategy: None,
115            summary_metrics: BTreeSet::new(),
116            summary: false,
117            dataloader_train,
118            dataloader_valid,
119            progress_logger: None,
120            label: None,
121        }
122    }
123}
124
125impl<M: LearnerModel> SupervisedTraining<M> {
126    /// Replace the default training strategy (SingleDeviceTrainingStrategy) with the provided one.
127    ///
128    /// # Arguments
129    ///
130    /// * `training_strategy` - The training strategy.
131    pub fn with_training_strategy(mut self, training_strategy: TrainingStrategy<M>) -> Self {
132        self.training_strategy = Some(training_strategy);
133        self
134    }
135
136    /// Replace the default metric loggers with the provided ones.
137    ///
138    /// # Arguments
139    ///
140    /// * `logger` - The training logger.
141    pub fn with_metric_logger<ML>(mut self, logger: ML) -> Self
142    where
143        ML: MetricLogger + 'static,
144    {
145        self.event_store.register_logger(logger);
146        self
147    }
148
149    /// Register a progress logger to track and store training progress.
150    ///
151    /// # Example
152    ///
153    /// ```ignore
154    /// // `MyTrainingProgressLogger` is a user-defined type that implements
155    /// // `burn_train::logger::TrainingProgressLogger`.
156    /// let learner = SupervisedTraining::new(...)
157    ///     .with_progress_logger(MyTrainingProgressLogger);
158    /// ```
159    pub fn with_progress_logger<PL>(mut self, logger: PL) -> Self
160    where
161        PL: TrainingProgressLogger + 'static,
162    {
163        self.progress_logger = Some(Box::new(logger));
164        self
165    }
166
167    /// Update the checkpointing_strategy.
168    pub fn with_checkpointing_strategy<CS: CheckpointingStrategy + 'static>(
169        mut self,
170        strategy: CS,
171    ) -> Self {
172        self.checkpointer_strategy = Box::new(strategy);
173        self
174    }
175
176    /// Replace the default CLI renderer with a custom one.
177    ///
178    /// # Arguments
179    ///
180    /// * `renderer` - The custom renderer.
181    pub fn renderer(mut self, renderer: Box<dyn MetricsRenderer + 'static>) -> Self {
182        self.renderer = Some(renderer);
183        self
184    }
185
186    /// Register all metrics as numeric for the training and validation set.
187    pub fn metrics<Me: MetricRegistration<M>>(self, metrics: Me) -> Self {
188        metrics.register(self)
189    }
190
191    /// Register all metrics as text for the training and validation set.
192    pub fn metrics_text<Me: TextMetricRegistration<M>>(self, metrics: Me) -> Self {
193        metrics.register(self)
194    }
195
196    /// Register a training metric.
197    pub fn metric_train<Me: Metric + 'static>(mut self, metric: Me) -> Self
198    where
199        TrainingModelOutput<M>: Adaptor<Me::Input>,
200    {
201        self.metrics.register_train_metric(metric);
202        self
203    }
204
205    /// Register a validation metric.
206    pub fn metric_valid<Me: Metric + 'static>(mut self, metric: Me) -> Self
207    where
208        InferenceModelOutput<M>: Adaptor<Me::Input>,
209    {
210        self.metrics.register_valid_metric(metric);
211        self
212    }
213
214    /// Enable gradients accumulation.
215    ///
216    /// # Notes
217    ///
218    /// When you enable gradients accumulation, the gradients object used by the optimizer will be
219    /// the sum of all gradients generated by each backward pass. It might be a good idea to
220    /// reduce the learning to compensate.
221    ///
222    /// The effect is similar to increasing the `batch size` and the `learning rate` by the `accumulation`
223    /// amount.
224    ///
225    /// # Panics
226    ///
227    /// Panics if `accumulation` is zero.
228    pub fn grads_accumulation(mut self, accumulation: usize) -> Self {
229        assert!(
230            accumulation > 0,
231            "Gradient accumulation must be greater than zero."
232        );
233        self.grad_accumulation = Some(accumulation);
234        self
235    }
236
237    /// Enables autodiff checkpointing.
238    ///
239    /// # Notes
240    /// Gradient checkpointing recomputes activations during backpropagation for operations
241    /// marked as memory-bound, while compute-bound operations still cache their
242    /// output. This reduces peak memory usage at the cost of additional computation
243    /// for memory-bound ops.
244    pub fn gradient_checkpointing(mut self) -> Self {
245        self.grad_checkpointing = true;
246        self
247    }
248
249    /// Register a [numeric](crate::metric::Numeric) training [metric](Metric).
250    pub fn metric_train_numeric<Me>(mut self, metric: Me) -> Self
251    where
252        Me: Metric + Numeric + 'static,
253        TrainingModelOutput<M>: Adaptor<Me::Input>,
254    {
255        self.summary_metrics.insert(metric.name().to_string());
256        self.metrics.register_train_metric_numeric(metric);
257        self
258    }
259
260    /// Register a [numeric](crate::metric::Numeric) validation [metric](Metric).
261    pub fn metric_valid_numeric<Me: Metric + Numeric + 'static>(mut self, metric: Me) -> Self
262    where
263        InferenceModelOutput<M>: Adaptor<Me::Input>,
264    {
265        self.summary_metrics.insert(metric.name().to_string());
266        self.metrics.register_valid_metric_numeric(metric);
267        self
268    }
269
270    /// The number of epochs the training should last.
271    pub fn num_epochs(mut self, num_epochs: usize) -> Self {
272        self.num_epochs = num_epochs;
273        self
274    }
275
276    /// The epoch from which the training must resume.
277    pub fn checkpoint(mut self, checkpoint: usize) -> Self {
278        self.checkpoint = Some(checkpoint);
279        self
280    }
281
282    /// Provides a handle that can be used to interrupt training.
283    pub fn interrupter(&self) -> Interrupter {
284        self.interrupter.clone()
285    }
286
287    /// Override the handle for stopping training with an externally provided handle
288    pub fn with_interrupter(mut self, interrupter: Interrupter) -> Self {
289        self.interrupter = interrupter;
290        self
291    }
292
293    /// Register an [early stopping strategy](EarlyStoppingStrategy) to stop the training when the
294    /// conditions are meet.
295    pub fn early_stopping<Strategy>(mut self, strategy: Strategy) -> Self
296    where
297        Strategy: EarlyStoppingStrategy + Clone + Send + Sync + 'static,
298    {
299        self.early_stopping = Some(Box::new(strategy));
300        self
301    }
302
303    /// By default, Rust logs are captured and written into
304    /// `experiment.log`. If disabled, standard Rust log handling
305    /// will apply.
306    pub fn with_application_logger(
307        mut self,
308        logger: Option<Box<dyn ApplicationLoggerInstaller>>,
309    ) -> Self {
310        self.tracing_logger = logger;
311        self
312    }
313
314    /// Register a checkpointer that will save the [optimizer](burn_optim::ModuleOptimizer), the
315    /// [model](LearnerModel) and the [learning rate scheduler](burn_optim::lr_scheduler::module_lr_scheduler::ModuleLrScheduler) to separate burnpack files.
316    pub fn with_default_checkpointers(mut self) -> Self {
317        let checkpoint_dir = self.directory.join("checkpoint");
318        let checkpointer_model = FileCheckpointer::new(&checkpoint_dir, "model");
319        let checkpointer_optimizer = FileCheckpointer::new(&checkpoint_dir, "optim");
320        let checkpointer_scheduler = FileCheckpointer::new(&checkpoint_dir, "scheduler");
321
322        self.checkpointers = Some((
323            AsyncCheckpointer::new(checkpointer_model),
324            AsyncCheckpointer::new(checkpointer_optimizer),
325            AsyncCheckpointer::new(checkpointer_scheduler),
326        ));
327
328        self
329    }
330
331    /// Register your own checkpointers that will save the [optimizer](burn_optim::ModuleOptimizer), the
332    /// [model](LearnerModel) and the [learning rate scheduler](burn_optim::lr_scheduler::module_lr_scheduler::ModuleLrScheduler) to separate burnpack files.
333    pub fn with_custom_checkpointers<CM, CO, CL>(
334        mut self,
335        module_checkpointer: CM,
336        optimizer_checkpointer: CO,
337        lr_checkpointer: CL,
338    ) -> Self
339    where
340        CM: Checkpointer<ModuleRecord> + 'static,
341        CO: Checkpointer<OptimizerRecord> + 'static,
342        CL: Checkpointer<LrSchedulerRecord> + 'static,
343    {
344        self.checkpointers = Some((
345            AsyncCheckpointer::new(module_checkpointer),
346            AsyncCheckpointer::new(optimizer_checkpointer),
347            AsyncCheckpointer::new(lr_checkpointer),
348        ));
349
350        self
351    }
352
353    /// Enable the training summary report.
354    ///
355    /// The summary will be displayed after `.fit()`, when the renderer is dropped.
356    pub fn summary(mut self) -> Self {
357        self.summary = true;
358        self
359    }
360
361    /// Set a label for this training, making it easier to differentiate multiple runs.
362    pub fn label(mut self, label: &str) -> Self {
363        self.label = Some(label.to_string());
364        self
365    }
366}
367
368impl<M: LearnerModel> SupervisedTraining<M> {
369    /// Launch this training with the given [Learner](Learner).
370    pub fn launch(mut self, learner: Learner<M>) -> LearningResult<M> {
371        if self.tracing_logger.is_some()
372            && let Err(e) = self.tracing_logger.as_ref().unwrap().install()
373        {
374            log::warn!("Failed to install the experiment logger: {e}");
375        }
376        let renderer = self
377            .renderer
378            .unwrap_or_else(|| default_renderer(self.interrupter.clone(), self.checkpoint));
379
380        if !self.event_store.has_loggers() {
381            self.event_store
382                .register_logger(FileMetricLogger::new(self.directory.clone()));
383        }
384
385        let event_store = Arc::new(EventStoreClient::new(self.event_store));
386        let full_processor =
387            FullEventProcessorTraining::new(self.metrics, renderer, event_store.clone());
388        let full_processor = match self.progress_logger {
389            Some(logger) => full_processor.with_progress_logger(logger),
390            None => full_processor,
391        };
392        let event_processor = AsyncProcessorTraining::new(full_processor);
393
394        let checkpointer = self.checkpointers.map(|(model, optim, scheduler)| {
395            LearningCheckpointer::new(
396                model.with_interrupter(self.interrupter.clone()),
397                optim.with_interrupter(self.interrupter.clone()),
398                scheduler.with_interrupter(self.interrupter.clone()),
399                self.checkpointer_strategy,
400            )
401        });
402
403        let summary = if self.summary {
404            Some(LearnerSummaryConfig {
405                directory: self.directory,
406                metrics: self.summary_metrics.into_iter().collect::<Vec<_>>(),
407            })
408        } else {
409            None
410        };
411
412        let components = TrainingComponents {
413            checkpoint: self.checkpoint,
414            checkpointer,
415            interrupter: self.interrupter,
416            early_stopping: self.early_stopping,
417            event_processor,
418            event_store,
419            num_epochs: self.num_epochs,
420            grad_accumulation: self.grad_accumulation,
421            summary,
422            label: self.label,
423        };
424
425        // Default to single device based on model
426        let training_strategy = self.training_strategy.unwrap_or(TrainingStrategy::Default(
427            ExecutionStrategy::SingleDevice(autodiff_device(
428                learner.model.devices()[0].clone(),
429                self.grad_checkpointing,
430            )),
431        ));
432
433        let mut learner = learner;
434        if let Some(checkpoint) = components.checkpoint
435            && let Some(checkpointer) = &components.checkpointer
436        {
437            learner = checkpointer.load_checkpoint(learner, checkpoint);
438        }
439
440        match training_strategy {
441            TrainingStrategy::Custom(learning_paradigm) => learning_paradigm.train(
442                learner,
443                self.dataloader_train,
444                self.dataloader_valid,
445                components,
446            ),
447            TrainingStrategy::Default(strategy) => match strategy {
448                ExecutionStrategy::SingleDevice(device) => {
449                    let single_device = SingleDeviceTrainingStrategy::new(autodiff_device(
450                        device,
451                        self.grad_checkpointing,
452                    ));
453                    single_device.train(
454                        learner,
455                        self.dataloader_train,
456                        self.dataloader_valid,
457                        components,
458                    )
459                }
460                ExecutionStrategy::MultiDevice(devices, multi_device_optim) => {
461                    let strategy: Box<dyn SupervisedLearningStrategy<M>> = match devices.len() == 1
462                    {
463                        true => Box::new(SingleDeviceTrainingStrategy::new(autodiff_device(
464                            devices[0].clone(),
465                            self.grad_checkpointing,
466                        ))),
467                        false => Box::new(MultiDeviceLearningStrategy::new(
468                            devices
469                                .into_iter()
470                                .map(|d| autodiff_device(d, self.grad_checkpointing))
471                                .collect(),
472                            multi_device_optim,
473                        )),
474                    };
475                    strategy.train(
476                        learner,
477                        self.dataloader_train,
478                        self.dataloader_valid,
479                        components,
480                    )
481                }
482                ExecutionStrategy::DistributedDataParallel { devices, context } => {
483                    use crate::ddp::DdpTrainingStrategy;
484
485                    let ddp = DdpTrainingStrategy::new(
486                        devices
487                            .into_iter()
488                            .map(|d| autodiff_device(d, self.grad_checkpointing))
489                            .collect(),
490                        context,
491                    );
492                    ddp.train(
493                        learner,
494                        self.dataloader_train,
495                        self.dataloader_valid,
496                        components,
497                    )
498                }
499            },
500        }
501    }
502}
503
504// Validate the autodiff device property and enable grad checkpointing.
505fn autodiff_device(mut device: Device, grad_checkpointing: bool) -> Device {
506    if !device.is_autodiff() {
507        device = device.autodiff();
508    }
509
510    if grad_checkpointing {
511        device = device.gradient_checkpointing();
512    }
513
514    device
515}
516
517/// Trait to fake variadic generics.
518pub trait MetricRegistration<M: LearnerModel>: Sized {
519    /// Register the metrics.
520    fn register(self, builder: SupervisedTraining<M>) -> SupervisedTraining<M>;
521}
522
523/// Trait to fake variadic generics.
524pub trait TextMetricRegistration<M: LearnerModel>: Sized {
525    /// Register the metrics.
526    fn register(self, builder: SupervisedTraining<M>) -> SupervisedTraining<M>;
527}
528
529macro_rules! gen_tuple {
530    ($($M:ident),*) => {
531        impl<$($M,)* M: LearnerModel> TextMetricRegistration<M> for ($($M,)*)
532        where
533            $(TrainingModelOutput<M>: Adaptor<$M::Input>,)*
534            $(InferenceModelOutput<M>: Adaptor<$M::Input>,)*
535            $($M: Metric + 'static,)*
536        {
537            #[allow(non_snake_case)]
538            fn register(
539                self,
540                builder: SupervisedTraining<M>,
541            ) -> SupervisedTraining<M> {
542                let ($($M,)*) = self;
543                $(let builder = builder.metric_train($M.clone());)*
544                $(let builder = builder.metric_valid($M);)*
545                builder
546            }
547        }
548
549        impl<$($M,)* M: LearnerModel> MetricRegistration<M> for ($($M,)*)
550        where
551            $(TrainingModelOutput<M>: Adaptor<$M::Input>,)*
552            $(InferenceModelOutput<M>: Adaptor<$M::Input>,)*
553            $($M: Metric + Numeric + 'static,)*
554        {
555            #[allow(non_snake_case)]
556            fn register(
557                self,
558                builder: SupervisedTraining<M>,
559            ) -> SupervisedTraining<M> {
560                let ($($M,)*) = self;
561                $(let builder = builder.metric_train_numeric($M.clone());)*
562                $(let builder = builder.metric_valid_numeric($M);)*
563                builder
564            }
565        }
566    };
567}
568
569gen_tuple!(M1);
570gen_tuple!(M1, M2);
571gen_tuple!(M1, M2, M3);
572gen_tuple!(M1, M2, M3, M4);
573gen_tuple!(M1, M2, M3, M4, M5);
574gen_tuple!(M1, M2, M3, M4, M5, M6);
burn_train/learner/supervised/paradigm.rs

burn_train/learner/supervised/
paradigm.rs
