Pyro: Remove dependency on TracePosterior and deprecate utilities in AbstractInfer

[neerajprad]

While attempting the MCMC class refactoring (as originally suggested in #1725), I realized that it will be best to not subclass off of TracePosterior.

• For MCMC, storing traces leads to inefficient memory usage where the full trace can often consume 5X more memory. This also leads to our internal trace structure leaking into the API, and makes it harder to integrate with other libraries. e.g. NUTS integration with GPyTorch.
• For SVI, once we have the trained parameters in the param store, we can just use these each time we need to generate traces, do predictions etc. rather than maintaining a bag of traces. This has often led to confusion that users have surfaced through the forum. example.

TODOs:

Next Release:

• [x] Refactor MCMC module to not subclass off of TracePosterior, and bring the interface closer to NumPyro, deprecate existing API. #1913
• [x] Replacement for TracePredictive for MCMC.

Next Release:

• [x] Deprecate .run method in SVI module to plan for removal
  (Note: raise UserWarning rather than DeprecationWarning)
• [x] Replacement for TracePredictive for SVI (this additionally needs to handle data subsampling).

Following Release:

• [ ] Removing these code paths completely.

[fritzo]

@neerajprad is any part of this issue blocking the 0.4 release?

[neerajprad]

Deprecating TracePosterior/TracePredictive for SVI should be punted for later, since it will involve changes in other tutorials, OED module, and may need some discussion.

In the meantime, I will try to see if we can support SVI by making some minor changes to mcmc's predictive utility. As this utility becomes mature, that will make it easier for us to deprecate the old interface.

[neerajprad]

@fehiepsi - Since you volunteered helping with the predictive utility, I have added you to this task. Note that we already have a predictive utility in mcmc.util. We need to figure out a more general API so that it works with SVI too, and which will allow us to remove TracePredictive completely.

[fehiepsi]

Sure, I am happy to work on this. I'll try to make a PR which covers usage cases which I know, then we can discuss there how to support more general cases, which I believe that you have some experience with.

[neerajprad]

I think we need to allow for a guide argument and figure out predictions when data subsamples are differently sized between training and testing. @ahmadsalim has already handled this in TracePredictive and we can use that as a template. It doesn't necessarily need to be merged with mcmc.util.predictive as a single function, but I think the internals will both share a significant amount of code.

[fehiepsi]

@neerajprad @ahmadsalim I have thought about the subsample stuff deeply but couldn't figure out in which case the situation becomes complicated. Could you help me identify which test in test_abstract_infer which I need to pay attention to?

I have been using predictive utility in NumPyro which has plate with different sized between training and testing and saw no problem with it. Is the following the way we use predictive?

def model(X, y=None):
    pyro.plate("plate1", X.shape[0]):
        pyro.sample("obs", dist.Normal(), obs=y)

def guide(X, y):
   ...

svi = SVI(...)
svi.step(X, y)
posterior = predictive(guide, {}, X_new)
posterior_predictive = predictive(model, posterior, X_new, return_sites=["obs"])

[neerajprad]

I think we should be able to just run the guide forward on the new data without worrying about differences in plate sizes. IIRC, we only had to do subsampling when we were storing posterior data in exec_traces, so that we didn't need to run the guide again. @ahmadsalim - correct me if I am mistaken.

@fehiepsi - I think your simplified solution is better. Let us just go ahead with that instead. Could you add @ahmadsalim to the PR, who had a few models that were using TracePredictive and it will be nice if we can ensure that these continue to work well with the refactoring?

[ahmadsalim]

@fehiepsi Please, see https://github.com/pyro-ppl/pyro/blob/dev/tests/infer/test_abstract_infer.py#L102-L103 and https://github.com/pyro-ppl/pyro/blob/dev/tests/infer/test_abstract_infer.py#L47-L48 . You can see that the batch size changes, because we use num_trials[:3] which changes the size between inference time and prediction time 😄

[fehiepsi]

Yeah, I just noticed it. I'll make a PR soon then we can discuss those situations there. :)

[neerajprad]

Removing the 1.0 milestone since all tasks are done. We should prioritize removing deprecated code in a future release.