def test_sequential_consistent(monkeypatch): # test if there is no stuff left from the previous chain monkeypatch.setattr(torch.multiprocessing, 'cpu_count', lambda: 1) class FirstKernel(NUTS): def setup(self, warmup_steps, *args, **kwargs): self._chain_id = 0 if '_chain_id' not in self.__dict__ else 1 pyro.set_rng_seed(self._chain_id) super(FirstKernel, self).setup(warmup_steps, *args, **kwargs) class SecondKernel(NUTS): def setup(self, warmup_steps, *args, **kwargs): self._chain_id = 1 if '_chain_id' not in self.__dict__ else 0 pyro.set_rng_seed(self._chain_id) super(SecondKernel, self).setup(warmup_steps, *args, **kwargs) data = torch.tensor([1.0]) kernel = FirstKernel(normal_normal_model) mcmc = MCMC(kernel, num_samples=100, warmup_steps=100, num_chains=2) mcmc.run(data) samples1 = mcmc.get_samples(group_by_chain=True) kernel = SecondKernel(normal_normal_model) mcmc = MCMC(kernel, num_samples=100, warmup_steps=100, num_chains=2) mcmc.run(data) samples2 = mcmc.get_samples(group_by_chain=True) assert_close(samples1["y"][0], samples2["y"][1]) assert_close(samples1["y"][1], samples2["y"][0])
def test_potential_fn_initial_params(Kernel): target = torch.distributions.Normal(loc=torch.tensor([10., 0.]), scale=torch.tensor([1., 1.])) def potential_fn(z): z = z['points'] return -target.log_prob(z).sum(1)[None] initial_params = {'points': torch.tensor([[0., 0.]])} kernel = Kernel(potential_fn=potential_fn) mcmc = MCMC(kernel=kernel, warmup_steps=20, initial_params=initial_params, num_samples=10) mcmc.run() mcmc.get_samples()['points']
def nuts_sampling(self, x_data, y_data, num_samples, warmup_steps): nuts_kernel = NUTS(self.model, target_accept_prob=0.99) mcmc = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps) mcmc.run(x_data, y_data) self.posterior_samples = mcmc.get_samples()
def test_predictive(num_samples, parallel): model, data, true_probs = beta_bernoulli() init_params, potential_fn, transforms, _ = initialize_model( model, model_args=(data, )) nuts_kernel = NUTS(potential_fn=potential_fn, transforms=transforms) mcmc = MCMC(nuts_kernel, 100, initial_params=init_params, warmup_steps=100) mcmc.run(data) samples = mcmc.get_samples() with ignore_experimental_warning(): with optional(pytest.warns(UserWarning), num_samples not in (None, 100)): predictive_samples = predictive(model, samples, num_samples=num_samples, return_sites=["beta", "obs"], parallel=parallel) # check shapes assert predictive_samples["beta"].shape == (100, 5) assert predictive_samples["obs"].shape == (100, 1000, 5) # check sample mean assert_close(predictive_samples["obs"].reshape([-1, 5]).mean(0), true_probs, rtol=0.1)
def test_mcmc_interface(num_draws, group_by_chain, num_chains): num_samples = 2000 data = torch.tensor([1.0]) initial_params, _, transforms, _ = initialize_model(normal_normal_model, model_args=(data,), num_chains=num_chains) kernel = PriorKernel(normal_normal_model) mcmc = MCMC(kernel=kernel, num_samples=num_samples, warmup_steps=100, num_chains=num_chains, mp_context="spawn", initial_params=initial_params, transforms=transforms) mcmc.run(data) samples = mcmc.get_samples(num_draws, group_by_chain=group_by_chain) # test sample shape expected_samples = num_draws if num_draws is not None else num_samples if group_by_chain: expected_shape = (mcmc.num_chains, expected_samples, 1) elif num_draws is not None: # FIXME: what is the expected behavior of num_draw is not None and group_by_chain=False? expected_shape = (expected_samples, 1) else: expected_shape = (mcmc.num_chains * expected_samples, 1) assert samples['y'].shape == expected_shape # test sample stats if group_by_chain: samples = {k: v.reshape((-1,) + v.shape[2:]) for k, v in samples.items()} sample_mean = samples['y'].mean() sample_std = samples['y'].std() assert_close(sample_mean, torch.tensor(0.0), atol=0.05) assert_close(sample_std, torch.tensor(1.0), atol=0.05)
def init_guide(self): nuts_kernel = NUTS(self.model) num_samples = 3 mcmc = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=7) mcmc.run() hmc_samples = { k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items() } # Get latent coordinates samples from mcmc # gp_mean = torch.tensor(hmc_samples['gp_mean']).reshape( num_samples, self.V_net, self.H_dim, self.K_net, self.n_dir, self.n_w) gp_coord_demean = torch.tensor( hmc_samples[f'gp_coord_demean']).reshape(num_samples, self.V_net, self.H_dim, self.K_net, self.n_dir, self.n_w, self.T_net) # keep only last sample gp_mean_ini = gp_mean[num_samples - 1] gp_coord_demean_ini = gp_coord_demean[num_samples - 1] return (gp_mean_ini, gp_coord_demean_ini)
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_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_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.)) with pyro.plate("data", data.shape[0]): assignments = pyro.sample("assignments", dist.Categorical(mix_proportions)) pyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data) return cluster_means true_cluster_means = torch.tensor([1., 5., 10.]) 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_null_model_with_hook(kernel, model, jit, num_chains): num_warmup, num_samples = 10, 10 initial_params, potential_fn, transforms, _ = initialize_model( model, num_chains=num_chains) iters = [] hook = partial(_hook, iters) mp_context = "spawn" if "CUDA_TEST" in os.environ else None kern = kernel(potential_fn=potential_fn, transforms=transforms, jit_compile=jit) mcmc = MCMC(kern, num_samples=num_samples, warmup_steps=num_warmup, num_chains=num_chains, initial_params=initial_params, hook_fn=hook, mp_context=mp_context) mcmc.run() samples = mcmc.get_samples() assert samples == {} if num_chains == 1: expected = [("Warmup", i) for i in range(num_warmup)] + [("Sample", i) for i in range(num_samples)] assert iters == expected
def test_bernoulli_latent_model(jit): def model(data): y_prob = pyro.sample("y_prob", dist.Beta(1.0, 1.0)) y = pyro.sample("y", dist.Bernoulli(y_prob)) with pyro.plate("data", data.shape[0]): z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1)) pyro.sample("obs", dist.Normal(2.0 * z, 1.0), obs=data) pyro.sample("nuisance", dist.Bernoulli(0.3)) N = 2000 y_prob = torch.tensor(0.3) y = dist.Bernoulli(y_prob).sample(torch.Size((N, ))) z = dist.Bernoulli(0.65 * y + 0.1).sample() data = dist.Normal(2.0 * z, 1.0).sample() hmc_kernel = HMC( model, trajectory_length=1, max_plate_nesting=1, jit_compile=jit, ignore_jit_warnings=True, ) mcmc = MCMC(hmc_kernel, num_samples=600, warmup_steps=200) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples["y_prob"].mean(0), y_prob, prec=0.06)
def nuts(data, model, iter=None, warmup=None, num_chains=None): assert type(data) == dict assert type(model) == Model iter, warmup, num_chains = apply_default_hmc_args(iter, warmup, num_chains) nuts_kernel = NUTS(model.fn, jit_compile=False, adapt_step_size=True) mcmc = MCMC(nuts_kernel, num_samples=iter, warmup_steps=warmup, num_chains=num_chains) mcmc.run(**data) samples = mcmc.get_samples() transformed_samples = run_model_on_samples_and_data( model.fn, samples, data) def loc(d): # Optimization: For the data used for inference, values for # `mu` are already computed and available from # `transformed_samples`. if d == data: return transformed_samples['mu'] else: # TODO: This computes more than is necessary. (i.e. It # build additional tensors we immediately throw away.) # This is minor, but might be worth addressing eventually. return run_model_on_samples_and_data(model.fn, samples, d)['mu'] all_samples = dict(samples, **transformed_samples) return Samples(all_samples, lambda name: all_samples[name], loc)
def test_logistic_regression(step_size, trajectory_length, num_steps, adapt_step_size, adapt_mass_matrix, full_mass): dim = 3 data = torch.randn(2000, dim) true_coefs = torch.arange(1., dim + 1.) labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample() def model(data): coefs_mean = pyro.param('coefs_mean', torch.zeros(dim)) coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim))) y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels) return y hmc_kernel = HMC(model, step_size=step_size, trajectory_length=trajectory_length, num_steps=num_steps, adapt_step_size=adapt_step_size, adapt_mass_matrix=adapt_mass_matrix, full_mass=full_mass) mcmc = MCMC(hmc_kernel, num_samples=500, warmup_steps=100, disable_progbar=True) mcmc.run(data) samples = mcmc.get_samples()['beta'] assert_equal(rmse(true_coefs, samples.mean(0)).item(), 0.0, prec=0.1)
def hmc(matrix_factorization_poisson, num_samples=1000, warmup_steps=200): nuts_kernel = NUTS(matrix_factorization_poisson) mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=200) mcmc.run(anime_matrix_train.values, k=k) hmc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()} return hmc_samples
def run_default_mcmc( data, kernel, num_samples, warmup_steps=None, initial_params=None, num_chains=1, hook_fn=None, mp_context=None, transforms=None, num_draws=None, group_by_chain=False, ): mcmc = MCMC( kernel=kernel, num_samples=num_samples, warmup_steps=warmup_steps, initial_params=initial_params, num_chains=num_chains, hook_fn=hook_fn, mp_context=mp_context, transforms=transforms, ) mcmc.run(data) return mcmc.get_samples(num_draws, group_by_chain=group_by_chain), mcmc.num_chains
def test_gamma_poisson(hyperpriors): def model(data): with pyro.plate("latent_dim", 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.]) 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., 10.]) 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 pyro_mcmc( self, num_samples: int, potential_function: Callable, context: Tensor, mcmc_method: str = "slice", thin: int = 10, warmup_steps: int = 200, num_chains: Optional[int] = 1, ): """Return samples obtained using Pyro's HMC, NUTS or slice kernels. Args: num_samples: desired number of samples potential_fn: defining the potential function as a callable **class** makes it picklable for Pyro's MCMC to use it across chains in parallel, even if the potential function requires evaluating a neural network. context: conditioning observation mcmc_method: one of "hmc", "nuts" or "slice" (default "slice") thin: thinning (subsampling) factor (default 10) warmup_steps: initial samples to discard (defaults to 200) num_chains: whether to sample in parallel. If None, will use all CPUs except one (default 1) Returns: tensor of shape num_samples x parameter dimension """ if num_chains is None: num_chains = mp.cpu_count - 1 # XXX move outside function, and assert inside; remember return to train # Always sample in eval mode. self.neural_net.eval() kernels = dict(slice=Slice, hmc=HMC, nuts=NUTS) try: kernel = kernels[mcmc_method](potential_fn=potential_function) except KeyError: raise ValueError("`mcmc_method` not one of 'slice', 'hmc', 'nuts'.") initial_params = self._prior.sample((num_chains,)) sampler = MCMC( kernel=kernel, num_samples=(thin * num_samples) // num_chains + num_chains, warmup_steps=warmup_steps, initial_params={"": initial_params}, num_chains=num_chains, mp_context="fork", ) sampler.run() samples = next(iter(sampler.get_samples().values())).reshape( -1, len(self._prior.mean) # len(prior.mean) = dim of theta ) samples = samples[::thin][:num_samples] assert samples.shape[0] == num_samples return samples
def init_guide(self): nuts_kernel = NUTS(self.model) mcmc = MCMC(nuts_kernel, num_samples=1, warmup_steps=30) mcmc.run() hmc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()} # Get latent coordinates from mcmc # gp_coord_mcmc = np.array([[hmc_samples[f'latent_coord_{v}_{h}'] for h in range(self.H_dim)] for v in range(self.V_net)]).transpose(0,1,3,2) self.gp_coord_mcmc = torch.from_numpy(gp_coord_mcmc)
def _pyro_mcmc( self, num_samples: int, potential_function: Callable, mcmc_method: str = "slice", thin: int = 10, warmup_steps: int = 200, num_chains: Optional[int] = 1, show_progress_bars: bool = True, ): r"""Return samples obtained using Pyro HMC, NUTS or slice kernels. Args: num_samples: Desired number of samples. potential_function: A callable **class**. A class, but not a function, is picklable for Pyro MCMC to use it across chains in parallel, even when the potential function requires evaluating a neural network. mcmc_method: One of `hmc`, `nuts` or `slice`. thin: Thinning (subsampling) factor. warmup_steps: Initial number of samples to discard. num_chains: Whether to sample in parallel. If None, use all but one CPU. show_progress_bars: Whether to show a progressbar during sampling. Returns: Tensor of shape (num_samples, shape_of_single_theta). """ num_chains = mp.cpu_count - 1 if num_chains is None else num_chains # Go into eval mode for evaluating during sampling self.net.eval() kernels = dict(slice=Slice, hmc=HMC, nuts=NUTS) sampler = MCMC( kernel=kernels[mcmc_method](potential_fn=potential_function), num_samples=(thin * num_samples) // num_chains + num_chains, warmup_steps=warmup_steps, initial_params={"": self._mcmc_init_params}, num_chains=num_chains, mp_context="fork", disable_progbar=not show_progress_bars, ) sampler.run() samples = next(iter(sampler.get_samples().values())).reshape( -1, len(self._prior.mean) # len(prior.mean) = dim of theta ) samples = samples[::thin][:num_samples] assert samples.shape[0] == num_samples # Back to training mode self.net.train(True) return samples
def sample_posterior(self, num_samples, thin=1): """ Samples from posterior for true observation q(theta | x0) ~ q(x0 | theta) p(theta) using most recent likelihood estimate q(x0 | theta) with MCMC. :param num_samples: Number of samples to generate. :param thin: Generate (num_samples * thin) samples in total, then select every 'thin' sample. :return: torch.Tensor of shape [num_samples, parameter_dim] """ # Always sample in eval mode. self._neural_likelihood.eval() if self._mcmc_method == "slice-np": self.posterior_sampler.gen(20) samples = torch.Tensor(self.posterior_sampler.gen(num_samples)) else: if self._mcmc_method == "slice": kernel = Slice(potential_function=self._potential_function) elif self._mcmc_method == "hmc": kernel = HMC(potential_fn=self._potential_function) elif self._mcmc_method == "nuts": kernel = NUTS(potential_fn=self._potential_function) else: raise ValueError( "'mcmc_method' must be one of ['slice', 'hmc', 'nuts']." ) num_chains = mp.cpu_count() - 1 # TODO: decide on way to initialize chain initial_params = self._prior.sample((num_chains,)) sampler = MCMC( kernel=kernel, num_samples=num_samples // num_chains + num_chains, warmup_steps=200, initial_params={"": initial_params}, num_chains=num_chains, ) sampler.run() samples = next(iter(sampler.get_samples().values())).reshape( -1, self._simulator.parameter_dim ) samples = samples[:num_samples].to(device) assert samples.shape[0] == num_samples # Back to training mode. self._neural_likelihood.train() return samples
def run_inference(data, gen_model, ode_model, method, iterations=10000, num_particles=1, num_samples=1000, warmup_steps=500, init_scale=0.1, seed=12, lr=0.5, return_sites="_RETURN"): torch_data = torch.tensor(data, dtype=torch.float) if isinstance(ode_model, ForwardSensManualJacobians) or \ isinstance(ode_model, ForwardSensTorchJacobians): ode_op = ForwardSensOp elif isinstance(ode_model, AdjointSensManualJacobians) or \ isinstance(ode_model, AdjointSensTorchJacobians): ode_op = AdjointSensOp else: raise ValueError('Unknown sensitivity solver: Use "Forward" or "Adjoint"') model = gen_model(ode_op, ode_model) pyro.set_rng_seed(seed) pyro.clear_param_store() if method == 'VI': guide = AutoMultivariateNormal(model, init_scale=init_scale) optim = AdagradRMSProp({"eta": lr}) if num_particles == 1: svi = SVI(model, guide, optim, loss=Trace_ELBO()) else: svi = SVI(model, guide, optim, loss=Trace_ELBO(num_particles=num_particles, vectorize_particles=True)) loss_trace = [] t0 = timer.time() for j in range(iterations): loss = svi.step(torch_data) loss_trace.append(loss) if j % 500 == 0: print("[iteration %04d] loss: %.4f" % (j + 1, np.mean(loss_trace[max(0, j - 1000):j + 1]))) t1 = timer.time() print('VI time: ', t1 - t0) predictive = Predictive(model, guide=guide, num_samples=num_samples, return_sites=return_sites) # "ode_params", "scale", vb_samples = predictive(torch_data) return vb_samples elif method == 'NUTS': nuts_kernel = NUTS(model, adapt_step_size=True, init_strategy=init_to_median) # mcmc = MCMC(nuts_kernel, num_samples=iterations, warmup_steps=warmup_steps, num_chains=2) mcmc = MCMC(nuts_kernel, num_samples=iterations, warmup_steps=warmup_steps, num_chains=1) t0 = timer.time() mcmc.run(torch_data) t1 = timer.time() print('NUTS time: ', t1 - t0) hmc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()} return hmc_samples else: raise ValueError('Unknown method: Use "NUTS" or "VI"')
def sample_posterior_mcmc(self, num_samples, thin=10): """ Samples from posterior for true observation q(theta | x0) using MCMC. :param num_samples: Number of samples to generate. :param mcmc_method: Which MCMC method to use ['metropolis-hastings', 'slice', 'hmc', 'nuts'] :param thin: Generate (num_samples * thin) samples in total, then select every 'thin' sample. :return: torch.Tensor of shape [num_samples, parameter_dim] """ # Always sample in eval mode. self._neural_posterior.eval() if self._mcmc_method == "slice-np": self.posterior_sampler.gen(20) # Burn-in for 200 samples samples = torch.Tensor(self.posterior_sampler.gen(num_samples)) else: if self._mcmc_method == "slice": kernel = Slice(potential_function=self._potential_function) elif self._mcmc_method == "hmc": kernel = HMC(potential_fn=self._potential_function) elif self._mcmc_method == "nuts": kernel = NUTS(potential_fn=self._potential_function) else: raise ValueError( "'mcmc_method' must be one of ['slice', 'hmc', 'nuts']." ) num_chains = mp.cpu_count() - 1 initial_params = self._prior.sample((num_chains,)) sampler = MCMC( kernel=kernel, num_samples=(thin * num_samples) // num_chains + num_chains, warmup_steps=200, initial_params={"": initial_params}, num_chains=num_chains, mp_context="spawn", ) sampler.run() samples = next(iter(sampler.get_samples().values())).reshape( -1, self._simulator.parameter_dim ) samples = samples[::thin][:num_samples] assert samples.shape[0] == num_samples # Back to training mode. self._neural_posterior.train() return samples
class bay_reg_mcmc: def __init__(self,model, num_samples=1000, warmup_steps=100, cuda=False): self.model=model self.num_samples=num_samples self.warmup_steps=warmup_steps self.cuda = cuda def fit(self,x_train,y_train, verbose=False): if verbose: disable_progbar=False else: disable_progbar=True if self.cuda == True: x_train=torch.tensor(x_train).cuda() y_train=torch.tensor(y_train).cuda() torch.set_default_tensor_type(torch.cuda.FloatTensor) else: x_train=torch.tensor(x_train) y_train=torch.tensor(y_train) self.posterior=MCMC(kernel=NUTS(self.model), num_samples=self.num_samples, warmup_steps=self.warmup_steps, disable_progbar=disable_progbar) self.posterior=self.posterior self.posterior.run(x_train,y_train) def predict(self, x_test, num_samples=300, percentiles=(5.0, 50.0, 95.0)): x_test=torch.tensor(x_test) samples=self.posterior.get_samples() #posterior_predictive=predictive(self.model, samples, x_test, num_samples=num_samples,return_sites=("w","b","_RETURN")) posterior_predictive= Predictive(self.model, posterior_samples=samples, return_sites=("_RETURN",'obs')).forward(x_test) #confidence intervall convidence_intervalls=posterior_predictive['_RETURN'].detach().cpu().numpy() y_pred=np.percentile(convidence_intervalls, percentiles, axis=0).T y_median_conv=y_pred[:,1].reshape(-1,1) y_lower_upper_quantil_conv=np.concatenate((y_pred[:,1].reshape(-1,1)-y_pred[:,0].reshape(-1,1),y_pred[:,2].reshape(-1,1)-y_pred[:,1].reshape(-1,1)),axis=1) #predictive intervall predictive_intervalls=posterior_predictive['obs'].detach().cpu().numpy() y_pred=np.percentile(predictive_intervalls, percentiles, axis=0).T y_median_pred=y_pred[:,1].reshape(-1,1) y_lower_upper_quantil_pred=np.concatenate((y_pred[:,1].reshape(-1,1)-y_pred[:,0].reshape(-1,1),y_pred[:,2].reshape(-1,1)-y_pred[:,1].reshape(-1,1)),axis=1) return(y_median_conv, y_lower_upper_quantil_conv, y_median_pred, y_lower_upper_quantil_pred)
def test_(): # if torch.cuda.is_available(): # device = torch.device("cuda") # torch.set_default_tensor_type("torch.cuda.FloatTensor") # else: # input("CUDA not available, do you wish to continue?") # device = torch.device("cpu") # torch.set_default_tensor_type("torch.FloatTensor") loc = torch.Tensor([0, 0]) covariance_matrix = torch.Tensor([[1, 0.99], [0.99, 1]]) likelihood = distributions.MultivariateNormal( loc=loc, covariance_matrix=covariance_matrix) bound = 1.5 low, high = -bound * torch.ones(2), bound * torch.ones(2) prior = distributions.Uniform(low=low, high=high) # def potential_function(inputs_dict): # parameters = next(iter(inputs_dict.values())) # return -(likelihood.log_prob(parameters) + prior.log_prob(parameters).sum()) prior = distributions.Uniform(low=-5 * torch.ones(4), high=2 * torch.ones(4)) from nsf import distributions as distributions_ likelihood = distributions_.LotkaVolterraOscillating() potential_function = PotentialFunction(likelihood, prior) # kernel = Slice(potential_function=potential_function) from pyro.infer.mcmc import HMC, NUTS # kernel = HMC(potential_fn=potential_function) kernel = NUTS(potential_fn=potential_function) num_chains = 3 sampler = MCMC( kernel=kernel, num_samples=10000 // num_chains, warmup_steps=200, initial_params={"": torch.zeros(num_chains, 4)}, num_chains=num_chains, ) sampler.run() samples = next(iter(sampler.get_samples().values())) utils.plot_hist_marginals(utils.tensor2numpy(samples), ground_truth=utils.tensor2numpy(loc), lims=[-6, 3]) # plt.show() plt.savefig("/home/conor/Dropbox/phd/projects/lfi/out/mcmc.pdf") plt.close()
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,)))) hmc_kernel = HMC(model, trajectory_length=1, jit_compile=jit, ignore_jit_warnings=True) mcmc = MCMC(hmc_kernel, num_samples=200, warmup_steps=100) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples['p_latent'].mean(0), true_probs, prec=0.02)
def test_gamma_normal(): def model(data): rate = torch.tensor([1.0, 1.0]) concentration = torch.tensor([1.0, 1.0]) p_latent = pyro.sample("p_latent", dist.Gamma(rate, concentration)) pyro.sample("obs", dist.Normal(3, p_latent), obs=data) return p_latent true_std = torch.tensor([0.5, 2]) data = dist.Normal(3, true_std).sample(sample_shape=(torch.Size((2000, )))) hmc_kernel = HMC(model, num_steps=15, step_size=0.01, adapt_step_size=True) mcmc = MCMC(hmc_kernel, num_samples=200, warmup_steps=200) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples["p_latent"].mean(0), true_std, prec=0.05)
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 test_beta_bernoulli(step_size, adapt_step_size, adapt_mass_matrix, full_mass): def model(data): alpha = torch.tensor([1.1, 1.1]) beta = torch.tensor([1.1, 1.1]) p_latent = pyro.sample("p_latent", dist.Beta(alpha, beta)) pyro.sample("obs", dist.Bernoulli(p_latent), obs=data) return p_latent true_probs = torch.tensor([0.9, 0.1]) data = dist.Bernoulli(true_probs).sample(sample_shape=(torch.Size((1000,)))) nuts_kernel = NUTS(model, step_size=step_size, adapt_step_size=adapt_step_size, adapt_mass_matrix=adapt_mass_matrix, full_mass=full_mass) mcmc = MCMC(nuts_kernel, num_samples=400, warmup_steps=200) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples["p_latent"].mean(0), true_probs, prec=0.02)
def test_hmc(model_class, X, y, kernel, likelihood): if model_class is SparseGPRegression or model_class is VariationalSparseGP: gp = model_class(X, y, kernel, X.clone(), likelihood) else: gp = model_class(X, y, kernel, likelihood) kernel.variance = PyroSample(dist.Uniform(torch.tensor(0.5), torch.tensor(1.5))) kernel.lengthscale = PyroSample(dist.Uniform(torch.tensor(1.0), torch.tensor(3.0))) hmc_kernel = HMC(gp.model, step_size=1) mcmc = MCMC(hmc_kernel, num_samples=10) mcmc.run() for name, param in mcmc.get_samples().items(): param_mean = torch.mean(param, 0) logger.info("Posterior mean - {}".format(name)) logger.info(param_mean)
def test_beta_bernoulli(jit): def model(data): alpha = torch.tensor([1.1, 1.1]) beta = torch.tensor([1.1, 1.1]) p_latent = pyro.sample('p_latent', dist.Beta(alpha, beta)) with pyro.plate("data", data.shape[0], dim=-2): pyro.sample('obs', dist.Bernoulli(p_latent), obs=data) return p_latent true_probs = torch.tensor([0.9, 0.1]) data = dist.Bernoulli(true_probs).sample(sample_shape=(torch.Size((1000,)))) hmc_kernel = HMC(model, trajectory_length=1, max_plate_nesting=2, jit_compile=jit, ignore_jit_warnings=True) mcmc = MCMC(hmc_kernel, num_samples=800, warmup_steps=500) mcmc.run(data) samples = mcmc.get_samples() assert_equal(samples['p_latent'].mean(0), true_probs, prec=0.05)