Repository: pyro-ppl/numpyro

Issue:
Simplify training phase of SVI
Currently, to train SVI, users need to know either
+ `jax.lax.fori_loop`
+ `jax.lax.scan`
+ or `jax.jit`

Either way, users code will look like
```
svi = SVI(model, guide, ...)
svi_state = svi.init(rng_key)
def body_fn:
    ...
svi_state = fori_loop/scan/python loop
params = svi.get_params(svi_state)
```

To me, this is unfortunate and in most of the cases (given the fact that we already support subsampling via `numpyro.subsample`), the model input will be constant. I think it is better to support a high-level api for the above job. For example, we can add a method `svi.run` (which mimics `mcmc.run`) which do all the job, kind of
```
svi = SVI(model, guide, ...)
svi.run(rng_key, num_steps, *args, progress_bar=True, **kwargs)
params = svi.get_params()
losses = svi.get_losses()
```

With this, we
+ need not to know what svi_state is
+ need not to know how to use those jax transforms
+ switch between with progress_bar and no progress_bar easily
+ can see that the code is more intuitively


Discussion:

Resolving patch:
diff --git a/examples/neutra.py b/examples/neutra.py
--- a/examples/neutra.py
+++ b/examples/neutra.py
@@ -21,7 +21,7 @@
 import matplotlib.pyplot as plt
 import seaborn as sns
 
-from jax import lax, random
+from jax import random
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
 
@@ -67,17 +67,15 @@ def main(args):
 
     guide = AutoBNAFNormal(dual_moon_model, hidden_factors=[args.hidden_factor, args.hidden_factor])
     svi = SVI(dual_moon_model, guide, optim.Adam(0.003), Trace_ELBO())
-    svi_state = svi.init(random.PRNGKey(1))
 
     print("Start training guide...")
-    last_state, losses = lax.scan(lambda state, i: svi.update(state), svi_state, jnp.zeros(args.num_iters))
-    params = svi.get_params(last_state)
+    svi_result = svi.run(random.PRNGKey(1), args.num_iters)
     print("Finish training guide. Extract samples...")
-    guide_samples = guide.sample_posterior(random.PRNGKey(2), params,
+    guide_samples = guide.sample_posterior(random.PRNGKey(2), svi_result.params,
                                            sample_shape=(args.num_samples,))['x'].copy()
 
     print("\nStart NeuTra HMC...")
-    neutra = NeuTraReparam(guide, params)
+    neutra = NeuTraReparam(guide, svi_result.params)
     neutra_model = neutra.reparam(dual_moon_model)
     nuts_kernel = NUTS(neutra_model)
     mcmc = MCMC(nuts_kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
@@ -111,7 +109,7 @@ def main(args):
     ax5 = fig.add_subplot(gs[0, 2])
     ax6 = fig.add_subplot(gs[1, 2])
 
-    ax1.plot(losses[1000:])
+    ax1.plot(svi_result.losses[1000:])
     ax1.set_title('Autoguide training loss\n(after 1000 steps)')
 
     ax2.contourf(X1, X2, P, cmap='OrRd')
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
--- a/numpyro/contrib/module.py
+++ b/numpyro/contrib/module.py
@@ -202,19 +202,13 @@ def apply(self, x):
         >>> guide = autoguide.AutoNormal(model, init_loc_fn=init_to_feasible)
         >>> svi = SVI(model, guide, numpyro.optim.Adam(5e-3), TraceMeanField_ELBO())
         >>>
-        >>> batch_size = 256
         >>> n_iterations = 3000
-        >>> svi_state = svi.init(random.PRNGKey(0), x_train, y_train, batch_size=batch_size)
-        >>> update_fn = jit(svi.update, static_argnums=(3,))
-        >>> for i in tqdm.trange(n_iterations):
-        ...     svi_state, loss = update_fn(svi_state, x_train, y_train, batch_size)
-        >>>
-        >>> params = svi.get_params(svi_state)
+        >>> params, losses = svi.run(random.PRNGKey(0), n_iterations, x_train, y_train, batch_size=256)
         >>> n_test_data = 100
         >>> x_test, y_test = generate_data(n_test_data)
         >>> predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
         >>> y_pred = predictive(random.PRNGKey(1), x_test[:100])["obs"].copy()
-        >>> assert loss < 3000
+        >>> assert losses[-1] < 3000
         >>> assert np.sqrt(np.mean(np.square(y_test - y_pred))) < 1
     """
     nn = flax_module(name, nn_module, input_shape=input_shape)
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -3,7 +3,9 @@
 
 from functools import namedtuple, partial
 
-from jax import random
+from jax import jit, lax, random
+import jax.numpy as jnp
+import tqdm
 
 from numpyro.distributions import constraints
 from numpyro.distributions.transforms import biject_to
@@ -18,6 +20,14 @@
 """
 
 
+SVIRunResult = namedtuple('SVIRunResult', ['params', 'losses'])
+"""
+A :func:`~collections.namedtuple` consisting of the following fields:
+ - **params** - the optimized parameters.
+ - **losses** - the losses collected at every step.
+"""
+
+
 def _apply_loss_fn(loss_fn, rng_key, constrain_fn, model, guide,
                    args, kwargs, static_kwargs, params):
     return loss_fn(rng_key, constrain_fn(params), model, guide, *args, **kwargs, **static_kwargs)
@@ -36,7 +46,7 @@ class SVI(object):
 
     .. doctest::
 
-        >>> from jax import lax, random
+        >>> from jax import random
         >>> import jax.numpy as jnp
         >>> import numpyro
         >>> import numpyro.distributions as dist
@@ -56,11 +66,8 @@ class SVI(object):
         >>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])
         >>> optimizer = numpyro.optim.Adam(step_size=0.0005)
         >>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
