def test_inference_data_no_posterior(self, data, eight_schools_params): posterior_samples = data.obj.get_samples() model = data.obj.kernel.model posterior_predictive = Predictive(model, posterior_samples)( eight_schools_params["J"], torch.from_numpy(eight_schools_params["sigma"]).float() ) prior = Predictive(model, num_samples=500)( eight_schools_params["J"], torch.from_numpy(eight_schools_params["sigma"]).float() ) idata = from_pyro( prior=prior, posterior_predictive=posterior_predictive, coords={"school": np.arange(eight_schools_params["J"])}, dims={"theta": ["school"], "eta": ["school"]}, ) test_dict = {"posterior_predictive": ["obs"], "prior": ["mu", "tau", "eta", "obs"]} fails = check_multiple_attrs(test_dict, idata) assert not fails
def predict(Xpred=None, data_pars={}, compute_pars=None, out_pars={}, **kw): global model, session compute_pars2 = model.compute_pars if compute_pars is None else compute_pars num_samples = compute_pars2.get('num_samples', 300) ###### Data load if Xpred is None: Xpred = get_dataset(data_pars, task_type="predict") cols_Xpred = list(Xpred.columns) max_size = compute_pars2.get('max_size', len(Xpred)) Xpred = Xpred.iloc[:max_size, :] Xpred_ = torch.tensor(Xpred.values, dtype=torch.float) ###### Post processing normalization post_process_fun = model.model_pars.get('post_process_fun', None) if post_process_fun is None: def post_process_fun(y): return y from pyro.infer import Predictive def summary(samples): site_stats = {} for k, v in samples.items(): site_stats[k] = { "mean": torch.mean(v, 0), "std": torch.std(v, 0), } return site_stats # If the model is loaded, it drops the guide param if it's None guide = getattr(model, "guide", None) predictive = Predictive(model.model, guide=guide, num_samples=num_samples, return_sites=("linear.weight", "obs", "_RETURN")) pred_samples = predictive(Xpred_) pred_summary = summary(pred_samples) mu = pred_summary["_RETURN"] y = pred_summary["obs"] dd = { "mu_mean": post_process_fun(mu["mean"].detach().numpy()), "y_mean": post_process_fun(y["mean"].detach().numpy()), } for i, col in enumerate(cols_Xpred): dd[col] = Xpred[col].values # "major_PHYSICS": x_data[:, -8], ypred_mean = pd.DataFrame(dd) model.pred_summary = {'pred_mean': ypred_mean, 'pred_summary': pred_summary, 'pred_samples': pred_samples} print('stored in model.pred_summary') ypred_proba = None ### No proba if compute_pars.get("probability", False): ypred_proba = model.model.predict_proba(Xpred) return dd['y_mean'], ypred_proba
def test_posterior_predictive_svi_manual_guide(parallel): 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}) svi = SVI(conditioned_model, beta_guide, optim.Adam(dict(lr=1.0)), Trace_ELBO()) for i in range(1000): svi.step(num_trials) posterior_predictive = Predictive(model, guide=beta_guide, num_samples=10000, parallel=parallel, return_sites=["_RETURN"]) marginal_return_vals = posterior_predictive.get_samples( num_trials)["_RETURN"] assert_close(marginal_return_vals.mean(dim=0), torch.ones(5) * 700, rtol=0.05)
def sample_prior(model, num_samples, sites=None): return { k: v.detach().numpy() for k, v in Predictive( model, {}, return_sites=sites, num_samples=num_samples )().items() }
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 marginal(guide, num_samples=25): posterior_predictive = Predictive(guide, num_samples=num_samples) posterior_samples = posterior_predictive.forward(data) mu_mean = posterior_samples['mu'].detach().mean(dim=0) prec_mean = posterior_samples['prec'].detach().mean(dim=0) corr_mean = torch.zeros(T, D, D) for t in range(T): corr_mean[t, ...] = posterior_samples['corr_chol_{}'.format( t)].detach().mean(dim=0) beta_mean = posterior_samples['beta'].detach().mean(dim=0) weights_mean = mix_weights(beta_mean) centers, sigmas, _ = truncate(alpha, mu_mean, prec_mean, corr_mean, weights_mean) return centers, sigmas
def predict_samples(inputs, digits, pre_trained_cvae, epoch_frac): predictive = Predictive(pre_trained_cvae.model, guide=pre_trained_cvae.guide, num_samples=1) preds = predictive(inputs) y_loc = preds['y'].squeeze().detach().cpu().numpy() dfs = pd.DataFrame(data=y_loc) dfs['digit'] = digits.numpy() dfs['epoch'] = epoch_frac return dfs
def get_inference_data(self, data, eight_schools_params): posterior_samples = data.obj.get_samples() model = data.obj.kernel.model posterior_predictive = Predictive( model, posterior_samples).get_samples( eight_schools_params["J"], torch.from_numpy(eight_schools_params["sigma"]).float()) prior = Predictive(model, num_samples=500).get_samples( eight_schools_params["J"], torch.from_numpy(eight_schools_params["sigma"]).float()) return from_pyro( posterior=data.obj, prior=prior, posterior_predictive=posterior_predictive, coords={"school": np.arange(eight_schools_params["J"])}, dims={ "theta": ["school"], "eta": ["school"] }, )
def evaluate_pointwise_pred_density(model, posterior_samples, baseball_dataset): """ Evaluate the log probability density of observing the unseen data (season hits) given a model and posterior distribution over the parameters. """ _, test, player_names = train_test_split(baseball_dataset) at_bats_season, hits_season = test[:, 0], test[:, 1] trace = Predictive(model, posterior_samples).get_vectorized_trace( at_bats_season, hits_season) # Use LogSumExp trick to evaluate $log(1/num_samples \sum_i p(new_data | \theta^{i})) $, # where $\theta^{i}$ are parameter samples from the model's posterior. trace.compute_log_prob() post_loglik = trace.nodes["obs"]["log_prob"] # computes expected log predictive density at each data point exp_log_density = (post_loglik.logsumexp(0) - math.log(post_loglik.shape[0])).sum() logging.info("\nLog pointwise predictive density") logging.info("--------------------------------") logging.info("{:.4f}\n".format(exp_log_density))
def test_end_to_end(model): # Test training. model = AutoReparam()(model) guide = AutoNormal(model) svi = SVI(model, guide, Adam({"lr": 1e-9}), Trace_ELBO()) for step in range(3): svi.step() # Test prediction. predictive = Predictive(model, guide=guide, num_samples=2) samples = predictive() assert set("abc").issubset(samples.keys())
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 predict(self, X): data = X observed = data.columns if 'class' in data.columns: observed = observed.drop(['class']) served_model = Predictive( model=model, guide=guide, return_sites=('class',), num_samples=1) predictions = served_model(data, self.graph, observed, self.n_categories) return predictions['class'].squeeze(0)
def predict(self, n_samples): predictive = Predictive(self.model, guide=self.guide, num_samples=n_samples, return_sites=("linear.weight", "obs", "_RETURN")) samples = predictive(self.X_test) for k, v in samples.items(): self.pred[k] = { "mean": torch.mean(v, 0), "std": torch.std(v, 0), } return
def predict(x, model, guide, num_samples=30): predictive = Predictive(model, guide=guide, num_samples=num_samples) # for a single image, output a mean and sd for each multivariate answer? yhats = predictive(x)["obs"].double() # yhats[0] seems to be integers 0 to 9, len 256 # prediction for one model, for all items in batch # 20, 256 mean = torch.mean(yhats, axis=0) std = torch.std(yhats.float(), 0).cpu().numpy() # yhats outputs a batch size number of predictions for 20 models # yhats seem to be a dictionary of weights return mean, std
def sample_posterior(self, num_samples = 200, attempts = 5): for i in range(attempts): try: samples = Predictive(self.model, guide=self.guide, num_samples=num_samples, return_sites=self.model.get_varnames())() return {varname.split('_')[-1] : samples[varname].detach().numpy() for varname in self.model.get_varnames()} except ValueError: pass raise ValueError('Posterior contains improper values')
def test_broken_plates_smoke(backend): def model(): with pyro.plate("i", 2): a = pyro.sample("a", dist.Normal(0, 1)) pyro.sample("b", dist.Normal(a.mean(-1), 1), obs=torch.tensor(0.0)) guide = AutoGaussian(model, backend=backend) svi = SVI(model, guide, ClippedAdam({"lr": 1e-8}), Trace_ELBO()) for step in range(2): with xfail_if_not_implemented(): svi.step() guide() predictive = Predictive(model, guide=guide, num_samples=2) predictive()
def summarize_posterior(self, raw_expr, encoded_expr, read_depth, num_samples=200, attempts=5): logging.info('Sampling posterior ...') self.posterior_predictive = Predictive(self.model, guide=self.guide, num_samples=num_samples, return_sites=self.var_names) trace = {} for i, batch in enumerate( self.epoch_batch(raw_expr, encoded_expr, read_depth, batch_size=512)): samples = self.posterior_predictive(*batch) if i == 0: for varname in self.global_vars: new_samples = samples[varname].cpu().detach().numpy() trace[varname] = new_samples for varname in self.local_vars: new_samples = samples[varname].cpu().detach().numpy() if not varname in trace: trace[varname] = [] trace[varname].append(new_samples) logging.info('Done {} batches.'.format(str(i + 1))) for varname in self.local_vars: trace[varname] = np.concatenate(trace[varname], axis=1) for varname in self.var_names: self.__setattr__(varname, np.mean(trace[varname], axis=0)) self.beta = self.get_beta() self.gamma = self.get_gamma() self.bias = self.get_bias() self.bn_mean = self.get_bn_mean() self.bn_var = self.get_bn_var() return self
def sample_all1(self, init='init_1', batch_size: int = 10): predictive = Predictive(self.model, guide=self.guide_i[init], num_samples=batch_size) post_samples = { k: v.detach().cpu().numpy() for k, v in self.step_predictive(predictive, self.x_data, self.extra_data_train).items() if k != "data_target" } return (post_samples)
def sample_node1(self, node, init, batch_size: int = 10): predictive = Predictive(self.model, guide=self.guide_i[init], num_samples=batch_size) post_samples = { k: v.detach().cpu().numpy() for k, v in self.step_predictive(predictive, self.x_data, self.extra_data_train).items() if k == node } return (post_samples[node])
def custom_l2_loss(self, model, guide, *args, **kwargs): # run the guide and trace its execution X, num_samples, pred_X_mean, pred_X_var = args predictive = Predictive(self, guide = self.guide, num_samples = num_samples, return_sites = ("obs", "_RETURN")) samples = predictive(X) pred_summary = self.summary(samples) mu = pred_summary["_RETURN"] y = pred_summary["obs"] mu_mean = mu["mean"] mu_std = mu["std"] return l2_loss(pred_X_mean, pred_X_var, mu_mean, mu_std)
def predict(self, data=None): self.latest_data = data print("evaluate model") predictor = Predictive(self.model, guide=self.guide, num_samples=5000) # print("predictor", predictor) if data is None: data = self.boards_test_data prediction = predictor(data) self.latest_means = prediction['obs'].T.detach().numpy().mean(axis=1) self.latest_stds = prediction['obs'].T.detach().numpy().std(axis=1) return self.latest_means
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_pyrocov_smoke(model, Guide, backend): T, P, S, F = 3, 4, 5, 6 dataset = { "features": torch.randn(S, F), "local_time": torch.randn(T, P), "weekly_strains": torch.randn(T, P, S).exp().round(), } guide = Guide(model, backend=backend) svi = SVI(model, guide, ClippedAdam({"lr": 1e-8}), Trace_ELBO()) for step in range(2): with xfail_if_not_implemented(): svi.step(dataset) guide(dataset) predictive = Predictive(model, guide=guide, num_samples=2) predictive(dataset)
def sample(self, x_test, num_samples = 128): self.guide.requires_grad_(False) predictive = Predictive(self, guide = self.guide, num_samples = num_samples, return_sites = ("obs", "_RETURN")) samples = predictive(x_test) pred_summary = self.summary(samples) mu = pred_summary["_RETURN"] y = pred_summary["obs"] mu_mean = mu["mean"] mu_std = mu["std"] mu_var = mu_std.pow(2) y_mean = y["mean"] y_std = y["std"] y_var = y_std.pow(2)
def test_predictive(auto_class): N, D = 3, 2 class RandomLinear(nn.Linear, PyroModule): def __init__(self, in_features, out_features): super().__init__(in_features, out_features) self.weight = PyroSample( dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2)) self.bias = PyroSample( dist.Normal(0., 10.).expand([out_features]).to_event(1)) class LinearRegression(PyroModule): def __init__(self): super().__init__() self.linear = RandomLinear(D, 1) def forward(self, x, y=None): mean = self.linear(x).squeeze(-1) sigma = pyro.sample("sigma", dist.LogNormal(0., 1.)) with pyro.plate('plate', N): return pyro.sample('obs', dist.Normal(mean, sigma), obs=y) x, y = torch.randn(N, D), torch.randn(N) model = LinearRegression() guide = auto_class(model) # XXX: Record `y` as observed in the prototype trace # Is there a better pattern to follow? guide(x, y=y) # Test predictive module model_trace = poutine.trace(model).get_trace(x, y=None) predictive = Predictive(model, guide=guide, num_samples=10) pyro.set_rng_seed(0) samples = predictive(x) for site in prune_subsample_sites(model_trace).stochastic_nodes: assert site in samples with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=torch.jit.TracerWarning) traced_predictive = torch.jit.trace_module(predictive, {"call": (x, )}) f = io.BytesIO() torch.jit.save(traced_predictive, f) f.seek(0) predictive_deser = torch.jit.load(f) pyro.set_rng_seed(0) samples_deser = predictive_deser.call(x) # Note that the site values are different in the serialized guide assert len(samples) == len(samples_deser)
def predict(x, model, guide, transform=False): predictive = Predictive(model, guide=guide, num_samples=20) if transform != False: x = transform(x) # for a single image, output a mean and sd for category yhats = predictive(x)["obs"].double() # yhats[0] seems to be integers 0 to 9, len 256 # prediction for one model, for all items in batch #yhats shape is 20, 256 # predictions for each item in batch and each model # take a mean across 20 models # Doesnt make sense to take mean, should take mode mode = torch.mode(yhats, axis=0) std = torch.std(yhats.float(), 0).numpy() # yhats outputs a batch size number of predictions for 20 models # yhats seem to be a dictionary of weights return mode
def fit_params_estimate(self, data): if self.inference_method == "svi": data = torch.tensor(data).float() predictive = Predictive(self.model, guide=self.multi_norm_guide(), num_samples=1000) svi_samples = { k: v.reshape(1000).detach().numpy() for k, v in predictive(data).items() if not k.startswith("obs") } return { "mean": np.array([v.mean() for v in svi_samples.values()]), "std": np.array([v.std() for v in svi_samples.values()]) } elif self.inference_method == "mcmc": return {"mean": self.mcmc.get_samples()["fit_params"].mean(0)} else: raise NotImplementedError
def predict(self, x_test, num_samples=500, percentiles=(5.0, 50.0, 95.0)): x_test=torch.tensor(x_test) posterior_predictive= Predictive(self.model, guide=self.guide, num_samples=800, 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_intractable_smoke(backend): def model(): i_plate = pyro.plate("i", 2, dim=-1) j_plate = pyro.plate("j", 3, dim=-2) with i_plate: a = pyro.sample("a", dist.Normal(0, 1)) with j_plate: b = pyro.sample("b", dist.Normal(0, 1)) with i_plate, j_plate: c = pyro.sample("c", dist.Normal(a + b, 1)) pyro.sample("d", dist.Normal(c, 1), obs=torch.zeros(3, 2)) guide = AutoGaussian(model, backend=backend) svi = SVI(model, guide, ClippedAdam({"lr": 1e-8}), Trace_ELBO()) for step in range(2): with xfail_if_not_implemented(): svi.step() guide() predictive = Predictive(model, guide=guide, num_samples=2) predictive()
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