burn_train/renderer/base.rs
1use std::sync::Arc;
2
3use crate::{
4 LearnerSummary,
5 logger::{EvaluationProgressLogger, TrainingProgressLogger},
6 metric::{MetricDefinition, MetricEntry, NumericEntry},
7};
8
9/// Trait for rendering metrics.
10pub trait MetricsRendererTraining: Send + Sync {
11 /// Updates the training metric state.
12 ///
13 /// # Arguments
14 ///
15 /// * `state` - The metric state.
16 fn update_train(&mut self, state: MetricState);
17
18 /// Updates the validation metric state.
19 ///
20 /// # Arguments
21 ///
22 /// * `state` - The metric state.
23 fn update_valid(&mut self, state: MetricState);
24
25 /// Callback method invoked when training ends, whether it
26 /// completed successfully or was interrupted.
27 ///
28 /// # Returns
29 ///
30 /// A result indicating whether the end-of-training actions were successful.
31 fn on_train_end(
32 &mut self,
33 summary: Option<LearnerSummary>,
34 ) -> Result<(), Box<dyn core::error::Error>> {
35 default_summary_action(summary);
36 Ok(())
37 }
38}
39
40/// A renderer that can be used for both training and evaluation.
41pub trait MetricsRenderer:
42 MetricsRendererEvaluation
43 + MetricsRendererTraining
44 + TrainingProgressLogger
45 + EvaluationProgressLogger
46{
47 /// Keep the renderer from automatically closing, requiring manual action to close it.
48 fn manual_close(&mut self);
49 /// Register a new metric.
50 fn register_metric(&mut self, definition: MetricDefinition);
51}
52
53#[derive(Clone)]
54/// The name of an evaluation.
55///
56/// This is going to group metrics together for easier analysis.
57pub struct EvaluationName {
58 pub(crate) name: Arc<String>,
59}
60
61impl EvaluationName {
62 /// Creates a new evaluation name.
63 pub fn new<S: core::fmt::Display>(s: S) -> Self {
64 Self {
65 name: Arc::new(format!("{s}")),
66 }
67 }
68
69 /// Returns the evaluation name.
70 pub fn as_str(&self) -> &str {
71 &self.name
72 }
73}
74
75impl core::fmt::Display for EvaluationName {
76 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
77 f.write_str(&self.name)
78 }
79}
80
81/// Trait for rendering metrics.
82pub trait MetricsRendererEvaluation: Send + Sync {
83 /// Updates the testing metric state.
84 ///
85 /// # Arguments
86 ///
87 /// * `state` - The metric state.
88 fn update_test(&mut self, name: EvaluationName, state: MetricState);
89
90 /// Callback method invoked when testing ends, whether it
91 /// completed successfully or was interrupted.
92 ///
93 /// # Returns
94 ///
95 /// A result indicating whether the end-of-testing actions were successful.
96 fn on_test_end(
97 &mut self,
98 summary: Option<LearnerSummary>,
99 ) -> Result<(), Box<dyn core::error::Error>> {
100 default_summary_action(summary);
101 Ok(())
102 }
103}
104
105/// The state of a metric.
106#[derive(Debug)]
107pub enum MetricState {
108 /// A generic metric.
109 Generic(MetricEntry),
110 /// A numeric metric.
111 ///
112 /// - `MetricEntry`: ID, formatted text string ("epoch N/A - batch N/A").
113 /// - `Option<NumericEntry>`:
114 /// - `Some(entry)`: A plottable point.
115 /// - `None`: A valid step tick on the X-axis, but with no drawable Y-coordinate point.
116 Numeric(MetricEntry, Option<NumericEntry>),
117}
118
119fn default_summary_action(summary: Option<LearnerSummary>) {
120 if let Some(summary) = summary {
121 println!("{summary}");
122 }
123}