fpm_rs/callbacks/
early_stop.rs

1use crate::{Result, diagnostics::DiagnosticRequest};
2
3use super::{Callback, CallbackAction, CallbackHook, StepContext};
4
5pub struct StopOnPlateau {
6    patience: usize,
7    minimum_improvement: f64,
8    best: f64,
9    stale_iterations: usize,
10}
11
12impl StopOnPlateau {
13    pub fn new(patience: usize, minimum_improvement: f64) -> Self {
14        Self {
15            patience: patience.max(1),
16            minimum_improvement: minimum_improvement.max(0.0),
17            best: f64::INFINITY,
18            stale_iterations: 0,
19        }
20    }
21}
22
23impl Callback for StopOnPlateau {
24    fn requires(&self) -> Vec<DiagnosticRequest> {
25        vec![DiagnosticRequest::Loss]
26    }
27
28    fn requires_for(&self, hook: CallbackHook, _iteration: usize) -> Vec<DiagnosticRequest> {
29        if hook == CallbackHook::IterationEnd {
30            self.requires()
31        } else {
32            Vec::new()
33        }
34    }
35
36    fn on_start(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
37        self.best = f64::INFINITY;
38        self.stale_iterations = 0;
39        for record in &context.history.iterations {
40            if self.best - record.loss > self.minimum_improvement {
41                self.best = record.loss;
42                self.stale_iterations = 0;
43            } else {
44                self.stale_iterations += 1;
45            }
46        }
47        Ok(CallbackAction::Continue)
48    }
49
50    fn on_iteration_end(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
51        if let Some(loss) = context.diagnostics.loss {
52            if self.best - loss > self.minimum_improvement {
53                self.best = loss;
54                self.stale_iterations = 0;
55            } else {
56                self.stale_iterations += 1;
57            }
58        }
59        Ok(if self.stale_iterations >= self.patience {
60            CallbackAction::Stop
61        } else {
62            CallbackAction::Continue
63        })
64    }
65}
fpm_rs/callbacks/early_stop.rs

fpm_rs/callbacks/
early_stop.rs
