def test_unpack_latent():
def model():
return pyro.sample('x', dist.LKJCorrCholesky(2, torch.tensor(1.)))
guide = AutoDiagonalNormal(model)
assert guide()['x'].shape == model().shape
latent = guide.sample_latent()
assert list(guide._unpack_latent(latent))[0][1].shape == (1,)
def main(args):
# load data
print('loading training data...')
dataset_directory = get_data_directory(__file__)
dataset_path = os.path.join(dataset_directory, 'faces_training.csv')
if not os.path.exists(dataset_path):
try:
os.makedirs(dataset_directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
pass
wget.download(
'https://d2hg8soec8ck9v.cloudfront.net/datasets/faces_training.csv',
dataset_path)
data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float()
sparse_gamma_def = SparseGammaDEF()
# Due to the special logic in the custom guide (e.g. parameter clipping), the custom guide
# seems to be more amenable to higher learning rates.
# Nevertheless, the easy guide performs the best (presumably because of numerical instabilities
# related to the gamma distribution in the custom guide).
learning_rate = 0.2 if args.guide in ['auto', 'easy'] else 4.5
momentum = 0.05 if args.guide in ['auto', 'easy'] else 0.1
opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum})
# use one of our three different guide types
if args.guide == 'auto':
guide = AutoDiagonalNormal(sparse_gamma_def.model,
init_loc_fn=init_to_feasible)
elif args.guide == 'easy':
guide = MyEasyGuide(sparse_gamma_def.model)
else:
guide = sparse_gamma_def.guide
# this is the svi object we use during training; we use TraceMeanField_ELBO to
# get analytic KL divergences
svi = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO())
# we use svi_eval during evaluation; since we took care to write down our model in
# a fully vectorized way, this computation can be done efficiently with large tensor ops
svi_eval = SVI(sparse_gamma_def.model,
guide,
opt,
loss=TraceMeanField_ELBO(num_particles=args.eval_particles,
vectorize_particles=True))
print('\nbeginning training with %s guide...' % args.guide)
# the training loop
for k in range(args.num_epochs):
loss = svi.step(data)
# for the custom guide we clip parameters after each gradient step
if args.guide == 'custom':
clip_params()
if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1:
loss = svi_eval.evaluate_loss(data)
print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def test_shapes(parallel):
num_samples = 10
def model():
x = pyro.sample("x", dist.Normal(0, 1).expand([2]).to_event(1))
with pyro.plate("plate", 5):
loc, log_scale = x.unbind(-1)
y = pyro.sample("y", dist.Normal(loc, log_scale.exp()))
return dict(x=x, y=y)
guide = AutoDiagonalNormal(model)
# Compute by hand.
vectorize = pyro.plate("_vectorize", num_samples, dim=-2)
trace = poutine.trace(vectorize(guide)).get_trace()
expected = poutine.replay(vectorize(model), trace)()
# Use Predictive.
predictive = Predictive(
model,
guide=guide,
return_sites=["x", "y"],
num_samples=num_samples,
parallel=parallel,
)
actual = predictive.get_samples()
assert set(actual) == set(expected)
assert actual["x"].shape == expected["x"].shape
assert actual["y"].shape == expected["y"].shape
def test_deterministic(with_plate, event_shape):
def model(y=None):
with pyro.util.optional(pyro.plate("plate", 3), with_plate):
x = pyro.sample("x", dist.Normal(0, 1).expand(event_shape).to_event())
x2 = pyro.deterministic("x2", x**2, event_dim=len(event_shape))
pyro.deterministic("x3", x2)
return pyro.sample("obs", dist.Normal(x2, 0.1).to_event(), obs=y)
y = torch.tensor(4.0)
guide = AutoDiagonalNormal(model)
svi = SVI(model, guide, optim.Adam(dict(lr=0.1)), Trace_ELBO())
for i in range(100):
svi.step(y)
actual = Predictive(
model, guide=guide, return_sites=["x2", "x3"], num_samples=1000
)()
x2_batch_shape = (3,) if with_plate else ()
assert actual["x2"].shape == (1000,) + x2_batch_shape + event_shape
# x3 shape is prepended 1 to match Pyro shape semantics
x3_batch_shape = (1, 3) if with_plate else ()
assert actual["x3"].shape == (1000,) + x3_batch_shape + event_shape
assert_close(actual["x2"].mean(), y, rtol=0.1)
assert_close(actual["x3"].mean(), y, rtol=0.1)
def __init__(self, in_features, num_hidden, learn_var = False):
super().__init__()
self.linear1 = PyroModule[nn.Linear](in_features, num_hidden)
self.linear1.weight = PyroSample(dist.Normal(0., 1.).expand([num_hidden, in_features]).to_event(2))
self.linear1.bias = PyroSample(dist.Normal(0., 10.).expand([num_hidden]).to_event(1))
self.learn_var = learn_var
if not self.learn_var:
self.linear2 = PyroModule[nn.Linear](num_hidden, 2)
self.linear2.weight = PyroSample(dist.Normal(0., 1.).expand([2, num_hidden]).to_event(2))
self.linear2.bias = PyroSample(dist.Normal(0., 10.).expand([2]).to_event(1))
else:
self.linear2 = PyroModule[nn.Linear](num_hidden, 1)
self.linear2.weight = PyroSample(dist.Normal(0., 1.).expand([1, num_hidden]).to_event(2))
self.linear2.bias = PyroSample(dist.Normal(0., 10.).expand([1]).to_event(1))
self.guide = AutoDiagonalNormal(self)
def test_empty_model_error():
def model():
pass
guide = AutoDiagonalNormal(model)
with pytest.raises(RuntimeError):
guide()
def create(self, split=True, model=NeuralNetwork(), load=False, guide=None, n=-1):
self.isSplit = split
self.model = model
if load:
self.load_model_and_guide()
else:
if guide is None:
guide = AutoDiagonalNormal(model)
self.create_model_and_guide(guide, n)
def __init__(self, in_features, out_features,
std_weights = std_weights_default,
std_bias = std_bias_default,
mean_prior = mean_prior_default,
hidden_features = hidden_features_default,
n_layers = n_layers_default,
device = device_default):
super().__init__(in_features, out_features, std_weights, std_bias, mean_prior, hidden_features, n_layers, device)
self.guide = AutoDiagonalNormal(self)
def train(model):
pyro.clear_param_store()
model(x)
guide = AutoDiagonalNormal(model)
guide(x)
adam = optim.Adam({"lr": 0.1})
svi = SVI(model, guide, adam, loss=Trace_ELBO())
for _ in range(1000):
svi.step(x)
import matplotlib.pyplot as plt
med = guide.median()['trans_score.weight'].detach()
plt.imshow(med)
return med
def __init__(self,
item_dim,
num_items,
num_slates,
item_group=None,
hidden_dim=None,
device="cpu"):
super().__init__()
self.item_dim = item_dim
self.hidden_dim = item_dim if hidden_dim == None else hidden_dim
self.num_items = num_items
self.num_slates = num_slates
self.item_group = item_group
self.device = device
#### item model
# Group-vec:
if item_group is None:
item_group = torch.zeros((self.num_items, )).long()
self.num_group = len(item_group.unique())
self.groupvec = PyroModule[nn.Embedding](num_embeddings=self.num_group,
embedding_dim=self.item_dim)
self.groupvec.weight = PyroSample(
dist.Normal(torch.zeros_like(self.groupvec.weight),
torch.ones_like(self.groupvec.weight)).to_event(2))
self.groupscale = PyroModule[nn.Embedding](
num_embeddings=self.num_group, embedding_dim=self.item_dim)
self.groupscale.weight = PyroSample(
dist.Uniform(0.01 * torch.ones_like(self.groupscale.weight),
torch.ones_like(self.groupscale.weight)).to_event(2))
# Item vec based on group vec hier prior:
self.itemvec = PyroModule[nn.Embedding](num_embeddings=num_items,
embedding_dim=item_dim)
self.itemvec.weight = PyroSample(lambda x: dist.Normal(
self.groupvec(self.item_group), self.groupscale(self.item_group)).
to_event(2))
# user model
self.gru = PyroModule[nn.GRU](input_size=self.item_dim,
hidden_size=self.hidden_dim,
bias=False,
num_layers=1,
batch_first=True)
self.gru.flatten_parameters()
set_noninform_prior(self.gru)
# Initial user state
self.z0 = PyroSample(
dist.Normal(torch.zeros((self.item_dim, )),
torch.ones((self.item_dim, ))).to_event(1))
self.guide = AutoDiagonalNormal(self.model)
def auto_guide_callable(model):
def guide_x():
x_loc = pyro.param("x_loc", torch.tensor(1.))
x_scale = pyro.param("x_scale", torch.tensor(.1), constraint=constraints.positive)
pyro.sample("x", dist.Normal(x_loc, x_scale))
def median_x():
return {"x": pyro.param("x_loc", torch.tensor(1.))}
guide = AutoGuideList(model)
guide.append(AutoCallable(model, guide_x, median_x))
guide.append(AutoDiagonalNormal(poutine.block(model, hide=["x"])))
return guide
def train(cls, *init_args, learning_rate = 0.01, iters = int(1e4), **init_kwargs):
model = cls(*init_args, **init_kwargs)
guide = AutoDiagonalNormal(model, init_loc_fn = init_to_mean)
adam = pyro.optim.Adam({"lr": learning_rate})
svi = SVI(model, guide, adam, loss=Trace_ELBO())
pyro.clear_param_store()
for j in tqdm.tqdm(range(iters//100)):
for i in range(100):
loss = svi.step()
return model, guide
def test_posterior_predictive_svi_auto_diag_normal_guide(return_trace):
true_probs = torch.ones(5) * 0.7
num_trials = torch.ones(5) * 1000
num_success = dist.Binomial(num_trials, true_probs).sample()
conditioned_model = poutine.condition(model, data={"obs": num_success})
guide = AutoDiagonalNormal(conditioned_model)
svi = SVI(conditioned_model, guide, optim.Adam(dict(lr=0.1)), Trace_ELBO())
for i in range(1000):
svi.step(num_trials)
posterior_predictive = Predictive(model, guide=guide, num_samples=10000, parallel=True)
if return_trace:
marginal_return_vals = posterior_predictive.get_vectorized_trace(num_trials).nodes["obs"]["value"]
else:
marginal_return_vals = posterior_predictive.get_samples(num_trials)["obs"]
assert_close(marginal_return_vals.mean(dim=0), torch.ones(5) * 700, rtol=0.05)
def __init__(self,model, guide=None, epochs=500, lr=0.05, cuda=False):
self.model=model
self.cuda=cuda
if guide != None:
self.guide=guide
else:
self.guide=AutoDiagonalNormal(model)
self.optimizer=optim.Adam({"lr": lr})
self.svi=SVI(self.model,
self.guide,
self.optimizer,
loss=Trace_ELBO())
self.epochs=epochs
def auto_guide_module_callable(model):
class GuideX(AutoGuide):
def __init__(self, model):
super().__init__(model)
self.x_loc = nn.Parameter(torch.tensor(1.))
self.x_scale = PyroParam(torch.tensor(.1), constraint=constraints.positive)
def forward(self, *args, **kwargs):
return {"x": pyro.sample("x", dist.Normal(self.x_loc, self.x_scale))}
def median(self, *args, **kwargs):
return {"x": self.x_loc.detach()}
guide = AutoGuideList(model)
guide.custom = GuideX(model)
guide.diagnorm = AutoDiagonalNormal(poutine.block(model, hide=["x"]))
return guide
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 test_posterior_predictive_svi_auto_diag_normal_guide():
true_probs = torch.ones(5) * 0.7
num_trials = torch.ones(5) * 1000
num_success = dist.Binomial(num_trials, true_probs).sample()
conditioned_model = poutine.condition(model, data={"obs": num_success})
opt = optim.Adam(dict(lr=0.1))
loss = Trace_ELBO()
guide = AutoDiagonalNormal(conditioned_model)
svi_run = SVI(conditioned_model,
guide,
opt,
loss,
num_steps=1000,
num_samples=100).run(num_trials)
posterior_predictive = TracePredictive(model, svi_run,
num_samples=10000).run(num_trials)
marginal_return_vals = posterior_predictive.marginal().empirical["_RETURN"]
assert_close(marginal_return_vals.mean, torch.ones(5) * 700, rtol=0.05)
def load(self, path):
"""Load model from path
Args:
load_path (string): Path to saved model
"""
model_path = path + '_model'
opt_path = path + '_opt'
guide_path = path + '_guide'
self.net = torch.load(model_path)
pyro.get_param_store().load(path + '_params')
self.optim = Adam({"lr": self.learning_rate})
self.optim.load(opt_path)
self.guide = AutoDiagonalNormal(self.model)
self.guide = torch.load(guide_path)
self.svi = SVI(self.net, self.guide, self.optim, loss=Trace_ELBO())
def __init__(self, net_enc, net_dec, predict_samples=None):
super().__init__()
self.encoder = net_enc
self.decoder = net_dec
self.decoder_guide = AutoDiagonalNormal(
poutine.block(self.decoder, hide=["obs"]))
self.optim = Adam({"lr": cfg.BAYESIAN.lr})
self.svi = SVI(self.decoder,
self.decoder_guide,
self.optim,
loss=Trace_ELBO())
predict_samples = predict_samples or cfg.BAYESIAN.predict_samples
self.predictive = Predictive(self.decoder,
guide=self.decoder_guide,
num_samples=predict_samples,
return_sites=['_RETURN'])
# Freeze encoder
self.encoder.eval()
for p in self.encoder.parameters():
p.requires_grad = False
def fit(data_pars=None, compute_pars=None, out_pars=None, **kw):
"""
"""
global model, session
session = None # Session type for compute
Xtrain, ytrain, Xtest, ytest = get_dataset(data_pars, task_type="train")
Xtrain = torch.tensor(Xtrain.values, dtype=torch.float)
Xtest = torch.tensor(Xtest.values, dtype=torch.float)
ytrain = torch.tensor(ytrain.values, dtype=torch.float)
ytest = torch.tensor(ytest.values, dtype=torch.float)
if VERBOSE: log(Xtrain, model.model)
###############################################################
compute_pars2 = compute_pars.get('compute_pars', {})
n_iter = compute_pars2.get('n_iter', 1000)
lr = compute_pars2.get('learning_rate', 0.01)
method = compute_pars2.get('method', 'svi_elbo')
guide = AutoDiagonalNormal(model.model)
adam = pyro.optim.Adam({"lr": lr})
### SVI + Elbo is faster than HMC
svi = SVI(model.model, guide, adam, loss=Trace_ELBO())
pyro.clear_param_store()
losses = []
for j in range(n_iter):
# calculate the loss and take a gradient step
loss = svi.step(Xtrain, ytrain)
losses.append({'loss': loss, 'iteration': j})
if j % 100 == 0:
log("[iteration %04d] loss: %.4f" % (j + 1, loss / len(Xtrain)))
model.guide = guide
df_loss = pd.DataFrame(losses)
df_loss['loss'].plot()
return df_loss
def train(self, data_loader, n_epochs, num_particles=1, lr=1e-3, log_per=5, show_smooth=True, save_per=10):
pyro.clear_param_store()
guide = AutoDiagonalNormal(self, init_loc_fn=init_loc, init_scale=1e-5)
self.guide = guide
svi = SVI(self, guide, Adam({"lr": lr}), TraceMeanField_ELBO(num_particles=num_particles))
losses = []
fig = None
pyro.clear_param_store()
pp = ProgressPlotter(losses, "$-ELBO$", log_per, show_smooth)
pp.start()
for epoch in range(n_epochs):
total_loss = 0.
for x, y in data_loader:
x, y = x.float().to(self.device), y.float().to(self.device)
loss = svi.step(x, y) / y.numel()
total_loss += loss
total_loss /= len(data_loader)
losses.append(total_loss)
fig = pp.update(epoch)
if epoch % save_per == 1:
self.save('MultilayerBayesian.pth')
return pp.fig, losses
#X_train_scaled = my_x_scaler.fit_transform(X_train)
#
#y_max = y_train.max()
#y_train_scaled = y_train/y_max
#
# Convert the data into tensors
X_train_torch = torch.tensor(X_train_scaled)
y_train_torch = torch.tensor(y_train_scaled)
pyro.clear_param_store()
# Provide a guide which fits a pre-defined distribution over each
# hidden parameter. The AutoDiagonalNormal guide fits a normal
# distribution over each coefficient and our rate parameter
my_guide = AutoDiagonalNormal(model_gamma)
# Initialize the SVI optimzation class
my_svi = SVI(model=model_gamma,
guide= my_guide,
optim=ClippedAdam({"lr": 0.01, 'clip_norm': 1.0}),
loss=Trace_ELBO())
losses = []
start_time = time.time()
# Perform optimization
for i in range(5000):
def main(args):
pyro.set_rng_seed(0)
pyro.clear_param_store()
pyro.enable_validation(__debug__)
# load data
if args.dataset == "dipper":
capture_history_file = os.path.dirname(
os.path.abspath(__file__)) + '/dipper_capture_history.csv'
elif args.dataset == "vole":
capture_history_file = os.path.dirname(
os.path.abspath(__file__)) + '/meadow_voles_capture_history.csv'
else:
raise ValueError("Available datasets are \'dipper\' and \'vole\'.")
capture_history = torch.tensor(
np.genfromtxt(capture_history_file, delimiter=',')).float()[:, 1:]
N, T = capture_history.shape
print(
"Loaded {} capture history for {} individuals collected over {} time periods."
.format(args.dataset, N, T))
if args.dataset == "dipper" and args.model in ["4", "5"]:
sex_file = os.path.dirname(
os.path.abspath(__file__)) + '/dipper_sex.csv'
sex = torch.tensor(np.genfromtxt(sex_file, delimiter=',')).float()[:,
1]
print("Loaded dipper sex data.")
elif args.dataset == "vole" and args.model in ["4", "5"]:
raise ValueError(
"Cannot run model_{} on meadow voles data, since we lack sex " +
"information for these animals.".format(args.model))
else:
sex = None
model = models[args.model]
# we use poutine.block to only expose the continuous latent variables
# in the models to AutoDiagonalNormal (all of which begin with 'phi'
# or 'rho')
def expose_fn(msg):
return msg["name"][0:3] in ['phi', 'rho']
# we use a mean field diagonal normal variational distributions (i.e. guide)
# for the continuous latent variables.
guide = AutoDiagonalNormal(poutine.block(model, expose_fn=expose_fn))
# since we enumerate the discrete random variables,
# we need to use TraceEnum_ELBO or TraceTMC_ELBO.
optim = Adam({'lr': args.learning_rate})
if args.tmc:
elbo = TraceTMC_ELBO(max_plate_nesting=1)
tmc_model = poutine.infer_config(model, lambda msg: {
"num_samples": args.tmc_num_samples,
"expand": False
} if msg["infer"].get("enumerate", None) == "parallel" else {}
) # noqa: E501
svi = SVI(tmc_model, guide, optim, elbo)
else:
elbo = TraceEnum_ELBO(max_plate_nesting=1,
num_particles=20,
vectorize_particles=True)
svi = SVI(model, guide, optim, elbo)
losses = []
print(
"Beginning training of model_{} with Stochastic Variational Inference."
.format(args.model))
for step in range(args.num_steps):
loss = svi.step(capture_history, sex)
losses.append(loss)
if step % 20 == 0 and step > 0 or step == args.num_steps - 1:
print("[iteration %03d] loss: %.3f" %
(step, np.mean(losses[-20:])))
# evaluate final trained model
elbo_eval = TraceEnum_ELBO(max_plate_nesting=1,
num_particles=2000,
vectorize_particles=True)
svi_eval = SVI(model, guide, optim, elbo_eval)
print("Final loss: %.4f" % svi_eval.evaluate_loss(capture_history, sex))
def guide_autodiagnorm(self):
self.guide = AutoDiagonalNormal(poutine.block(self.model,
expose=['weights',
'locs',
'scale']))
def auto_guide_list_x(model):
guide = AutoGuideList(model)
guide.append(AutoDelta(poutine.block(model, expose=["x"])))
guide.append(AutoDiagonalNormal(poutine.block(model, hide=["x"])))
return guide
def build_guide(self):
self.guide = AutoDiagonalNormal(self.model)
import torch
import pyro
from bnn import ClassifierBnn, train_bnn, validate_model
from pyro.infer import Predictive
from load_mnist import setup_data_loaders
from vae import VAE
train_loader, test_loader = setup_data_loaders()
vae = VAE()
encoder = vae.encoder
encoder.load_state_dict(torch.load("encoder.checkpoint"))
model = ClassifierBnn(num_in=100)
pyro.enable_validation(True)
from pyro.infer.autoguide import AutoDiagonalNormal
guide = AutoDiagonalNormal(model, init_scale=1e-1)
pyro.clear_param_store()
validate_model(train_loader, transform=transform)
train_bnn(100, train_loader, test_loader, model, guide, encoder=encoder)
def forward(self, x, y=None):
sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
mean = self.linear(x).squeeze(-1)
with pyro.plate("data", x.shape[0]):
obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
return mean
from pyro.infer.autoguide import AutoDiagonalNormal
model = BayesianRegression(3, 1)
guide = AutoDiagonalNormal(model)
from pyro.infer import SVI, Trace_ELBO
adam = pyro.optim.Adam({"lr": 0.03})
svi = SVI(model, guide, adam, loss=Trace_ELBO())
def main(args):
pyro.set_rng_seed(args.rng_seed)
fig = plt.figure(figsize=(8, 16), constrained_layout=True)
gs = GridSpec(4, 2, figure=fig)
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[1, 0])
ax4 = fig.add_subplot(gs[2, 0])
ax5 = fig.add_subplot(gs[3, 0])
ax6 = fig.add_subplot(gs[1, 1])
ax7 = fig.add_subplot(gs[2, 1])
ax8 = fig.add_subplot(gs[3, 1])
xlim = tuple(int(x) for x in args.x_lim.strip().split(','))
ylim = tuple(int(x) for x in args.y_lim.strip().split(','))
assert len(xlim) == 2
assert len(ylim) == 2
# 1. Plot samples drawn from BananaShaped distribution
x1, x2 = torch.meshgrid(
[torch.linspace(*xlim, 100),
torch.linspace(*ylim, 100)])
d = BananaShaped(args.param_a, args.param_b)
p = torch.exp(d.log_prob(torch.stack([x1, x2], dim=-1)))
ax1.contourf(
x1,
x2,
p,
cmap='OrRd',
)
ax1.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='BananaShaped distribution: \nlog density')
# 2. Run vanilla HMC
logging.info('\nDrawing samples using vanilla HMC ...')
mcmc = run_hmc(args, model)
vanilla_samples = mcmc.get_samples()['x'].cpu().numpy()
ax2.contourf(x1, x2, p, cmap='OrRd')
ax2.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='Posterior \n(vanilla HMC)')
sns.kdeplot(vanilla_samples[:, 0], vanilla_samples[:, 1], ax=ax2)
# 3(a). Fit a diagonal normal autoguide
logging.info('\nFitting a DiagNormal autoguide ...')
guide = AutoDiagonalNormal(model, init_scale=0.05)
fit_guide(guide, args)
with pyro.plate('N', args.num_samples):
guide_samples = guide()['x'].detach().cpu().numpy()
ax3.contourf(x1, x2, p, cmap='OrRd')
ax3.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='Posterior \n(DiagNormal autoguide)')
sns.kdeplot(guide_samples[:, 0], guide_samples[:, 1], ax=ax3)
# 3(b). Draw samples using NeuTra HMC
logging.info(
'\nDrawing samples using DiagNormal autoguide + NeuTra HMC ...')
neutra = NeuTraReparam(guide.requires_grad_(False))
neutra_model = poutine.reparam(model, config=lambda _: neutra)
mcmc = run_hmc(args, neutra_model)
zs = mcmc.get_samples()['x_shared_latent']
sns.scatterplot(zs[:, 0], zs[:, 1], alpha=0.2, ax=ax4)
ax4.set(xlabel='x0',
ylabel='x1',
title='Posterior (warped) samples \n(DiagNormal + NeuTra HMC)')
samples = neutra.transform_sample(zs)
samples = samples['x'].cpu().numpy()
ax5.contourf(x1, x2, p, cmap='OrRd')
ax5.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='Posterior (transformed) \n(DiagNormal + NeuTra HMC)')
sns.kdeplot(samples[:, 0], samples[:, 1], ax=ax5)
# 4(a). Fit a BNAF autoguide
logging.info('\nFitting a BNAF autoguide ...')
guide = AutoNormalizingFlow(
model, partial(iterated, args.num_flows, block_autoregressive))
fit_guide(guide, args)
with pyro.plate('N', args.num_samples):
guide_samples = guide()['x'].detach().cpu().numpy()
ax6.contourf(x1, x2, p, cmap='OrRd')
ax6.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='Posterior \n(BNAF autoguide)')
sns.kdeplot(guide_samples[:, 0], guide_samples[:, 1], ax=ax6)
# 4(b). Draw samples using NeuTra HMC
logging.info('\nDrawing samples using BNAF autoguide + NeuTra HMC ...')
neutra = NeuTraReparam(guide.requires_grad_(False))
neutra_model = poutine.reparam(model, config=lambda _: neutra)
mcmc = run_hmc(args, neutra_model)
zs = mcmc.get_samples()['x_shared_latent']
sns.scatterplot(zs[:, 0], zs[:, 1], alpha=0.2, ax=ax7)
ax7.set(xlabel='x0',
ylabel='x1',
title='Posterior (warped) samples \n(BNAF + NeuTra HMC)')
samples = neutra.transform_sample(zs)
samples = samples['x'].cpu().numpy()
ax8.contourf(x1, x2, p, cmap='OrRd')
ax8.set(xlabel='x0',
ylabel='x1',
xlim=xlim,
ylim=ylim,
title='Posterior (transformed) \n(BNAF + NeuTra HMC)')
sns.kdeplot(samples[:, 0], samples[:, 1], ax=ax8)
plt.savefig(os.path.join(os.path.dirname(__file__), 'neutra.pdf'))
def __init__(self, in_features, out_features):
super().__init__()
self.model = BayesianRegressionPyroModel(in_features, out_features)
self.guide = AutoDiagonalNormal(self.model)
