burn_train/checkpoint/
base.rs

1use burn_core::store::{ModuleRecord, RecordError};
2use burn_optim::OptimizerRecord;
3use burn_optim::lr_scheduler::LrSchedulerRecord;
4use burn_std::Bytes;
5use std::path::PathBuf;
6use thiserror::Error;
7
8/// The error type for checkpointer.
9#[derive(Error, Debug)]
10pub enum CheckpointerError {
11    /// IO error.
12    #[error("I/O Error: `{0}`")]
13    IOError(std::io::Error),
14
15    /// Record (burnpack) error.
16    #[error("Record error: `{0}`")]
17    Record(RecordError),
18
19    /// Other errors.
20    #[error("Unknown error: `{0}`")]
21    Unknown(String),
22}
23
24/// A record that can be saved to and loaded from a burnpack file.
25///
26/// Implemented for the burnpack record types used during training: the module
27/// ([`ModuleRecord`]), the optimizer ([`OptimizerRecord`]) and the learning rate scheduler
28/// ([`LrSchedulerRecord`]).
29///
30/// Records are device-free: a checkpoint is just file-backed bytes. Device placement is decided
31/// when a record is applied (the module keeps its existing parameter device; optimizer state
32/// migrates to each parameter's device on the next step), not when the checkpoint is loaded.
33pub trait Checkpoint: Sized + Send + 'static {
34    /// Save the record to `path`.
35    fn save(self, path: PathBuf) -> Result<(), CheckpointerError>;
36    /// Load the record from `path`.
37    fn load(path: PathBuf) -> Result<Self, CheckpointerError>;
38    /// Creates a checkpoint from bytes
39    fn checkpoint_from_bytes(bytes: Bytes) -> Result<Self, RecordError>;
40    /// Transforms a checkpoint into bytes
41    fn checkpoint_into_bytes(self) -> Result<Bytes, RecordError>;
42}
43
44/// A stateless record: nothing to save or load.
45impl Checkpoint for () {
46    fn save(self, _path: PathBuf) -> Result<(), CheckpointerError> {
47        Ok(())
48    }
49    fn load(_path: PathBuf) -> Result<Self, CheckpointerError> {
50        Ok(())
51    }
52    fn checkpoint_from_bytes(_bytes: Bytes) -> Result<Self, RecordError> {
53        Ok(())
54    }
55    fn checkpoint_into_bytes(self) -> Result<Bytes, RecordError> {
56        Ok(Bytes::from_bytes_vec(vec![0]))
57    }
58}
59
60impl Checkpoint for ModuleRecord {
61    fn save(self, path: PathBuf) -> Result<(), CheckpointerError> {
62        ModuleRecord::save(self, path).map_err(CheckpointerError::Record)
63    }
64    fn load(path: PathBuf) -> Result<Self, CheckpointerError> {
65        ModuleRecord::load(path).map_err(CheckpointerError::Record)
66    }
67    fn checkpoint_into_bytes(self) -> Result<Bytes, RecordError> {
68        self.into_bytes()
69    }
70    fn checkpoint_from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
71        ModuleRecord::from_bytes(bytes)
72    }
73}
74
75impl Checkpoint for OptimizerRecord {
76    fn save(self, path: PathBuf) -> Result<(), CheckpointerError> {
77        OptimizerRecord::save(self, path).map_err(CheckpointerError::Record)
78    }
79    fn load(path: PathBuf) -> Result<Self, CheckpointerError> {
80        OptimizerRecord::load(path).map_err(CheckpointerError::Record)
81    }
82    fn checkpoint_from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
83        OptimizerRecord::from_bytes(bytes)
84    }
85    fn checkpoint_into_bytes(self) -> Result<Bytes, RecordError> {
86        self.into_bytes()
87    }
88}
89
90impl Checkpoint for LrSchedulerRecord {
91    fn save(self, path: PathBuf) -> Result<(), CheckpointerError> {
92        LrSchedulerRecord::save(self, path).map_err(CheckpointerError::Record)
93    }
94    fn load(path: PathBuf) -> Result<Self, CheckpointerError> {
95        LrSchedulerRecord::load(path).map_err(CheckpointerError::Record)
96    }
97    fn checkpoint_from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
98        LrSchedulerRecord::from_bytes(bytes)
99    }
100    fn checkpoint_into_bytes(self) -> Result<Bytes, RecordError> {
101        self.into_bytes()
102    }
103}
104
105/// The trait for checkpointer.
106pub trait Checkpointer<R>: Send + Sync
107where
108    R: Checkpoint,
109{
110    /// Save the record.
111    ///
112    /// # Arguments
113    ///
114    /// * `epoch` - The epoch.
115    /// * `record` - The record.
116    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError>;
117
118    /// Delete the checkpoint saved at the given epoch if present.
119    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError>;
120
121    /// Restore the record from the checkpoint saved at the given epoch.
122    fn restore(&self, epoch: usize) -> Result<R, CheckpointerError>;
123}
burn_train/checkpoint/base.rs

burn_train/checkpoint/
base.rs
