burn_train/checkpoint/strategy/
metric.rs

1use super::CheckpointingStrategy;
2use crate::{
3    checkpoint::CheckpointingAction,
4    metric::{
5        Metric, MetricName,
6        store::{Aggregate, Direction, EventStoreClient, Split},
7    },
8};
9
10/// Keep the best checkpoint based on a metric.
11pub struct MetricCheckpointingStrategy {
12    current: Option<usize>,
13    aggregate: Aggregate,
14    direction: Direction,
15    split: Split,
16    name: MetricName,
17}
18
19impl MetricCheckpointingStrategy {
20    /// Create a new metric checkpointing strategy.
21    pub fn new<M>(metric: &M, aggregate: Aggregate, direction: Direction, split: Split) -> Self
22    where
23        M: Metric,
24    {
25        Self {
26            current: None,
27            name: metric.name(),
28            aggregate,
29            direction,
30            split,
31        }
32    }
33}
34
35impl CheckpointingStrategy for MetricCheckpointingStrategy {
36    fn checkpointing(
37        &mut self,
38        epoch: usize,
39        store: &EventStoreClient,
40    ) -> Vec<CheckpointingAction> {
41        let best_epoch =
42            match store.find_epoch(&self.name, self.aggregate, self.direction, &self.split) {
43                Some(epoch_best) => epoch_best,
44                None => epoch,
45            };
46
47        let mut actions = Vec::new();
48
49        if let Some(current) = self.current
50            && current != best_epoch
51        {
52            actions.push(CheckpointingAction::Delete(current));
53        }
54
55        if best_epoch == epoch {
56            actions.push(CheckpointingAction::Save);
57        }
58
59        self.current = Some(best_epoch);
60
61        actions
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use crate::{
68        EventProcessorTraining,
69        logger::InMemoryMetricLogger,
70        metric::{
71            LossMetric,
72            processor::{
73                MetricsTraining, MinimalEventProcessor,
74                test_utils::{end_epoch, process_train},
75            },
76            store::LogEventStore,
77        },
78        test_utils::start_epoch,
79    };
80
81    use super::*;
82    use std::sync::Arc;
83
84    #[test]
85    fn always_keep_the_best_epoch() {
86        let loss = LossMetric::new();
87        let mut store = LogEventStore::default();
88        let mut strategy = MetricCheckpointingStrategy::new(
89            &loss,
90            Aggregate::Mean,
91            Direction::Lowest,
92            Split::Train,
93        );
94        let mut metrics = MetricsTraining::<f64, f64>::default();
95        // Register an in memory logger.
96        store.register_logger(InMemoryMetricLogger::default());
97        // Register the loss metric.
98        metrics.register_train_metric_numeric(loss);
99        let store = Arc::new(EventStoreClient::new(store));
100        let mut processor = MinimalEventProcessor::new(metrics, store.clone());
101        processor
102            .process_train(crate::LearnerEvent::Start {
103                total_epochs: 0,
104                starting_epoch: 0,
105                label: None,
106            })
107            .unwrap();
108
109        // Two points for the first epoch. Mean 0.75
110        let mut epoch = 1;
111        start_epoch(&mut processor, epoch, 2);
112        process_train(&mut processor, 1.0, epoch);
113        process_train(&mut processor, 0.5, epoch);
114        end_epoch(&mut processor, epoch);
115
116        // Should save the current record.
117        assert_eq!(
118            vec![CheckpointingAction::Save],
119            strategy.checkpointing(epoch, &store)
120        );
121
122        // Two points for the second epoch. Mean 0.4
123        epoch += 1;
124        start_epoch(&mut processor, epoch, 2);
125        process_train(&mut processor, 0.5, epoch);
126        process_train(&mut processor, 0.3, epoch);
127        end_epoch(&mut processor, epoch);
128
129        // Should save the current record and delete the previous one.
130        assert_eq!(
131            vec![CheckpointingAction::Delete(1), CheckpointingAction::Save],
132            strategy.checkpointing(epoch, &store)
133        );
134
135        // Two points for the last epoch. Mean 2.0
136        epoch += 1;
137        start_epoch(&mut processor, epoch, 2);
138        process_train(&mut processor, 1.0, epoch);
139        process_train(&mut processor, 3.0, epoch);
140        end_epoch(&mut processor, epoch);
141
142        // Should not delete the previous record, since it's the best one, and should not save a
143        // new one.
144        assert!(strategy.checkpointing(epoch, &store).is_empty());
145    }
146}
burn_train/checkpoint/strategy/metric.rs

burn_train/checkpoint/strategy/
metric.rs
