Training with Variational Inference

In this example, we train a Bayesian neural network using Variational Inference.

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)

Our neural network will approximate the function \(f(x)=0.4 \sin(4x) + 0.5 \cos(12x) + \epsilon\) over the domain \(x \in [-1, 1]\). \(\epsilon\) represents the noise in our measurement defined as the Gaussian random variable \(\epsilon \sim N(0, 0.05)\).

Below we define the dataset by subclassing 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]

Next, we define our model architecture using UQpy’s BayesianLinear layers and train the model. The model is trained using the Bayes-by-backprop implementation in BBBTrainer.

width = 20
network = nn.Sequential(
    sml.BayesianLinear(1, width),
    nn.ReLU(),
    sml.BayesianLinear(width, width),
    nn.ReLU(),
    sml.BayesianLinear(width, width),
    nn.ReLU(),
    sml.BayesianLinear(width, width),
    nn.ReLU(),
    sml.BayesianLinear(width, 1),
)
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.BBBTrainer(model, optimizer)
print("Starting Training...", end="")
trainer.run(train_data=train_data, epochs=5_000, beta=1e-6, num_samples=10)
print("done")

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()