def test_gaussian_hmm_elbo(batch_shape, num_steps, hidden_dim, obs_dim): init_dist = random_mvn(batch_shape, hidden_dim) trans_mat = torch.randn(batch_shape + (num_steps, hidden_dim, hidden_dim), requires_grad=True) trans_dist = random_mvn(batch_shape + (num_steps, ), hidden_dim) obs_mat = torch.randn(batch_shape + (num_steps, hidden_dim, obs_dim), requires_grad=True) obs_dist = random_mvn(batch_shape + (num_steps, ), obs_dim) data = obs_dist.sample() assert data.shape == batch_shape + (num_steps, obs_dim) prior = dist.GaussianHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist) likelihood = dist.Normal(data, 1).to_event(2) posterior, log_normalizer = prior.conjugate_update(likelihood) def model(data): with pyro.plate_stack("plates", batch_shape): z = pyro.sample("z", prior) pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data) def guide(data): with pyro.plate_stack("plates", batch_shape): pyro.sample("z", posterior) reparam_model = poutine.reparam(model, {"z": ConjugateReparam(likelihood)}) def reparam_guide(data): pass elbo = Trace_ELBO(num_particles=1000, vectorize_particles=True) expected_loss = elbo.differentiable_loss(model, guide, data) actual_loss = elbo.differentiable_loss(reparam_model, reparam_guide, data) assert_close(actual_loss, expected_loss, atol=0.01) params = [trans_mat, obs_mat] expected_grads = torch.autograd.grad(expected_loss, params, retain_graph=True) actual_grads = torch.autograd.grad(actual_loss, params, retain_graph=True) for a, e in zip(actual_grads, expected_grads): assert_close(a, e, rtol=0.01)
def test_beta_binomial_elbo(): total = 10 counts = dist.Binomial(total, 0.3).sample() concentration1 = torch.tensor(0.5, requires_grad=True) concentration0 = torch.tensor(1.5, requires_grad=True) prior = dist.Beta(concentration1, concentration0) likelihood = dist.Beta(1 + counts, 1 + total - counts) posterior = dist.Beta(concentration1 + counts, concentration0 + total - counts) def model(): prob = pyro.sample("prob", prior) pyro.sample("counts", dist.Binomial(total, prob), obs=counts) def guide(): pyro.sample("prob", posterior) reparam_model = poutine.reparam(model, {"prob": ConjugateReparam(likelihood)}) def reparam_guide(): pass elbo = Trace_ELBO(num_particles=10000, vectorize_particles=True, max_plate_nesting=0) expected_loss = elbo.differentiable_loss(model, guide) actual_loss = elbo.differentiable_loss(reparam_model, reparam_guide) assert_close(actual_loss, expected_loss, atol=0.01) params = [concentration1, concentration0] expected_grads = torch.autograd.grad(expected_loss, params, retain_graph=True) actual_grads = torch.autograd.grad(actual_loss, params, retain_graph=True) for a, e in zip(actual_grads, expected_grads): assert_close(a, e, rtol=0.01)
def test_stable_hmm_smoke(batch_shape, num_steps, hidden_dim, obs_dim): init_dist = random_stable(batch_shape + (hidden_dim, )).to_event(1) trans_mat = torch.randn(batch_shape + (num_steps, hidden_dim, hidden_dim), requires_grad=True) trans_dist = random_stable(batch_shape + (num_steps, hidden_dim)).to_event(1) obs_mat = torch.randn(batch_shape + (num_steps, hidden_dim, obs_dim), requires_grad=True) obs_dist = random_stable(batch_shape + (num_steps, obs_dim)).to_event(1) data = obs_dist.sample() assert data.shape == batch_shape + (num_steps, obs_dim) def model(data): hmm = dist.LinearHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=num_steps) with pyro.plate_stack("plates", batch_shape): z = pyro.sample("z", hmm) pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data) # Test that we can combine these two reparameterizers. reparam_model = poutine.reparam( model, { "z": LinearHMMReparam(StableReparam(), StableReparam(), StableReparam()), }, ) reparam_model = poutine.reparam( reparam_model, { "z": ConjugateReparam(dist.Normal(data, 1).to_event(2)), }, ) reparam_guide = AutoDiagonalNormal( reparam_model) # Models auxiliary variables. # Smoke test only. elbo = Trace_ELBO(num_particles=5, vectorize_particles=True) loss = elbo.differentiable_loss(reparam_model, reparam_guide, data) params = [trans_mat, obs_mat] torch.autograd.grad(loss, params, retain_graph=True)
def check_backends_agree(model): guide1 = AutoGaussian(model, backend="dense") guide2 = AutoGaussian(model, backend="funsor") guide1() with xfail_if_not_implemented(): guide2() # Inject random noise into all unconstrained parameters. params1 = dict(guide1.named_parameters()) params2 = dict(guide2.named_parameters()) assert set(params1) == set(params2) for k, v in params1.items(): v.data.add_(torch.zeros_like(v).normal_()) params2[k].data.copy_(v.data) names = sorted(params1) # Check densities agree between backends. with torch.no_grad(), poutine.trace() as tr: aux = guide2._sample_aux_values(temperature=1.0) flat = guide1._dense_flatten(aux) tr.trace.compute_log_prob() log_prob_funsor = tr.trace.nodes["_AutoGaussianFunsor_latent"]["log_prob"] with torch.no_grad(), poutine.trace() as tr: with poutine.condition(data={"_AutoGaussianDense_latent": flat}): guide1._sample_aux_values(temperature=1.0) tr.trace.compute_log_prob() log_prob_dense = tr.trace.nodes["_AutoGaussianDense_latent"]["log_prob"] assert_equal(log_prob_funsor, log_prob_dense) # Check Monte Carlo estimate of entropy. entropy1 = guide1._dense_get_mvn().entropy() with pyro.plate("particle", 100000, dim=-3), poutine.trace() as tr: guide2._sample_aux_values(temperature=1.0) tr.trace.compute_log_prob() entropy2 = -tr.trace.nodes["_AutoGaussianFunsor_latent"]["log_prob"].mean() assert_close(entropy1, entropy2, atol=1e-2) grads1 = torch.autograd.grad( entropy1, [params1[k] for k in names], allow_unused=True ) grads2 = torch.autograd.grad( entropy2, [params2[k] for k in names], allow_unused=True ) for name, grad1, grad2 in zip(names, grads1, grads2): # Gradients should agree to very high precision. if grad1 is None and grad2 is not None: grad1 = torch.zeros_like(grad2) elif grad2 is None and grad1 is not None: grad2 = torch.zeros_like(grad1) assert_close(grad1, grad2, msg=f"{name}:\n{grad1} vs {grad2}") # Check elbos agree between backends. elbo = Trace_ELBO(num_particles=1000000, vectorize_particles=True) loss1 = elbo.differentiable_loss(model, guide1) loss2 = elbo.differentiable_loss(model, guide2) assert_close(loss1, loss2, atol=1e-2, rtol=0.05) grads1 = torch.autograd.grad(loss1, [params1[k] for k in names], allow_unused=True) grads2 = torch.autograd.grad(loss2, [params2[k] for k in names], allow_unused=True) for name, grad1, grad2 in zip(names, grads1, grads2): assert_close( grad1, grad2, atol=0.05, rtol=0.05, msg=f"{name}:\n{grad1} vs {grad2}" )