search

coreylowman / dfdx

Deep learning in Rust, with shape checked tensors and neural networks
Other
1.71k stars 98 forks source link

Optimizer serialization #101

Open coreylowman opened 2 years ago

coreylowman commented 2 years ago

A common feature to use when building training frameworks is to save the optimizer state along with network state.

There's already a way to convert a D::Vec<E> into a rust Vec<E>, which should make it easy to save Gradients object. However a pain point is that gradients are currently keyed off of UniqueId, which may not be persistent across saving/loading depending on how model is initialized.

Options:

  1. Move optimizers away from Gradients and towards using HashMap<String, D::Vec<E>> where String would be the full path to the tensor. This would make it easy to use with TensorCollection, and also require minimal changes to implement this serialization.
  2. The API for serializing optimizers could go through a model, so optimizer gets full paths to tensors on save.
kurnevsky commented 1 year ago

It's also good to have a way to clone in memory the state of optimizer. It's useful if you want to pit 2 networks to see if the trained network performs better than the one before training (which is part of alpha go algorithm).

coreylowman commented 1 year ago

@kurnevsky thanks for bringing that up, cloning optimizers is actually really easy with some changes recently. Addressed in the above PR 👍