Running my model with parameters from trace, after HMC inference

I have implemented a Bayesian Neural Network similar to this: bnn-model.py · GitHub

Then I sample a trace with NUTS and MCMC, like so:
nuts_kernel = NUTS(model, adapt_step_size=True)
mcmc_run = MCMC(nuts_kernel, num_samples=100, warmup_steps=30).run(inp, outp)

Now that I have a trace of sampled weights for my BNN, I would like to run my model on new data, with parameters from my trace. This feels like something that should be easy to do, but I can’t quite find how to do this in the documentation. What is the canonical way?

I have figured out a way of doing it, but I don’t know if this unnecessarily complicated.

  1. Unlike the linked example, have your model-function return a value with .sample().
  2. Create a TracePredictive object and run that on your data (X below). This seems to create a new type of trace.
  3. Create an EmpiricalMarginal from your new trace.
  4. Sample from that
    tp = TracePredictive(model, mcmc_run, 100)
    em = EmpiricalMarginal(tp.run(X, None))
    em.sample()

It seems to me that TracePredictive, which samples N times from the MCMC trace, is an odd step, but perhaps this is the right way?

You should be able to sample from the TracePredictive directly as:

tp = TracePredictive(model, mcmc_run, num_samples=100)
sample = tp()

The resampling is needed since you may have K samples from MCMC, and might want to draw N samples from your predictive distribution where N could more more than K. I will be refactoring the TracePredictive class, so please feel free to add any suggestions on improving the interface in Make AbstractInfer classes lazily consume traces · Issue #1725 · pyro-ppl/pyro · GitHub.