initial commit of project files

Author: coursekevin

README.md

@@ -0,0 +1,15 @@
+# variationalsparsebayes
+
+This library provides a PyTorch implementation for learning sparse models with with half-cauchy priors using stochastic variational inference.
+
+# Features
+
+The main features of the library are methods for performing:
+
+- [sparse polynomial regression](https://github.com/coursekevin/variationalsparsebayes/blob/main/examples/sparse_poly_regression.py)
+- sparse learning with [precomputed features](https://github.com/coursekevin/variationalsparsebayes/blob/main/examples/support_vectors.py)
+- sparse learning of [Bayesian neural networks](https://github.com/coursekevin/variationalsparsebayes/blob/main/examples/sparse_bnn_regression.py).
+
+To implement your own custom features, you can inherit from the [SparseFeaturesLibrary](https://github.com/coursekevin/variationalsparsebayes/blob/main/variationalsparsebayes/sparse_glm.py) class.
+
+More generally you can use the [SVIHalfCauchyPrior](https://github.com/coursekevin/variationalsparsebayes/blob/main/variationalsparsebayes/svi_half_cauchy.py) class to perform sparse regression with _any_ parameterized model. To do so you need to define a method which takes reparameterized sample weights and computes the expected log-likelihood of your data using these weights.

examples/sparse_bnn_regression.py

@@ -0,0 +1,73 @@
+import torch
+import matplotlib.pyplot as plt
+import torch.optim as optim
+from torch.optim import optimizer
+from variationalsparsebayes import *
+import numpy as np
+import math
+
+plt.style.use("default")
+
+torch.set_default_dtype(torch.float64)
+
+plt.style.use("default")
+
+n_data = 1000
+train_x = torch.linspace(0, 1, 1000)
+noise = 0.2
+train_y = torch.stack(
+    [
+        torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * noise,
+        torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * noise,
+        torch.sin(train_x * (2 * math.pi))
+        + 2 * torch.cos(train_x * (2 * math.pi))
+        + torch.randn(train_x.size()) * noise,
+        -torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * noise,
+    ],
+    -1,
+)
+n_tasks = train_y.size(-1)
+
+
+def sampler():
+    idx = np.random.choice(n_data, size=(64,), replace=False)
+    return (train_x[idx].unsqueeze(-1), train_y[idx, :])
+
+
+if __name__ == "__main__":
+    d = 1
+    model = SparseBNN(
+        in_features=1, out_features=n_tasks, n_hidden=30, n_layers=2, n_reparam=16
+    )
+    num_epochs = 1000
+    beta_warmup_iters = 250
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
+    fig, axs = plt.subplots(1, n_tasks, figsize=(4 * n_tasks, 3))
+    beta = 0.0
+    for j in range(1, num_epochs + 1):
+        optimizer.zero_grad()
+        x_batch, y_batch = sampler()
+        beta = min(1.0, (1.0 * j) / (beta_warmup_iters))
+        model.reparam_sample()
+        loss = -model.elbo(x_batch, y_batch, n_data, beta=beta)
+        loss.backward()
+        optimizer.step()
+
+        if j % 100 == 0:
+            print(f"iter: {j:05d} | loss: {loss.item():.2f}")
+            # plotting
+            x_t = torch.linspace(0, 1, 300)
+            y_pred = model(x_t.unsqueeze(-1))
+            mu = y_pred.mean(0)
+            var = y_pred.var(0) + model.sigma.pow(2)
+            lb = mu - 2 * var.sqrt()
+            ub = mu + 2 * var.sqrt()
+            for j, ax in enumerate(axs):
+                ax.cla()
+                ax.plot(train_x, train_y[:, j], "C7o", alpha=0.3)
+                with torch.no_grad():
+                    ax.plot(x_t, mu[:, j], "C0")
+                    ax.fill_between(x_t, lb[:, j], ub[:, j], color="C0", alpha=0.3)
+            fig.tight_layout()
+            plt.pause(1 / 60)
+    plt.show()

examples/sparse_poly_regression.py

@@ -0,0 +1,45 @@
+import torch
+import matplotlib.pyplot as plt
+import torch.optim as optim
+from torch.optim import optimizer
+from variationalsparsebayes import *
+import numpy as np
+
+torch.set_default_dtype(torch.float64)
+
+plt.style.use("default")
+
+n_data = 1000
+x = torch.linspace(-3, 3, n_data)
+noise = 1e-1
+y = -0.1 * x + 2.0 * x.pow(3)
+y_data = y + torch.randn_like(x) * noise
+
+
+def sampler():
+    idx = np.random.choice(np.arange(n_data), 128, replace=False)
+    return (x[idx].unsqueeze(-1), y_data[idx])
+
+
+if __name__ == "__main__":
+    d = 1
+    features = SparsePolynomialFeatures(d, 8, include_bias=True, input_labels=["x"])
+    model = SparseGLMGaussianLikelihood(d, features, noise=noise, learn_noise=False)
+    opt_summary = model.optimize(
+        data_sampler=sampler, n_data_total=n_data, print_progress=True
+    )
+    # prune basis
+    model.prune_basis()
+    print(model.features)
+    # plotting
+    x_t = x  # torch.linspace(-2, 2, 500)
+    mu, cov = model(x_t.unsqueeze(-1))
+    lb = mu - 2 * torch.sqrt(cov.diag())
+    ub = mu + 2 * torch.sqrt(cov.diag())
+
+    plt.plot(x, y_data, "C7o", alpha=0.3)
+    with torch.no_grad():
+        plt.plot(x_t, mu, "C0")
+        plt.fill_between(x_t, lb, ub, color="C0", alpha=0.3)
+    plt.plot(x, y, "C7--", alpha=1.0)
+    plt.show()

examples/support_vectors.py

@@ -0,0 +1,88 @@
+from variationalsparsebayes.sparse_glm import SparsePrecomputedFeatures
+
+# import sparsebayes
+import matplotlib.pyplot as plt
+import torch
+import numpy as np
+from variationalsparsebayes import *
+import math
+
+import time
+
+torch.set_default_dtype(torch.float64)
+torch.manual_seed(2021)
+
+
+plt.style.use("default")
+
+# n_data = 1000
+n_data = 100
+x = torch.linspace(-10, 10, n_data)
+
+noise = 1e-1  # 0.2
+y = torch.sinc(x / math.pi)
+# y_data = y + (torch.rand(n_data) * (0.2 + 0.2) - 0.2)
+y_data = y + torch.randn(n_data) * noise
+
+
+def rbf(x_in):
+    x_in = x_in.unsqueeze(-1)
+    xn = x.unsqueeze(0)
+    d = (x_in - xn).pow(2)
+    l = 3.0
+    return torch.exp(-d / (l**2))
+
+
+phi = rbf(x)
+
+
+def data_sampler():
+    return (phi, y_data)
+
+
+if __name__ == "__main__":
+    x_test = torch.linspace(-10, 10, n_data * 2)
+    phi_test = rbf(x_test)
+    features = SparsePrecomputedFeatures(n_data)
+    start_time = time.time()
+    model = SparseGLMGaussianLikelihood(
+        n_data, features, noise=noise, learn_noise=False, tau=1.0
+    )
+    opt_summary = model.optimize(
+        data_sampler=data_sampler,
+        n_data_total=n_data,
+        lr=1e-1,
+        beta_warmup_iters=3000,
+        max_iter=20000,
+        n_reparams=20,
+        print_progress=True,
+    )
+    print("Time elapsed: {}".format(time.time() - start_time))
+    # prune basis
+    model.prune_basis()
+    axs = [0]
+    f, axs[0] = plt.subplots(1, 1, sharey=True)
+    mu, cov = model(phi_test)
+    with torch.no_grad():
+        axs[0].plot(x, y_data, "k.", alpha=0.3)
+        axs[0].plot(
+            x_test,
+            mu,
+            "--",
+            label="half-cauchy, {} support vecs.".format(model.num_sparse_features),
+        )
+        lb = mu - 2 * torch.sqrt(cov.diag())
+        ub = mu + 2 * torch.sqrt(cov.diag())
+        axs[0].fill_between(x_test, lb, ub, alpha=0.3)
+        axs[0].axis("off")
+        axs[0].legend()
+        axs[0].plot(
+            x[model.sparse_index],
+            y_data[model.sparse_index],
+            "ro",
+            fillstyle="none",
+            linewidth=1.0,
+            markersize=10,
+        )
+        axs[0].plot(x, y, "k-")
+    plt.show()