samuelinferences/transformers-can-do-bayesian-inference

1import gradio as gr2import numpy as np3import matplotlib.pyplot as plt4import gpytorch5import torch6import sys7 8import gpytorch9 10# We will use the simplest form of GP model, exact inference11class ExactGPModel(gpytorch.models.ExactGP):12    def __init__(self, train_x, train_y, likelihood):13        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)14        self.mean_module = gpytorch.means.ConstantMean()15        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())16 17    def forward(self, x):18        mean_x = self.mean_module(x)19        covar_x = self.covar_module(x)20        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)21 22def get_model(x, y, hyperparameters):23    likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1.e-9))24    model = ExactGPModel(x, y, likelihood)25    model.likelihood.noise = torch.ones_like(model.likelihood.noise) * hyperparameters["noise"]26    model.covar_module.outputscale = torch.ones_like(model.covar_module.outputscale) * hyperparameters["outputscale"]27    model.covar_module.base_kernel.lengthscale = torch.ones_like(model.covar_module.base_kernel.lengthscale) * \28                                                 hyperparameters["lengthscale"]29    return model, likelihood30 31 32 33excuse = "Please only specify numbers, x values should be in [0,1] and y values in [-1,1]."34excuse_max_examples = "This model is trained to work with up to 4 input points."35hyperparameters = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .1, 'fast_computations': (False,False,False)}36 37 38conf = .539 40def mean_and_bounds_for_gp(x,y,test_xs):41    gp_model, likelihood = get_model(x,y,hyperparameters)42    gp_model.eval()43    l = likelihood(gp_model(test_xs))44    means = l.mean.squeeze()45    varis = torch.diagonal(l.covariance_matrix.squeeze())46    stds = varis.sqrt()47    return means, means-stds, means+stds48 49 50def mean_and_bounds_for_pnf(x,y,test_xs, choice):51    sys.path.append('prior-fitting/')52    model = torch.load(f'onefeature_gp_ls.1_pnf_{choice}.pt')53 54    logits = model((torch.cat([x,test_xs],0).unsqueeze(1),y.unsqueeze(1)),single_eval_pos=len(x))55    bounds = model.criterion.quantile(logits,center_prob=.682).squeeze(1)56    return model.criterion.mean(logits).squeeze(1), bounds[:,0], bounds[:,1]57 58def plot_w_conf_interval(ax_or_plt, x, m, lb, ub, color, label_prefix):59    ax_or_plt.plot(x.squeeze(-1),m, color=color, label=label_prefix+' mean')60    ax_or_plt.fill_between(x.squeeze(-1), lb, ub, alpha=.1, color=color, label=label_prefix+' conf. interval')61 62 63 64 65@torch.no_grad()66def infer(table, choice):67    vfunc = np.vectorize(lambda s: len(s))68    non_empty_row_mask = (vfunc(table).sum(1) != 0)69    table = table[non_empty_row_mask]70 71    try:72        table = table.astype(np.float32)73    except ValueError:74        return excuse, None75    x = torch.tensor(table[:,0]).unsqueeze(1)76    y = torch.tensor(table[:,1])77    fig = plt.figure(figsize=(8,4),dpi=1000)78 79    if len(x) > 4:80        return excuse_max_examples, None81    if (x<0.).any() or (x>1.).any() or (y<-1).any() or (y>1).any():82        return excuse, None83 84    plt.scatter(x,y, color='black', label='Examples in given dataset')85 86 87    88    test_xs = torch.linspace(0,1,100).unsqueeze(1)89    90    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_gp(x,y,test_xs), 'green', 'GP')91    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_pnf(x,y,test_xs, choice), 'blue', 'PFN')92    93    plt.legend(ncol=2,bbox_to_anchor=[0.5,-.14],loc="upper center")94    plt.xlabel('x')95    plt.ylabel('y')96    plt.tight_layout()97 98    99    return 'There you go, your plot. 📈', plt.gcf()100 101iface = gr.Interface(fn=infer,102                     title='GP Posterior Approximation with Transformers',103                     description='''This is a demo of PFNs as we describe them in our recent paper (https://openreview.net/forum?id=KSugKcbNf9).104Lines represent means and shaded areas are the confidence interval (68.2% quantile). In green, we have the ground truth GP posterior and in blue we have our approximation.105We provide three models that are architecturally the same, but with different training budgets.106The GP (approximated) uses an RBF Kernel with a little noise (1e-4), 0 mean and a length scale of 0.1.107                     ''',108                     article="<p style='text-align: center'><a href='https://arxiv.org/abs/2112.10510'>Paper: Transformers Can Do Bayesian Inference</a></p>",109                     inputs=[110                         gr.inputs.Dataframe(headers=["x", "y"], datatype=["number", "number"], type='numpy', default=[['.25','.1'],['.75','.4']], col_count=2, label='The data: you can change this and increase the number of data points using the `enter` key.'),111                         gr.inputs.Radio(['160K','800K','4M'], type="value", default='4M', label='Number of Sampled Datasets in Training (Training Costs), higher values yield better results')112                     ], outputs=["text",gr.outputs.Plot(type="matplotlib")])113iface.launch()114 115 116 117