African countries regression on ruggeness of terrain: porting to MCMC

Hi,

I am working my way through the pyro tutorials and I am trying to use a MCMC method
instead of SVI in the linear regression example.
I use version 1.7.0 for pyro.

Just in case, i modified the original model to make it explicitely callable (not sure it is necessary).

class BayesianRegression(PyroModule):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = PyroModule[nn.Linear](in_features, out_features)
        self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
        self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))
    def forward(self, x, y=None):
        sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
        mean = self.linear(x).squeeze(-1)
        with pyro.plate("data", x.shape[0]):
            obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
        return mean
    def __call__(self, x, y=None):
        return self.forward(x,y)

I checked that it does run with the SVI algorithm.
Now, I am just calling the MCMC method:

import pyro.distributions as dist
from pyro.infer import MCMC, NUTS, HMC

#kernel = NUTS(BayesianRegression, jit_compile=True, ignore_jit_warnings=True, max_tree_depth=3)
kernel = NUTS(BayesianRegression, jit_compile=True)
posterior1 = MCMC(kernel, num_samples=7000, warmup_steps=800)
posterior1.run(x_data, y_data);
  # Traitement des résulats
hmc_samples = {k: v.detach().cpu().numpy() for k, v in posterior1.get_samples().items()}

I get the following error:
empty(): argument ‘size’ must be tuple of ints, but found element of type Tensor at pos 1
which I managed to work around thanks to the answer below.

Now I am trying to get the predictions:
from pyro.infer import Predictive

def summary_mcmc(samples):
    site_stats = {}
    for k, v in samples.items():
        site_stats[k] = {
            "mean": torch.mean(v, 0).detach().cpu().numpy(),
            "std": torch.std(v, 0).detach().cpu().numpy(),
            "5%": v.kthvalue(int(len(v) * 0.05), dim=0)[0].detach().cpu().numpy(),
            "95%": v.kthvalue(int(len(v) * 0.95), dim=0)[0].detach().cpu().numpy(),
        }
    return site_stats

hmc_samples_mcmc = posterior1.get_samples()
predictive_mcmc = Predictive(model, hmc_samples_mcmc, num_samples=300,
                        return_sites=("linear.weight", "obs", "_RETURN"))
samples = predictive_mcmc(x_data)
pred_summary = summary_mcmc(samples)

but I get an error on samples=predictive_mcmc(x_data)
shape ‘[300, 1, 1, 3]’ is invalid for input of size 153000
Any suggestion as how to work around this one ?

Here is a work around:

def model_mcmc(x_data, y_data):
    reg = BayesianRegression( x_data.shape[1], y_data.shape[0])
    return reg(x_data, y_data)

import pyro.distributions as dist
from pyro.infer import MCMC, NUTS, HMC


kernel = NUTS(model_mcmc, jit_compile=True)
posterior = MCMC(kernel, num_samples=400, warmup_steps=300)
posterior.run(x_data, y_data);

i recommend you write your model using sample statements like is done here. PyroModules are probably a better fit for optimization based inference like SVI.

This is a valuable insight. I shall follow your advice.