Spectral Normalized Gaussian Process (SNGP) Regression

%%capture
%pip install git+https://github.com/lightning-uq-box/lightning-uq-box.git

import os
import tempfile
from functools import partial

import matplotlib.pyplot as plt
import torch
from lightning import Trainer
from lightning.pytorch import seed_everything
from lightning.pytorch.loggers import CSVLogger
from torch.optim import Adam

from lightning_uq_box.datamodules import ToyDUE
from lightning_uq_box.models.fc_resnet import FCResNet
from lightning_uq_box.uq_methods import SNGPRegression
from lightning_uq_box.viz_utils import (
    plot_calibration_uq_toolbox,
    plot_predictions_regression,
    plot_toy_regression_data,
    plot_training_metrics,
)

plt.rcParams["figure.figsize"] = [14, 5]

%load_ext autoreload
%autoreload 2

# temporary directory for saving
my_temp_dir = tempfile.mkdtemp()

seed_everything(42)

Seed set to 42

42

Datamodule

dm = ToyDUE(batch_size=256, n_samples=1000)
X_train, Y_train, X_test, Y_test, train_loader, test_loader, X_gtext, Y_gtext = (
    dm.X_train,
    dm.Y_train,
    dm.X_test,
    dm.Y_test,
    dm.train_dataloader(),
    dm.val_dataloader(),
    dm.X_gtext,
    dm.Y_gtext,
)

fig = plot_toy_regression_data(X_train, Y_train, X_test, Y_test)

Model

feature_extractor = FCResNet(input_dim=1, features=64, depth=4)

sngp = SNGPRegression(
    feature_extractor=feature_extractor,
    loss_fn=torch.nn.MSELoss(),
    optimizer=partial(Adam, lr=3e-3),
)

Trainer

logger = CSVLogger(my_temp_dir)
trainer = Trainer(
    accelerator="cpu",
    max_epochs=100,  # number of epochs we want to train
    logger=logger,  # log training metrics for later evaluation
    log_every_n_steps=1,
    enable_checkpointing=False,
    enable_progress_bar=False,
    default_root_dir=my_temp_dir,
)

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

trainer.fit(sngp, dm)

  | Name              | Type                  | Params | Mode  | FLOPs
----------------------------------------------------------------------------
0 | feature_extractor | FCResNet              | 16.8 K | train | 0
1 | loss_fn           | MSELoss               | 0      | train | 0
2 | normalize         | LayerNorm             | 256    | train | 0
3 | rff               | RandomFourierFeatures | 0      | train | 0
4 | beta              | Linear                | 1.0 K  | train | 0
5 | train_metrics     | MetricCollection      | 0      | train | 0
6 | val_metrics       | MetricCollection      | 0      | train | 0
7 | test_metrics      | MetricCollection      | 0      | train | 0
----------------------------------------------------------------------------
18.0 K    Trainable params
0         Non-trainable params
18.0 K    Total params
0.072     Total estimated model params size (MB)
24        Modules in train mode
0         Modules in eval mode
0         Total Flops
/home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages/lightning/pytorch/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
`Trainer.fit` stopped: `max_epochs=100` reached.

fig = plot_training_metrics(
    os.path.join(my_temp_dir, "lightning_logs"), ["train_loss", "trainRMSE"]
)

Evaluate Predictions

The constructed Data Module contains two possible test variable. X_test are IID samples from the same noise distribution as the training data, while X_gtext (“X ground truth extended”) are dense inputs from the underlying “ground truth” function without any noise that also extends the input range to either side, so we can visualize the method’s UQ tendencies when extrapolating beyond the training data range. Thus, we will use X_gtext for visualization purposes, but use X_test to compute uncertainty and calibration metrics because we want to analyse how well the method has learned the noisy data distribution.

preds = sngp.predict_step(X_gtext.to(sngp.device))

fig = plot_predictions_regression(
    X_train,
    Y_train,
    X_gtext,
    Y_gtext,
    preds["pred"].squeeze(),
    preds["pred_uct"],
    epistemic=preds["epistemic_uct"],
    title="SNGP",
    show_bands=False,
)

preds = sngp.predict_step(X_test.to(sngp.device))

fig = plot_calibration_uq_toolbox(
    preds["pred"].cpu().numpy(),
    preds["pred_uct"].cpu().numpy(),
    Y_test.cpu().numpy(),
    X_test.cpu().numpy(),
)