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                burn_pack::Tensor::new(
84                    tensor.name,
85                    tensor.data.dtype,
86                    tensor.data.shape,
87                    tensor.param_id,
88                    tensor.data.bytes,
89                )
90            })
91            .collect();
92        let mut writer = Writer::new(tensors);
93        for (key, value) in &self.scalars {
94            writer = writer.with_scalar(key, *value);
95        }
96        for (key, value) in &self.paths {
97            writer = writer.with_metadata(key, value);
98        }
99        writer
100    }
101
102    fn from_reader(reader: Reader) -> Result<Self, RecordError> {
103        let scalars = reader.scalars().clone();
104        let paths = reader.metadata().clone();
105        let tensors = reader
106            .into_tensors()?
107            .into_iter()
108            .map(|tensor| {
109                let (name, dtype, shape, param_id, bytes) = tensor.into_parts()?;
110                Ok(RecordTensor {
111                    name,
112                    param_id,
113                    data: TensorData::from_bytes(bytes, shape, dtype),
114                })
115            })
116            .collect::<Result<Vec<_>, RecordError>>()?;
117        Ok(Self {
118            tensors,
119            scalars,
120            paths,
121        })
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use burn::tensor::TensorData;
129
130    fn record(data: TensorData) -> OptimizerRecord {
131        OptimizerRecord {
132            tensors: alloc::vec![RecordTensor {
133                name: "42.moment".into(),
134                param_id: Some(42),
135                data,
136            }],
137            scalars: Default::default(),
138            paths: Default::default(),
139        }
140    }
141
142    fn assert_record(decoded: &OptimizerRecord, expected: &TensorData) {
143        let tensor = &decoded.tensors[0];
144
145        assert_eq!(tensor.name, "42.moment");
146        assert_eq!(tensor.param_id, Some(42));
147        assert_eq!(tensor.data.dtype, expected.dtype);
148        assert_eq!(tensor.data.shape, expected.shape);
149        assert_eq!(&tensor.data.bytes[..], &expected.bytes[..]);
150    }
151
152    #[test]
153    fn tensor_data_stays_internal_across_burnpack_boundary() {
154        let expected = TensorData::from([1.0_f32, 2.0, 3.0]);
155        let decoded =
156            OptimizerRecord::from_bytes(record(expected.clone()).into_bytes().unwrap()).unwrap();
157
158        assert_record(&decoded, &expected);
159    }
160
161    #[cfg(feature = "std")]
162    #[test]
163    fn tensor_data_round_trips_through_burnpack_file() {
164        let expected = TensorData::from([1.0_f32, 2.0, 3.0]);
165        let path = std::env::temp_dir().join(format!(
166            "burn_optim_record_{}_tensor_data.bpk",
167            std::process::id()
168        ));
169
170        record(expected.clone()).save(&path).unwrap();
171        let decoded = OptimizerRecord::load(&path).unwrap();
172        assert_record(&decoded, &expected);
173
174        drop(decoded);
175        std::fs::remove_file(path).unwrap();
176    }
177}
burn_optim/optim/module/record/mod.rs

burn_optim/optim/module/record/
mod.rs
