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)