def main(args): baseball_dataset = pd.read_csv(DATA_URL, "\t") train, _, player_names = train_test_split(baseball_dataset) at_bats, hits = train[:, 0], train[:, 1] nuts_kernel = NUTS(conditioned_model, adapt_step_size=True) logging.info("Original Dataset:") logging.info(baseball_dataset) # (1) Full Pooling Model posterior_fully_pooled = MCMC(nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps) \ .run(fully_pooled, at_bats, hits) logging.info("\nModel: Fully Pooled") logging.info("===================") logging.info("\nphi:") logging.info( summary(posterior_fully_pooled, sites=["phi"], player_names=player_names)["phi"]) posterior_predictive = TracePredictive(fully_pooled, posterior_fully_pooled, num_samples=args.num_samples) sample_posterior_predictive(posterior_predictive, baseball_dataset) evaluate_log_predictive_density(fully_pooled, posterior_fully_pooled, baseball_dataset) # (2) No Pooling Model posterior_not_pooled = MCMC(nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps) \ .run(not_pooled, at_bats, hits) logging.info("\nModel: Not Pooled") logging.info("=================") logging.info("\nphi:") logging.info( summary(posterior_not_pooled, sites=["phi"], player_names=player_names)["phi"]) posterior_predictive = TracePredictive(not_pooled, posterior_not_pooled, num_samples=args.num_samples) sample_posterior_predictive(posterior_predictive, baseball_dataset) evaluate_log_predictive_density(not_pooled, posterior_not_pooled, baseball_dataset) # (3) Partially Pooled Model posterior_partially_pooled = MCMC(nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps) \ .run(partially_pooled, at_bats, hits) logging.info("\nModel: Partially Pooled") logging.info("=======================") logging.info("\nSigmoid(alpha):") logging.info( summary(posterior_partially_pooled, sites=["alpha"], player_names=player_names, transforms={"alpha": lambda x: 1. / (1 + np.exp(-x))})["alpha"]) posterior_predictive = TracePredictive(partially_pooled, posterior_partially_pooled, num_samples=args.num_samples) sample_posterior_predictive(posterior_predictive, baseball_dataset) evaluate_log_predictive_density(partially_pooled, posterior_partially_pooled, baseball_dataset)
def test_arrowhead_mass(): def model(prec): w = pyro.sample("w", dist.Normal(0, 1000).expand([2]).to_event(1)) x = pyro.sample("x", dist.Normal(0, 1000).expand([1]).to_event(1)) y = pyro.sample("y", dist.Normal(0, 1000).expand([1]).to_event(1)) z = pyro.sample("z", dist.Normal(0, 1000).expand([2]).to_event(1)) wyxz = torch.cat([w, y, x, z]) pyro.sample("obs", dist.MultivariateNormal(torch.zeros(6), precision_matrix=prec), obs=wyxz) A = torch.randn(6, 12) prec = A @ A.t() * 0.1 # smoke tests for dense_mass in [True, False]: kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True, full_mass=dense_mass) mcmc = MCMC(kernel, num_samples=1, warmup_steps=1) mcmc.run(prec) assert kernel.inverse_mass_matrix[("w", "x", "y", "z")].dim() == 1 + int(dense_mass) kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True, full_mass=[("w",), ("y", "x")]) kernel.mass_matrix_adapter = ArrowheadMassMatrix() mcmc = MCMC(kernel, num_samples=1, warmup_steps=1000) mcmc.run(prec) assert ("w", "y", "x", "z") in kernel.inverse_mass_matrix mass_matrix = kernel.mass_matrix_adapter.mass_matrix[("w", "y", "x", "z")] assert mass_matrix.top.shape == (4, 6) assert mass_matrix.bottom_diag.shape == (2,) assert_close(mass_matrix.top, prec[:4], atol=0.2, rtol=0.2) assert_close(mass_matrix.bottom_diag, prec.diag()[4:], atol=0.2, rtol=0.2)
def test_structured_mass(): def model(cov): w = pyro.sample("w", dist.Normal(0, 1000).expand([2]).to_event(1)) x = pyro.sample("x", dist.Normal(0, 1000).expand([1]).to_event(1)) y = pyro.sample("y", dist.Normal(0, 1000).expand([1]).to_event(1)) z = pyro.sample("z", dist.Normal(0, 1000).expand([1]).to_event(1)) wxyz = torch.cat([w, x, y, z]) pyro.sample("obs", dist.MultivariateNormal(torch.zeros(5), cov), obs=wxyz) w_cov = torch.tensor([[1.5, 0.5], [0.5, 1.5]]) xy_cov = torch.tensor([[2., 1.], [1., 3.]]) z_var = torch.tensor([2.5]) cov = torch.zeros(5, 5) cov[:2, :2] = w_cov cov[2:4, 2:4] = xy_cov cov[4, 4] = z_var # smoke tests for dense_mass in [True, False]: kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True, full_mass=dense_mass) mcmc = MCMC(kernel, num_samples=1, warmup_steps=1) mcmc.run(cov) assert kernel.inverse_mass_matrix[("w", "x", "y", "z")].dim() == 1 + int(dense_mass) kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True, full_mass=[("w",), ("x", "y")]) mcmc = MCMC(kernel, num_samples=1, warmup_steps=1000) mcmc.run(cov) assert_close(kernel.inverse_mass_matrix[("w",)], w_cov, atol=0.5, rtol=0.5) assert_close(kernel.inverse_mass_matrix[("x", "y")], xy_cov, atol=0.5, rtol=0.5) assert_close(kernel.inverse_mass_matrix[("z",)], z_var, atol=0.5, rtol=0.5)
def numpyro_schools_model(data, draws, chains): """Centered eight schools implementation in NumPyro.""" import jax import numpyro import numpyro.distributions as dist from numpyro.mcmc import MCMC, NUTS def model(): mu = numpyro.sample("mu", dist.Normal(0, 5)) tau = numpyro.sample("tau", dist.HalfCauchy(5)) # TODO: use numpyro.plate or `sample_shape` kwargs instead of # pylint: disable=fixme # multiplying with np.ones(J) in future versions of NumPyro theta = numpyro.sample("theta", dist.Normal(mu * np.ones(data["J"]), tau)) numpyro.sample("obs", dist.Normal(theta, data["sigma"]), obs=data["y"]) mcmc = MCMC( NUTS(model), num_warmup=draws, num_samples=draws, num_chains=chains, chain_method="sequential", ) mcmc.run(jax.random.PRNGKey(0), collect_fields=("z", "diverging")) return mcmc
def pyro_centered_schools(data, draws, chains): """Centered eight schools implementation in Pyro. Note there is not really a deterministic node in pyro, so I do not know how to do a non-centered implementation. """ import torch from pyro.infer.mcmc import MCMC, NUTS del chains y = torch.Tensor(data["y"]).type(torch.Tensor) sigma = torch.Tensor(data["sigma"]).type(torch.Tensor) nuts_kernel = NUTS(_pyro_conditioned_model, adapt_step_size=True) posterior = MCMC( # pylint:disable=not-callable nuts_kernel, num_samples=draws, warmup_steps=500 ).run(_pyro_centered_model, sigma, y) # This block lets the posterior be pickled for trace in posterior.exec_traces: for node in trace.nodes.values(): node.pop("fn", None) posterior.kernel = None posterior.run = None posterior.logger = None if hasattr(posterior, "sampler"): posterior.sampler = None return posterior
def fit(self, inputs, targets, *args, **kwargs): self.mcmc_run = MCMC(self.kernel, num_samples=self.num_samples, warmup_steps=100).run(inputs, targets) self.all_samples = torch.cat(list( self.mcmc_run.marginal(sites="t").support(flatten=True).values()), dim=-1)
def compute_posteriors(self, x_obs: torch.Tensor, mcmc_kwargs: dict = None): """ :param x_obs: :param mcmc_kwargs: By default: {num_samples=1000, warmup_steps=1000, num_chains=4) :return: """ if mcmc_kwargs is None: mcmc_kwargs = { "num_samples": 1000, "warmup_steps": 1000, "num_chains": 4 } kernel = NUTS( self.pyro_mdl, adapt_step_size=True, max_plate_nesting=1, jit_compile=True, target_accept_prob=0.6, ) mcmc_run = MCMC(kernel, **mcmc_kwargs).run(data=x_obs) marginals = mcmc_run.marginal(sites=["z"]) marginals_supp = marginals.support() z_x = marginals_supp["z"] return z_x, marginals
def test_nuts_conjugate_gaussian( fixture, num_samples, warmup_steps, expected_means, expected_precs, mean_tol, std_tol, ): pyro.get_param_store().clear() nuts_kernel = NUTS(fixture.model) mcmc = MCMC(nuts_kernel, num_samples, warmup_steps) mcmc.run(fixture.data) samples = mcmc.get_samples() for i in range(1, fixture.chain_len + 1): param_name = "loc_" + str(i) latent = samples[param_name] latent_loc = latent.mean(0) latent_std = latent.std(0) expected_mean = torch.ones(fixture.dim) * expected_means[i - 1] expected_std = 1 / torch.sqrt(torch.ones(fixture.dim) * expected_precs[i - 1]) # Actual vs expected posterior means for the latents logger.debug("Posterior mean (actual) - {}".format(param_name)) logger.debug(latent_loc) logger.debug("Posterior mean (expected) - {}".format(param_name)) logger.debug(expected_mean) assert_equal(rmse(latent_loc, expected_mean).item(), 0.0, prec=mean_tol) # Actual vs expected posterior precisions for the latents logger.debug("Posterior std (actual) - {}".format(param_name)) logger.debug(latent_std) logger.debug("Posterior std (expected) - {}".format(param_name)) logger.debug(expected_std) assert_equal(rmse(latent_std, expected_std).item(), 0.0, prec=std_tol)
def test_saasbo_sample(self): for use_saas, use_input_warping in product((False, True), repeat=2): with torch.random.fork_rng(): torch.manual_seed(0) X = torch.randn(3, 2) Y = torch.randn(3, 1) Yvar = torch.randn(3, 1) kernel = NUTS(pyro_model, max_tree_depth=1) mcmc = MCMC(kernel, warmup_steps=0, num_samples=1) mcmc.run( X, Y, Yvar, use_input_warping=use_input_warping, use_saas=use_saas, ) samples = mcmc.get_samples() if use_saas: self.assertTrue("kernel_tausq" in samples) self.assertTrue("_kernel_inv_length_sq" in samples) self.assertTrue("lengthscale" not in samples) else: self.assertTrue("kernel_tausq" not in samples) self.assertTrue("_kernel_inv_length_sq" not in samples) self.assertTrue("lengthscale" in samples) if use_input_warping: self.assertIn("c0", samples) self.assertIn("c1", samples) else: self.assertNotIn("c0", samples) self.assertNotIn("c1", samples)
def test_beta_binomial(hyperpriors): def model(data): with pyro.plate("plate_0", data.shape[-1]): alpha = pyro.sample( "alpha", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor( [1., 1.]) beta = pyro.sample( "beta", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor( [1., 1.]) beta_binom = BetaBinomialPair() with pyro.plate("plate_1", data.shape[-2]): probs = pyro.sample("probs", beta_binom.latent(alpha, beta)) with pyro.plate("data", data.shape[0]): pyro.sample("binomial", beta_binom.conditional( probs=probs, total_count=total_count), obs=data) true_probs = torch.tensor([[0.7, 0.4], [0.6, 0.4]]) total_count = torch.tensor([[1000, 600], [400, 800]]) num_samples = 80 data = dist.Binomial( total_count=total_count, probs=true_probs).sample(sample_shape=(torch.Size((10, )))) hmc_kernel = NUTS(collapse_conjugate(model), jit_compile=True, ignore_jit_warnings=True) mcmc = MCMC(hmc_kernel, num_samples=num_samples, warmup_steps=50) mcmc.run(data) samples = mcmc.get_samples() posterior = posterior_replay(model, samples, data, num_samples=num_samples) assert_equal(posterior["probs"].mean(0), true_probs, prec=0.05)
def test_gaussian_mixture_model(jit): K, N = 3, 1000 def gmm(data): mix_proportions = pyro.sample("phi", dist.Dirichlet(torch.ones(K))) with pyro.plate("num_clusters", K): cluster_means = pyro.sample( "cluster_means", dist.Normal(torch.arange(float(K)), 1.0) ) with pyro.plate("data", data.shape[0]): assignments = pyro.sample("assignments", dist.Categorical(mix_proportions)) pyro.sample("obs", dist.Normal(cluster_means[assignments], 1.0), obs=data) return cluster_means true_cluster_means = torch.tensor([1.0, 5.0, 10.0]) true_mix_proportions = torch.tensor([0.1, 0.3, 0.6]) cluster_assignments = dist.Categorical(true_mix_proportions).sample( torch.Size((N,)) ) data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample() nuts_kernel = NUTS( gmm, max_plate_nesting=1, jit_compile=jit, ignore_jit_warnings=True ) mcmc = MCMC(nuts_kernel, num_samples=300, warmup_steps=100) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples["phi"].mean(0).sort()[0], true_mix_proportions, prec=0.05) assert_equal( samples["cluster_means"].mean(0).sort()[0], true_cluster_means, prec=0.2 )
def test_gaussian_hmm(num_steps): dim = 4 def model(data): initialize = pyro.sample("initialize", dist.Dirichlet(torch.ones(dim))) with pyro.plate("states", dim): transition = pyro.sample("transition", dist.Dirichlet(torch.ones(dim, dim))) emission_loc = pyro.sample( "emission_loc", dist.Normal(torch.zeros(dim), torch.ones(dim)) ) emission_scale = pyro.sample( "emission_scale", dist.LogNormal(torch.zeros(dim), torch.ones(dim)) ) x = None with ignore_jit_warnings([("Iterating over a tensor", RuntimeWarning)]): for t, y in pyro.markov(enumerate(data)): x = pyro.sample( "x_{}".format(t), dist.Categorical(initialize if x is None else transition[x]), infer={"enumerate": "parallel"}, ) pyro.sample( "y_{}".format(t), dist.Normal(emission_loc[x], emission_scale[x]), obs=y, ) def _get_initial_trace(): guide = AutoDelta( poutine.block( model, expose_fn=lambda msg: not msg["name"].startswith("x") and not msg["name"].startswith("y"), ) ) elbo = TraceEnum_ELBO(max_plate_nesting=1) svi = SVI(model, guide, optim.Adam({"lr": 0.01}), elbo) for _ in range(100): svi.step(data) return poutine.trace(guide).get_trace(data) def _generate_data(): transition_probs = torch.rand(dim, dim) emissions_loc = torch.arange(dim, dtype=torch.Tensor().dtype) emissions_scale = 1.0 state = torch.tensor(1) obs = [dist.Normal(emissions_loc[state], emissions_scale).sample()] for _ in range(num_steps): state = dist.Categorical(transition_probs[state]).sample() obs.append(dist.Normal(emissions_loc[state], emissions_scale).sample()) return torch.stack(obs) data = _generate_data() nuts_kernel = NUTS( model, max_plate_nesting=1, jit_compile=True, ignore_jit_warnings=True ) if num_steps == 30: nuts_kernel.initial_trace = _get_initial_trace() mcmc = MCMC(nuts_kernel, num_samples=5, warmup_steps=5) mcmc.run(data)
def _main(self, ratings, sigma, args): self.n_thread = args.n_threads # split data to training & testing train_pct = 0.9 rand_state.shuffle(ratings) train_size = int(train_pct * ratings.shape[0]) train = ratings[:train_size] validation = ratings[train_size:] self.data = train nuts_kernel = NUTS(self._conditioned_model, jit_compile=args.jit,) posterior = MCMC(nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps, num_chains=args.num_chains, disable_progbar=False).run(self._model, sigma, train) sites = ['mu_item', 'mu_user'] + ['u_temp_feature' + str(user_id) for user_id in xrange(self.n_user)] + ['i_temp_feature' + str(item_id) for item_id in xrange(self.n_item)] marginal = posterior.marginal(sites=sites) marginal = torch.cat(list(marginal.support(flatten=True).values()), dim=-1).cpu().numpy() avg_mu_item = np.average(marginal[:, :self.n_feature], axis=0) avg_mu_user = np.average(marginal[:, self.n_feature:2*self.n_feature], axis=0) avg_u_temp_feature = np.average(marginal[:, 2*self.n_feature:2*self.n_feature + self.n_user*self.n_feature].reshape(-1,self.n_user,self.n_feature), axis=0) avg_i_temp_feature = np.average(marginal[:, 2*self.n_feature + self.n_user*self.n_feature:].reshape(-1,self.n_item, self.n_feature), axis=0) train_preds = self._predict(train[:, :2], True, avg_u_temp_feature, avg_i_temp_feature) train_rmse = RMSE(train_preds, train[:, 2]) val_preds = self._predict(validation[:, :2], True, avg_u_temp_feature, avg_i_temp_feature) val_rmse = RMSE(val_preds, validation[:, 2]) print("After %d iteration, train RMSE: %.6f, validation RMSE: %.6f" % (self.iter_, train_rmse, val_rmse))
def test_gamma_poisson(hyperpriors): def model(data): with pyro.plate("latent_dim", data.shape[1]): alpha = ( pyro.sample("alpha", dist.HalfCauchy(1.0)) if hyperpriors else torch.tensor([1.0, 1.0]) ) beta = ( pyro.sample("beta", dist.HalfCauchy(1.0)) if hyperpriors else torch.tensor([1.0, 1.0]) ) gamma_poisson = GammaPoissonPair() rate = pyro.sample("rate", gamma_poisson.latent(alpha, beta)) with pyro.plate("data", data.shape[0]): pyro.sample("obs", gamma_poisson.conditional(rate), obs=data) true_rate = torch.tensor([3.0, 10.0]) num_samples = 100 data = dist.Poisson(rate=true_rate).sample(sample_shape=(torch.Size((100,)))) hmc_kernel = NUTS( collapse_conjugate(model), jit_compile=True, ignore_jit_warnings=True ) mcmc = MCMC(hmc_kernel, num_samples=num_samples, warmup_steps=50) mcmc.run(data) samples = mcmc.get_samples() posterior = posterior_replay(model, samples, data, num_samples=num_samples) assert_equal(posterior["rate"].mean(0), true_rate, prec=0.3)
def sample(draws=500, model=None, warmup_steps=None, num_chains=1, kernel='nuts'): """Markov-chain Monte Carlo sampling. Sampling should be run within the context of a model or the model should be passed as an argument `model` explicitly. Number of samples is given by `draws` which defaults to `500`. Warm-up steps are assumed to be 30% of sample count. MCMC kernel can be selected by setting `kernel`. `hmc` and `nuts` are available. `pmpyro.inference.sample` returns a trace of samples """ # get model from context if model is None: model = Context.get_context() stfn = model.stfn # get stochastic function from model data = model.args # get data # make nuts kernel kernels = {'nuts': NUTS(stfn, adapt_step_size=True), 'hmc': HMC(stfn)} # if not num_chains: # figure out number of chains # num_chains = max(os.cpu_count() -1, 2) if not warmup_steps: # figure out warm-up steps warmup_steps = int(0.3 * draws) # run MCMC mcmc = MCMC(kernels[kernel], num_samples=draws, warmup_steps=warmup_steps, num_chains=num_chains) mcmc.run(*data) # get num samples num_samples = num_chains * draws return mcmc.get_samples()
def fit(self): pyro.clear_param_store() self.kernel = NUTS(self.model) self.mcmc = MCMC(self.kernel, warmup_steps=self.warmup_steps, num_samples=self.num_samples) self.mcmc.run() self._posterior_sample_df()
def train_nuts(model, data, num_warmup, num_samples, num_chains=1, **kwargs): _kwargs = dict(adapt_step_size=True, adapt_mass_matrix=True, jit_compile=True) _kwargs.update(kwargs) print(_kwargs) kernel = NUTS(model, **_kwargs) engine = MCMC(kernel, num_samples, num_warmup, num_chains=num_chains) engine.run(data, training=True) return engine
def nested(): 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}) nuts_kernel = NUTS(conditioned_model, adapt_step_size=True) mcmc_run = MCMC(nuts_kernel, num_samples=10, warmup_steps=2).run(num_trials) return mcmc_run
def run_inference( pyro_model: Callable[[Tensor, Tensor, Tensor, bool, str, float], None], X: Tensor, Y: Tensor, Yvar: Tensor, num_samples: int = 512, warmup_steps: int = 1024, thinning: int = 16, use_input_warping: bool = False, max_tree_depth: int = 6, use_saas: bool = False, disable_progbar: bool = False, ) -> Tensor: start = time.time() try: from pyro.infer.mcmc import NUTS, MCMC except ImportError: # pragma: no cover raise RuntimeError("Cannot call run_inference without pyro installed!") kernel = NUTS( pyro_model, jit_compile=True, full_mass=True, ignore_jit_warnings=True, max_tree_depth=max_tree_depth, ) mcmc = MCMC( kernel, warmup_steps=warmup_steps, num_samples=num_samples, disable_progbar=disable_progbar, ) mcmc.run( # there is an issue with jit-compilation and cuda # for now, we run MCMC on the CPU. X.cpu(), Y.cpu(), Yvar.cpu(), use_input_warping=use_input_warping, use_saas=use_saas, ) # this prints the summary orig_std_out = sys.stdout.write sys.stdout.write = logger.info mcmc.summary() sys.stdout.write = orig_std_out logger.info(f"MCMC elapsed time: {time.time() - start}") samples = mcmc.get_samples() if use_saas: # compute the lengthscale for saas and throw away everything else inv_length_sq = (samples["kernel_tausq"].unsqueeze(-1) * samples["_kernel_inv_length_sq"]) samples["lengthscale"] = (1.0 / inv_length_sq).sqrt() # pyre-ignore [16] del samples["kernel_tausq"], samples["_kernel_inv_length_sq"] # thin for k, v in samples.items(): # apply thinning and move back to X's device samples[k] = v[::thinning].to(device=X.device) return samples
def main(): train_x = torch.linspace(0, 1, 1000).double() train_y = torch.sin(train_x * (2 * math.pi)).double() + torch.randn(train_x.size()).double() * 0.1 # Use a positive constraint instead of usual GreaterThan(1e-4) so that LogNormal has support over full range. likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.Positive()) model = ExactGPModel(train_x, train_y, likelihood) # model.mean_module.register_prior("mean_prior", UniformPrior(-1, 1), "constant") # model.covar_module.base_kernel.register_prior("lengthscale_prior", UniformPrior(0.0, 9.0), "lengthscale") # # model.covar_module.base_kernel.register_prior("period_length_prior", UniformPrior(0.0, 4.0), "period_length") # model.covar_module.register_prior("outputscale_prior", UniformPrior(0, 4), "outputscale") # likelihood.register_prior("noise_prior", UniformPrior(0.0, 0.25), "noise") mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model) def pyro_model(x, y): priors= { 'covar_module.base_kernel.raw_lengthscale': Normal(0, 2).expand([1, 1]), 'covar_module.raw_outputscale': Normal(0, 2), 'likelihood.noise_covar.raw_noise': Normal(0, 2).expand([1]), 'mean_module.constant': Normal(0, 2), } fn = pyro.random_module("model", model, prior=priors) sampled_model = fn() output = sampled_model.likelihood(sampled_model(x)) pyro.sample("obs", output, obs=y) # model.mean_module.constant.data.fill_(0.0) # model.covar_module.outputscale = 0.5**2 # model.covar_module.base_kernel.lengthscale = 1 # model.likelihood.noise = 0.05**2 model.double() likelihood.double() nuts_kernel = NUTS(pyro_model, adapt_step_size=True, jit_compile=False) hmc_kernel = HMC(pyro_model, step_size=0.1, num_steps=10, adapt_step_size=True,\ init_strategy=pyro.infer.autoguide.initialization.init_to_median(num_samples=20)) mcmc_run = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps)#, initial_params=initial_params) return model, likelihood, mll, mcmc_run, train_x, train_y labels = ["c", "ls","os","noise"] fig, axes = plt.subplots(nrows=2, ncols=2) for i in range(4): if i == 0: samples = getattr(chain, labels[i]).reshape(-1) else: samples = 1/(1+np.exp(-1*getattr(chain, labels[i]).reshape(-1))) sns.distplot(samples, ax=axes[int(i/2), int(i%2)]) axes[int(i/2)][int(i%2)].legend([labels[i]]) plt.show() pickle.dump(chain, open("results/test_mcmc.pkl", "wb")) return
def fit_MCMC(self, models, items, responses, num_epochs): sites = ['theta', 'b'] nuts_kernel = NUTS(self.model_vague, adapt_step_size=True) hmc_posterior = MCMC(nuts_kernel, num_samples=1000, warmup_steps=100) \ .run(models, items, responses) theta_sum = self.summary(hmc_posterior, ['theta']).items() b_sum = self.summary(hmc_posterior, ['b']).items() print(theta_sum) print(b_sum)
def fit_MCMC(self, models, items, responses, num_epochs): """Fit the IRT model with MCMC""" sites = ["theta", "b"] nuts_kernel = NUTS(self.model_vague, adapt_step_size=True) hmc_posterior = MCMC(nuts_kernel, num_samples=1000, warmup_steps=100).run(models, items, responses) theta_sum = self.summary(hmc_posterior, ["theta"]).items() b_sum = self.summary(hmc_posterior, ["b"]).items() print(theta_sum) print(b_sum)
def test_posterior_predictive(): 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}) nuts_kernel = NUTS(conditioned_model, adapt_step_size=True) mcmc_run = MCMC(nuts_kernel, num_samples=1000, warmup_steps=200).run(num_trials) posterior_predictive = TracePredictive(model, mcmc_run, num_samples=10000).run(num_trials) marginal_return_vals = EmpiricalMarginal(posterior_predictive) assert_equal(marginal_return_vals.mean, torch.ones(5) * 700, prec=30)
def fit_fully_bayesian_model_nuts( model: SaasFullyBayesianSingleTaskGP, max_tree_depth: int = 6, warmup_steps: int = 512, num_samples: int = 256, thinning: int = 16, disable_progbar: bool = False, ) -> None: r"""Fit a fully Bayesian model using the No-U-Turn-Sampler (NUTS) Args: model: SaasFullyBayesianSingleTaskGP to be fitted. max_tree_depth: Maximum tree depth for NUTS warmup_steps: The number of burn-in steps for NUTS. num_samples: The number of MCMC samples. Note that with thinning, num_samples / thinning samples are retained. thinning: The amount of thinning. Every nth sample is retained. disable_progbar: A boolean indicating whether to print the progress bar and diagnostics during MCMC. Example: >>> gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y) >>> fit_fully_bayesian_model_nuts(gp) """ model.train() # Do inference with NUTS nuts = NUTS( model.pyro_model.sample, jit_compile=True, full_mass=True, ignore_jit_warnings=True, max_tree_depth=max_tree_depth, ) mcmc = MCMC( nuts, warmup_steps=warmup_steps, num_samples=num_samples, disable_progbar=disable_progbar, ) mcmc.run() # Get final MCMC samples from the Pyro model mcmc_samples = model.pyro_model.postprocess_mcmc_samples( mcmc_samples=mcmc.get_samples()) for k, v in mcmc_samples.items(): mcmc_samples[k] = v[::thinning] # Load the MCMC samples back into the BoTorch model model.load_mcmc_samples(mcmc_samples) model.eval()
def test_gp_kernels(self): torch.manual_seed(0) X = torch.randn(3, 2) Y = torch.randn(3, 1) Yvar = torch.randn(3, 1) kernel = NUTS(single_task_pyro_model, max_tree_depth=1) with self.assertRaises(ValueError): mcmc = MCMC(kernel, warmup_steps=0, num_samples=1) mcmc.run( X, Y, Yvar, gp_kernel="some_kernel_we_dont_support", )
def test_pyro_sampling(self): try: import pyro # noqa from pyro.infer.mcmc import NUTS, MCMC except ImportError: return train_x, test_x, train_y, test_y = self._get_data(cuda=False) likelihood = GaussianLikelihood( noise_constraint=gpytorch.constraints.Positive()) gp_model = ExactGPModel(train_x, train_y, likelihood) # Register normal GPyTorch priors gp_model.mean_module.register_prior("mean_prior", UniformPrior(-1, 1), "constant") gp_model.covar_module.base_kernel.register_prior( "lengthscale_prior", UniformPrior(0.01, 0.5), "lengthscale") gp_model.covar_module.register_prior("outputscale_prior", UniformPrior(1, 2), "outputscale") likelihood.register_prior("noise_prior", UniformPrior(0.05, 0.3), "noise") def pyro_model(x, y): with gpytorch.settings.fast_computations(False, False, False): sampled_model = gp_model.pyro_sample_from_prior() output = sampled_model.likelihood(sampled_model(x)) pyro.sample("obs", output, obs=y) return y nuts_kernel = NUTS(pyro_model, adapt_step_size=True) mcmc_run = MCMC(nuts_kernel, num_samples=3, warmup_steps=20, disable_progbar=True) mcmc_run.run(train_x, train_y) gp_model.pyro_load_from_samples(mcmc_run.get_samples()) gp_model.eval() expanded_test_x = test_x.unsqueeze(-1).repeat(3, 1, 1) output = gp_model(expanded_test_x) self.assertEqual(output.mean.size(0), 3) # All 3 samples should do reasonably well on a noiseless dataset. self.assertLess( torch.norm(output.mean[0] - test_y) / test_y.norm(), 0.2) self.assertLess( torch.norm(output.mean[1] - test_y) / test_y.norm(), 0.2) self.assertLess( torch.norm(output.mean[2] - test_y) / test_y.norm(), 0.2)
def test_gamma_beta(jit): def model(data): alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=1., rate=1.)) beta_prior = pyro.sample('beta', dist.Gamma(concentration=1., rate=1.)) pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data) true_alpha = torch.tensor(5.) true_beta = torch.tensor(1.) data = dist.Beta(concentration1=true_alpha, concentration0=true_beta).sample(torch.Size((5000,))) nuts_kernel = NUTS(model, jit_compile=jit, ignore_jit_warnings=True) mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=200) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples["alpha"].mean(0), true_alpha, prec=0.08) assert_equal(samples["beta"].mean(0), true_beta, prec=0.05)
def nuts_fit(model, data, num_samples=None, warmup_steps=None, var_names=None): """ Use No U-turn Sampler for inferring latent variables """ num_samples = int(0.8 * len(data)) if num_samples is None else num_samples warmup_steps = len( data) - num_samples if warmup_steps is None else warmup_steps nuts_kernel = NUTS(model, adapt_step_size=True) trace = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps).run(data) if var_names is None: posteriors = None else: posteriors = trace.marginal(var_names).empirical return trace, posteriors
def test_dirichlet_categorical(jit): def model(data): concentration = torch.tensor([1.0, 1.0, 1.0]) p_latent = pyro.sample("p_latent", dist.Dirichlet(concentration)) pyro.sample("obs", dist.Categorical(p_latent), obs=data) return p_latent true_probs = torch.tensor([0.1, 0.6, 0.3]) data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,)))) nuts_kernel = NUTS(model, jit_compile=jit, ignore_jit_warnings=True) mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=100) mcmc.run(data) samples = mcmc.get_samples() posterior = samples["p_latent"] assert_equal(posterior.mean(0), true_probs, prec=0.02)
def test_pyro_sampling(self): try: import pyro from pyro.infer.mcmc import NUTS, MCMC except: return train_x, test_x, train_y, test_y = self._get_data(cuda=False) likelihood = GaussianLikelihood( noise_constraint=gpytorch.constraints.Positive()) gp_model = ExactGPModel(train_x, train_y, likelihood) # Register normal GPyTorch priors gp_model.mean_module.register_prior("mean_prior", UniformPrior(-1, 1), "constant") gp_model.covar_module.base_kernel.register_prior( "lengthscale_prior", UniformPrior(0.01, 0.2), "lengthscale") gp_model.covar_module.register_prior("outputscale_prior", UniformPrior(1, 2), "outputscale") likelihood.register_prior("noise_prior", LogNormalPrior(-1.5, 0.1), "noise") mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model) def pyro_model(x, y): gp_model.pyro_sample_from_prior() output = gp_model(x) loss = mll.pyro_factor(output, y) return y nuts_kernel = NUTS(pyro_model, adapt_step_size=True) mcmc_run = MCMC(nuts_kernel, num_samples=3, warmup_steps=20) mcmc_run.run(train_x, train_y) gp_model.pyro_load_from_samples(mcmc_run.get_samples()) gp_model.eval() expanded_test_x = test_x.unsqueeze(-1).repeat(3, 1, 1) output = gp_model(expanded_test_x) self.assertEqual(output.mean.size(0), 3) # All 3 samples should do reasonably well on a noiseless dataset. self.assertLess( torch.norm(output.mean[0] - test_y) / test_y.norm(), 0.2) self.assertLess( torch.norm(output.mean[1] - test_y) / test_y.norm(), 0.2) self.assertLess( torch.norm(output.mean[2] - test_y) / test_y.norm(), 0.2)