search

unda-ml / unda

General purpose machine learning crate
Apache License 2.0
13 stars 4 forks source link

Model save/load #30

Open atlv24 opened 8 months ago

atlv24 commented 8 months ago

expose a way to save out a computation graph, accompanied by all constants parameters and model weights as relative-path hlo format files, xla-rs has utils for this and its the expected format for model/optimizer saving in XLA

Ebanflo42 commented 8 months ago

xla-rs has utils for saving a converted XLA computation as HLO and also utils for saving literals as numpy, i think it would be dope to have these things combined in the same file, so the compute graph is serialized to HLO and eventually we have a Model object which contains weights that get serialized to a dictionary of numpy arrays.

Ebanflo42 commented 7 months ago

Reading HLO into an unda context would be a neat trick. Might also be good to look into the ONNX format, I've never used this but apparently everyone is using it. Would boost our compatibility with existing APIs.

Ebanflo42 commented 3 months ago

@BradenEverson or @atlv24 this should be implemented sooner rather than later if either of you have the bandwidth. There is already utilities for this in xla-rs/src/wrappers/mod.rs (of all places). If one of you could look into saving and loading SupervisedModel with these utilities I will clean up the organization of xla-rs.