burn_optim/optim/module/record/
mod.rs

1use alloc::collections::BTreeMap;
2use alloc::string::String;
3use alloc::vec::Vec;
4
5use burn::store::RecordError;
6use burn::tensor::Bytes;
7use burn_core as burn;
8
9use burn_pack::{Reader, Scalar, Writer};
10
11/// A serialized optimizer state, stored in the [burnpack](burn_pack) format.
12///
13/// Unlike a module record (keyed by module path), an optimizer record is keyed per parameter:
14/// each parameter's state is decomposed into tensors named `"{param_id}.{field}"` (carrying the
15/// originating `param_id`) plus a few typed scalar entries kept in the burnpack scalar map.
16///
17/// Obtain one from a [`ModuleOptimizer`](crate::optim::ModuleOptimizer) with
18/// [`to_record`](crate::optim::ModuleOptimizer::to_record), then save it
19/// ([`save`](Self::save) / [`into_bytes`](Self::into_bytes)) or apply it back with
20/// [`load_record`](crate::optim::ModuleOptimizer::load_record).
21#[derive(Default)]
22pub struct OptimizerRecord {
23    pub(crate) tensors: Vec<burn_pack::Tensor>,
24    pub(crate) scalars: BTreeMap<String, Scalar>,
25    pub(crate) paths: BTreeMap<String, String>,
26}
27
28impl core::fmt::Debug for OptimizerRecord {
29    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
30        f.debug_struct("OptimizerRecord")
31            .field("num_tensors", &self.tensors.len())
32            .field("num_scalars", &self.scalars.len())
33            .finish()
34    }
35}
36
37impl OptimizerRecord {
38    /// The number of tensors in the record.
39    pub fn len(&self) -> usize {
40        self.tensors.len()
41    }
42
43    /// Whether the record holds no tensors.
44    pub fn is_empty(&self) -> bool {
45        self.tensors.is_empty()
46    }
47
48    /// Serialize the record to an in-memory burnpack byte buffer.
49    pub fn into_bytes(self) -> Result<Bytes, RecordError> {
50        Ok(self.into_writer().into_bytes()?)
51    }
52
53    /// Reconstruct a record from an in-memory burnpack byte buffer.
54    pub fn from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
55        Self::from_reader(Reader::from_bytes(bytes)?)
56    }
57
58    /// Save the record to a burnpack file on disk.
59    #[cfg(feature = "std")]
60    pub fn save<P: AsRef<std::path::Path>>(self, path: P) -> Result<(), RecordError> {
61        self.into_writer().write_to_file(path)?;
62        Ok(())
63    }
64
65    /// Load a record from a burnpack file on disk.
66    #[cfg(feature = "std")]
67    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, RecordError> {
68        Self::from_reader(Reader::from_file(path)?)
69    }
70
71    fn into_writer(self) -> Writer {
72        let mut writer = Writer::new(self.tensors);
73        for (key, value) in &self.scalars {
74            writer = writer.with_scalar(key, *value);
75        }
76        for (key, value) in &self.paths {
77            writer = writer.with_metadata(key, value);
78        }
79        writer
80    }
81
82    fn from_reader(reader: Reader) -> Result<Self, RecordError> {
83        let scalars = reader.scalars().clone();
84        let paths = reader.metadata().clone();
85        let tensors = reader.into_tensors()?;
86        Ok(Self {
87            tensors,
88            scalars,
89            paths,
90        })
91    }
92}
burn_optim/optim/module/record/mod.rs

burn_optim/optim/module/record/
mod.rs
