Error in multiprocessing MCMC with potential_fn (Windows)

Misc.
1

Hi,

I’ve recently tested the option to run more than 1 chain and I am getting TypeError: can't pickle weakref objects in ForkingPickler(file, protocol).dump(obj), when trying to pickle the multiprocessing.context.Process class. I am only getting this error if I pass to the MCMC kernel the potential_fn. Has this been tested?

As a minimal working example, I have used the fully_pooled model in examples/baseball.py. The difference is that I convert the model to a potential function first.

init_params, potential_fn, transforms, model_trace = pyro.infer.mcmc.util.initialize_model(fully_pooled, model_args=(at_bats, hits))
nuts_kernel = NUTS(None, potential_fn=potential_fn, transforms=transforms)
nuts_kernel.initial_params = init_params
posterior_fully_pooled = MCMC(nuts_kernel, num_samples=200, warmup_steps=100, num_chains=2).run()

It works fine if num_chains=1 or if I use fully_pooled directly in NUTS.

2

What OS are you on? Note that multiprocessing support on Windows is limited. This seems to work fine for me on mac or linux.

3

Yes, I confirm that it works on Linux. At the time I was testing this on a Windows machine.
I found a hack to make it run on Windows, but I haven’t yet made sure this completely fixes the problem. It basically consisted in replacing the pickle package by the dill package in the multiprocessing module that pickles objects. It’s possible that there is a better solution, I just don’t know what it would be.