Gpytorch: [Docs] Need help in loading a Full Bayesian model that uses Pyro for MCMC

[natarajanmolecule]

I am trying to build a fully Bayesian GP with MCMC sampled parameters using Pyro based on the documentation. When I try to save and then load the model (using the save/load model documentation) later, I get an error about size mismatch:

RuntimeError: Error(s) in loading state_dict for ExactGPModel:
size mismatch for likelihood.noise_covar.raw_noise: copying a param with shape torch.Size([100, 1]) from checkpoint, the shape in current model is torch.Size([1]).
size mismatch for mean_module.constant: copying a param with shape torch.Size([100, 1]) from checkpoint, the shape in current model is torch.Size([1]).
size mismatch for covar_module.base_kernel.raw_lengthscale: copying a param with shape torch.Size([100, 1, 13]) from checkpoint, the shape in current model is torch.Size([1, 13]).

In this case, my dataset has 13 features, and I have run a NUTS sampler for 100 samples when i originally built the model.
I think this is because when I instantiate ExactGPModel as indicated in the documentation, it is allocating only one sample for the different hyperparameters. Any thoughts on how to get around this error? It will also be great to include such an example in the saving and loading documentation.

[jacobrgardner]

@natarajanmolecule Ah, yeah. This is definitely an oversight on my part -- didn't even think about saving and loading when I implemented the sampling :-).

I'll put up a PR today or tomorrow that fixes this, and add an example saving and loading at the end of the sampling notebook!

[natarajanmolecule]

That will be awesome @jacobrgardner ! Thank you for the help. I will look forward to the PR.

[natarajanmolecule]

Hi @jacobrgardner I wanted to follow up with you on this issue. Please keep me posted and thanks so much for your help.

[jacobrgardner]

Hi @natarajanmolecule Sorry for the wait -- I'll try to get the PR that fixes this up by the end of the day.

[jacobrgardner]

@natarajanmolecule Actually I might need some more info from you. I finally got around to playing with this today, and the following works for me (in the basic fully Bayesian example notebook):

mcmc_run.run(train_x, train_y)
model.pyro_load_from_samples(mcmc_run.get_samples())
model.load_state_dict(model.state_dict())

Are you doing anything in particular different from the example notebook?

[natarajanmolecule]

Yes, this works for me too. But when I save the model using the model.save(), and then in a new session (where I dont have the mcmc_run object) when I instantiate the model, and try to load the dict I get the error.

[jacobrgardner]

Got it. It's because pyro_load_from_samples correctly updates the batch_shape and size of each parameter, but when you reinstantiate the model and try to load the parameters are batched (e.g., num_samples batch) but the current model parameters are not.

[natarajanmolecule]

I see - what should the solution be then? How shall i instantiate the model with relevant samples beforehand or based on the loaded state dictionary?

[jacobrgardner]

@natarajanmolecule This isn't a problem on your end, it's just some slight complication in fixing it on ours. We are targeting to release a new minor version of GPyTorch by the end of the week that will include a fix for this.

If you absolutely need this to work before the end of the week, when you create your GP the second time (when you intend on loading) you could pass in batch_shape=torch.Size([num_samples]) to your base kernel, your scale kernel and your likelihood.

[natarajanmolecule]

Hi @jacobrgardner I wanted to follow up on this again. Actually I tried the batch_shape option, but that also resulted in an error. so it will be very helpful if you can help me figure this out.

[jacobrgardner]

Finally got around to putting up a PR for this. Sorry for the delay -- been travelling a lot recently.

Once it's merged in, everything should work on master and I believe we are planning to cut a new release soon.