Setting ADVI start state

I would like to continue an ADVI inference right where it ended. Pickling the python object is not an option as it consumes too much memory. Additionally, I would like to switch from float32 GPU to float64 CPU calculations and am not sure If that would work once the inference was started on GPU.

I belive the state of the inference consists of three vectors:

  • parameter means
  • parameter standard deviations stds
  • stochastic optimizer state opst

During the first round of inference I get

means = approx.bij.rmap(approx.mean.get_value())
stds = approx.bij.rmap(approx.std.eval())

and save them with some means (in this case with h5py). I do this during a callback where approx is the first argument passed to the callback function. I do not know how to obtain opts yet.

To start the second round of inference I load means and stds and start the process with

advi = pm.fit(model=model, method='advi', start=means)

However, how do I initialize the inference with stds and `opts?

See Saving ADVI results and reloading :wink:

1 Like

Oh, great! Thanks!

In case somone wonders:

import numpy as np
rho = {var_name:np.log(np.expm1(values))
       for var_name, values in stds.items()}
1 Like