Prior definition and accessing individual ELBO loss terms

Hello community,

I am new in using Pyro and I have a few doubts I would like to clarify. I am refracting my code which implements SVI and ELBO manually (in PyTorch) and I would like to adopt Pyro as it offers a unified interface to many variation inference methods. I have a few “conceptual” questions:

  1. The prior over latent variable(s) is defined in the model and it shares the same name as the variational posterior distribution implemented in the guide. Pyro automatically fetches the prior by looking at name shared between model and guide. Is this correct?

  2. Can I access to individual loss terms inside the ELBO (Trace_ELBO) loss, that is fitting term and regularisation term?

  3. Can I save some user defined attributes when the model is run to save some metrics that are irrelevant to Pyro but important for the user such as training accuracy…? This avoids me to run again do an additional inference after svi.step().

Thanks in advance for clarifications and help,

  1. Correct, Pyro matches the guide and model sites by site name.
  2. It is nontrivial to access individual loss terms in the Trace_ELBO. To do so you’d want to inspect the model_trace: regularization terms are attached to latent sample sites and fitting terms are attached to observed sites (site["is_observed"] is True). But the actual terms are quite tricky due to reparametrization trick, sticking the landing, etc.

You might take a look at minipyro to get a sense of the loss terms.

  1. Yes you can save user defined attributes when the model is run. However beware that Pyro uses nonstandard to run the model in multiple different ways (e.g. running vectorized, running with enumeration, sometimes running multiple times per svi step to determine tensor shapes). So do take care to inspect any data you save each iteration; you might want to save an svi step with each iteration and check some of the tensor shapes to get comfortable with the nonstandard evaluation going on.

Good luck!