burn_train/metric/
learning_rate.rs

1use std::sync::Arc;
2
3use super::{
4    MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry,
5    state::{FormatOptions, NumericMetricState},
6};
7use crate::metric::{Metric, MetricName, Numeric, SerializedEntry};
8use burn_core::tensor::TensorReadError;
9
10/// Track the learning rate across iterations.
11#[derive(Clone)]
12pub struct LearningRateMetric {
13    name: MetricName,
14    state: NumericMetricState,
15}
16
17impl LearningRateMetric {
18    /// Creates a new learning rate metric.
19    pub fn new() -> Self {
20        Self {
21            name: Arc::new("Learning Rate".to_string()),
22            state: NumericMetricState::new(),
23        }
24    }
25}
26
27impl Default for LearningRateMetric {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl Metric for LearningRateMetric {
34    type Input = ();
35
36    fn update(
37        &mut self,
38        _item: &(),
39        metadata: &MetricMetadata,
40    ) -> Result<SerializedEntry, TensorReadError> {
41        // TODO: We only log the default learning rate. Yet another motivation to introduce metric groups.
42        let lr = metadata.lr.as_ref().map(|val| val.base()).unwrap_or(0.0);
43
44        self.state.update(lr, 1);
45        Ok(self
46            .state
47            .compute_update(FormatOptions::new(self.name()).precision(2)))
48    }
49
50    fn compute(&mut self) -> Result<SerializedEntry, TensorReadError> {
51        Ok(self
52            .state
53            .compute_final(FormatOptions::new(self.name()).precision(2)))
54    }
55
56    fn clear(&mut self) {
57        self.state.reset()
58    }
59
60    fn name(&self) -> MetricName {
61        self.name.clone()
62    }
63
64    fn attributes(&self) -> MetricAttributes {
65        NumericAttributes {
66            unit: None,
67            higher_is_better: false,
68        }
69        .into()
70    }
71}
72
73// TODO: LR should probably just report the current value, the aggregated values don't make as much sense esp. for visualization
74impl Numeric for LearningRateMetric {
75    fn value(&self) -> Option<NumericEntry> {
76        Some(self.state.current_value())
77    }
78
79    fn running_value(&self) -> Option<NumericEntry> {
80        Some(self.state.running_value())
81    }
82
83    fn final_value(&self) -> NumericEntry {
84        self.state.final_value()
85    }
86}
burn_train/metric/learning_rate.rs

burn_train/metric/
learning_rate.rs
