burn_train/evaluator/
base.rs

1use crate::{
2    AsyncProcessorEvaluation, EvaluationItem, FullEventProcessorEvaluation, InferenceStep,
3    Interrupter, LearnerSummaryConfig,
4    evaluator::components::EvaluatorComponentTypes,
5    metric::processor::{EvaluatorEvent, EventProcessorEvaluation},
6    renderer::{EvaluationName, MetricsRenderer},
7};
8use burn_core::{data::dataloader::DataLoader, module::Module};
9use std::sync::Arc;
10
11pub(crate) type TestInput<EC> = <<EC as EvaluatorComponentTypes>::Model as InferenceStep>::Input;
12pub(crate) type TestOutput<EC> = <<EC as EvaluatorComponentTypes>::Model as InferenceStep>::Output;
13
14pub(crate) type TestLoader<EC> = Arc<dyn DataLoader<TestInput<EC>>>;
15
16/// The result of an evaluation.
17pub struct EvaluationResult {
18    /// The renderer that can be used for follow up training and evaluation.
19    pub renderer: Box<dyn MetricsRenderer>,
20    /// The stop that ended evaluation early, if
21    /// [`Interrupter::stop`](crate::Interrupter::stop) was called and no error happened.
22    pub interrupted: Option<crate::Interruption>,
23    /// The error that ended evaluation early, if it hit one.
24    ///
25    /// Both are `None` when every dataset was evaluated.
26    pub error: Option<Arc<crate::TrainingError>>,
27}
28
29/// Evaluates a model on a specific dataset.
30pub struct Evaluator<EC: EvaluatorComponentTypes> {
31    pub(crate) model: EC::Model,
32    pub(crate) interrupter: Interrupter,
33    pub(crate) event_processor:
34        AsyncProcessorEvaluation<FullEventProcessorEvaluation<TestOutput<EC>>>,
35    /// Config for creating a summary of the evaluation
36    pub summary: Option<LearnerSummaryConfig>,
37}
38
39impl<EC: EvaluatorComponentTypes> Evaluator<EC> {
40    /// Run the evaluation on the given dataset.
41    ///
42    /// The data will be stored and displayed under the provided name.
43    pub fn eval<S: core::fmt::Display>(
44        self,
45        name: S,
46        dataloader: TestLoader<EC>,
47    ) -> EvaluationResult {
48        self.eval_all([(name, dataloader)])
49    }
50
51    /// Run the evaluation on multiple named datasets sequentially.
52    ///
53    /// Prefer this over calling [`eval`](Self::eval) in a loop — the progress logger
54    /// receives the correct `total_tests` count and `end_test` is called between splits.
55    pub fn eval_all<S: core::fmt::Display>(
56        mut self,
57        splits: impl IntoIterator<Item = (S, TestLoader<EC>)>,
58    ) -> EvaluationResult {
59        let splits: Vec<_> = splits.into_iter().collect();
60        let total_tests = splits.len();
61
62        self.interrupter.fail_on_error(
63            self.event_processor
64                .process_test(EvaluatorEvent::Start { total_tests }),
65        );
66
67        for (name, dataloader) in splits {
68            let dataloader = dataloader.to_device(self.model.devices().first().unwrap());
69            let name = EvaluationName::new(name);
70            let total_items = dataloader.num_items();
71            let mut iterator = dataloader.iter();
72            let mut iteration = 0;
73
74            self.interrupter.fail_on_error(
75                self.event_processor
76                    .process_test(EvaluatorEvent::StartTest(name.clone(), total_items)),
77            );
78
79            while let Some(item) = iterator.next() {
80                let item = match item {
81                    Ok(item) => item,
82                    Err(err) => {
83                        self.interrupter.fail(err);
84                        break;
85                    }
86                };
87                let progress = iterator.progress();
88                iteration += 1;
89
90                let item = self.model.step(item);
91                let item = EvaluationItem::new(item, progress, Some(iteration));
92
93                self.interrupter.fail_on_error(
94                    self.event_processor
95                        .process_test(EvaluatorEvent::ProcessedItem(name.clone(), item)),
96                );
97
98                if self.interrupter.should_stop() {
99                    log::info!("Testing interrupted.");
100                    break;
101                }
102            }
103
104            self.interrupter
105                .fail_on_error(self.event_processor.process_test(EvaluatorEvent::EndTest));
106
107            // The remaining datasets are not evaluated once interrupted.
108            if self.interrupter.should_stop() {
109                break;
110            }
111        }
112
113        let summary = self.summary.and_then(|summary| {
114            summary
115                .init()
116                .map(|summary| summary.with_model(self.model.to_string()))
117                .ok()
118        });
119
120        self.interrupter.fail_on_error(
121            self.event_processor
122                .process_test(EvaluatorEvent::End(summary)),
123        );
124
125        // Finish processing remaining events.
126        self.interrupter.fail_on_error(self.event_processor.flush());
127
128        EvaluationResult {
129            renderer: self.event_processor.renderer(),
130            interrupted: self.interrupter.interruption(),
131            error: self.interrupter.error(),
132        }
133    }
134}
burn_train/evaluator/base.rs

burn_train/evaluator/
base.rs
