burn-train 0.22.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
123
124
125
126
127
128
129
130
131
132
133
134

use burn_core::data::dataloader::Progress;
use burn_optim::lr_scheduler::module_lr_scheduler::ModuleLearningRate;

use crate::{
    LearnerSummary,
    renderer::{EvaluationName, MetricsRenderer},
};

/// Event happening during the training/validation process.
pub enum LearnerEvent<T> {
    /// Signal the start of the process (e.g., training start).
    Start {
        /// The total number of training epochs.
        total_epochs: usize,
        /// The starting epoch.
        starting_epoch: usize,
    },
    /// Signal that an item have been processed.
    ProcessedItem(TrainingItem<T>),
    /// Signal the start of a split.
    StartSplit {
        /// The epoch number.
        epoch_number: usize,
        /// The total number of items to be processed during this split.
        total_items: usize,
    },
    /// Signal the end of a split, carrying the current epoch number.
    EndSplit(usize),
    /// Signal the end of a full epoch.
    EndEpoch(usize),
    /// Signal the end of the process (e.g., training end).
    End(Option<LearnerSummary>),
}

/// Event happening during the evaluation process.
pub enum EvaluatorEvent<T> {
    /// Signal the start of the process (e.g., evaluation start)
    Start {
        /// The total number of items to evaluate.
        total_tests: usize,
    },
    /// Signal the start of a test split, carrying the split name and total number of items.
    StartTest(EvaluationName, usize),
    /// Signal that an item have been processed.
    ProcessedItem(EvaluationName, EvaluationItem<T>),
    /// Signal the end of a single test split.
    EndTest,
    /// Signal the end of the process (e.g., evaluation end).
    End(Option<LearnerSummary>),
}

/// Items that are lazy are not ready to be processed by metrics.
///
/// We want to sync them on a different thread to avoid blocking training.
pub trait ItemLazy: Send {
    /// Sync the item.
    fn sync(self) -> Self;
}

/// Process events happening during training and validation.
pub trait EventProcessorTraining<TrainEvent, ValidEvent>: Send {
    /// Collect a training event.
    fn process_train(&mut self, event: TrainEvent);
    /// Collect a validation event.
    fn process_valid(&mut self, event: ValidEvent);
    /// Returns the renderer used for training.
    fn renderer(self) -> Box<dyn MetricsRenderer>;
}

/// Process events happening during evaluation.
pub trait EventProcessorEvaluation: Send {
    /// The test item.
    type ItemTest: ItemLazy;

    /// Collect a test event.
    fn process_test(&mut self, event: EvaluatorEvent<Self::ItemTest>);

    /// Returns the renderer used for evaluation.
    fn renderer(self) -> Box<dyn MetricsRenderer>;
}

/// A learner item.
#[derive(new)]
pub struct TrainingItem<T> {
    /// The item.
    pub item: T,

    /// The progress.
    pub progress: Progress,

    /// The iteration, if it it different from the items processed.
    pub iteration: Option<usize>,

    /// The learning rate for a module's parameters.
    pub lr: Option<ModuleLearningRate>,
}

impl<T: ItemLazy> ItemLazy for TrainingItem<T> {
    fn sync(self) -> Self {
        TrainingItem {
            item: self.item.sync(),
            progress: self.progress,
            iteration: self.iteration,
            lr: self.lr,
        }
    }
}

/// An evaluation item.
#[derive(new)]
pub struct EvaluationItem<T> {
    /// The item.
    pub item: T,

    /// The progress.
    pub progress: Progress,

    /// The iteration, if it it different from the items processed.
    pub iteration: Option<usize>,
}

impl<T: ItemLazy> ItemLazy for EvaluationItem<T> {
    fn sync(self) -> Self {
        EvaluationItem {
            item: self.item.sync(),
            progress: self.progress,
            iteration: self.iteration,
        }
    }
}

impl ItemLazy for () {
    fn sync(self) -> Self {}
}