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}"
)
