diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py index f85d6f854..f9756dcc3 100644 --- a/numpyro/infer/svi.py +++ b/numpyro/infer/svi.py @@ -403,12 +403,16 @@ def body_fn(svi_state, _): svi_state = init_state if progress_bar: losses = [] + jitted_body_fn = jit(body_fn) with tqdm.trange(1, num_steps + 1) as t: batch = max(num_steps // 20, 1) for i in t: - svi_state, loss = jit(body_fn)(svi_state, None) - losses.append(jax.device_get(loss)) + svi_state, loss = jitted_body_fn(svi_state, None) + losses.append(loss) if i % batch == 0: + # transfer the whole batch to the host at once instead of + # blocking on the device at every step + losses[i - batch :] = jax.device_get(losses[i - batch :]) if stable_update: valid_losses = [x for x in losses[i - batch :] if x == x] num_valid = len(valid_losses) diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py index daacc200d..c1318bfe8 100644 --- a/test/infer/test_svi.py +++ b/test/infer/test_svi.py @@ -282,6 +282,33 @@ def guide(data): ) +def test_run_progress_bar_matches_scan_path(): + data = jnp.array([1.0] * 8 + [0.0] * 2) + + def model(data): + f = numpyro.sample("beta", dist.Beta(1.0, 1.0)) + with numpyro.plate("N", len(data)): + 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()) + # use a step count that is not a multiple of the progress-bar update batch + # so that the last losses are collected outside a batch boundary + result_pbar = svi.run(random.key(1), 123, data, progress_bar=True) + result_scan = svi.run(random.key(1), 123, data, progress_bar=False) + assert result_pbar.losses.shape == result_scan.losses.shape == (123,) + assert result_pbar.losses.dtype == result_scan.losses.dtype + # the two paths run differently compiled programs (jitted python loop vs + # lax.scan), so they can differ by a few ulps + assert_allclose(result_pbar.losses, result_scan.losses, rtol=1e-5) + for name in result_scan.params: + assert_allclose(result_pbar.params[name], result_scan.params[name], rtol=1e-5) + + def test_jitted_update_fn(): data = jnp.array([1.0] * 8 + [0.0] * 2)