How to save fitted ADVI Result?

Edit:, originally a github issue, then posted here, now it’s back to being a github issue #2615

If I have a (lengthy) advi optimization, whats the best way to store the result for a future notebook?

with pymc_model:

    approx = pm.fit(
        method='advi', n=50000, random_seed=1, callbacks=cb,
        obj_optimizer=pm.adagrad_window(learning_rate=.5, n_win=100))

    trace = approx.sample(1000)

Pickling the trace works,

import pickle

with open('trace.p', 'wb') as f:
    pickle.dump(trace, f)
    
with open('trace.p', 'rb') as f:
    test1 = pickle.load(f)

But pickling the approx class fails when I try to load it.

with open('approx.p', 'wb') as f:
    pickle.dump(approx, f)
    
with open('approx.p', 'rb') as f:
    test2 = pickle.load(f)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-31-54dd7d55c9e6> in <module>()
      3 
      4 with open('approx.p', 'rb') as f:
----> 5     test2 = pickle.load(f)

TypeError: __new__() missing 1 required positional argument: 'group'

Versions and main components

  • PyMC3 Version: 3.1
  • Theano Version: 0.9.0
  • Python Version: 3.6
  • Operating system: Mac OS X
  • How did you install PyMC3: from source

One way to go around it is to initialize the approximation, then do the fit, and then saved the approximation parameters. For example:

with model:
    approx=pm.ADVI()
    approx.fit()
saveparam = [(param.name, param.eval()) for param in approx.approx.params]

and load it for later:

with model:
    approx=pm.ADVI()
for i, param in enumerate(saveparam):
    approx.approx.params[i].set_value(param[1])

This is not the safe way to do as showed in Saving ADVI results and reloading
A better solution is below:

The safe way to do is to save the parameters as a dictionary, and map it back to a vector whenever the new model is created.

...
# save inference
bij = inference.approx.groups[0].bij
saveparam = {param.name: bij.rmap(param.eval())
	 for param in inference.approx.params}

# new model
model2 = construct_neural_network(
  X, Y, Yerr, neuronsPerHiddenlayer, ninputs, noutputs, ndata)
with model2:
  inference2 = pm.ADVI()

# load inference
bij2 = inference2.approx.groups[0].bij
inference2.approx.params[0].set_value(bij2.map(saveparam['mu']))
inference2.approx.params[1].set_value(bij2.map(saveparam['rho']))

It’s also now possible to pickle ADVI objects and reload them, thanks to a PR from ferrine: #2619

1 Like