The correct approach to save the SVI results

numpyro
6

I think you need to install the latest version (which seems to fix some pickle issues). We have tests to cover pickling autoguides but we might miss something. For example, we don’t need dill for the above model

optimizer = optax.adam(1e-3)
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
svi_result = svi.run(rng_key, 10000, x, y)

import pickle

with open("result.pkl", "wb") as f:
    pickle.dump((model, guide, svi_result), f)

with open("result.pkl", "rb") as g:
    model, guide, svi_result = pickle.load(g)

# draw some posterior predictive
# Posterior Predictive
num_sample_dist = 200
posterior_predictive = Predictive(model=model, guide=guide,
                                  params=svi_result.params,
                                  num_samples=num_sample_dist)
y_pred = posterior_predictive(rng_key_predict, x=x)['obs']
1 Like