burn_train/metric/
state.rs

1use std::sync::Arc;
2
3use burn_core::tensor::{Bool, Tensor, TensorReadError};
4
5use crate::metric::{ClassReduction, MetricName, NumericEntry, SerializedEntry, format_float};
6
7/// Useful utility to implement numeric metrics.
8///
9/// # Notes
10///
11/// The numeric metric store values inside floats.
12/// Even if some metric are integers, their mean are floats.
13#[derive(Clone)]
14pub struct NumericMetricState {
15    sum: f64,
16    count: usize,
17    current: f64,
18    current_count: usize,
19}
20
21/// Accumulates raw predictions and targets across batches.
22///
23/// Used by rank-based metrics (AUROC, AUC-PR) that must recompute over the
24/// whole epoch. Buffers are freed on [`reset`](Self::reset).
25#[derive(Clone)]
26pub struct PredictionAccumulatorState {
27    predictions: Vec<Tensor<2>>,
28    targets: Vec<Tensor<2, Bool>>,
29    current: Option<f64>,
30}
31
32/// Formatting options for the [numeric metric state](NumericMetricState).
33pub struct FormatOptions {
34    name: Arc<String>,
35    unit: Option<String>,
36    precision: Option<usize>,
37}
38
39impl PredictionAccumulatorState {
40    /// Create a new [prediction accumulator state](PredictionAccumulatorState).
41    pub fn new() -> Self {
42        Self {
43            predictions: vec![],
44            targets: vec![],
45            current: None,
46        }
47    }
48
49    /// Accumulate a batch of predictions and targets.
50    pub fn accumulate(&mut self, preds: Tensor<2>, targets: Tensor<2, Bool>) {
51        self.predictions.push(preds);
52        self.targets.push(targets);
53    }
54
55    /// All accumulated predictions and targets, concatenated along the samples.
56    pub fn tensors(&self) -> (Tensor<2>, Tensor<2, Bool>) {
57        (
58            Tensor::cat(self.predictions.clone(), 0),
59            Tensor::cat(self.targets.clone(), 0),
60        )
61    }
62
63    /// Record the value computed over the accumulated set and return the entry
64    /// to log.
65    pub fn compute(&mut self, value: f64, format: FormatOptions) -> SerializedEntry {
66        self.current = Some(value);
67
68        let serialized = NumericEntry::Value(value).serialize();
69
70        let formatted_value = match format.precision {
71            Some(precision) => format_float(value, precision),
72            None => format!("{value}"),
73        };
74
75        // Rank-based metrics have no mathematically valid "batch" slice, so we render N/A
76        let formatted = match format.unit {
77            Some(unit) => format!("epoch {formatted_value} {unit} - batch N/A {unit}"),
78            None => format!("epoch {formatted_value} - batch N/A"),
79        };
80
81        SerializedEntry::new(formatted, serialized)
82    }
83
84    /// Get the current metric value, when available.
85    pub fn value(&self) -> Option<NumericEntry> {
86        self.current.map(NumericEntry::Value)
87    }
88
89    /// Reset the state, freeing the accumulated tensors.
90    pub fn reset(&mut self) {
91        self.predictions.clear();
92        self.targets.clear();
93        self.current = None;
94    }
95
96    pub(crate) fn is_empty(&self) -> bool {
97        self.predictions.is_empty()
98    }
99
100    pub(crate) fn serialize_placeholder(&self, format: FormatOptions) -> SerializedEntry {
101        SerializedEntry::not_available(format.unit.as_deref())
102    }
103}
104
105impl Default for PredictionAccumulatorState {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111impl FormatOptions {
112    /// Create the [formatting options](FormatOptions) with a name.
113    pub fn new(name: MetricName) -> Self {
114        Self {
115            name: name.clone(),
116            unit: None,
117            precision: None,
118        }
119    }
120
121    /// Specify the metric unit.
122    pub fn unit(mut self, unit: &str) -> Self {
123        self.unit = Some(unit.to_string());
124        self
125    }
126
127    /// Specify the floating point precision.
128    pub fn precision(mut self, precision: usize) -> Self {
129        self.precision = Some(precision);
130        self
131    }
132
133    /// Get the metric name.
134    pub fn name(&self) -> &Arc<String> {
135        &self.name
136    }
137
138    /// Get the metric unit.
139    pub fn unit_value(&self) -> &Option<String> {
140        &self.unit
141    }
142
143    /// Get the precision.
144    pub fn precision_value(&self) -> Option<usize> {
145        self.precision
146    }
147}
148
149impl NumericMetricState {
150    /// Create a new [numeric metric state](NumericMetricState).
151    pub fn new() -> Self {
152        Self {
153            sum: 0.0,
154            count: 0,
155            current: f64::NAN,
156            current_count: 0,
157        }
158    }
159
160    /// Reset the state.
161    pub fn reset(&mut self) {
162        self.sum = 0.0;
163        self.count = 0;
164        self.current = f64::NAN;
165        self.current_count = 0;
166    }
167
168    /// Update the state.
169    ///
170    /// `count` is the number of underlying units `value` is averaged over.
171    /// This is typically the batch size (e.g. for accuracy, loss), but it can be any
172    /// unit the metric's `value` is implicitly a rate over.
173    /// A zero-count update records the current value without changing the running aggregate.
174    pub fn update(&mut self, value: f64, count: usize) {
175        self.current = value;
176        self.current_count = count;
177
178        if count > 0 {
179            self.sum += value * count as f64;
180            self.count += count;
181        }
182    }
183
184    /// Compute the metric for the current update.
185    pub fn compute_update(&self, format: FormatOptions) -> SerializedEntry {
186        self.compute(format, false)
187    }
188
189    /// Compute the final metric for the accumulated global state.
190    pub fn compute_final(&self, format: FormatOptions) -> SerializedEntry {
191        self.compute(format, true)
192    }
193
194    fn compute(&self, format: FormatOptions, final_entry: bool) -> SerializedEntry {
195        let value_current = self.current;
196        let count = self.current_count;
197        let value_running = self.sum / self.count as f64;
198        // Numeric metric state is an aggregated value
199        let serialized = if final_entry {
200            NumericEntry::Final(value_running).serialize()
201        } else {
202            NumericEntry::Aggregated {
203                aggregated_value: value_current,
204                count,
205            }
206            .serialize()
207        };
208
209        let (formatted_current, formatted_running) = match format.precision {
210            Some(precision) => (
211                format_float(value_current, precision),
212                format_float(value_running, precision),
213            ),
214            None => (format!("{value_current}"), format!("{value_running}")),
215        };
216
217        // TODO: naming inconsistent with RL.
218        let formatted = match format.unit {
219            Some(unit) => {
220                format!("epoch {formatted_running} {unit} - batch {formatted_current} {unit}")
221            }
222            None => format!("epoch {formatted_running} - batch {formatted_current}"),
223        };
224
225        SerializedEntry::new(formatted, serialized)
226    }
227
228    /// Get the numeric value.
229    pub fn current_value(&self) -> NumericEntry {
230        NumericEntry::Aggregated {
231            aggregated_value: self.current,
232            count: self.current_count,
233        }
234    }
235
236    /// Get the running aggregated value.
237    pub fn running_value(&self) -> NumericEntry {
238        NumericEntry::Aggregated {
239            aggregated_value: self.sum / self.count as f64,
240            count: self.count,
241        }
242    }
243
244    /// Get the final aggregated value.
245    pub fn final_value(&self) -> NumericEntry {
246        NumericEntry::Final(self.sum / self.count as f64)
247    }
248}
249
250impl Default for NumericMetricState {
251    fn default() -> Self {
252        Self::new()
253    }
254}
255
256/// Accumulates confusion-matrix counts (TP, FP, FN) across an epoch.
257///
258/// Used by metrics derived from confusion-matrix ratios (precision, recall, etc.)
259/// which must accumulate raw counts and compute the ratio over the summed counts.
260/// Averaging the per-batch ratios (using [NumericMetricState]) is not equivalent;
261/// it is biased whenever per-batch class support varies.
262/// Accumulates confusion-matrix counts (TP, FP, FN) across an epoch using native tensors.
263#[derive(Clone)]
264pub struct ConfusionStatsState {
265    // Epoch-level accumulated totals
266    true_positive: Option<Tensor<1>>,
267    false_positive: Option<Tensor<1>>,
268    false_negative: Option<Tensor<1>>,
269    running_count: usize,
270
271    // Most recent batch snapshots (retained for the compute split)
272    current_tp: Option<Tensor<1>>,
273    current_fp: Option<Tensor<1>>,
274    current_fn: Option<Tensor<1>>,
275    current_count: usize,
276
277    // Already computed values
278    current_value: f64,
279    running_value: f64,
280}
281
282impl Default for ConfusionStatsState {
283    fn default() -> Self {
284        Self::new()
285    }
286}
287
288impl ConfusionStatsState {
289    /// Create a new [ConfusionStatsState].
290    pub fn new() -> Self {
291        Self {
292            true_positive: None,
293            false_positive: None,
294            false_negative: None,
295            running_count: 0,
296            current_tp: None,
297            current_fp: None,
298            current_fn: None,
299            current_count: 0,
300            current_value: f64::NAN,
301            running_value: f64::NAN,
302        }
303    }
304
305    /// Reset the state.
306    pub fn reset(&mut self) {
307        self.true_positive = None;
308        self.false_positive = None;
309        self.false_negative = None;
310        self.running_count = 0;
311        self.current_tp = None;
312        self.current_fp = None;
313        self.current_fn = None;
314        self.current_count = 0;
315        self.current_value = f64::NAN;
316        self.running_value = f64::NAN;
317    }
318
319    fn accumulate(running: &mut Option<Tensor<1>>, batch: Option<Tensor<1>>) {
320        if let Some(b) = batch {
321            *running = Some(match running.take() {
322                Some(acc) => acc + b,
323                None => b,
324            });
325        }
326    }
327
328    /// Accumulate running totals and capture the batch snapshot.
329    pub fn update(
330        &mut self,
331        tp: Option<Tensor<1>>,
332        fp: Option<Tensor<1>>,
333        fn_: Option<Tensor<1>>,
334        sample_size: usize,
335    ) {
336        // Accumulate running totals for global metric value
337        Self::accumulate(&mut self.true_positive, tp.clone());
338        Self::accumulate(&mut self.false_positive, fp.clone());
339        Self::accumulate(&mut self.false_negative, fn_.clone());
340        self.running_count += sample_size;
341
342        // Retain current batch tensors for the subsequent compute pass
343        self.current_tp = tp;
344        self.current_fp = fp;
345        self.current_fn = fn_;
346        self.current_count = sample_size;
347    }
348
349    /// Compute the batch-level value and the running epoch-level value (from all accumulated counts)
350    ///
351    /// # Errors
352    ///
353    /// Returns a [`TensorReadError`] if the computed metric tensors cannot be read back.
354    pub fn compute_update(
355        &mut self,
356        class_reduction: ClassReduction,
357        format: FormatOptions,
358        compute_fn: impl Fn(Option<Tensor<1>>, Option<Tensor<1>>, Option<Tensor<1>>) -> Tensor<1>,
359    ) -> Result<SerializedEntry, TensorReadError> {
360        let current_tp = self.current_tp.take();
361        let current_fp = self.current_fp.take();
362        let current_fn = self.current_fn.take();
363
364        // Compute batch-level value
365        let batch_metric = compute_fn(current_tp, current_fp, current_fn);
366        self.current_value = Self::class_average(batch_metric, class_reduction)?;
367
368        // Compute epoch-level value from running totals
369        let total_tp = self.true_positive.clone();
370        let total_fp = self.false_positive.clone();
371        let total_fn = self.false_negative.clone();
372
373        let epoch_metric = compute_fn(total_tp, total_fp, total_fn);
374        self.running_value = Self::class_average(epoch_metric, class_reduction)?;
375
376        // Serialize and format
377        let serialized = NumericEntry::Aggregated {
378            aggregated_value: self.current_value,
379            count: self.current_count,
380        }
381        .serialize();
382        Ok(self.serialized_entry(format, serialized))
383    }
384
385    /// Compute the final metric for the accumulated global state.
386    pub fn compute_final(&mut self, format: FormatOptions) -> SerializedEntry {
387        self.serialized_entry(format, NumericEntry::Final(self.running_value).serialize())
388    }
389
390    fn serialized_entry(&mut self, format: FormatOptions, serialized: String) -> SerializedEntry {
391        let (formatted_current, formatted_running) = match format.precision {
392            Some(p) => (
393                format_float(self.current_value, p),
394                format_float(self.running_value, p),
395            ),
396            None => (
397                format!("{}", self.current_value),
398                format!("{}", self.running_value),
399            ),
400        };
401
402        let formatted = match format.unit {
403            Some(unit) => {
404                format!("epoch {formatted_running} {unit} - batch {formatted_current} {unit}")
405            }
406            None => format!("epoch {formatted_running} - batch {formatted_current}"),
407        };
408
409        SerializedEntry::new(formatted, serialized)
410    }
411
412    fn class_average(
413        mut metric: Tensor<1>,
414        class_reduction: ClassReduction,
415    ) -> Result<f64, TensorReadError> {
416        use ClassReduction::{Macro, Micro};
417        let avg = match class_reduction {
418            Micro => metric,
419            Macro => {
420                if metric
421                    .clone()
422                    .contains_nan()
423                    .any()
424                    .try_into_scalar::<bool>()?
425                {
426                    let mask = metric.clone().is_nan();
427                    metric = metric
428                        .clone()
429                        .select(0, mask.bool_not().argwhere().squeeze_dim(1));
430                }
431                metric.mean()
432            }
433        };
434        avg.try_into_scalar::<f64>()
435    }
436
437    /// Get the current batch value.
438    pub fn current_value(&self) -> Option<NumericEntry> {
439        (!self.current_value.is_nan()).then_some(NumericEntry::Aggregated {
440            aggregated_value: self.current_value,
441            count: self.current_count,
442        })
443    }
444
445    /// Get the running aggregated value.
446    pub fn running_value(&self) -> Option<NumericEntry> {
447        (!self.running_value.is_nan()).then_some(NumericEntry::Aggregated {
448            aggregated_value: self.running_value,
449            count: self.running_count,
450        })
451    }
452
453    /// Get the final value of the metric.
454    pub fn final_value(&self) -> NumericEntry {
455        // The running value holds the epoch-level value from accumulated totals, which should hold
456        // the correct final value after all batches have been processed
457        NumericEntry::Final(self.running_value)
458    }
459}
burn_train/metric/state.rs

burn_train/metric/
state.rs
