Skip to main content

burn_train/learner/
error.rs

1use burn_core::data::dataset::DatasetError;
2
3use crate::{
4    EventProcessorError,
5    checkpoint::CheckpointerError,
6    train::{MultiDeviceStepError, WorkerFailure, fmt_worker_failures},
7};
8
9/// An error that stopped training or evaluation.
10#[derive(Debug)]
11pub enum TrainingError {
12    /// Error while processing events.
13    EventProcessor(EventProcessorError),
14    /// Error while loading data.
15    Dataset(DatasetError),
16    /// The training step panicked on one or more workers of a multi-device strategy.
17    Workers(Vec<WorkerFailure>),
18    /// Error while saving, restoring or deleting a checkpoint.
19    Checkpoint(CheckpointerError),
20}
21
22impl TrainingError {
23    /// Whether the error comes from a poisoned device. In this case, every later operation
24    /// will fail until the process is restarted.
25    pub fn is_device_poisoned(&self) -> bool {
26        match self {
27            Self::EventProcessor(err) => err.is_device_poisoned(),
28            _ => false,
29        }
30    }
31}
32
33impl core::fmt::Display for TrainingError {
34    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
35        match self {
36            Self::EventProcessor(err) => write!(f, "{err}"),
37            Self::Dataset(err) => write!(f, "Dataset error: {err}"),
38            Self::Workers(failures) => fmt_worker_failures(failures, f),
39            Self::Checkpoint(err) => write!(f, "Checkpoint error: {err}"),
40        }
41    }
42}
43
44impl core::error::Error for TrainingError {
45    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
46        match self {
47            Self::EventProcessor(err) => Some(err),
48            Self::Dataset(err) => Some(err),
49            Self::Workers(_) => None,
50            Self::Checkpoint(err) => Some(err),
51        }
52    }
53}
54
55impl From<EventProcessorError> for TrainingError {
56    fn from(err: EventProcessorError) -> Self {
57        Self::EventProcessor(err)
58    }
59}
60
61impl From<DatasetError> for TrainingError {
62    fn from(err: DatasetError) -> Self {
63        Self::Dataset(err)
64    }
65}
66
67impl From<CheckpointerError> for TrainingError {
68    fn from(err: CheckpointerError) -> Self {
69        Self::Checkpoint(err)
70    }
71}
72
73impl From<MultiDeviceStepError> for TrainingError {
74    fn from(err: MultiDeviceStepError) -> Self {
75        match err {
76            MultiDeviceStepError::Dataset(err) => Self::Dataset(err),
77            MultiDeviceStepError::Workers(failures) => Self::Workers(failures),
78        }
79    }
80}