Summary
A training run cannot be resumed. nn::serialize::save writes named_parameters only, so BatchNorm2d's running statistics are lost (documented in LIMITATIONS) and Adam's moments and step count are lost (not documented).
Where (main at 853a224)
crates/oxmera-nn/src/serialize.rs:20-48 — save iterates module.named_parameters("").crates/oxmera-nn/src/norm.rs:80-83 — running_mean: Mutex<Tensor>, running_var: Mutex<Tensor> are not parameters.crates/oxmera-optim/src/lib.rs:157-164 — AdamCore { step, m, v } is private with no accessor.
Why it matters
A model saved after training with BatchNorm loads with running_mean = 0, running_var = 1 and evaluates wrongly; any interrupted run restarts its optimizer cold. Both are table stakes for a framework that describes a training dashboard.
Fix
Module::named_buffers(&self, prefix) -> Vec<(String, Tensor)> with a default empty implementation, implemented by BatchNorm2d; save/load include buffers.Optimizer::state_dict() -> Vec<(String, Tensor)> / load_state_dict(..) keyed by parameter name, stored as safetensors tensors (step count as a one-element tensor).- Document the file layout in
nn::serialize.
Done when
A test trains two steps, saves model + optimizer, loads both into fresh objects and asserts the third step's parameters equal an uninterrupted three-step run bit for bit on the CPU.
Summary
A training run cannot be resumed.
nn::serialize::savewritesnamed_parametersonly, soBatchNorm2d's running statistics are lost (documented in LIMITATIONS) and Adam's moments and step count are lost (not documented).Where (main at 853a224)
crates/oxmera-nn/src/serialize.rs:20-48—saveiteratesmodule.named_parameters("").crates/oxmera-nn/src/norm.rs:80-83—running_mean: Mutex<Tensor>, running_var: Mutex<Tensor>are not parameters.crates/oxmera-optim/src/lib.rs:157-164—AdamCore { step, m, v }is private with no accessor.Why it matters
A model saved after training with BatchNorm loads with
running_mean = 0,running_var = 1and evaluates wrongly; any interrupted run restarts its optimizer cold. Both are table stakes for a framework that describes a training dashboard.Fix
Module::named_buffers(&self, prefix) -> Vec<(String, Tensor)>with a default empty implementation, implemented byBatchNorm2d;save/loadinclude buffers.Optimizer::state_dict() -> Vec<(String, Tensor)>/load_state_dict(..)keyed by parameter name, stored as safetensors tensors (step count as a one-element tensor).nn::serialize.Done when
A test trains two steps, saves model + optimizer, loads both into fresh objects and asserts the third step's parameters equal an uninterrupted three-step run bit for bit on the CPU.