Make Approximation serializeable

[pstjohn]

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

[fonnesbeck]

Please use our Discourse page for general questions.

https://discourse.pymc.io

[pstjohn]

I suppose this was a bug report disguised as a general question 😄
you've defined a custom __new__ method for Group, but haven't provided a corresponding __getnewargs__. As a result it will always be impossible to pickle that class.

[fonnesbeck]

OK, I will reopen and change the title to clarify.

[pstjohn]

I took a bit of a closer look, this is definitely a question for @ferrine.
It is possible to have group construction proceed without overloading __new__? I.e., just a function make_group(group, vfam=None, params=None, *args, **kwargs)? That would certainly make the code easier to understand, but would be a more involved fix.

I tried the simple fix of just adding a __getnewargs__ function to Group, but that didn't improve things:

    def __getnewargs__(self):
        return (self.group, self._vfam, self.user_params)

Traceback (most recent call last):
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  [Previous line repeated 991 more times]
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 539, in __getattr__
    def __getattr__(self, item):
  File "/Users/pstjohn/anaconda3/envs/pymc/lib/python3.6/bdb.py", line 50, in trace_dispatch
    return self.dispatch_call(frame, arg)
  File "/Users/pstjohn/anaconda3/envs/pymc/lib/python3.6/bdb.py", line 76, in dispatch_call
    if not (self.stop_here(frame) or self.break_anywhere(frame)):
  File "/Users/pstjohn/anaconda3/envs/pymc/lib/python3.6/bdb.py", line 173, in break_anywhere
    return self.canonic(frame.f_code.co_filename) in self.breaks
  File "/Users/pstjohn/anaconda3/envs/pymc/lib/python3.6/bdb.py", line 29, in canonic
    if filename == "<" + filename[1:-1] + ">":
RecursionError: maximum recursion depth exceeded in comparison

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  [Previous line repeated 996 more times]
RecursionError: maximum recursion depth exceeded while calling a Python object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  File "/Users/pstjohn/Packages/pymc3/pymc3/variational/approximations.py", line 540, in __getattr__
    return getattr(self.groups[0], item)
  [Previous line repeated 996 more times]
RecursionError: maximum recursion depth exceeded while calling a Python object

Anyways, here's a test case we should add to test_variational_inference.py when we have a fix:

def test_pickle_approx(three_var_approx):
    import pickle
    dump = pickle.dumps(three_var_approx)
    pickle.loads(dump)

[junpenglao]

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])

[pstjohn]

So the recursion issue isn't fixed by deleting __new__ and replacing calls to Group() with a separate function,

def create_group(group, vfam=None, params=None, *args, **kwargs):
    if vfam is not None and params is not None:
        raise TypeError('Cannot call Group with both `vfam` and `params` provided')
    elif vfam is not None:
        return group_for_short_name(vfam)(*args, **kwargs)
    elif params is not None:
        return group_for_params(params)(*args, **kwargs)
    else:
        raise TypeError('Need to call Group with either `vfam` or `params` provided')

So those might be separate issues. (and the __getnewargs__ solution might work otherwise). It seems to be implying that a SingleGroupApproximation is ending up as its own group, causing the recursion error here:
https://github.com/pymc-devs/pymc3/blob/master/pymc3/variational/approximations.py#L539-L540

I'm not sure why that would only happen on unpickling though.

[ferrine]

I'll investigate this issue today