How to properly serialize an AutoMultivariateNormal's state?

numpyro
1

I am trying to fit an SVI model using an AutoMultivariateNormal guide in one file and trying to serialize the results, so that I can generate predictions in a different file. However, I’m finding that just serializing the SVI params is not enough. AutoMultivariateNormal seems to contain important state that AutoNormal does not.

When I generate predictions using a newly generated guide object, the predictions are very bad. But the predictions are great if I use the original guide object that I used to fit the model. I had erroneously assumed that guides were stateless(?)

What’s the proper way to serialize / load this state? I don’t want to resort to pickle because I’d like to share these files later for publication?

2

I think you might need to use ParamStoreDict.save and load methods.

3

Should have included it in the original post, but I’m using numpyro and I don’t think it has a ParamStoreDict

4

Could you include minimal code to illustrate the issue?