mixture model failure with njobs > 1

[DeliciousHair]

Greetings!

I'm trying to speed up a rather lengthy sampling process on a mixture model, first thing to try was simply using the njobs= flag in sample(). This is not working...

As a minimal example, if I alter the example provided here from:

    tr = pm.sample(10000, step=[step1, step2])

to

   tr = pm.sample(10000, step=[step1, step2], njobs=3)

then I get:

Applied stickbreaking-transform to p and added transformed p_stickbreaking to model.
Applied interval-transform to sd and added transformed sd_interval to model.
---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/usr/local/lib/python3.4/dist-packages/theano/compile/function_module.py", line 859, in __call__
    outputs = self.fn()
IndexError: index 3 is out of bounds for size 3

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.4/multiprocessing/pool.py", line 119, in worker
    result = (True, func(*args, **kwds))
  File "/usr/lib/python3.4/multiprocessing/pool.py", line 44, in mapstar
    return list(map(*args))
  File "/usr/local/lib/python3.4/dist-packages/pymc3/sampling.py", line 296, in argsample
    return _sample(*args)
  File "/usr/local/lib/python3.4/dist-packages/pymc3/sampling.py", line 165, in _sample
    for i, strace in enumerate(sampling):
  File "/usr/local/lib/python3.4/dist-packages/pymc3/sampling.py", line 247, in _iter_sample
    point = step.step(point)
  File "/usr/local/lib/python3.4/dist-packages/pymc3/step_methods/compound.py", line 14, in step
    point = method.step(point)
  File "/usr/local/lib/python3.4/dist-packages/pymc3/step_methods/compound.py", line 14, in step
    point = method.step(point)
  File "/usr/local/lib/python3.4/dist-packages/pymc3/step_methods/compound.py", line 14, in step
    point = method.step(point)
  File "/usr/local/lib/python3.4/dist-packages/pymc3/step_methods/arraystep.py", line 118, in step
    apoint = self.astep(bij.map(point))
  File "/usr/local/lib/python3.4/dist-packages/pymc3/step_methods/metropolis.py", line 125, in astep
    q_new = metrop_select(self.delta_logp(q,q0), q, q0)
  File "/usr/local/lib/python3.4/dist-packages/theano/compile/function_module.py", line 871, in __call__
    storage_map=getattr(self.fn, 'storage_map', None))
  File "/usr/local/lib/python3.4/dist-packages/theano/gof/link.py", line 314, in raise_with_op
    reraise(exc_type, exc_value, exc_trace)
  File "/usr/local/lib/python3.4/dist-packages/six.py", line 685, in reraise
    raise value.with_traceback(tb)
  File "/usr/local/lib/python3.4/dist-packages/theano/compile/function_module.py", line 859, in __call__
    outputs = self.fn()
IndexError: index 3 is out of bounds for size 3
Apply node that caused the error: AdvancedSubtensor1(p, Subtensor{int64:int64:}.0)
Toposort index: 28
Inputs types: [TensorType(float64, vector), TensorType(int64, vector)]
Inputs shapes: [(3,), (2000,)]
Inputs strides: [(8,), (8,)]
Inputs values: [array([ 0.33333333,  0.33333333,  0.33333333]), 'not shown']
Outputs clients: [[Elemwise{Composite{Switch((GE(i0, i1) * LE(i0, i2) * i3), log(i4), i5)}}[(0, 4)](Subtensor{int64:int64:}.0, TensorConstant{(1,) of 0}, Elemwise{Add}[(0, 1)].0, Elemwise{Composite{LE(Abs((i0 + i1)), i2)}}.0, AdvancedSubtensor1.0, TensorConstant{(1,) of -inf})]]

