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

use burn_core::data::dataset::DatasetError;

use crate::{
    EventProcessorError,
    checkpoint::CheckpointerError,
    train::{MultiDeviceStepError, WorkerFailure, fmt_worker_failures},
};

/// An error that stopped training or evaluation.
#[derive(Debug)]
pub enum TrainingError {
    /// Error while processing events.
    EventProcessor(EventProcessorError),
    /// Error while loading data.
    Dataset(DatasetError),
    /// The training step panicked on one or more workers of a multi-device strategy.
    Workers(Vec<WorkerFailure>),
    /// Error while saving, restoring or deleting a checkpoint.
    Checkpoint(CheckpointerError),
}

impl TrainingError {
    /// Whether the error comes from a poisoned device. In this case, every later operation
    /// will fail until the process is restarted.
    pub fn is_device_poisoned(&self) -> bool {
        match self {
            Self::EventProcessor(err) => err.is_device_poisoned(),
            _ => false,
        }
    }
}

impl core::fmt::Display for TrainingError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::EventProcessor(err) => write!(f, "{err}"),
            Self::Dataset(err) => write!(f, "Dataset error: {err}"),
            Self::Workers(failures) => fmt_worker_failures(failures, f),
            Self::Checkpoint(err) => write!(f, "Checkpoint error: {err}"),
        }
    }
}

impl core::error::Error for TrainingError {
    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
        match self {
            Self::EventProcessor(err) => Some(err),
            Self::Dataset(err) => Some(err),
            Self::Workers(_) => None,
            Self::Checkpoint(err) => Some(err),
        }
    }
}

impl From<EventProcessorError> for TrainingError {
    fn from(err: EventProcessorError) -> Self {
        Self::EventProcessor(err)
    }
}

impl From<DatasetError> for TrainingError {
    fn from(err: DatasetError) -> Self {
        Self::Dataset(err)
    }
}

impl From<CheckpointerError> for TrainingError {
    fn from(err: CheckpointerError) -> Self {
        Self::Checkpoint(err)
    }
}

impl From<MultiDeviceStepError> for TrainingError {
    fn from(err: MultiDeviceStepError) -> Self {
        match err {
            MultiDeviceStepError::Dataset(err) => Self::Dataset(err),
            MultiDeviceStepError::Workers(failures) => Self::Workers(failures),
        }
    }
}