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, TensorData};
7use burn_core as burn;
8
9use burn_pack::{Reader, Scalar, Writer};
10
11/// A single optimizer state tensor and its optional originating parameter id.
12pub(crate) struct RecordTensor {
13    pub(crate) name: String,
14    pub(crate) param_id: Option<u64>,
15    pub(crate) data: TensorData,
16}
17
18/// A serialized optimizer state, stored in the [burnpack](burn_pack) format.
19///
20/// Unlike a module record (keyed by module path), an optimizer record is keyed per parameter:
21/// each parameter's state is decomposed into tensors named `"{param_id}.{field}"` (carrying the
22/// originating `param_id`) plus a few typed scalar entries kept in the burnpack scalar map.
23///
24/// Obtain one from a [`ModuleOptimizer`](crate::optim::ModuleOptimizer) with
25/// [`to_record`](crate::optim::ModuleOptimizer::to_record), then save it
26/// ([`save`](Self::save) / [`into_bytes`](Self::into_bytes)) or apply it back with
27/// [`load_record`](crate::optim::ModuleOptimizer::load_record).
28#[derive(Default)]
29pub struct OptimizerRecord {
30    pub(crate) tensors: Vec<RecordTensor>,
31    pub(crate) scalars: BTreeMap<String, Scalar>,
32    pub(crate) paths: BTreeMap<String, String>,
33}
34
35impl core::fmt::Debug for OptimizerRecord {
36    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
37        f.debug_struct("OptimizerRecord")
38            .field("num_tensors", &self.tensors.len())
39            .field("num_scalars", &self.scalars.len())
40            .finish()
41    }
42}
43
44impl OptimizerRecord {
45    /// The number of tensors in the record.
46    pub fn len(&self) -> usize {
47        self.tensors.len()
48    }
49
50    /// Whether the record holds no tensors.
51    pub fn is_empty(&self) -> bool {
52        self.tensors.is_empty()
53    }
54
55    /// Serialize the record to an in-memory burnpack byte buffer.
56    pub fn into_bytes(self) -> Result<Bytes, RecordError> {
57        Ok(self.into_writer().into_bytes()?)
58    }
59
60    /// Reconstruct a record from an in-memory burnpack byte buffer.
61    pub fn from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
62        Self::from_reader(Reader::from_bytes(bytes)?)
63    }
64
65    /// Save the record to a burnpack file on disk.
66    #[cfg(feature = "std")]
67    pub fn save<P: AsRef<std::path::Path>>(self, path: P) -> Result<(), RecordError> {
68        self.into_writer().write_to_file(path)?;
69        Ok(())
70    }
71
72    /// Load a record from a burnpack file on disk.
73    #[cfg(feature = "std")]
74    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, RecordError> {
75        Self::from_reader(Reader::from_file(path)?)
76    }
77
78    fn into_writer(self) -> Writer {
79        let tensors = self
80            .tensors
81            .into_iter()
82            .map(|tensor| {
83                let (bytes, shape, dtype) = tensor.data.into_parts();
84                burn_pack::Tensor::new(tensor.name, dtype, shape, tensor.param_id, bytes)
85            })
86            .collect();
87        let mut writer = Writer::new(tensors);
88        for (key, value) in &self.scalars {
89            writer = writer.with_scalar(key, *value);
90        }
91        for (key, value) in &self.paths {
92            writer = writer.with_metadata(key, value);
93        }
94        writer
95    }
96
97    fn from_reader(reader: Reader) -> Result<Self, RecordError> {
98        let scalars = reader.scalars().clone();
99        let paths = reader.metadata().clone();
100        let tensors = reader
101            .into_tensors()?
102            .into_iter()
103            .map(|tensor| {
104                let (name, dtype, shape, param_id, bytes) = tensor.into_parts()?;
105                Ok(RecordTensor {
106                    name,
107                    param_id,
108                    data: TensorData::from_bytes(bytes, shape, dtype),
109                })
110            })
111            .collect::<Result<Vec<_>, RecordError>>()?;
112        Ok(Self {
113            tensors,
114            scalars,
115            paths,
116        })
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use burn::tensor::TensorData;
124
125    fn record(data: TensorData) -> OptimizerRecord {
126        OptimizerRecord {
127            tensors: alloc::vec![RecordTensor {
128                name: "42.moment".into(),
129                param_id: Some(42),
130                data,
131            }],
132            scalars: Default::default(),
133            paths: Default::default(),
134        }
135    }
136
137    fn assert_record(decoded: &OptimizerRecord, expected: &TensorData) {
138        let tensor = &decoded.tensors[0];
139
140        assert_eq!(tensor.name, "42.moment");
141        assert_eq!(tensor.param_id, Some(42));
142        assert_eq!(tensor.data.dtype(), expected.dtype());
143        assert_eq!(tensor.data.shape(), expected.shape());
144        assert_eq!(tensor.data.as_bytes(), expected.as_bytes());
145    }
146
147    #[test]
148    fn tensor_data_stays_internal_across_burnpack_boundary() {
149        let expected = TensorData::from([1.0_f32, 2.0, 3.0]);
150        let decoded =
151            OptimizerRecord::from_bytes(record(expected.clone()).into_bytes().unwrap()).unwrap();
152
153        assert_record(&decoded, &expected);
154    }
155
156    #[cfg(feature = "std")]
157    #[test]
158    fn tensor_data_round_trips_through_burnpack_file() {
159        let expected = TensorData::from([1.0_f32, 2.0, 3.0]);
160        let path = std::env::temp_dir().join(format!(
161            "burn_optim_record_{}_tensor_data.bpk",
162            std::process::id()
163        ));
164
165        record(expected.clone()).save(&path).unwrap();
166        let decoded = OptimizerRecord::load(&path).unwrap();
167        assert_record(&decoded, &expected);
168
169        drop(decoded);
170        std::fs::remove_file(path).unwrap();
171    }
172}
burn_optim/optim/module/record/mod.rs

burn_optim/optim/module/record/
mod.rs
