burn_train/checkpoint/
file.rs

1use std::path::{Path, PathBuf};
2
3use super::{Checkpoint, Checkpointer, CheckpointerError};
4
5/// The file checkpointer.
6///
7/// Saves each record as a [burnpack](burn_core::store) file in the given directory.
8pub struct FileCheckpointer {
9    directory: PathBuf,
10    name: String,
11}
12
13impl FileCheckpointer {
14    /// Creates a new file checkpointer.
15    ///
16    /// # Arguments
17    ///
18    /// * `directory` - The directory to save the checkpoints.
19    /// * `name` - The name of the checkpoint.
20    pub fn new(directory: impl AsRef<Path>, name: &str) -> Self {
21        let directory = directory.as_ref();
22        std::fs::create_dir_all(directory).ok();
23
24        Self {
25            directory: directory.to_path_buf(),
26            name: name.to_string(),
27        }
28    }
29
30    fn path_for_epoch(&self, epoch: usize) -> PathBuf {
31        self.directory.join(format!("{}-{}.bpk", self.name, epoch))
32    }
33}
34
35impl<R> Checkpointer<R> for FileCheckpointer
36where
37    R: Checkpoint,
38{
39    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> {
40        let file_path = self.path_for_epoch(epoch);
41        log::trace!("Saving checkpoint {} to {}", epoch, file_path.display());
42
43        record.save(file_path)
44    }
45
46    fn restore(&self, epoch: usize) -> Result<R, CheckpointerError> {
47        let file_path = self.path_for_epoch(epoch);
48        log::info!(
49            "Restoring checkpoint {} from {}",
50            epoch,
51            file_path.display()
52        );
53
54        R::load(file_path)
55    }
56
57    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> {
58        let file_path = self.path_for_epoch(epoch);
59
60        if file_path.exists() {
61            log::trace!("Removing checkpoint {}", file_path.display());
62            std::fs::remove_file(file_path).map_err(CheckpointerError::IOError)?;
63        }
64
65        Ok(())
66    }
67}
burn_train/checkpoint/file.rs

burn_train/checkpoint/
file.rs
