SVI methods for Bayesian neural networks to deal with trainable sigma

docs/code/scientific_machine_learning/bbb_trainer/svitrainer_trig.py

@@ -0,0 +1,156 @@
+"""
+Training with Variational Inference
+=============================================================
+
+In this example, we train a Bayesian neural network using Variational Inference.
+
+"""
+
+# %% md
+# First, we import the necessary modules and, optionally, set UQpy to print logs to the console.
+
+# %%
+
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import matplotlib.pyplot as plt
+import UQpy.scientific_machine_learning as sml
+
+logger = logging.getLogger("UQpy")  # Optional, display UQpy logs to console
+logger.setLevel(logging.INFO)
+
+# %% md
+# Our neural network will approximate the function :math:`f(x)=0.4 \sin(4x) + 0.5 \cos(12x) + \epsilon` over the domain
+# :math:`x \in [-1, 1]`. :math:`\epsilon` represents the noise in our measurement defined as the Gaussian random
+# variable :math:`\epsilon \sim N(0, 0.05)`.
+#
+# Below we define the dataset by subclassing :py:class:`torch.utils.data.Dataset`.
+
+# %%
+
+
+class SinusoidalDataset(Dataset):
+    def __init__(self, n_samples=20, noise_std=0.05):
+        self.n_samples = n_samples
+        self.noise_std = noise_std
+        self.x = torch.linspace(-1, 1, n_samples).reshape(-1, 1)
+        self.y = torch.tensor(
+            0.4 * np.sin(4 * self.x)
+            + 0.5 * np.cos(12 * self.x)
+            + np.random.normal(0, self.noise_std, self.x.shape),
+            dtype=torch.float,
+        )
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, item):
+        return self.x[item], self.y[item]
+
+
+# %% md
+# Next, we define our model architecture using UQpy's :py:class:`BayesianLinear` layers and train the model.
+# The model is trained using the Bayes-by-backprop implementation in :py:class:`BBBTrainer`.
+
+# %%
+beta = 1e-6  # weight on the KL divergence term in the loss function scaled as
+n_samples = 20 # The same parameter as in bbbtrainer_trig.py
+
+width = 20
+class BayesianNeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layer1 = sml.BayesianLinear(1, width)
+        self.layer2 = sml.BayesianLinear(width, width)
+        self.layer3 = sml.BayesianLinear(width, width)
+        self.layer4 = sml.BayesianLinear(width, width)
+        self.layer5 = sml.BayesianLinear(width, 1)
+        self.activation = nn.ReLU()
+        self.sigma = nn.Parameter(torch.tensor((n_samples * beta/2)**0.5), requires_grad=False)
+
+    def forward(self, x, return_log_prob: bool = False):
+        x = self.activation(self.layer1(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer2(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer3(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer4(x, return_log_prob=return_log_prob))
+        x = self.layer5(x, return_log_prob=return_log_prob)
+        return x
+
+network = BayesianNeuralNetwork()
+model = sml.FeedForwardNeuralNetwork(network)
+
+dataset = SinusoidalDataset()
+train_data = DataLoader(dataset, batch_size=20, shuffle=True)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+trainer = sml.SVITrainer(model, optimizer)
+print("Starting Training...", end="")
+trainer.run(train_data=train_data, epochs=5_000, num_samples=10)
+print("done")
+
+# %% md
+# That's the hard part done! We defined our dataset, our model, and then fit the model to the data.
+# The rest of this example plots the model predictions to compare them to the exact solution.
+
+# %%
+
+# Plot training history
+fig, ax = plt.subplots()
+ax.semilogy(trainer.history["train_loss"])
+ax.set_title("Bayes By Backpropagation Training Loss")
+ax.set(xlabel="Epoch", ylabel="Loss")
+
+# Plot model predictions
+x_noisy = dataset.x
+y_noisy = dataset.y
+x_exact = torch.linspace(-1, 1, 1000).reshape(-1, 1)
+y_exact = 0.4 * torch.sin(4 * x_exact) + 0.5 * torch.cos(12 * x_exact)
+
+# compute mean prediction from model
+model.eval()
+model.sample(False)
+print("BNN is deterministic:", model.is_deterministic())
+with torch.no_grad():
+    mean_prediction = model(x_exact)
+
+# compute stochastic prediction from model
+model.sample(True)
+print("BNN is deterministic:", model.is_deterministic())
+n = 1_000
+samples = torch.zeros(len(x_exact), n)
+with torch.no_grad():
+    for i in range(n):
+        samples[:, i] = model(x_exact).squeeze()
+standard_deviation = torch.std(samples, dim=1)
+
+# convert tensors to numpy arrays for matplotlib
+x_exact = x_exact.squeeze().detach().numpy()
+y_exact = y_exact.squeeze().detach().numpy()
+mean_prediction = mean_prediction.squeeze().detach().numpy()
+standard_deviation = standard_deviation.squeeze().detach().numpy()
+
+fig, ax = plt.subplots()
+ax.scatter(x_noisy, y_noisy, label="Training Data", color="black")
+ax.plot(
+    x_exact,
+    y_exact,
+    label="Exact",
+    color="black",
+    linestyle="dashed",
+)
+ax.plot(x_exact, mean_prediction, label="Model $\mu$", color="tab:blue")
+ax.fill_between(
+    x_exact,
+    mean_prediction - (3 * standard_deviation),
+    mean_prediction + (3 * standard_deviation),
+    label="$\mu \pm 3\sigma$,",
+    color="tab:blue",
+    alpha=0.3,
+)
+ax.set_title("Bayesian Neural Network Predictions")
+ax.set(xlabel="x", ylabel="f(x)")
+ax.legend()
+
+plt.show()

docs/code/scientific_machine_learning/bbb_trainer/svitrainer_trig_sigma.py

@@ -0,0 +1,196 @@
+"""
+Training with Variational Inference
+=============================================================
+
+In this example, we train a Bayesian neural network using Variational Inference.
+
+"""
+
+# %% md
+# First, we import the necessary modules and, optionally, set UQpy to print logs to the console.
+
+# %%
+
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import matplotlib.pyplot as plt
+import UQpy.scientific_machine_learning as sml
+from torch.nn import functional as F
+
+logger = logging.getLogger("UQpy")  # Optional, display UQpy logs to console
+logger.setLevel(logging.INFO)
+
+# %% md
+# Our neural network will approximate the function :math:`f(x)=0.4 \sin(4x) + 0.5 \cos(12x) + \epsilon` over the domain
+# :math:`x \in [-1, 1]`. :math:`\epsilon` represents the noise in our measurement defined as the Gaussian random
+# variable :math:`\epsilon \sim N(0, 0.05)`.
+#
+# Below we define the dataset by subclassing :py:class:`torch.utils.data.Dataset`.
+
+# %%
+
+
+class SinusoidalDataset(Dataset):
+    def __init__(self, n_samples=20, noise_std=0.05):
+        self.n_samples = n_samples
+        self.noise_std = noise_std
+        self.x = torch.linspace(-1, 1, n_samples).reshape(-1, 1)
+        self.y = torch.tensor(
+            0.4 * np.sin(4 * self.x)
+            + 0.5 * np.cos(12 * self.x)
+            + np.random.normal(0, self.noise_std, self.x.shape),
+            dtype=torch.float,
+        )
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, item):
+        return self.x[item], self.y[item]
+
+
+# %% md
+# Next, we define our model architecture using UQpy's :py:class:`BayesianLinear` layers and train the model.
+# The model is trained using the Bayes-by-backprop implementation in :py:class:`BBBTrainer`.
+
+# %%
+
+width = 20
+# class Sigma(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.prior_concentration = 0.1
+#         self.prior_rate = 4.0
+#         self.prior_distribution = torch.distributions.Gamma(self.prior_concentration, self.prior_rate)
+#         self.posterior_concentration = nn.Parameter(torch.tensor(1.0))
+#         self.posterior_rate = nn.Parameter(torch.tensor(4.0))
+#     def sample(self):
+#         gamma_dist = torch.distributions.Gamma(self.posterior_concentration, self.posterior_rate)
+#         gamma_sample = (gamma_dist.rsample()).clamp(min=1e-6)
+#         log_pdf_posterior = gamma_dist.log_prob(gamma_sample)
+#         log_pdf_prior = self.prior_distribution.log_prob(gamma_sample)
+#         self.qs = log_pdf_posterior
+#         self.ps = log_pdf_prior
+#         return gamma_sample
+
+
+class Sigma(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Define prior distribution as Gamma
+        self.prior_concentration = 0.1
+        self.prior_rate = 4.0
+        self.prior_distribution = torch.distributions.Gamma(self.prior_concentration, self.prior_rate)
+        # Assume variational distribution as Normal
+        self.loc = nn.Parameter(torch.tensor(-4.0))
+        self.log_scale = nn.Parameter(torch.tensor(-1.0)) 
+    def sample(self):
+        scale = F.softplus(self.log_scale)
+        base_dist = torch.distributions.Normal(self.loc, scale)
+        z = base_dist.rsample()
+        sigma_sample = z.exp().clamp(min=1e-6)
+
+        self.qs = base_dist.log_prob(z) - torch.log(sigma_sample)  # q_sigma(sigma) = q_z(z) * |dz/dsigma| = q_z(z) * |d log(sigma)/dsigma| = q_z(z) * 1/sigma
+        self.ps = self.prior_distribution.log_prob(sigma_sample)
+        return sigma_sample
+
+class BayesianNeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layer1 = sml.BayesianLinear(1, width)
+        self.layer2 = sml.BayesianLinear(width, width)
+        self.layer3 = sml.BayesianLinear(width, width)
+        self.layer4 = sml.BayesianLinear(width, width)
+        self.layer5 = sml.BayesianLinear(width, 1)
+        self.activation = nn.ReLU()
+        self.sigma_class = Sigma()
+
+    def forward(self, x, return_log_prob: bool = False):
+        x = self.activation(self.layer1(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer2(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer3(x, return_log_prob=return_log_prob))
+        x = self.activation(self.layer4(x, return_log_prob=return_log_prob))
+        x = self.layer5(x, return_log_prob=return_log_prob)
+        self.sigma = self.sigma_class.sample()
+        return x
+
+network = BayesianNeuralNetwork()
+model = sml.FeedForwardNeuralNetwork(network)
+
+dataset = SinusoidalDataset()
+train_data = DataLoader(dataset, batch_size=20, shuffle=True)
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+trainer = sml.SVITrainer(model, optimizer)
+print("Starting Training...", end="")
+trainer.run(train_data=train_data, epochs=5_000, num_samples=10, tolerance=-10000000.0)
+print("done")
+
+# %% md
+# That's the hard part done! We defined our dataset, our model, and then fit the model to the data.
+# The rest of this example plots the model predictions to compare them to the exact solution.
+
+# %%
+
+# Plot training history
+fig, ax = plt.subplots()
+ax.semilogy(trainer.history["train_loss"])
+ax.set_title("Bayes By Backpropagation Training Loss")
+ax.set(xlabel="Epoch", ylabel="Loss")
+
+# Plot model predictions
+x_noisy = dataset.x
+y_noisy = dataset.y
+x_exact = torch.linspace(-1, 1, 1000).reshape(-1, 1)
+y_exact = 0.4 * torch.sin(4 * x_exact) + 0.5 * torch.cos(12 * x_exact)
+
+# compute mean prediction from model
+model.eval()
+model.sample(False)
+print("BNN is deterministic:", model.is_deterministic())
+with torch.no_grad():
+    mean_prediction = model(x_exact)
+
+# compute stochastic prediction from model
+model.sample(True)
+print("BNN is deterministic:", model.is_deterministic())
+n = 1_000
+samples = torch.zeros(len(x_exact), n)
+with torch.no_grad():
+    for i in range(n):
+        samples[:, i] = model(x_exact).squeeze()
+standard_deviation = torch.std(samples, dim=1)
+
+# convert tensors to numpy arrays for matplotlib
+x_exact = x_exact.squeeze().detach().numpy()
+y_exact = y_exact.squeeze().detach().numpy()
+mean_prediction = mean_prediction.squeeze().detach().numpy()
+standard_deviation = standard_deviation.squeeze().detach().numpy()
+
+fig, ax = plt.subplots()
+ax.scatter(x_noisy, y_noisy, label="Training Data", color="black")
+ax.plot(
+    x_exact,
+    y_exact,
+    label="Exact",
+    color="black",
+    linestyle="dashed",
+)
+ax.plot(x_exact, mean_prediction, label="Model $\mu$", color="tab:blue")
+ax.fill_between(
+    x_exact,
+    mean_prediction - (3 * standard_deviation),
+    mean_prediction + (3 * standard_deviation),
+    label="$\mu \pm 3\sigma$,",
+    color="tab:blue",
+    alpha=0.3,
+)
+ax.set_title("Bayesian Neural Network Predictions")
+ax.set(xlabel="x", ylabel="f(x)")
+ax.legend()
+
+plt.show()
+
+# %%

src/UQpy/scientific_machine_learning/baseclass/NormalBayesianLayer.py

@@ -5,6 +5,7 @@
 from typing import Union
 from beartype import beartype
 from UQpy.utilities.ValidationTypes import PositiveFloat
+import math
 
 
 @beartype
@@ -92,7 +93,7 @@ def reset_parameters(self):
             rho = getattr(self, f"{name}_rho")
             rho.data.normal_(*self.posterior_rho_initial)
 
-    def get_bayesian_weights(self) -> tuple:
+    def get_bayesian_weights(self, return_log_prob: bool = False) -> tuple:
         """Get the weights for the Bayesian layer.
 
         If ``sampling`` is ``True``, then sample weights from their respective distributions.
@@ -102,6 +103,8 @@ def get_bayesian_weights(self) -> tuple:
         """
         if self.sampling:
             weights = []
+            log_qs = [] if return_log_prob else None
+            log_ps = [] if return_log_prob else None
             for name in self.parameter_shapes:
                 if self.parameter_shapes[name] is None:
                     weights.append(None)
@@ -113,8 +116,16 @@ def get_bayesian_weights(self) -> tuple:
                 sigma = torch.log1p(torch.exp(rho))
                 weight = mu + (epsilon * sigma)
                 weights.append(weight)
+                if return_log_prob:
+                    log_pdf_weights = -(0.5 * epsilon.pow(2) + sigma.log()).sum()  # Log PDF of weights (without constants)
+                    log_pdf_prior = -0.5 * ((mu + epsilon * sigma - self.prior_mu) / self.prior_sigma).pow(2).sum() - weight.numel() * math.log(self.prior_sigma)
+                    log_qs.append(log_pdf_weights)
+                    log_ps.append(log_pdf_prior)
         else:
             weights = (getattr(self, f"{name}_mu") for name in self.parameter_shapes)
+        if return_log_prob:
+            self.qs = torch.stack(log_qs).sum()
+            self.ps = torch.stack(log_ps).sum()
         return tuple(weights)
 
     def sample(self, mode: bool = True):