Saving and loading trained estimator

[vagechirkov]

I’d like to save and load trained estimators and/or posterior objects — basically anything that takes a while to compute. Ideally, I’d be able to persist all the expensive parts of the inference workflow and reload them later without issues.

I’ve looked at the current docs, especially this FAQ. But I find it a bit hard to follow and not super clear on best practices.

I think this would be a great thing to include in the “How-to” guide — maybe with a simple example showing how to save/load a posterior and a trained network, and what needs to be saved alongside it (e.g. prior, simulator config, etc.). @michaeldeistler

This might also be related to the logger issue #1524, since having saved weights with relevant metadata would probably be the most useful setup.

If there’s working code for this already, I’d be happy to help turn it into a short guide or snippet for the docs. Otherwise, I think this would be a really useful thing to have in the documentation.

Let me know if this is discussed somewhere else or if I missed an example!

[danielmk]

I don't want to hijack this issue and I can take this somewhere else if that's not what you are looking for but I've been thinking a lot about adding a .save/.load interface to the NeuralInference object. Especially if we are looking towards a mature release in the future, that interface could handle compatibility issues more gracefully than the generic pickle can.

That said I've used pickle a lot on the inference object and it works generally well. For my recent publication I plan to provide the inference object containing the trained density estimators as pickle files for reproducibility purposes. I don't think anything other than the density estimator needs to be saved to get the posterior and sample from it. If one wanted to take already trained weights and do more training on them, that could prove more complicated though but I never tried that.

I can make my scripts that save/load the pickles available when I'm back in the office.

[vagechirkov]

Thanks @danielmk, that’s super helpful!

Yeah, having a .save() / .load() interface on NeuralInference sounds like a great idea — especially for making things more user-friendly and future-proof. Pickle works for now, but I agree that something more built-in would make it easier to share and reuse results.

Would love to see your scripts — I’d be happy to help turn them into a small example and update the docs 🙂

[jorobledo]

I would suggest to use dill over pickle for a more robust approach.

I've had problems on saving instances with pickle in pymc MCMC chains, which I could save using dill.

From their documentation:

dill provides the user the same interface as the pickle module, and also includes some additional features. In addition to pickling Python objects, dill provides the ability to save the state of an interpreter session in a single command. Hence, it would be feasible to save an interpreter session, close the interpreter, ship the pickled file to another computer, open a new interpreter, unpickle the session and thus continue from the ‘saved’ state of the original interpreter session.