burn-train 0.23.0-pre.1

Training crate for the Burn framework
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use burn_core::data::dataloader::Progress;

/// Trait for logging training progress at each step and end of epoch.
///
/// # Call sequence
///
/// Implementors can expect the following sequence of calls for a complete training run:
///
/// ```text
/// start(total_epochs, starting_epoch, total_items, label)
///   for each epoch:
///     start_split("train", progress)
///       update_split(progress)  // called once per batch
///       ...
///     end_split(progress)
///     start_split("valid", progress)
///       update_split(progress)  // called once per batch
///       ...
///     end_split(progress)
///     update_epoch(epoch)
/// end()
/// ```
///
/// `end()` is called whether training completes normally or is interrupted early.
///
/// Split updates carry the current count, expected total, and unit.
/// Final snapshots include any total resolved at exhaustion.
pub trait TrainingProgressLogger: Send {
    /// Called once at the start of training, providing the total number of epochs and the starting epoch.
    ///
    /// The total number of items of the training can optionally be provided if it is known.
    fn start(
        &mut self,
        total_epochs: usize,
        starting_epoch: usize,
        total_items: Option<usize>,
        label: Option<&str>,
    );

    /// Called at the end of each full epoch (after both train and valid splits complete).
    fn update_epoch(&mut self, epoch: usize);

    /// Called at the start of a training split, providing the split name and initial progress.
    fn start_split(&mut self, split: &str, progress: &Progress);

    /// Log a progress snapshot for the current split.
    fn update_split(&mut self, progress: &Progress);

    /// Called once at the end of a training split with final measurements.
    ///
    /// This also applies to empty, interrupted, and failed splits. A final snapshot does not
    /// produce an iteration event or an additional `update_split` call.
    fn end_split(&mut self, progress: &Progress);

    /// Called at the end of training, whether it completed successfully or was interrupted.
    fn end(&mut self);

    /// Log a custom named event that falls outside the standard training lifecycle callbacks.
    fn log_event_training(&mut self, event: String);
}

/// Trait for logging evaluation progress at each step and end of evaluation.
///
/// # Call sequence
///
/// Implementors can expect the following sequence of calls for a complete evaluation run:
///
/// ```text
/// start_global_progress(total_tests)
///   for each test split:
///     start_test(name, progress)
///       update_test_progress(progress)  // called once per batch
///       ...
///     end_test(progress)
/// end_global_progress()
/// ```
///
/// `end_global_progress()` is called whether evaluation completes normally or is interrupted early.
///
/// Test updates carry the current count, expected total, and unit.
/// Final snapshots include any total resolved at exhaustion.
pub trait EvaluationProgressLogger: Send {
    /// Called once at the start of evaluation, providing the total number of test splits.
    fn start_global_progress(&mut self, total_tests: usize);

    /// Called at the start of a test split, providing the split name and initial progress.
    fn start_test(&mut self, name: &str, progress: &Progress);

    /// Log a progress snapshot for the current test split.
    fn update_test_progress(&mut self, progress: &Progress);

    /// Called once at the end of a test split with final measurements.
    ///
    /// This also applies to empty, interrupted, and failed splits. A final snapshot does not
    /// produce an iteration event or an additional `update_test_progress` call.
    fn end_test(&mut self, progress: &Progress);

    /// Called at the end of evaluation.
    fn end_global_progress(&mut self);

    /// Log a custom named event that falls outside the standard evaluation lifecycle callbacks.
    fn log_event_evaluation(&mut self, event: String);
}

/// Two-level progress snapshot combining run-level and phase-level tracking.
///
/// `global_progress` spans the full training run (e.g., epochs completed out of total),
/// while `split_progress` tracks the current phase (e.g., batches within the current epoch).
#[derive(Debug, Clone)]
pub struct ProgressSnapshot {
    /// Progress across the entire training run.
    pub global: Progress,
    /// Progress within the current phase (epoch or evaluation split).
    pub split: Progress,
}

impl ProgressSnapshot {
    /// Create a new overall progress snapshot.
    pub fn new(global: Progress, split: Progress) -> Self {
        Self { global, split }
    }
}