burn_train/
lib.rs

1#![warn(missing_docs)]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3
4//! Training and evaluation for Burn models.
5//!
6//! Applications enable the `train` feature of `burn` and use this crate as `burn::train`.
7//!
8//! - [`Learner`] bundles a model, optimizer and learning rate scheduler, and [`SupervisedTraining`]
9//!   runs it over training and validation data loaders on one or several devices.
10//! - [`metric`]: loss, accuracy, precision and recall, F-scores, AUROC, BLEU, CER, WER, perplexity,
11//!   system usage and more; with the `vision` feature, image metrics such as PSNR, SSIM, LPIPS and
12//!   FID.
13//! - [`renderer`]: a terminal dashboard (`tui` feature) or plain CLI output for training progress.
14//! - [`checkpoint`]: periodic and metric-based checkpointing, and resuming from a checkpoint.
15//! - Early stopping, interruption and an [`Evaluator`] for test sets.
16//!
17//! Failures on a device during training are returned as errors rather than panics.
18//!
19//! # Feature flags
20//!
21//! - `tui` (default): terminal dashboard.
22//! - `sys-metrics` (default): CPU, memory and GPU usage metrics.
23//! - `vision`: image quality metrics.
24//! - `rl`: reinforcement learning training through `burn-rl`.
25//! - `tracing`: instrument operations with the `tracing` crate.
26
27#[macro_use]
28extern crate derive_new;
29
30/// The checkpoint module.
31pub mod checkpoint;
32
33pub(crate) mod components;
34
35/// Renderer modules to display metrics and training information.
36pub mod renderer;
37
38/// The logger module.
39pub mod logger;
40
41/// The metric module.
42pub mod metric;
43
44pub use metric::processor::*;
45
46mod learner;
47
48pub use learner::*;
49
50mod evaluator;
51
52pub use evaluator::*;
53
54pub use components::*;
55
56#[cfg(test)]
57pub(crate) mod tests {
58    use burn_core::{prelude::Tensor, tensor::Bool};
59    use std::default::Default;
60
61    /// Probability of tp before adding errors
62    pub const THRESHOLD: f64 = 0.5;
63
64    #[derive(Debug, Default)]
65    pub enum ClassificationType {
66        #[default]
67        Binary,
68        Multiclass,
69        Multilabel,
70    }
71
72    /// Sample x Class shaped matrix for use in
73    /// classification metrics testing
74    pub fn dummy_classification_input(
75        classification_type: &ClassificationType,
76    ) -> (Tensor<2>, Tensor<2, Bool>) {
77        match classification_type {
78            ClassificationType::Binary => {
79                (
80                    Tensor::from_data([[0.3], [0.2], [0.7], [0.1], [0.55]], &Default::default()),
81                    // targets
82                    Tensor::from_data([[0], [1], [0], [0], [1]], &Default::default()),
83                    // predictions @ threshold=0.5
84                    //                     [[0], [0], [1], [0], [1]]
85                )
86            }
87            ClassificationType::Multiclass => {
88                (
89                    Tensor::from_data(
90                        [
91                            [0.2, 0.8, 0.0],
92                            [0.3, 0.6, 0.1],
93                            [0.7, 0.25, 0.05],
94                            [0.1, 0.15, 0.8],
95                            [0.9, 0.03, 0.07],
96                        ],
97                        &Default::default(),
98                    ),
99                    Tensor::from_data(
100                        // targets
101                        [[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [1, 0, 0]],
102                        // predictions @ top_k=1
103                        //   [[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1], [1, 0,  0]]
104                        // predictions @ top_k=2
105                        //   [[1, 1, 0], [1, 1, 0], [1, 1, 0], [0, 1, 1], [1, 0,  1]]
106                        &Default::default(),
107                    ),
108                )
109            }
110            ClassificationType::Multilabel => {
111                (
112                    Tensor::from_data(
113                        [
114                            [0.1, 0.7, 0.6],
115                            [0.3, 0.9, 0.05],
116                            [0.8, 0.9, 0.4],
117                            [0.7, 0.5, 0.9],
118                            [1.0, 0.3, 0.2],
119                        ],
120                        &Default::default(),
121                    ),
122                    // targets
123                    Tensor::from_data(
124                        [[1, 1, 0], [1, 0, 1], [1, 1, 1], [0, 0, 1], [1, 0, 0]],
125                        // predictions @ threshold=0.5
126                        //   [[0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 1], [1, 0, 0]]
127                        &Default::default(),
128                    ),
129                )
130            }
131        }
132    }
133}
burn_train/lib.rs

burn_train/
lib.rs