-        >>> init_state = svi.init(random.PRNGKey(0), data)
-        >>> state = lax.fori_loop(0, 2000, lambda i, state: svi.update(state, data)[0], init_state)
-        >>> # or to collect losses during the loop
-        >>> # state, losses = lax.scan(lambda state, i: svi.update(state, data), init_state, jnp.arange(2000))
-        >>> params = svi.get_params(state)
+        >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)
+        >>> params = svi_result.params
         >>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])
 
     :param model: Python callable with Pyro primitives for the model.
@@ -82,15 +89,14 @@ def __init__(self, model, guide, optim, loss, **static_kwargs):
 
     def init(self, rng_key, *args, **kwargs):
         """
+        Gets the initial SVI state.
 
         :param jax.random.PRNGKey rng_key: random number generator seed.
         :param args: arguments to the model / guide (these can possibly vary during
             the course of fitting).
         :param kwargs: keyword arguments to the model / guide (these can possibly vary
             during the course of fitting).
-        :return: tuple containing initial :data:`SVIState`, and `get_params`, a callable
-            that transforms unconstrained parameter values from the optimizer to the
-            specified constrained domain
+        :return: the initial :data:`SVIState`
         """
         rng_key, model_seed, guide_seed = random.split(rng_key, 3)
         model_init = seed(self.model, model_seed)
@@ -114,7 +120,8 @@ def get_params(self, svi_state):
         """
         Gets values at `param` sites of the `model` and `guide`.
 
-        :param svi_state: current state of the optimizer.
+        :param svi_state: current state of SVI.
+        :return: the corresponding parameters
         """
         params = self.constrain_fn(self.optim.get_params(svi_state.optim_state))
         return params
@@ -137,6 +144,51 @@ def update(self, svi_state, *args, **kwargs):
         loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)
         return SVIState(optim_state, rng_key), loss_val
 
+    def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):
+        """
+        (EXPERIMENTAL INTERFACE) Run SVI with `num_steps` iterations, then return
+        the optimized parameters and the stacked losses at every step. If `num_steps`
+        is large, setting `progress_bar=False` can make the run faster.
+
+        .. note:: For a complex training process (e.g. the one requires early stopping,
+            epoch training, varying args/kwargs,...), we recommend to use the more
+            flexible methods :meth:`init`, :meth:`update`, :meth:`evaluate` to
+            customize your training procedure.
+
+        :param jax.random.PRNGKey rng_key: random number generator seed.
+        :param int num_steps: the number of optimization steps.
+        :param args: arguments to the model / guide
+        :param bool progress_bar: Whether to enable progress bar updates. Defaults to
+            ``True``.
+        :param kwargs: keyword arguments to the model / guide
+        :return: a namedtuple with fields `params` and `losses` where `params`
+            holds the optimized values at :class:`numpyro.param` sites,
+            and `losses` is the collected loss during the process.
+        :rtype: SVIRunResult
+        """
+        def body_fn(svi_state, carry):
+            svi_state, loss = self.update(svi_state, *args, **kwargs)
+            return svi_state, loss
+
+        svi_state = self.init(rng_key, *args, **kwargs)
+        if progress_bar:
+            losses = []
+            with tqdm.trange(1, num_steps + 1) as t:
+                batch = num_steps // 20
+                for i in t:
+                    svi_state, loss = jit(body_fn)(svi_state, None)
+                    if i % batch == 0:
+                        avg_loss = sum(losses[i-batch:]) / batch
+                        t.set_postfix_str("init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}"
+                                          .format(losses[0], i - batch + 1, i, avg_loss),
+                                          refresh=False)
+                    losses.append(loss)
+            losses = jnp.stack(losses)
+        else:
+            svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)
+
+        return SVIRunResult(self.get_params(svi_state), losses)
+
     def evaluate(self, svi_state, *args, **kwargs):
         """
         Take a single step of SVI (possibly on a batch / minibatch of data).


Test patch:
diff --git a/test/test_svi.py b/test/test_svi.py
--- a/test/test_svi.py
+++ b/test/test_svi.py
@@ -70,6 +70,27 @@ def body_fn(i, val):
     assert_allclose(params['alpha_q'] / (params['alpha_q'] + params['beta_q']), 0.8, atol=0.05, rtol=0.05)
 
 
+@pytest.mark.parametrize('progress_bar', [True, False])
+def test_run(progress_bar):
+    data = jnp.array([1.0] * 8 + [0.0] * 2)
+
+    def model(data):
+        f = numpyro.sample("beta", dist.Beta(1., 1.))
+        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+
+    def guide(data):
+        alpha_q = numpyro.param("alpha_q", 1.0,
+                                constraint=constraints.positive)
+        beta_q = numpyro.param("beta_q", 1.0,
+                               constraint=constraints.positive)
+        numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
+
+    svi = SVI(model, guide, optim.Adam(0.05), Trace_ELBO())
+    params, losses = svi.run(random.PRNGKey(1), 1000, data, progress_bar=progress_bar)
+    assert losses.shape == (1000,)
+    assert_allclose(params['alpha_q'] / (params['alpha_q'] + params['beta_q']), 0.8, atol=0.05, rtol=0.05)
+
+
 def test_jitted_update_fn():
     data = jnp.array([1.0] * 8 + [0.0] * 2)