Better support for nested inference with SVI

pyro-ppl / pyro

Deep universal probabilistic programming with Python and PyTorch

http://pyro.ai

Apache License 2.0

8.5k stars 982 forks source link

Better support for nested inference with SVI #1227

Open eb8680 opened 6 years ago

eb8680 commented 6 years ago

The internals of our SVI code need a bit of cleanup to better support nesting SVI inside other inference algorithms, a necessary component of several projects including contrib.oed.

Tasks:

[x] Make infer.SVI a TracePosterior instance again so that it works with EmpiricalMarginal and can be interchanged with other methods like MCMC
[ ] Have step update the parameter store with new values instead of updating parameter tensors in place as pyro.optim.PyroOptim optimizers do

fritzo commented 6 years ago

Note regarding #628 and https://forum.pyro.ai/t/checkpoint-in-pyro/303: There are some disadvantages of using grad instead of .backward() (at least in Trace_ELBO):

grad has less jit support at the moment
grad has no checkpointing support at the moment
grad is slightly more expensive for tiny tensors (since it traverses the graph twice)

However these are minor and could be outweighed by code simplification.