Skip to main content

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}