Hello,
I am relatively new to PyTorch and its ecosystem. I have been porting previous code over to this framework and so far I have been using GPyTorch for Bayesian Optimization. At some point, I get a callable that is differentiable and represents an unnormalized pdf from which I wish to sample. I figured that I could use the Pyro MCMC samplers to tackle this problem.
The Pyro samplers require the specification of a callable containing Pyro primitives (Pyro model) and it was not obvious how to do such thing. I was able to solve this issue using the hacks suggested here.
Regarding reproducibility, I would like to be able to obtain the same traces upon specification of an initial seed and an initial sample (MAP estimate). After inspecting the code for pyro.infer.mcmc.mcmc, I see that I have to redefine the init method of _Worker so as to replace torch.initial_seed(). This would work for the _ParallelSampler. How about the _SingleSampler? I guess I would have to include pyro.set_rng_seed(my_seed) somewhere, possibly in the _traces method? Regarding my initial sample (initial trace), after initializing the kernel (hmc or nuts), I think I only need to set the property initial_trace in kernel. Any tips?
We are currently working on a more general interface to HMC (see Improvements to the HMC interface · Issue #1816 · pyro-ppl/pyro · GitHub) so that you will be able to generate samples from your target distribution by simply specifying the potential_energy function (negative unnormalized log density in your case), instead of having to pass in a callable that contains Pyro primitives.
I guess I would have to include pyro.set_rng_seed(my_seed) somewhere, possibly in the _traces method? Regarding my initial sample (initial trace), after initializing the kernel (hmc or nuts), I think I only need to set the property initial_trace in kernel. Any tips?
I think calling pyro.set_rng_seed(x) at the beginning of your script is a great idea in general, and it should work for your use case as well. To be safe, you can just call this right before you call the .run method. You should get deterministic traces from MCMC if you do so. There shouldn’t be a need to place this inside the MCMC class. The reason why _ParallelSampler needs to do this inside the class is because the global variables are not shared across processes, and this needs to be done do get deterministic results with multiprocessing.
I see that the interface has been changed and now allows the specification of potential_fn as an alternative to model. I am however still passing a model and in this case my initial_params dict that I specify before running the MCMC is replaced by some other dict that is created in _get_init_params inside initialize_model. Wouldn’t it be more reasonable to only set this as the default if initial_params has not yet been set in the HMC kernel? Is there a way to set it without having to modify the base code (specifically HMC._initialize_model_properties?
1 Like
Good point, thanks for pointing this out! We are going to be actively refactoring MCMC classes over the next several weeks, during which the interface will continue to change. e.g. the model argument is still there to preserve backwards compatibility, but there is complexity involved in supporting both potential_fn and model and what you experienced is a resulting bug. We’ll fix this shortly. If you are looking for a more well tested interface, my suggestion would be to use the latest 0.3.3 release which should have all the functionality. If you are willing to be a bit more adventurous and help us find bugs like this, you are of course more than welcome to work off of the dev branch! 
Sorry, I have not kept backward compatibility for setting initial_params with model. For the time being, I guess you can do
class FixedHMC(HMC):
initial_params = None
then work with FixedHMC instead of HMC.
I managed to fix this by explicitly calling pyro.infer.mcmc.util.initialize_model. Something like:
_, potential_fn, transforms, model_trace = initialize_model(my_pyro_model)
mcmc_kernel = NUTS(None, potential_fn=potential_fn, transforms=transforms)
mcmc_kernel.initial_params = my_init_params
mcmc_run = MCMC(mcmc_kernel, num_samples=my_num_samples).run()
The difference is that now the mcmc_run.exec_traces is a bit different and I do have to explicitly transform the parameters back to the constrained space.
Yea, your fix is the practice we would recommend the next release. You are right that the result will be different w.r.t. manual setting initial_trace because I haven’t thought about this case. But we have verified that the algorithm is the same as before. I hope that the difference due to the order of setting manual initial_params is not important for you.
About transforms, we should transform back to normal space automatically. Thank you so much for catching that! Let me make a quick fix for it.
p/s: If you have further suggestions, please let us know. About the reproducibility which you suggested in the main topic, we will address them for the next release.