burn_train/learner/supervised/strategies/
base.rs

1use crate::{
2    EarlyStoppingStrategyRef, Interrupter, Learner, LearnerModel, LearnerSummaryConfig,
3    LearningCheckpointer, LearningResult, SupervisedTrainingEventProcessor, TrainLoader,
4    ValidLoader,
5    metric::{
6        processor::{EventProcessorTraining, LearnerEvent},
7        store::EventStoreClient,
8    },
9};
10use burn_core::prelude::Device;
11use burn_core::tensor::distributed::{DistributedConfig, DistributedContext};
12use std::sync::Arc;
13
14/// A reference to an implementation of SupervisedLearningStrategy.
15pub type CustomLearningStrategy<M> = Arc<dyn SupervisedLearningStrategy<M>>;
16
17#[derive(Clone, Copy, Debug)]
18/// Determine how the optimization is performed when training with multiple devices.
19pub enum MultiDeviceOptim {
20    /// The optimization is done on an elected device.
21    OptimMainDevice,
22    /// The optimization is sharded across all devices.
23    OptimSharded,
24}
25
26/// Describes where training runs.
27pub enum ExecutionStrategy {
28    /// Training on one device
29    SingleDevice(Device),
30    /// Performs data-parallel distributed training where the optimization is
31    /// done on an elected master device.
32    MultiDevice(Vec<Device>, MultiDeviceOptim),
33    /// Training with input distributed across devices, each device has its own copy of the model.
34    /// Collective ops are used to sync the gradients after each pass.
35    DistributedDataParallel {
36        /// Devices on this node for the DDP
37        devices: Vec<Device>,
38        /// The distributed runtime.
39        context: DistributedContext,
40    },
41}
42
43impl ExecutionStrategy {
44    /// Returns the primary device responsible for coordination.
45    pub fn main_device(&self) -> &Device {
46        match self {
47            ExecutionStrategy::SingleDevice(device) => device,
48            ExecutionStrategy::MultiDevice(devices, _optim) => &devices[0],
49            ExecutionStrategy::DistributedDataParallel {
50                devices,
51                context: _,
52            } => &devices[0],
53        }
54    }
55
56    /// Creates a strategy for a single device.
57    pub fn single(device: Device) -> Self {
58        Self::SingleDevice(device)
59    }
60
61    /// Creates a multi-device strategy.
62    pub fn multi(devices: Vec<Device>, optim: MultiDeviceOptim) -> Self {
63        Self::MultiDevice(devices, optim)
64    }
65}
66
67impl ExecutionStrategy {
68    /// Creates a distributed data parallel (DDP) strategy.
69    pub fn ddp(devices: Vec<Device>, config: DistributedConfig) -> Self {
70        let context = DistributedContext::init(devices.clone(), config);
71        Self::DistributedDataParallel { devices, context }
72    }
73}
74
75/// How should the learner run the learning for the model
76pub enum TrainingStrategy<M: LearnerModel> {
77    /// Default training loop with specified device strategy.
78    Default(ExecutionStrategy),
79    /// Training using a custom learning strategy
80    Custom(CustomLearningStrategy<M>),
81}
82
83impl<M: LearnerModel> From<ExecutionStrategy> for TrainingStrategy<M> {
84    fn from(value: ExecutionStrategy) -> Self {
85        Self::Default(value)
86    }
87}
88
89impl<M: LearnerModel> Default for TrainingStrategy<M> {
90    fn default() -> Self {
91        Self::Default(ExecutionStrategy::SingleDevice(Default::default()))
92    }
93}
94
95/// Struct to minimise parameters passed to [SupervisedLearningStrategy::train].
96/// These components are used during training.
97pub struct TrainingComponents<M: LearnerModel> {
98    /// An optional label for this training.
99    pub label: Option<String>,
100    /// The total number of epochs
101    pub num_epochs: usize,
102    /// The epoch number from which to continue the training.
103    pub checkpoint: Option<usize>,
104    /// A checkpointer used to load and save learner checkpoints.
105    pub checkpointer: Option<LearningCheckpointer<M>>,
106    /// Enables gradients accumulation.
107    pub grad_accumulation: Option<usize>,
108    /// An [Interupter](Interrupter) that allows aborting the training/evaluation process early.
109    pub interrupter: Interrupter,
110    /// Cloneable reference to an early stopping strategy.
111    pub early_stopping: Option<EarlyStoppingStrategyRef>,
112    /// An [EventProcessor](crate::EventProcessorTraining) that processes events happening during training and validation.
113    pub event_processor: SupervisedTrainingEventProcessor<M>,
114    /// A reference to an [EventStoreClient](EventStoreClient).
115    pub event_store: Arc<EventStoreClient>,
116    /// Config for creating a summary of the learning
117    pub summary: Option<LearnerSummaryConfig>,
118}
119
120/// Provides the `fit` function for any learning strategy
121pub trait SupervisedLearningStrategy<M: LearnerModel> {
122    /// Train the learner's model with this strategy.
123    fn train(
124        &self,
125        learner: Learner<M>,
126        dataloader_train: TrainLoader<M>,
127        dataloader_valid: ValidLoader<M>,
128        mut training_components: TrainingComponents<M>,
129    ) -> LearningResult<M> {
130        let starting_epoch = training_components.checkpoint.unwrap_or(0) + 1;
131        let summary_config = training_components.summary.clone();
132        let interrupter = training_components.interrupter.clone();
133
134        // Event processor start training
135        interrupter.fail_on_error(training_components.event_processor.process_train(
136            LearnerEvent::Start {
137                total_epochs: training_components.num_epochs,
138                starting_epoch,
139                label: training_components.label.clone(),
140            },
141        ));
142        // Training loop
143        let (model, mut event_processor) = self.fit(
144            training_components,
145            learner,
146            dataloader_train,
147            dataloader_valid,
148            starting_epoch,
149        );
150
151        let summary = summary_config.and_then(|summary| {
152            summary
153                .init()
154                .map(|summary| summary.with_model(model.to_string()))
155                .ok()
156        });
157
158        // Signal training end. For the TUI renderer, this handles the exit & return to main screen.
159        interrupter.fail_on_error(event_processor.process_train(LearnerEvent::End(summary)));
160
161        let model = model.valid();
162        // Finish processing remaining events.
163        interrupter.fail_on_error(event_processor.flush());
164        let renderer = event_processor.renderer();
165
166        LearningResult::<M> {
167            model,
168            renderer,
169            interrupted: interrupter.interruption(),
170            error: interrupter.error(),
171        }
172    }
173
174    /// Training loop for this strategy
175    fn fit(
176        &self,
177        training_components: TrainingComponents<M>,
178        learner: Learner<M>,
179        dataloader_train: TrainLoader<M>,
180        dataloader_valid: ValidLoader<M>,
181        starting_epoch: usize,
182    ) -> (M, SupervisedTrainingEventProcessor<M>);
183}
burn_train/learner/supervised/strategies/base.rs

burn_train/learner/supervised/strategies/
base.rs
