burn_telemetry 0.1.0

burn training logger
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

//! Helper factories for constructing metric loggers based on enabled features.

use burn::train::{
    logger::MetricLogger,
    metric::{MetricEntry, NumericEntry},
};

#[cfg(feature = "rerun")]
use crate::{
    metric::{metric_identifier, numeric_value},
    rerun::RerunLogger,
};
#[cfg(feature = "rerun")]
use std::sync::Arc;

#[cfg(feature = "rerun")]
pub struct RerunMetricLogger {
    split: String,
    logger: Arc<RerunLogger>,
    step: usize,
    current_epoch: usize,
    last_epoch_step: Option<usize>,
}

#[cfg(feature = "rerun")]
impl RerunMetricLogger {
    pub fn new(split: impl Into<String>, logger: Arc<RerunLogger>) -> Self {
        Self {
            split: split.into(),
            logger,
            step: 0,
            current_epoch: 1,
            last_epoch_step: None,
        }
    }
}

#[cfg(feature = "rerun")]
impl MetricLogger for RerunMetricLogger {
    fn log(&mut self, item: &MetricEntry) {
        if self.last_epoch_step != Some(self.step) {
            let entity = metrics_epoch_entity_path(&self.split);
            self.logger
                .log_scalar(&entity, self.current_epoch as f64, self.step);
            self.last_epoch_step = Some(self.step);
        }

        let entity = metric_entity_path(&self.split, item.name.as_ref());
        if let Some(value) = numeric_value(item) {
            self.logger.log_scalar(&entity, value, self.step);
        } else {
            self.logger
                .log_text(&format!("{entity}/text"), &item.formatted, self.step);
        }
        self.step += 1;
    }

    fn end_epoch(&mut self, epoch: usize) {
        self.current_epoch = epoch + 1;
    }

    fn read_numeric(&mut self, _name: &str, _epoch: usize) -> Result<Vec<NumericEntry>, String> {
        Err("Rerun metric logger does not support reading metrics".into())
    }
}

#[cfg(feature = "rerun")]
pub fn rerun_metric_logger_pair(
    logger: Arc<RerunLogger>,
) -> (RerunMetricLogger, RerunMetricLogger) {
    (
        RerunMetricLogger::new("train", logger.clone()),
        RerunMetricLogger::new("valid", logger),
    )
}

#[cfg(feature = "rerun")]
fn metric_entity_path(split: &str, name: &str) -> String {
    format!("metrics/{split}/{}", metric_identifier(name))
}

#[cfg(feature = "rerun")]
fn metrics_epoch_entity_path(split: &str) -> String {
    format!("metrics/{split}/epoch")
}