Backtrace when the node is created:
  File "/usr/local/lib/python3.4/dist-packages/pymc3/distributions/discrete.py", line 379, in logp
    return bound(T.log(p[value]),

HINT: Use the Theano flag 'exception_verbosity=high' for a debugprint and storage map footprint of this apply node.
"""

The above exception was the direct cause of the following exception:

IndexError                                Traceback (most recent call last)
<ipython-input-15-6ee43e24a0b0> in <module>()
     32     step1 = pm.Metropolis(vars=[p, sd, means])
     33     step2 = pm.ElemwiseCategoricalStep(vars=[category], values=[0, 1, 2])
---> 34     tr = pm.sample(2000, step=[step1, step2], njobs=3)

/usr/local/lib/python3.4/dist-packages/pymc3/sampling.py in sample(draws, step, start, trace, chain, njobs, tune, progressbar, model, random_seed)
    154         sample_args = [draws, step, start, trace, chain,
    155                        tune, progressbar, model, random_seed]
--> 156     return sample_func(*sample_args)
    157 
    158 

/usr/local/lib/python3.4/dist-packages/pymc3/sampling.py in _mp_sample(njobs, args)
    275 def _mp_sample(njobs, args):
    276     p = mp.Pool(njobs)
--> 277     traces = p.map(argsample, args)
    278     p.close()
    279     return merge_traces(traces)

/usr/lib/python3.4/multiprocessing/pool.py in map(self, func, iterable, chunksize)
    258         in a list that is returned.
    259         '''
--> 260         return self._map_async(func, iterable, mapstar, chunksize).get()
    261 
    262     def starmap(self, func, iterable, chunksize=None):

/usr/lib/python3.4/multiprocessing/pool.py in get(self, timeout)
    597             return self._value
    598         else:
--> 599             raise self._value
    600 
    601     def _set(self, i, obj):

IndexError: index 3 is out of bounds for size 3
Apply node that caused the error: AdvancedSubtensor1(p, Subtensor{int64:int64:}.0)
Toposort index: 28
Inputs types: [TensorType(float64, vector), TensorType(int64, vector)]
Inputs shapes: [(3,), (2000,)]
Inputs strides: [(8,), (8,)]
Inputs values: [array([ 0.33333333,  0.33333333,  0.33333333]), 'not shown']
Outputs clients: [[Elemwise{Composite{Switch((GE(i0, i1) * LE(i0, i2) * i3), log(i4), i5)}}[(0, 4)](Subtensor{int64:int64:}.0, TensorConstant{(1,) of 0}, Elemwise{Add}[(0, 1)].0, Elemwise{Composite{LE(Abs((i0 + i1)), i2)}}.0, AdvancedSubtensor1.0, TensorConstant{(1,) of -inf})]]

Backtrace when the node is created:
  File "/usr/local/lib/python3.4/dist-packages/pymc3/distributions/discrete.py", line 379, in logp
    return bound(T.log(p[value]),

HINT: Use the Theano flag 'exception_verbosity=high' for a debugprint and storage map footprint of this apply node.

That last hint is not very helpful to me BTW :-)

Any ideas?

[fonnesbeck]

The last hint is a Theano message, not one generated by PyMC. It is just a way of generating more verbose output.

Are you using the current master branch, or something earlier?

[abojchevski]

I have the exactly same error, and I think it's not Theano's fault.
With some debugging what I've discovered so far is:

• with njobs=1 (default) in CompoundStep in self.methods all of the steps are as you have assigned them in your model construction
• but with njobs>1 in CompoundStep in self.methods each step is of the type Metropolis for some reason

still debugging to try and figure out why this happens, but maybe you have a clue.

[twiecki]

Interesting, so we're not communicating the step-method assignment to the individual workers?

[abojchevski]

It appears so. In my case I have 4 variables, 2 of them with Metropolis step and 2 of them with a custom SequentialScanDiscreteMetropolis suggested on another thread.

With njobs=1 the methods list in the outermost CompundStep contains:

• 1 CompundStep which in turn contains the two Metropolis steps inside
• 2 separate SequentialScanDiscreteMetropolis steps
• which makes list of 3 steps in total

but with njobs>1 the methods list in the outermost CompundStep contains:

• 4 separate Metropolis steps
• which makes list of 4 steps in total.

[abojchevski]

Even more so, I tried a different configuration of steps:

• Metropolis
• 2 x BinaryGibbsMetropolis
• NUTS

the CompoundStep ends up having 4 separate Metropolis steps regardless.

[abojchevski]

Additional information:

It seems that in BlockedStep when __new__ is called for the first variable
this else part get executed:

 else: # Assume all model variables
            vars = model.vars

and then we go into if not blocked and len(vars) > 1: and to every single var we assign the step of the first var, which in my case happens to be Metropolis

[fonnesbeck]

OK, I thought joblib took care of that, but apparently not. I will have a look. Thanks for reporting

[benavente]

Hi. Is there any workaround for this?
Thanks!

[benavente]

I found the problem. When a BlockedStep is unpickled after being transmitted to the new process spawned, pickle initializes the class like so (https://docs.python.org/3.5/library/pickle.html#pickling-class-instances):

def load(cls, attributes):
    obj = cls.__new__(cls)
    obj.__dict__.update(attributes)
    return obj

That is, pickle calls the __new__ method again and with no arguments, so both args and kwargs is empty and all the model variables are assumed like @abojchevski pointed out.

I did a very dirty workaround (https://github.com/benavente/pymc3/commit/8d67245da83ee8a9a4323289a69633567a336389) by saving the arguments in an instance variable and returning them through the __getnewargs_ex__ method (https://docs.python.org/3.5/library/pickle.html#object.__getnewargs_ex__). However, this only works in Python 3, because the Python 2 pickle protocol does not allow specifying keyword arguments in __getnewargs__, and it's not an ideal solution...

[twiecki]

@benavente Great you figured it out. Can you think of a better way to fix it? If not, the code diff doesn't look too bad for the problem it solves. And who uses python2 anyhow? ;)

[benavente]

@twiecki I created a pull request #1127 if you want to add it.

[twiecki]

Thanks!

[abojchevski]

I can confirm that the workaround works for me too. Good job @benavente.

There is however one caveat (related to njobs but not to the workaround) which might need to be documented. That is if you have a custom distubution defined with a lambda function then the pickling does not work:

e.g.

        A = pm.DensityDist(
            name='A',
            logp=lambda value: pm.Bernoulli.dist(p=0.5).logp(value))

And you would get an error like this one:

AttributeError: Can't pickle local object 'usbm.<locals>.<lambda>'

Also, I am not sure if running multiple njobs is supposed to be faster (i.e. draw your same X number samples just in parallel) or just run multiple separate chains.

[twiecki]

Thanks for confirming @abojchevski. Can we catch that error somewhere smart and display a better error message?

[abojchevski]

Perhaps in _mp_sample?