burn_train/metric/processor/
metrics.rs

1use std::collections::HashMap;
2
3use burn_core::tensor::TensorReadError;
4
5use super::{ItemLazy, MetricError, TrainingItem};
6use crate::{
7    EvaluationItem,
8    metric::{
9        Adaptor, Metric, MetricDefinition, MetricEntry, MetricId, MetricMetadata, MetricName,
10        Numeric,
11        store::{MetricsUpdate, NumericMetricUpdate, Split},
12    },
13};
14
15pub(crate) struct MetricsTraining<T: ItemLazy, V: ItemLazy> {
16    train: Vec<Box<dyn MetricUpdater<T>>>,
17    valid: Vec<Box<dyn MetricUpdater<V>>>,
18    train_numeric: Vec<Box<dyn NumericMetricUpdater<T>>>,
19    valid_numeric: Vec<Box<dyn NumericMetricUpdater<V>>>,
20    // Vec preserves metrics registration order; reflected in TUI tabs order
21    metric_definitions: Vec<MetricDefinition>,
22}
23
24pub(crate) struct MetricsEvaluation<T: ItemLazy> {
25    test: Vec<Box<dyn MetricUpdater<T>>>,
26    test_numeric: Vec<Box<dyn NumericMetricUpdater<T>>>,
27    metric_definitions: HashMap<MetricId, MetricDefinition>,
28}
29
30impl<T: ItemLazy> Default for MetricsEvaluation<T> {
31    fn default() -> Self {
32        Self {
33            test: Default::default(),
34            test_numeric: Default::default(),
35            metric_definitions: HashMap::default(),
36        }
37    }
38}
39
40impl<T: ItemLazy, V: ItemLazy> Default for MetricsTraining<T, V> {
41    fn default() -> Self {
42        Self {
43            train: Vec::default(),
44            valid: Vec::default(),
45            train_numeric: Vec::default(),
46            valid_numeric: Vec::default(),
47            metric_definitions: Vec::default(),
48        }
49    }
50}
51
52impl<T: ItemLazy> MetricsEvaluation<T> {
53    /// Register a testing metric.
54    pub(crate) fn register_test_metric<Me: Metric + 'static>(&mut self, metric: Me)
55    where
56        T: Adaptor<Me::Input> + 'static,
57    {
58        let metric = MetricWrapper::new(metric);
59        self.register_definition(&metric);
60        self.test.push(Box::new(metric))
61    }
62
63    /// Register a numeric testing metric.
64    pub(crate) fn register_test_metric_numeric<Me: Metric + Numeric + 'static>(
65        &mut self,
66        metric: Me,
67    ) where
68        T: Adaptor<Me::Input> + 'static,
69    {
70        let metric = MetricWrapper::new(metric);
71        self.register_definition(&metric);
72        self.test_numeric.push(Box::new(metric))
73    }
74
75    fn register_definition<Me: Metric>(&mut self, metric: &MetricWrapper<Me>) {
76        self.metric_definitions.insert(
77            metric.id.clone(),
78            MetricDefinition::new(metric.id.clone(), &metric.metric),
79        );
80    }
81
82    /// Get metric definitions.
83    pub(crate) fn metric_definitions(&mut self) -> Vec<MetricDefinition> {
84        self.metric_definitions.values().cloned().collect()
85    }
86
87    /// Update the testing information from the testing item.
88    ///
89    /// Every metric is updated, even after one fails: returns the entries of those that
90    /// succeeded and the failures of the others.
91    pub(crate) fn update_test(
92        &mut self,
93        item: &EvaluationItem<T>,
94        metadata: &MetricMetadata,
95        split: Split,
96    ) -> (MetricsUpdate, Vec<MetricError>) {
97        update_metrics(
98            &mut self.test,
99            &mut self.test_numeric,
100            &item.item,
101            metadata,
102            split,
103        )
104    }
105}
106
107impl<T: ItemLazy, V: ItemLazy> MetricsTraining<T, V> {
108    /// Register a training metric.
109    pub(crate) fn register_train_metric<Me: Metric + 'static>(&mut self, metric: Me)
110    where
111        T: Adaptor<Me::Input> + 'static,
112    {
113        let metric = MetricWrapper::new(metric);
114        self.register_definition(&metric);
115        self.train.push(Box::new(metric))
116    }
117
118    /// Register a validation metric.
119    pub(crate) fn register_valid_metric<Me: Metric + 'static>(&mut self, metric: Me)
120    where
121        V: Adaptor<Me::Input> + 'static,
122    {
123        let metric = MetricWrapper::new(metric);
124        self.register_definition(&metric);
125        self.valid.push(Box::new(metric))
126    }
127
128    /// Register a numeric training metric.
129    pub(crate) fn register_train_metric_numeric<Me: Metric + Numeric + 'static>(
130        &mut self,
131        metric: Me,
132    ) where
133        T: Adaptor<Me::Input> + 'static,
134    {
135        let metric = MetricWrapper::new(metric);
136        self.register_definition(&metric);
137        self.train_numeric.push(Box::new(metric))
138    }
139
140    /// Register a numeric validation metric.
141    pub(crate) fn register_valid_metric_numeric<Me>(&mut self, metric: Me)
142    where
143        V: Adaptor<Me::Input> + 'static,
144        Me: Metric + Numeric + 'static,
145    {
146        let metric = MetricWrapper::new(metric);
147        self.register_definition(&metric);
148        self.valid_numeric.push(Box::new(metric))
149    }
150
151    fn register_definition<Me: Metric>(&mut self, metric: &MetricWrapper<Me>) {
152        // Avoid duplicate definitions if the same metric is registered for both train and valid
153        if !self
154            .metric_definitions
155            .iter()
156            .any(|def| def.metric_id == metric.id)
157        {
158            self.metric_definitions
159                .push(MetricDefinition::new(metric.id.clone(), &metric.metric));
160        }
161    }
162
163    /// Get metric definitions for all splits
164    pub(crate) fn metric_definitions(&mut self) -> Vec<MetricDefinition> {
165        self.metric_definitions.clone()
166    }
167
168    /// Update the training information from the training item.
169    pub(crate) fn update_train(
170        &mut self,
171        item: &TrainingItem<T>,
172        metadata: &MetricMetadata,
173    ) -> (MetricsUpdate, Vec<MetricError>) {
174        update_metrics(
175            &mut self.train,
176            &mut self.train_numeric,
177            &item.item,
178            metadata,
179            Split::Train,
180        )
181    }
182
183    /// Update the training information from the validation item.
184    pub(crate) fn update_valid(
185        &mut self,
186        item: &TrainingItem<V>,
187        metadata: &MetricMetadata,
188    ) -> (MetricsUpdate, Vec<MetricError>) {
189        update_metrics(
190            &mut self.valid,
191            &mut self.valid_numeric,
192            &item.item,
193            metadata,
194            Split::Valid,
195        )
196    }
197
198    /// Signal the end of a training epoch.
199    /// Returns the final metric entries for the epoch.
200    ///
201    /// Every metric is computed and cleared, even after one fails, so the next epoch starts
202    /// clean: returns the entries of those that succeeded and the failures of the others.
203    pub(crate) fn end_epoch_train(&mut self) -> (MetricsUpdate, Vec<MetricError>) {
204        end_epoch_metrics(&mut self.train, &mut self.train_numeric, Split::Train)
205    }
206
207    /// Signal the end of a validation epoch.
208    /// Returns the final metric entries for the epoch.
209    ///
210    /// Every metric is computed and cleared, even after one fails, so the next epoch starts
211    /// clean: returns the entries of those that succeeded and the failures of the others.
212    pub(crate) fn end_epoch_valid(&mut self) -> (MetricsUpdate, Vec<MetricError>) {
213        end_epoch_metrics(&mut self.valid, &mut self.valid_numeric, Split::Valid)
214    }
215}
216
217pub(crate) fn update_metrics<I>(
218    metrics: &mut [Box<dyn MetricUpdater<I>>],
219    metrics_numeric: &mut [Box<dyn NumericMetricUpdater<I>>],
220    item: &I,
221    metadata: &MetricMetadata,
222    split: Split,
223) -> (MetricsUpdate, Vec<MetricError>) {
224    let mut entries = Vec::with_capacity(metrics.len());
225    let mut entries_numeric = Vec::with_capacity(metrics_numeric.len());
226    let mut failures = Vec::new();
227    let mut fail = |metric: MetricName, source| {
228        failures.push(MetricError {
229            metric,
230            split: split.clone(),
231            source,
232        })
233    };
234
235    for metric in metrics.iter_mut() {
236        match metric.update(item, metadata) {
237            Ok(entry) => entries.push(entry),
238            Err(source) => fail(metric.name(), source),
239        }
240    }
241    for metric in metrics_numeric.iter_mut() {
242        match metric.update(item, metadata) {
243            Ok(entry) => entries_numeric.push(entry),
244            Err(source) => fail(metric.name(), source),
245        }
246    }
247
248    (MetricsUpdate::new(entries, entries_numeric), failures)
249}
250
251pub(crate) fn end_epoch_metrics<I>(
252    metrics: &mut [Box<dyn MetricUpdater<I>>],
253    metrics_numeric: &mut [Box<dyn NumericMetricUpdater<I>>],
254    split: Split,
255) -> (MetricsUpdate, Vec<MetricError>) {
256    let mut entries = Vec::with_capacity(metrics.len());
257    let mut entries_numeric = Vec::with_capacity(metrics_numeric.len());
258    let mut failures = Vec::new();
259    let mut fail = |metric: MetricName, source| {
260        failures.push(MetricError {
261            metric,
262            split: split.clone(),
263            source,
264        })
265    };
266
267    for metric in metrics.iter_mut() {
268        match metric.compute() {
269            Ok(entry) => entries.push(entry),
270            Err(source) => fail(metric.name(), source),
271        }
272        metric.clear();
273    }
274    for metric in metrics_numeric.iter_mut() {
275        match metric.compute() {
276            Ok(entry) => entries_numeric.push(entry),
277            Err(source) => fail(metric.name(), source),
278        }
279        metric.clear();
280    }
281
282    (MetricsUpdate::new(entries, entries_numeric), failures)
283}
284
285impl<T> From<&TrainingItem<T>> for MetricMetadata {
286    fn from(item: &TrainingItem<T>) -> Self {
287        Self {
288            progress: item.progress.clone(),
289            iteration: item.iteration,
290            lr: item.lr.clone(),
291        }
292    }
293}
294
295impl<T> From<&EvaluationItem<T>> for MetricMetadata {
296    fn from(item: &EvaluationItem<T>) -> Self {
297        Self {
298            progress: item.progress.clone(),
299            iteration: item.iteration,
300            lr: None,
301        }
302    }
303}
304
305pub(crate) trait NumericMetricUpdater<T>: Send + Sync {
306    fn update(
307        &mut self,
308        item: &T,
309        metadata: &MetricMetadata,
310    ) -> Result<NumericMetricUpdate, TensorReadError>;
311    fn compute(&mut self) -> Result<NumericMetricUpdate, TensorReadError>;
312    fn clear(&mut self);
313    fn name(&self) -> MetricName;
314}
315
316pub(crate) trait MetricUpdater<T>: Send + Sync {
317    fn update(
318        &mut self,
319        item: &T,
320        metadata: &MetricMetadata,
321    ) -> Result<MetricEntry, TensorReadError>;
322    fn compute(&mut self) -> Result<MetricEntry, TensorReadError>;
323    fn clear(&mut self);
324    fn name(&self) -> MetricName;
325}
326
327pub(crate) struct MetricWrapper<M> {
328    pub id: MetricId,
329    pub metric: M,
330}
331
332impl<M: Metric> MetricWrapper<M> {
333    pub fn new(metric: M) -> Self {
334        Self {
335            id: MetricId::new(metric.name()),
336            metric,
337        }
338    }
339}
340
341impl<T, M> NumericMetricUpdater<T> for MetricWrapper<M>
342where
343    T: 'static,
344    M: Metric + Numeric + 'static,
345    T: Adaptor<M::Input>,
346{
347    fn update(
348        &mut self,
349        item: &T,
350        metadata: &MetricMetadata,
351    ) -> Result<NumericMetricUpdate, TensorReadError> {
352        let serialized_entry = self.metric.update(&item.adapt(), metadata)?;
353        let update = MetricEntry::new(self.id.clone(), serialized_entry);
354        let numeric = self.metric.value();
355        let running = self.metric.running_value();
356
357        Ok(NumericMetricUpdate {
358            entry: update,
359            numeric_entry: numeric,
360            running_entry: running,
361        })
362    }
363
364    fn compute(&mut self) -> Result<NumericMetricUpdate, TensorReadError> {
365        let serialized_entry = self.metric.compute()?;
366        let update = MetricEntry::new(self.id.clone(), serialized_entry);
367        let final_entry = self.metric.final_value();
368
369        Ok(NumericMetricUpdate {
370            entry: update,
371            // Running entry is not applicable. This is the final epoch-level value computed.
372            numeric_entry: Some(final_entry),
373            running_entry: None,
374        })
375    }
376
377    fn clear(&mut self) {
378        self.metric.clear()
379    }
380
381    fn name(&self) -> MetricName {
382        self.metric.name()
383    }
384}
385
386impl<T, M> MetricUpdater<T> for MetricWrapper<M>
387where
388    T: 'static,
389    M: Metric + 'static,
390    T: Adaptor<M::Input>,
391{
392    fn update(
393        &mut self,
394        item: &T,
395        metadata: &MetricMetadata,
396    ) -> Result<MetricEntry, TensorReadError> {
397        let serialized_entry = self.metric.update(&item.adapt(), metadata)?;
398        Ok(MetricEntry::new(self.id.clone(), serialized_entry))
399    }
400
401    fn compute(&mut self) -> Result<MetricEntry, TensorReadError> {
402        let serialized_entry = self.metric.compute()?;
403        Ok(MetricEntry::new(self.id.clone(), serialized_entry))
404    }
405
406    fn clear(&mut self) {
407        self.metric.clear()
408    }
409
410    fn name(&self) -> MetricName {
411        self.metric.name()
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418    use crate::metric::{SerializedEntry, processor::EventProcessorError};
419    use burn_std::ExecutionError;
420    use std::sync::Arc;
421
422    /// A metric that fails its update, its compute, or neither, and counts its clears.
423    #[derive(Clone)]
424    struct TestMetric {
425        name: &'static str,
426        fail_update: bool,
427        fail_compute: bool,
428        cleared: Arc<std::sync::atomic::AtomicUsize>,
429    }
430
431    impl TestMetric {
432        fn new(name: &'static str, fail_update: bool, fail_compute: bool) -> Self {
433            Self {
434                name,
435                fail_update,
436                fail_compute,
437                cleared: Default::default(),
438            }
439        }
440
441        fn failure(&self) -> TensorReadError {
442            ExecutionError::with_context(format!("{} could not read", self.name)).into()
443        }
444    }
445
446    impl Metric for TestMetric {
447        type Input = ();
448
449        fn name(&self) -> MetricName {
450            Arc::new(self.name.to_string())
451        }
452
453        fn update(
454            &mut self,
455            _item: &(),
456            _metadata: &MetricMetadata,
457        ) -> Result<SerializedEntry, TensorReadError> {
458            match self.fail_update {
459                true => Err(self.failure()),
460                false => Ok(SerializedEntry::new(self.name.into(), self.name.into())),
461            }
462        }
463
464        fn compute(&mut self) -> Result<SerializedEntry, TensorReadError> {
465            match self.fail_compute {
466                true => Err(self.failure()),
467                false => Ok(SerializedEntry::new(self.name.into(), self.name.into())),
468            }
469        }
470
471        fn clear(&mut self) {
472            self.cleared
473                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
474        }
475    }
476
477    fn metrics(list: &[TestMetric]) -> Vec<Box<dyn MetricUpdater<()>>> {
478        list.iter()
479            .map(|metric| Box::new(MetricWrapper::new(metric.clone())) as Box<_>)
480            .collect()
481    }
482
483    fn names(failures: &[MetricError]) -> Vec<String> {
484        failures.iter().map(|err| err.to_string()).collect()
485    }
486
487    fn metadata() -> MetricMetadata {
488        MetricMetadata {
489            progress: burn_core::data::dataloader::Progress {
490                items_processed: 1,
491                items_total: 1,
492                unit: None,
493            },
494            iteration: None,
495            lr: None,
496        }
497    }
498
499    #[test]
500    fn update_reports_failing_metric_and_keeps_others() {
501        let list = [
502            TestMetric::new("first", true, false),
503            TestMetric::new("working", false, false),
504            TestMetric::new("last", true, false),
505        ];
506
507        let (update, failures) =
508            update_metrics(&mut metrics(&list), &mut [], &(), &metadata(), Split::Train);
509
510        assert_eq!(
511            update.entries.len(),
512            1,
513            "the working metric's entry is kept"
514        );
515        assert_eq!(
516            names(&failures),
517            [
518                "train/first: An error happened during execution\nCaused by:\n  first could not read",
519                "train/last: An error happened during execution\nCaused by:\n  last could not read",
520            ]
521        );
522    }
523
524    #[test]
525    fn epoch_end_clears_metrics_and_reports_failures() {
526        let list = [
527            TestMetric::new("first", false, true),
528            TestMetric::new("working", false, false),
529            TestMetric::new("last", false, true),
530        ];
531
532        let (update, failures) = end_epoch_metrics(&mut metrics(&list), &mut [], Split::Valid);
533
534        assert_eq!(
535            update.entries.len(),
536            1,
537            "the working metric's entry is kept"
538        );
539        let failed: Vec<_> = failures.iter().map(|err| err.metric.as_str()).collect();
540        assert_eq!(failed, ["first", "last"]);
541        assert!(failures.iter().all(|err| err.split == Split::Valid));
542        for metric in &list {
543            assert_eq!(
544                metric.cleared.load(std::sync::atomic::Ordering::SeqCst),
545                1,
546                "{} must be cleared so the next epoch starts clean",
547                metric.name
548            );
549        }
550    }
551
552    #[test]
553    fn metrics_error_reports_failures_and_poisoned_device() {
554        let failure = |metric: &str, source: ExecutionError| MetricError {
555            metric: Arc::new(metric.to_string()),
556            split: Split::Test(Some(Arc::new("holdout".to_string()))),
557            source: source.into(),
558        };
559        let error = EventProcessorError::from_errors(vec![
560            failure("Accuracy", ExecutionError::with_context("refused")),
561            failure("Loss", ExecutionError::device_poisoned("status 700")),
562        ])
563        .unwrap_err();
564
565        assert!(error.is_device_poisoned());
566        let message = error.to_string();
567        assert!(
568            message.starts_with("Event processing failed 2 times:"),
569            "{message}"
570        );
571        assert!(message.contains("test/holdout/Accuracy"), "{message}");
572        assert!(message.contains("test/holdout/Loss"), "{message}");
573
574        assert!(EventProcessorError::from_errors(Vec::new()).is_ok());
575    }
576
577    #[test]
578    fn a_sync_failure_is_one_failure_naming_its_split() {
579        let error = EventProcessorError::sync(
580            Split::Test(Some(Arc::new("holdout".to_string()))),
581            ExecutionError::device_poisoned("status 700"),
582        );
583
584        assert!(error.is_device_poisoned());
585        assert_eq!(
586            error.failures().len(),
587            1,
588            "one failure, whatever the metrics"
589        );
590        let message = error.to_string();
591        assert!(
592            message.contains("test/holdout: the event could not be synced"),
593            "{message}"
594        );
595    }
596}
burn_train/metric/processor/metrics.rs

burn_train/metric/processor/
metrics.rs
