lightning_uq_box.uq_methods

Single Forward Pass Methods

Mean Variance Estimation

Mean Variance Base

class lightning_uq_box.uq_methods.mean_variance_estimation.MVEBase(model: ~torch.nn.modules.module.Module, burnin_epochs: int, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Mean Variance Estimation Network Base Class.

If you use this model in your research, please cite the following paper:

• https://ieeexplore.ieee.org/document/374138

__init__(model: ~torch.nn.modules.module.Module, burnin_epochs: int, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of Deterministic Gaussian Model.

Parameters:

• model – pytorch model

• burnin_epochs – number of burnin epochs before switiching to NLL

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

Mean Variance Regression

class lightning_uq_box.uq_methods.mean_variance_estimation.MVERegression(model: ~torch.nn.modules.module.Module, burnin_epochs: int, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Mean Variance Estimation Model for Regression that is trained with NLL.

If you use this model in your research, please cite the following paper:

• https://ieeexplore.ieee.org/document/374138

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, burnin_epochs: int, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of Mean Variance Estimation Model for Regression.

Parameters:

• model – pytorch model

• burnin_epochs – number of burnin epochs before switiching to NLL

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Quantile Regression

Quantile Regression Base

class lightning_uq_box.uq_methods.quantile_regression.QuantileRegressionBase(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Quantile Regression Base Module.

If you use this model in your research, please cite the following paper:

• https://www.jstor.org/stable/1913643

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of Quantile Regression Model.

Parameters:

• model – pytorch model

• loss_fn – loss function

• quantiles – quantiles to compute

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

Quantile Regression

class lightning_uq_box.uq_methods.quantile_regression.QuantileRegression(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Quantile Regression Module for Regression.

If you use this model in your research, please cite the following paper:

• https://www.jstor.org/stable/1913643

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of Quantile Regression Model.

Parameters:

• model – pytorch model

• optimizer – optimizer used for training

• loss_fn – loss function

• quantiles – quantiles to compute

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

Parameters:

out – output from self.forward() [batch_size x num_outputs]

Returns:

extracted mean used for metric computation [batch_size x 1]

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with Quantile Regression.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

predicted uncertainties

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Quantile Pixelwise Regression

class lightning_uq_box.uq_methods.quantile_regression.QuantilePxRegression(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False)

Quantile Regression for Pixelwise Regression.

Added in version 0.2.0.

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, quantiles: list[float] = [0.1, 0.5, 0.9], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False) → None

Initialize a new instance of Quantile Regression Model.

Parameters:

• model – pytorch model

• optimizer – optimizer used for training

• loss_fn – loss function

• quantiles – quantiles to compute

• freeze_backbone – whether to freeze the backbone

• freeze_decoder – whether to freeze the decoder

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

Parameters:

out – output from self.forward() [batch_size x num_outputs x height x width]

Returns:

extracted mean used for metric computation [batch_size x 1 x height x width]

on_test_start() → None

Create logging directory and initialize metrics.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

Parameters:

• batch – batch of testing data

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with Quantile Regression.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

predicted uncertainties

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Deep Evidential Regression

Deep Evidential Regression

class lightning_uq_box.uq_methods.deep_evidential_regression.DER(model: ~torch.nn.modules.module.Module, coeff: float = 0.01, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deep Evidential Regression Model.

Following the suggested implementation of the Unreasonable Effectiveness of Deep Evidential Regression

If you use this model in your work, please cite:

• https://arxiv.org/abs/2205.10060

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, coeff: float = 0.01, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Base Model.

Parameters:

• model – pytorch model

• coeff – coefficient for the DER loss from the predictive distribution

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

forward(X: Tensor) → Any

Forward pass of the model.

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction Step Deep Evidential Regression.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

dictionary with predictions and uncertainty measures

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

Parameters:

out – output from self.forward() [batch_size x 4]

Returns:

extracted mean used for metric computation [batch_size x 1]

compute_aleatoric_uct(beta: Tensor, alpha: Tensor, nu: Tensor) → Tensor

Compute the aleatoric uncertainty for DER model.

Equation 10 of the paper

Parameters:

• beta – beta output DER model

• alpha – alpha output DER model

• nu – nu output DER model

Returns:

Aleatoric Uncertainty

compute_epistemic_uct(nu: Tensor) → Tensor

Compute the aleatoric uncertainty for DER model.

Equation 10: of the paper

Parameters:

nu – nu output DER model

Returns:

Epistemic Uncertainty

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Deep Evidential Pixelwise Regression

class lightning_uq_box.uq_methods.deep_evidential_regression.DERPxRegression(model: ~torch.nn.modules.module.Module, coeff: float = 0.01, freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False)

Deep Evidential Regression Model for Pixelwise Regression with NLL.

Added in version 0.2.0.

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, coeff: float = 0.01, freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False) → None

Initialize a new instance of the DER for Pixelwise Regression.

Parameters:

• model – pytorch model

• coeff – coefficient for the DER loss from the predictive distribution

• freeze_backbone – whether to freeze the model backbone

• freeze_decoder – whether to freeze the model decoder

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

Parameters:

out – output from self.forward() [batch_size x 4]

Returns:

extracted mean used for metric computation [batch_size x 1]

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Zig Zag

Zig Zag Base

class lightning_uq_box.uq_methods.zigzag.ZigZagBase(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, blank_const: int = -100, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

ZigZag Uncertainty Quantification Base.

If you use this method in your work, please cite:

• https://openreview.net/forum?id=QSvb6jBXML

Added in version 0.2.

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, blank_const: int = -100, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None) → None

Initialize a new instance of ZigZag.

Parameters:

• model – PyTorch model.

• loss_fn – Loss function.

• blank_const – constant for the blank zig zag input, should be a value far from training targets

• freeze_backbone – Whether or not to freeze the backbone.

• optimizer – Optimizer.

• lr_scheduler – Learning rate scheduler.

check_input_layer() → None

Check whether the input layer is linear or convolutional.

This has an effect on how the inputs are concatenated for ZigZag.

forward(x: Tensor, y: Tensor | None = None, training: bool = False) → Tensor

Forward pass of Zig Zag method.

Parameters:

• x – Input tensor.

• y – Target tensor.

• training – Whether or not the model is in training mode, which affects the Zig Zag operation for conv input layers

Returns:

Output of model with Zig Zag operation.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

Zig Zag Regression

class lightning_uq_box.uq_methods.zigzag.ZigZagRegression(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, blank_const: int = -100, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Zig Zag Uncertainty Estimation for Regression.

If you use this method in your work, please cite:

• https://openreview.net/forum?id=QSvb6jBXML

Added in version 0.2.

pred_file_name = 'preds.csv'

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict Step.

Conducts two forward passes. One with the input, and a second one with the input and the output of the first forward pass.

Parameters:

• X – prediction input tensor

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Zig Zag Classification

class lightning_uq_box.uq_methods.zigzag.ZigZagClassification(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, blank_const: int = -100, task: str = 'multiclass', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Zig Zag Uncertainty Estimation for Classification.

If you use this method in your work, please cite:

• https://openreview.net/forum?id=QSvb6jBXML

Added in version 0.2.

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, blank_const: int = -100, task: str = 'multiclass', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None) → None

Initialize a new instance of ZigZag for classification.

Parameters:

• model – PyTorch model.

• loss_fn – Loss function.

• blank_const – constant for the blank zig zag input, should be a value far from training targets

• task – Task type. One of “binary”, “multiclass”, “multilabel”.

• freeze_backbone – Whether or not to freeze the backbone.

• optimizer – Optimizer.

• lr_scheduler – Learning rate scheduler.

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Mixture Density Networks

Mixture Density Regression

Approximate Bayesian Methods

Monte Carlo Dropout

MC-Dropout Base

class lightning_uq_box.uq_methods.mc_dropout.MCDropoutBase(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

MC-Dropout Base class.

If you use this model in your research, please cite the following paper:

• https://proceedings.mlr.press/v48/gal16.html

__init__(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of MCDropoutModel.

Parameters:

• model – pytorch model with dropout layers

• num_mc_samples – number of MC samples during prediction

• loss_fn – loss function

• dropout_layer_names – names of dropout layers to activate during prediction

• freeze_backbone – freeze backbone during training

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

activate_dropout() → None

Activate dropout layers.

MC-Dropout Regression

class lightning_uq_box.uq_methods.mc_dropout.MCDropoutRegression(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, burnin_epochs: int = 0, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

MC-Dropout Model for Regression.

If you use this model in your research, please cite the following paper:

• https://proceedings.mlr.press/v48/gal16.html

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, burnin_epochs: int = 0, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of MC-Dropout Model for Regression.

Parameters:

• model – pytorch model with dropout layers

• num_mc_samples – number of MC samples during prediction

• loss_fn – loss function

• burnin_epochs – number of burnin epochs before using the loss_fn

• dropout_layer_names – names of dropout layers to activate during prediction

• freeze_backbone – freeze backbone during training

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler from the predictive distribution

setup_task() → None

Set up task specific attributes.

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation..

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict steps via Monte Carlo Sampling.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

mean and standard deviation of MC predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

MC-Dropout Classification

class lightning_uq_box.uq_methods.mc_dropout.MCDropoutClassification(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, task: str = 'multiclass', dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

MC-Dropout Model for Classification.

If you use this model in your research, please cite the following paper:

• https://proceedings.mlr.press/v48/gal16.html

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, task: str = 'multiclass', dropout_layer_names: list[str] = [], freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of MC-Dropout Model for Classification.

Parameters:

• model – pytorch model with dropout layers

• num_mc_samples – number of MC samples during prediction

• loss_fn – loss function

• task – classification task, one of [‘binary’, ‘multiclass’, ‘multilabel’]

• dropout_layer_names – names of dropout layers to activate during prediction

• freeze_backbone – freeze backbone during training

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Extract mean output from model.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict steps via Monte Carlo Sampling.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

mean and standard deviation of MC predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

MC-Dropout Segmentation

class lightning_uq_box.uq_methods.mc_dropout.MCDropoutSegmentation(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, task: str = 'multiclass', dropout_layer_names: list[str] = [], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False)

MC-Dropout Model for Segmentation.

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, task: str = 'multiclass', dropout_layer_names: list[str] = [], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False) → None

Initialize a new instance of MC-Dropout Model for Segmentation.

Parameters:

• model – pytorch model with dropout layers

• num_mc_samples – number of MC samples during prediction

• loss_fn – loss function

• task – classification task, one of [‘binary’, ‘multiclass’, ‘multilabel’]

• dropout_layer_names – names of dropout layers to activate during prediction

• freeze_backbone – whether to freeze the model backbone, by default this is supported for torchseg Unet models

• freeze_decoder – whether to freeze the model decoder, by default this is supported for torchseg Unet models

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

setup_task() → None

Set up task specific attributes for segmentation.

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict steps via Monte Carlo Sampling.

Parameters:

• X – prediction batch of shape [batch_size x num_channels x height x width]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

mean and standard deviation of MC predictions

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

MC-Dropout Pixelwise Regression

class lightning_uq_box.uq_methods.mc_dropout.MCDropoutPxRegression(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, burnin_epochs: int = 0, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False)

MC-Dropout Model for Pixel-wise Regression.

Added in version 0.2.0.

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, num_mc_samples: int, loss_fn: ~torch.nn.modules.module.Module, burnin_epochs: int = 0, dropout_layer_names: list[str] = [], freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False) → None

Initialize a new instance of MC-Dropout Model for Pixel-wise Regression.

Parameters:

• model – pytorch model with dropout layers

• num_mc_samples – number of MC samples during prediction

• loss_fn – loss function

• burnin_epochs – number of burnin epochs before using the loss_fn

• dropout_layer_names – names of dropout layers to activate during prediction

• freeze_backbone – freeze backbone during training

• freeze_decoder – freeze decoder during training

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation..

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Laplace Approximation

Laplace Base

class lightning_uq_box.uq_methods.laplace_model.LaplaceBase(laplace_model: Laplace, pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None)

Laplace Approximation Method.

This is a lightning module wrapper for the Laplace library. # noqa: E501

If you use this model in your research, please cite the following papers:

• https://arxiv.org/abs/2106.14806

pred_file_name = 'preds.csv'

__init__(laplace_model: Laplace, pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None) → None

Initialize a new instance of Laplace Model Wrapper.

Parameters:

• laplace_model – initialized Laplace model

• pred_type – prediction type, one of [‘glm’, ‘nn’]

• link_approx – link function approximation, one of [‘mc’, ‘probit’, ‘bridge’] for pred_type=’nn’ only ‘mc’ is supported

• num_samples – number of samples for prediction, if specified will call predictive_samples instead of predictive method in Laplace library

Changed in version 0.2: Add ‘pred_type’ and ‘link_approx’ arguments.

setup_task() → None

Set up task.

property num_input_features: int

Retrieve input dimension to the model.

Returns:

number of input dimension to the model

property num_outputs: int

Retrieve output dimension to the model.

Returns:

number of output dimension to model

on_test_start() → None

Fit the Laplace approximation before testing.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test step.

on_test_epoch_end()

Log epoch-level test metrics.

Laplace Regression

class lightning_uq_box.uq_methods.laplace_model.LaplaceRegression(laplace_model: Laplace, pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None, tune_prior_precision: bool = True, tune_sigma_noise: bool = False, tuning_lr: float = 0.001, n_epochs_tuning: int = 100)

Laplace Approximation Wrapper for regression.

This is a lightning module wrapper for the Laplace library.

If you use this model in your research, please cite the following paper:

• https://arxiv.org/abs/2106.14806

__init__(laplace_model: Laplace, pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None, tune_prior_precision: bool = True, tune_sigma_noise: bool = False, tuning_lr: float = 0.001, n_epochs_tuning: int = 100) → None

Initialize a new instance of Laplace Model Wrapper for regression.

Parameters:

• laplace_model – initialized Laplace model

• pred_type – prediction type, one of [‘glm’, ‘nn’]

• link_approx – link function approximation, one of [‘mc’, ‘probit’, ‘bridge’] for pred_type=’nn’ only ‘mc’ is supported

• num_samples – number of samples for prediction, if specified will call predictive_samples instead of predictive method in Laplace library

• tune_prior_precision – whether to tune prior precision

• tune_sigma_noise – whether to tune sigma noise

• tuning_lr – learning rate for tuning prior precision and sigma

• n_epochs_tuning – number of epochs to tune prior precision and sigma

on_test_start() → None

Fit the Laplace approximation before testing.

setup_task() → None

Set up task specific attributes.

forward(X: Tensor) → dict[str, Tensor]

Fitted Laplace Model Forward Pass.

Parameters:

X – tensor of data to run through the model [batch_size, input_dim]

Returns:

output from the laplace model

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with Laplace Approximation.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Laplace Classification

class lightning_uq_box.uq_methods.laplace_model.LaplaceClassification(laplace_model: Laplace, task: str = 'multiclass', pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None)

Laplace Approximation Wrapper for classification.

This is a lightning module wrapper for the Laplace library.

If you use this model in your research, please cite the following paper:

• https://arxiv.org/abs/2106.14806

valid_tasks = ['binary', 'multiclass']

__init__(laplace_model: Laplace, task: str = 'multiclass', pred_type: str = 'glm', link_approx: str = 'probit', num_samples: int | None = None) → None

Initialize a new instance of Laplace Model Wrapper for Classification.

Parameters:

• laplace_model – initialized Laplace model

• task – classification task, one of [‘binary’, ‘multiclass’]

• pred_type – prediction type, one of [‘glm’, ‘nn’]

• link_approx – link function approximation, one of [‘mc’, ‘probit’, ‘bridge’] for pred_type=’nn’ only ‘mc’ is supported

• num_samples – number of samples for prediction, if specified will call predictive_samples instead of predictive method in Laplace library

forward(X: Tensor, **kwargs: Any) → dict[str, Tensor]

Fitted Laplace Model Forward Pass.

Parameters:

• X – tensor of data to run through the model [batch_size, input_dim]

• kwargs – additional arguments for laplace forward pass

Returns:

output from the laplace model

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with Laplace Approximation.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Bayesian Neural Networks ELBO

BNN ELBO Base

class lightning_uq_box.uq_methods.bnn_vi_elbo.BNN_VI_ELBO_Base(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayes By Backprop Base with Variational Inference (VI).

If you use this model in your work, please cite:

• https://arxiv.org/abs/1505.05424

__init__(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Model instance.

Parameters:

• model – pytorch model that will be converted into a BNN

• criterion – loss function used for optimization

• beta – beta factor for negative elbo loss computation, should be number of weights and biases

• num_mc_samples_train – number of MC samples during training when computing the negative ELBO loss. When setting num_mc_samples_train=1, this is just Bayes by Backprop.

• num_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• bayesian_layer_type – flipout or reparameterization

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if num_mc_samples_train is not positive.

• AssertionError – if num_mc_samples_test is not positive.

setup_task() → None

Set up task.

forward(X: Tensor) → Tensor

Forward pass BNN+VI.

Parameters:

X – input data

Returns:

bnn output

on_fit_start() → None

Before fitting compute number of training points.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute validation loss and log example predictions.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

compute_elbo_loss(X: Tensor, y: Tensor) → tuple[Tensor]

Compute the ELBO loss with mse/nll.

Parameters:

• X – input data

• y – target

Returns:

negative elbo loss and mean model output [batch_size] for logging

compute_task_loss(X: Tensor, y: Tensor) → Tensor

Compute the loss for the respective task for a single sampling iteration.

Parameters:

• X – input data

• y – target

Returns:

nll loss for the task

exclude_from_wt_decay(named_params, weight_decay: float, skip_list: list[str] = ('mu', 'rho'))

Exclude non VI parameters from weight_decay optimization.

Parameters:

• named_params – named parameters of the model

• weight_decay – weight decay factor

• skip_list – list of strings that if found in parameter name excludes the parameter from weight decay

Returns:

split parameter groups for optimization with and without weight_decay

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

BNN ELBO Regression

class lightning_uq_box.uq_methods.bnn_vi_elbo.BNN_VI_ELBO_Regression(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, burnin_epochs: int, beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayes By Backprop Model with Variational Inference (VI) for Regression.

If you use this model in your work, please cite:

• https://arxiv.org/abs/1505.05424

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, burnin_epochs: int, beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Model instance.

Parameters:

• model – pytorch model that will be converted into a BNN

• criterion – loss function used for optimization

• burnin_epochs – number of epochs to train before switching to nll loss

• beta – beta factor for negative elbo loss computation, should be number of weights and biases

• num_mc_samples_train – number of MC samples during training when computing the negative ELBO loss. When setting num_mc_samples_train=1, this is just Bayes by Backprop.

• num_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• bayesian_layer_type – flipout or reparameterization

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if num_mc_samples_train is not positive.

• AssertionError – if num_mc_samples_test is not positive.

setup_task() → None

Set up task specific attributes.

compute_task_loss(pred: Tensor, y: Tensor) → Tensor

Compute the loss for the respective task for a single sampling iteration.

Parameters:

• pred – model_prediction

• y – target

Returns:

nll loss for the task

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

BNN ELBO Classification

class lightning_uq_box.uq_methods.bnn_vi_elbo.BNN_VI_ELBO_Classification(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, task: str = 'multiclass', beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayes By Backprop Model with Variational Inference (VI) for Classification.

If you use this model in your work, please cite:

• https://arxiv.org/abs/1505.05424

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, task: str = 'multiclass', beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Model instance.

Parameters:

• model – pytorch model that will be converted into a BNN

• criterion – loss function used for optimization

• task – classification task, one of binary, multiclass, multilabel

• beta – beta factor for negative elbo loss computation, should be number of weights and biases

• num_mc_samples_train – number of MC samples during training when computing the negative ELBO loss. When setting num_mc_samples_train=1, this is just Bayes by Backprop.

• num_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• bayesian_layer_type – flipout or reparameterization

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• lr_scheduler – learning rate scheduler

• optimizer – optimizer used for training

Raises:

• AssertionError – if num_mc_samples_train is not positive.

• AssertionError – if num_mc_samples_test is not positive.

setup_task() → None

Set up task specific attributes.

compute_task_loss(pred: Tensor, y: Tensor) → Tensor

Compute the loss for the respective task for a single sampling iteration.

Parameters:

• pred – model_prediction

• y – target

Returns:

nll loss for the task

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

BNN ELBO Segmentation

class lightning_uq_box.uq_methods.bnn_vi_elbo.BNN_VI_ELBO_Segmentation(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, task: str = 'multiclass', beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False)

Bayes By Backprop Model with Variational Inference (VI) for Segmentation.

If you use this model in your work, please cite:

• https://arxiv.org/abs/1505.05424

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, criterion: ~torch.nn.modules.module.Module, task: str = 'multiclass', beta: float = 100, num_mc_samples_train: int = 10, num_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, freeze_decoder: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None, save_preds: bool = False) → None

Initialize a new BNN VI ELBO Segmentation instance.

Parameters:

• model – pytorch model that will be converted into a BNN

• criterion – loss function used for optimization

• task – classification task, one of binary, multiclass, multilabel

• beta – beta factor for negative elbo loss computation, should be number of weights and biases

• num_mc_samples_train – number of MC samples during training when computing the negative ELBO loss. When setting num_mc_samples_train=1, this is just Bayes by Backprop.

• num_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• bayesian_layer_type – flipout or reparameterization

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the model backbone, by default this is supported for torchseg Unet models

• freeze_decoder – whether to freeze the model decoder, by default this is supported for torchseg Unet models

• lr_scheduler – learning rate scheduler

• optimizer – optimizer used for training

• save_preds – whether to save predictions

freeze_model() → None

Freeze model backbone.

By default, assumes a timm model with a backbone and head. Alternatively, selected the last layer with parameters to freeze.

setup_task() → None

Set up task specific attributes for segmentation.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step for segmentation.

Parameters:

• X – prediction batch of shape [batch_size x num_channels x height x width]

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Bayesian Neural Networks with Alpha Divergence

BNN VI Base

class lightning_uq_box.uq_methods.bnn_vi.BNN_VI_Base(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, alpha: float = 1.0, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayesian Neural Network (BNN) with VI.

Trained with (VI) Variational Inferece and energy loss.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, alpha: float = 1.0, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of BNN VI.

Parameters:

• model – pytorch model that will be converted into a BNN

• n_mc_samples_train – number of MC samples during training when computing the energy loss

• n_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• alpha – alpha divergence parameter

• bayesian_layer_type – flipout or reparameterization

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if n_mc_samples_train is not positive.

• AssertionError – if n_mc_samples_test is not positive.

setup_task() → None

Set up task specific attributes.

forward(X: Tensor) → Tensor

Forward pass BNN+LI.

Parameters:

X – input data

Returns:

bnn output

on_fit_start() → None

Before fitting compute number of training points.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the data loader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute validation loss and log example predictions.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the data loader

Returns:

validation loss

freeze_layers() → None

Freeze BNN Layers to fix the stochasticity over forward passes.

unfreeze_layers() → None

Unfreeze BNN Layers to make them fully stochastic.

exclude_from_wt_decay(named_params, weight_decay: float, skip_list: list[str] = ('mu', 'rho'))

Exclude non VI parameters from weight_decay optimization.

Parameters:

• named_params – named parameters of the model

• weight_decay – weight decay factor

• skip_list – list of strings that if found in parameter name excludes the parameter from weight decay

Returns:

split parameter groups for optimization with and without weight_decay

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation – https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers

BNN VI Regression

class lightning_uq_box.uq_methods.bnn_vi.BNN_VI_Regression(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayesian Neural Network (BNN) with VI.

Trained with (VI) Variational Inferece and energy loss.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[int] | list[str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of BNN VI Regression.

Parameters:

• model – pytorch model that will be converted into a BNN

• optimizer – optimizer used for training

• n_mc_samples_train – number of MC samples during training when computing the energy loss

• n_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• alpha – alpha divergence parameter

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• bayesian_layer_type – reparameterization layer type, “reparametrization” or “flipout”

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if n_mc_samples_train is not positive.

• AssertionError – if n_mc_samples_test is not positive.

setup_task() → None

Set up task specific attributes.

compute_energy_loss(X: Tensor, y: Tensor) → None

Compute the loss for BNN with alpha divergence.

Parameters:

• X – input tensor

• y – target tensor

Returns:

mean output

over samples, dim [n_mc_samples_train, output_dim]

Return type:

energy loss and mean output for logging mean_out

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

BNN VI Batched Regression

class lightning_uq_box.uq_methods.bnn_vi.BNN_VI_BatchedRegression(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Batched sampling version of BNN_VI.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

__init__(model: ~torch.nn.modules.module.Module, n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of BNN VI Batched.

Parameters:

• model – pytorch model that will be converted into a BNN

• n_mc_samples_train – number of MC samples during training when computing the energy loss

• n_mc_samples_test – number of MC samples during test and prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• alpha – alpha divergence parameter

• bayesian_layer_type – reparameterization layer type, “reparametrization” or “flipout”

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• lr_scheduler – learning rate scheduler

• optimizer – optimizer used for training

Raises:

• AssertionError – if n_mc_samples_train is not positive.

• AssertionError – if n_mc_samples_test is not positive.

forward(X: Tensor, n_samples: int) → Tensor

Forward pass BNN+LI.

Parameters:

• X – input data

• n_samples – number of samples to compute

Returns:

bnn output of shape [num_samples, batch_size, num_outputs]

compute_energy_loss(X: Tensor, y: Tensor) → tuple[Tensor]

Compute the loss for BNN with alpha divergence.

Parameters:

• X – input tensor

• y – target tensor

Returns:

mean output

over samples, dim [n_mc_samples_train, output_dim]

Return type:

energy loss and mean output for logging mean_out

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

freeze_layers(n_samples: int) → None

Freeze BNN Layers to fix the stochasticity over forward passes.

Parameters:

n_samples – number of samples used in frozen layers

Bayesian Neural Networks with Latent Variables (BNN-LV)

BNN LV VI Base

class lightning_uq_box.uq_methods.bnn_lv_vi.BNN_LV_VI_Base(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, alpha: float = 1.0, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0.0, lv_prior_std: float = 1.0, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayesian Neural Network (BNN) with Latent Variables (LV).

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

lv_intro_options = ['first', 'last']

__init__(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, alpha: float = 1.0, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0.0, lv_prior_std: float = 1.0, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of BNN+LV.

Parameters:

• model – pytorch model that will be converted into a BNN

• latent_net – latent variable network

• num_training_points – num of data points contained in the training dataset

• prediction_head – prediction head that will be attached to the model

• latent_variable_intro – whether to introduce the latent variable at the first or last layer of the model

• n_mc_samples_train – number of MC samples during training when computing the negative ELBO loss

• n_mc_samples_test – number of MC samples during test and prediction

• n_mc_samples_epistemic – number of epistemic samples during prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• alpha – alpha divergence parameter

• bayesian_layer_type – flipout or reparameterization

• lv_prior_mu – prior mean for latent variable network

• lv_prior_std – prior std for latent variable network

• lv_latent_dim – number of latent dimension

• init_scaling – init scaling factor for q(z) in latent variable network

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if n_mc_samples_train is not positive

• AssertionError – if n_mc_samples_test is not positive

setup_task() → None

Set up task.

forward(X: Tensor, y: Tensor | None = None, training: bool = True) → Tensor

Forward pass BNN LV.

Parameters:

• X – input data

• y – target

• training – if yes, smple from lv posterior, else use sample from prior or provide z

Returns:

bnn output of size [batch_size, output_dim]

sample_latent_variable_prior(X: Tensor) → Tensor

Sample the latent variable prior during inference.

Parameters:

X – inference tensor that gets concatenated with z

Returns:

sampled latent variable of shape [batch_size, lv_latent_dim]

compute_energy_loss(X: Tensor, y: Tensor) → tuple[Tensor]

Compute the loss for BNN with alpha divergence.

Parameters:

• X – input tensor

• y – target tensor

Returns:

mean output

over samples, dim [n_mc_samples_train, output_dim]

Return type:

energy loss and mean output for logging mean_out

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary

freeze_layers() → None

Freeze BNN Layers to fix the stochasticity over forward passes.

unfreeze_layers() → None

Unfreeze BNN Layers to make them fully stochastic.

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

class lightning_uq_box.uq_methods.bnn_lv_vi.BNN_LV_VI_Batched_Base(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0, lv_prior_std: float = 1, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Batched sampling version of BNN_LV_VI.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

__init__(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0, lv_prior_std: float = 1, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instace of BNN+LV Batched.

Parameters:

• model – pytorch model that will be converted into a BNN

• latent_net – latent variable network

• num_training_points – number of data points contained in the training dataset

• num_training_points – number of data points contained in the training dataset

• prediction_head – prediction head that will be attached to the model

• latent_variable_intro – whether to introcde the LV at first or last layer

• n_mc_samples_train – number of MC samples during training when computing the negative ELBO loss

• n_mc_samples_test – number of MC samples during test and prediction

• n_mc_samples_epistemic – number of epistemic samples during prediction

• output_noise_scale – scale of predicted sigmas

• prior_mu – prior mean value for bayesian layer

• prior_sigma – prior variance value for bayesian layer

• posterior_mu_init – mean initialization value for approximate posterior

• posterior_rho_init – variance initialization value for approximate posterior through softplus σ = log(1 + exp(ρ))

• alpha – alpha divergence parameter

• bayesian_layer_type – flipout or reparameterization

• lv_prior_mu – prior mean for latent variable network

• lv_prior_std – prior std for latent variable network

• lv_latent_dim – number of latent dimension

• init_scaling – init scaling factor for q(z) in latent variable network

• stochastic_module_names – list of module names or indices that should be converted to variational layers

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Raises:

• AssertionError – if n_mc_samples_train is not positive

• AssertionError – if n_mc_samples_test is not positive

forward(X: Tensor, y: Tensor | None = None, n_samples: int = 5, training: bool = True) → Tensor

Forward pass BNN+LI.

Parameters:

• X – input data

• y – target

• n_samples – number of samples to compute

• training – if yes, sample from lv posterior, else use sample from prior or provide z

Returns:

bnn output [batch_size, output_dim, num_samples]

sample_latent_variable_prior(X: Tensor) → Tensor

Sample the latent variable prior during inference.

Parameters:

X – inference tensor that gets concatenated with z

Returns:

sampled latent variable of shape [batch_size, lv_latent_dim]

compute_energy_loss(X: Tensor, y: Tensor) → tuple[Tensor]

Compute the loss for BNN with alpha divergence.

Parameters:

• X – input tensor

• y – target tensor

Returns:

mean output

over samples, dim [n_mc_samples_train, output_dim]

Return type:

energy loss and mean output for logging mean_out

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0, n_samples_pred: int | None = None) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the data loader

• n_samples_pred – number of samples to use for prediction

Returns:

prediction dictionary

freeze_layers(n_samples: int) → None

Freeze BNN Layers to fix the stochasticity over forward passes.

Parameters:

n_samples – number of samples used in frozen layers

BNN LV VI Regression

class lightning_uq_box.uq_methods.bnn_lv_vi.BNN_LV_VI_Regression(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0.0, prior_sigma: float = 1.0, posterior_mu_init: float = 0.0, posterior_rho_init: float = -5.0, alpha: float = 1.0, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0.0, lv_prior_std: float = 1.0, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayesian Latent Variable Network with Variational Inference for Regression.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

nll_loss = GaussianNLLLoss()

pred_file_name = 'preds.csv'

setup_task() → None

Set up task.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

class lightning_uq_box.uq_methods.bnn_lv_vi.BNN_LV_VI_Batched_Regression(model: ~torch.nn.modules.module.Module, latent_net: ~torch.nn.modules.module.Module, num_training_points: int, prediction_head: ~torch.nn.modules.module.Module | None = None, latent_variable_intro: str = 'first', n_mc_samples_train: int = 25, n_mc_samples_test: int = 50, n_mc_samples_epistemic: int = 50, output_noise_scale: float = 1.3, prior_mu: float = 0, prior_sigma: float = 1, posterior_mu_init: float = 0, posterior_rho_init: float = -5, alpha: float = 1, bayesian_layer_type: str = 'reparameterization', lv_prior_mu: float = 0, lv_prior_std: float = 1, lv_latent_dim: int = 1, init_scaling: float = 0.1, stochastic_module_names: list[str | int] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Bayesian Latent Variable Network with VI Batched for Regression.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v80/depeweg18a

nll_loss = GaussianNLLLoss()

pred_file_name = 'preds.csv'

setup_task() → None

Set up task.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Stochastic Weight Averaging Gaussian (SWAG)

SWAG Base

class lightning_uq_box.uq_methods.swag.SWAGBase(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int | str] | None = None)

Stochastic Weight Averaging - Gaussian (SWAG).

If you use this model in your research, please cite the following paper:

• https://proceedings.neurips.cc/paper_files/paper/2019/hash/118921efba23fc329e6560b27861f0c2-Abstract.html # noqa: E501

__init__(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int | str] | None = None) → None

Initialize a new instance of SWAG Model Wrapper.

Parameters:

• model – pytorch model

• max_swag_snapshots – maximum number of snapshots to store

• snapshot_freq – frequency of snapshots

• num_mc_samples – number of MC samples during prediction

• swag_lr – learning rate for swag

• loss_fn – loss function

• stochastic_module_names – list of module names or indices that should be converted to variational layer

setup_task() → None

Set up task specific attributes.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Compute SWAG optimization step.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

on_train_epoch_end()

Do not Log epoch-level training metrics.

on_train_end() → None

After training stage is completed, swag is fitted.

validation_step(*args: Any, **kwargs: Any) → Tensor

Not intended to be used.

on_validation_epoch_end() → None

Do not log any validation metrics.

update_uncertainty_buffers()

Update the running average over weights.

sample_state()

Update the state with a sample.

sample_predictions(X: Tensor) → Tensor

Sample predictions.

Parameters:

X – input batch of shape [batch_size x input_dims]

Returns:

predictions of shape [batch_size x num_outputs x num_mc_samples]

configure_optimizers() → dict[str, Any]

Manually implemented SWAG optimization.

SWAG Regression

class lightning_uq_box.uq_methods.swag.SWAGRegression(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int] | list[str] | None = None)

SWAG Model for Regression.

If you use this model in your research, please cite the following paper:

• https://proceedings.neurips.cc/paper_files/paper/2019/hash/118921efba23fc329e6560b27861f0c2-Abstract.html

pred_file_name = 'preds.csv'

__init__(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int] | list[str] | None = None) → None

Initialize a new instance of SWAG Model for Regression.

Parameters:

• model – pytorch model

• num_swag_epochs – number of epochs to train swag

• max_swag_snapshots – maximum number of snapshots to store

• snapshot_freq – frequency of snapshots

• num_mc_samples – number of MC samples during prediction

• swag_lr – learning rate for swag

• loss_fn – loss function

• stochastic_module_names – names of modules that are partially stochastic

setup_task() → None

Set up task specific attributes.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step that with SWAG uncertainty.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary

SWAG Classification

class lightning_uq_box.uq_methods.swag.SWAGClassification(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, task: str = 'multiclass', stochastic_module_names: list[int] | list[str] | None = None)

SWAG Model for Classification.

If you use this model in your research, please cite the following paper:

• https://proceedings.neurips.cc/paper_files/paper/2019/hash/118921efba23fc329e6560b27861f0c2-Abstract.html

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, task: str = 'multiclass', stochastic_module_names: list[int] | list[str] | None = None) → None

Initialize a new instance of SWAG Model for Classification.

Parameters:

• model – pytorch model

• num_swag_epochs – number of epochs to train swag

• max_swag_snapshots – maximum number of snapshots to store

• snapshot_freq – frequency of snapshots

• num_mc_samples – number of MC samples during prediction

• swag_lr – learning rate for swag

• loss_fn – loss function

• task – classification task, one of [‘binary’, ‘multiclass’, ‘multilabel’]

• stochastic_module_names – names of modules that are partially stochastic

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt model output to be compatible for metric computation.

Parameters:

out – output from the model

Returns:

mean output

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step with SWAG uncertainty.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

SWAG Segmentation

class lightning_uq_box.uq_methods.swag.SWAGSegmentation(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, task: str = 'multiclass', stochastic_module_names: list[int] | list[str] | None = None, save_preds: bool = False)

SWAG Model for Segmentation.

pred_dir_name = 'preds'

__init__(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, task: str = 'multiclass', stochastic_module_names: list[int] | list[str] | None = None, save_preds: bool = False) → None

Initialize a new instance of SWAG Model for Segmentation.

Parameters:

• model – pytorch model

• num_swag_epochs – number of epochs to train swag

• max_swag_snapshots – maximum number of snapshots to store

• snapshot_freq – frequency of snapshots

• num_mc_samples – number of MC samples during prediction

• swag_lr – learning rate for swag

• loss_fn – loss function

• task – segmentation task, one of [‘binary’, ‘multiclass’]

• stochastic_module_names – names of modules that are partially stochastic

• save_preds – save predictions

setup_task() → None

Set up task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step with SWAG uncertainty.

Parameters:

• X – prediction batch of shape [batch_size x num_channels x height x width]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

SWAG Pixelwise Regression

class lightning_uq_box.uq_methods.swag.SWAGPxRegression(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int] | list[str] | None = None, save_preds: bool = False)

SWAG Model for Pixelwise Regression.

Added in version 0.2.0.

pred_dir_name = 'preds'

__init__(model: Module, max_swag_snapshots: int, snapshot_freq: int, num_mc_samples: int, swag_lr: float, loss_fn: Module, stochastic_module_names: list[int] | list[str] | None = None, save_preds: bool = False) → None

Initialize a new instance of SWAG Model for Pixelwise Regression.

Parameters:

• model – pytorch model

• num_swag_epochs – number of epochs to train swag

• max_swag_snapshots – maximum number of snapshots to store

• snapshot_freq – frequency of snapshots

• num_mc_samples – number of MC samples during prediction

• swag_lr – learning rate for swag

• loss_fn – loss function

• stochastic_module_names – names of modules that are partially stochastic

• save_preds – save predictions

setup_task() → None

Set up task specific attributes.

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Stochastic Gradient Langevin Dynamics (SGLD)

class lightning_uq_box.uq_methods.sgld.SGLD(params: Iterator[Parameter], lr: float, noise_factor: float, weight_decay: float = 0.0)

Stochastic Gradient Langevian Dynamics Optimzer.

If you use this optimizer in your research, please cite the following paper:

• https://www.stats.ox.ac.uk/~teh/research/compstats/WelTeh2011a.pdf

__init__(params: Iterator[Parameter], lr: float, noise_factor: float, weight_decay: float = 0.0) → None

Initialize new instance of SGLD Optimier.

Parameters:

• params – model parameters

• lr – initial learning rate

• noise_factor – parameter denoting how much noise to inject in the SGD update

• weight_decay – weight decay parameter for SGLD optimizer

step(closure: callable)

Perform a single optimization step.

Parameters:

closure – A closure that reevaluates the model

Returns:

updated loss

Variational Bayes Last Layer

VBLL Classification

class lightning_uq_box.uq_methods.vbll.VBLLClassification(model: ~torch.nn.modules.module.Module, regularization_weight: float, num_targets: int, replace_ll: bool = True, parameterization: str = 'dense', prior_scale: float = 1, wishart_scale: float = 0.01, dof: int = 1, layer_type: str = 'disc', freeze_backbone: bool = False, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Any | None = None)

Variational Bayes Last Layer (VBLL) for Classification.

If you use this method in your research, please cite the following paper:

• https://arxiv.org/abs/2404.11599

Added in version 0.2.

valid_layer_types = ['disc', 'gen']

__init__(model: ~torch.nn.modules.module.Module, regularization_weight: float, num_targets: int, replace_ll: bool = True, parameterization: str = 'dense', prior_scale: float = 1, wishart_scale: float = 0.01, dof: int = 1, layer_type: str = 'disc', freeze_backbone: bool = False, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Any | None = None) → None

Initialize a new instance of VBLL Classification.

Parameters:

• model – The backbone model

• regularization_weight – regularization weight term in ELBO, and should be 1 / (dataset size) by default. This term impacts the epistemic uncertainty estimate.

• num_targets – Number of targets

• replace_ll – If True, replace the last layer of the model with VBLL or add a new layer

• parameterization – Parameterization of covariance matrix. One of [‘dense’,’diagonal’]

• prior_scale – prior covariance matrix scale Scale of prior covariance matrix

• wishart_scale – Scale of Wishart prior on noise covariance. This term has an impact on the aleatoric uncertainty estimate.

• dof – Degrees of freedom of Wishart prior on noise covariance

• layer_type – The type of layer to use. One of [‘disc’, ‘gen’], a Discriminative or Generative layer

• freeze_backbone – If True, the backbone model will be frozen and only the VBBL layer will be trained

• task – The type of task. One of [‘binary’, ‘multiclass’]

• optimizer – The optimizer to use for training

• lr_scheduler – The learning rate scheduler to use for training

build_model() → None

Build Classification Model.

freeze_model() → None

Freeze model.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for metrics.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Training step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Validation step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Test step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

test loss

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Any]

Predict step with VBLL model.

Parameters:

• X – The input data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

prediction dictionary

configure_optimizers() → dict[str, Any]

Configure Optimizers.

VBLL Regression

class lightning_uq_box.uq_methods.vbll.VBLLRegression(model: ~torch.nn.modules.module.Module, regularization_weight, replace_ll: bool = True, num_targets: int = 1, parameterization: str = 'dense', prior_scale: float = 1.0, wishart_scale: float = 0.01, dof: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Variational Bayesian Last Layer (VBLL) for Regression.

If you use this model in your research, please cite the following paper:

• https://arxiv.org/abs/2404.11599

Added in version 0.2.

__init__(model: ~torch.nn.modules.module.Module, regularization_weight, replace_ll: bool = True, num_targets: int = 1, parameterization: str = 'dense', prior_scale: float = 1.0, wishart_scale: float = 0.01, dof: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None) → None

Initialize the VBLL regression model.

Parameters:

• model – The backbone model

• regularization_weight – regularization weight term in ELBO, and should be 1 / (dataset size) by default. This term impacts the epistemic uncertainty estimate.

• replace_ll – Whether to replace the last layer of the model with VBLL or add a new layer

• num_targets – Number of targets

• parameterization – Parameterization of covariance matrix. One of [‘dense’,’diagonal’]

• prior_scale – prior covariance matrix scale Scale of prior covariance matrix

• wishart_scale – Scale of Wishart prior on noise covariance. This term has an impact on the aleatoric uncertainty estimate.

• dof – Degrees of freedom of Wishart prior on noise covariance

• freeze_backbone – If True, the backbone model will be frozen and only the VBBL layer will be trained

• optimizer – The optimizer to use for training

• lr_scheduler – The learning rate scheduler to use for training

build_model() → None

Build model.

freeze_model() → None

Freeze model.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for metrics.

Parameters:

out – the output from the VBLL module

Returns:

the mean prediction

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Training step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Validation step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Test step.

Parameters:

• batch – The batch of data

• batch_idx – The index of the batch

• dataloader_idx – The index of the dataloader

Returns:

test loss

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step with VBLL model.

configure_optimizers() → dict[str, Any]

Configure Optimizers.

Spectral Normalized Gaussian Process (SNGP)

SNGPBase

class lightning_uq_box.uq_methods.sngp.SNGPBase(feature_extractor: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_targets: int = 1, num_gp_features: int = 128, num_random_features: int = 1024, normalize_gp_features: bool = True, feature_scale: int = 2, ridge_penalty: float = 1.0, coeff: float = 0.95, n_power_iterations: int = 1, input_size: int | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Specral Normalized Gaussian Process (SNGP).

If you use this code, please cite the following paper:

• https://arxiv.org/abs/2006.10108

pred_file_name = 'preds.csv'

__init__(feature_extractor: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_targets: int = 1, num_gp_features: int = 128, num_random_features: int = 1024, normalize_gp_features: bool = True, feature_scale: int = 2, ridge_penalty: float = 1.0, coeff: float = 0.95, n_power_iterations: int = 1, input_size: int | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new SNGP model.

Parameters:

• feature_extractor – Feature extractor model

• loss_fn – Loss function

• num_targets – Number of output units / targets

• num_gp_features – Number of GP features

• num_deep_features – Number of deep features

• num_random_features – Number of random features

• normalize_gp_features – Whether to normalize GP features

• feature_scale – Feature scale

• ridge_penalty – Ridge penalty

• coeff – soft normalization only when sigma larger than coeff, should be (0, 1)

• n_power_iterations – number of power iterations for spectral normalization

• input_size – image dimension input size needed for spectral normalization

• freeze_backbone – whether to freeze the feature extractor

• optimizer – Optimizer

• lr_scheduler – Learning rate scheduler

forward(x: Tensor) → tuple[Tensor]

Forward pass of the SNGP model.

Parameters:

x – Input tensor

Returns:

Output tensor after applying SNGP

reset_precision_matrix()

Reset the precision matrix to identity matrix.

recompute_covariance_matrix()

Recompute the covariance matrix.

on_fit_start() → None

Before fitting compute number of training points.

on_train_epoch_start() → None

Called when the train epoch begins.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

on_train_epoch_end()

Log epoch-level training metrics.

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

on_validation_epoch_end()

Log epoch-level validation metrics.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Compute and return the predictions.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

predictions

on_test_epoch_end()

Log epoch-level training metrics.

predict_step(X: Tensor) → dict[str, Tensor]

Predict the output for a batch of inputs.

Parameters:

X – Input tensor

Returns:

The predicted output

configure_optimizers()

Configure optimizers and learning rate schedulers.

SNGPRegression

class lightning_uq_box.uq_methods.sngp.SNGPRegression(feature_extractor: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_targets: int = 1, num_gp_features: int = 128, num_random_features: int = 1024, normalize_gp_features: bool = True, feature_scale: int = 2, ridge_penalty: float = 1.0, coeff: float = 0.95, n_power_iterations: int = 1, input_size: int | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

SNGP for regression.

If you use this code, please cite the following paper:

• https://arxiv.org/abs/2006.10108

setup_task() → None

Set up task.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

SNGPClassification

class lightning_uq_box.uq_methods.sngp.SNGPClassification(feature_extractor: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_targets: int = 1, num_gp_features: int = 128, num_random_features: int = 1024, normalize_gp_features: bool = True, feature_scale: int = 2, ridge_penalty: float = 1, coeff: float = 0.95, n_power_iterations: int = 1, input_size: int | None = None, mean_field_factor: float | None = 0.39269908169872414, task: str = 'multiclass', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

SNGP for classification.

If you use this code, please cite the following paper:

• https://arxiv.org/abs/2006.10108

valid_tasks = ['binary', 'multiclass']

__init__(feature_extractor: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_targets: int = 1, num_gp_features: int = 128, num_random_features: int = 1024, normalize_gp_features: bool = True, feature_scale: int = 2, ridge_penalty: float = 1, coeff: float = 0.95, n_power_iterations: int = 1, input_size: int | None = None, mean_field_factor: float | None = 0.39269908169872414, task: str = 'multiclass', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new SNGP model for classification.

Parameters:

• feature_extractor – Feature extractor model

• loss_fn – Loss function

• num_targets – Number of output units / targets

• num_gp_features – Number of GP features

• num_deep_features – Number of deep features

• num_random_features – Number of random features

• normalize_gp_features – Whether to normalize GP features

• feature_scale – Feature scale

• ridge_penalty – Ridge penalty

• coeff – soft normalization only when sigma larger than coeff, should be (0, 1)

• n_power_iterations – number of power iterations for spectral normalization

• input_size – image dimension input size needed for spectral normalization

• mean_field_factor – Mean field factor, required for classification problems

• task – classification task, one of [‘binary’, ‘multiclass’]

• freeze_backbone – whether to freeze the feature extractor

• optimizer – Optimizer

• lr_scheduler – Learning rate scheduler

setup_task() → None

Set up task.

mean_field_logits(logits: Tensor, pred_cov: Tensor) → Tensor

Applies the Mean-Field approximation to the provided logits.

Based on: https://arxiv.org/abs/2006.07584

Parameters:

• logits – The logits to be transformed

• pred_cov – The predicted covariance

Returns:

The transformed logits

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor) → dict[str, Tensor]

Predict the output for a batch of inputs.

Parameters:

X – Input tensor

Returns:

output dictionary

Deep Kernel Learning (DKL)

DKL Base

class lightning_uq_box.uq_methods.deep_kernel_learning.DKLBase(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, gp_kernel: str = 'RBF', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deep Kernel Learning Base Module.

If you use this model in your work, please cite:

• https://proceedings.mlr.press/v51/wilson16.html

kernel_choices = ['RBF', 'Matern12', 'Matern32', 'Matern52', 'RQ']

pred_file_name = 'preds.csv'

__init__(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, gp_kernel: str = 'RBF', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Deep Kernel Learning Model.

Initialize a new Deep Kernel Learning Model.

Parameters:

• feature_extractor – feature extractor model

• n_inducing_points – number of inducing points

• gp_kernel – kernel choice, supports one of [‘RBF’, ‘Matern12’, ‘Matern32’, ‘Matern52’, ‘RQ’]

• elbo_fn – gpytorch elbo function used for optimization

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

property num_input_features: int

Retrieve input dimension to the model.

Returns:

number of input dimension to the model

property num_outputs: int

Retrieve output dimension to the model.

Returns:

number of output dimension from model

setup_task() → None

Set up task specific attributes.

forward(X: Tensor) → MultivariateNormal

Forward pass through model.

Parameters:

X – input tensor to backbone

Returns:

output from GP

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

on_train_epoch_end()

Log epoch-level training metrics.

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute validation loss and log example predictions.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

on_validation_epoch_end() → None

Log epoch level validation metrics.

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

DKL Regression

class lightning_uq_box.uq_methods.deep_kernel_learning.DKLRegression(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, num_targets: int = 1, gp_kernel: str = 'RBF', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deep Kernel Learning Model.

If you use this model in your research, please cite the following papers:

• https://proceedings.mlr.press/v51/wilson16.html

• https://arxiv.org/abs/2102.11409

__init__(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, num_targets: int = 1, gp_kernel: str = 'RBF', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Deep Kernel Learning Model for Regression.

Parameters:

• feature_extractor – feature extractor model

• n_inducing_points – number of inducing points

• num_targets – number of targets

• gp_kernel – kernel choice, supports one of [‘RBF’, ‘Matern12’, ‘Matern32’, ‘Matern52’, ‘RQ’]

• elbo_fn – gpytorch elbo function used for optimization

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_epoch_end()

Log epoch-level test metrics.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

DKL Classification

class lightning_uq_box.uq_methods.deep_kernel_learning.DKLClassification(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, num_classes: int, task: str = 'multiclass', gp_kernel: str = 'RBF', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deep Kernel Learning for Classification.

If you use this model in your research, please cite the following papers:

• https://proceedings.mlr.press/v51/wilson16.html

• https://arxiv.org/abs/2102.11409

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, num_classes: int, task: str = 'multiclass', gp_kernel: str = 'RBF', freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Deep Kernel Learning Model for Classification.

Parameters:

• feature_extractor – feature extractor model

• n_inducing_points – number of inducing points

• gp_kernel – GP kernel choice, supports one of ‘RBF’, ‘Matern12’, ‘Matern32’, ‘Matern52’, ‘RQ’]

• num_classes – number of classes

• task – classification task, one of [‘binary’, ‘multiclass’, ‘multilabel’]

• freeze_backbone – whether to freeze the backbone

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

setup_task() → None

Set up task specific attributes.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute validation loss and log example predictions.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

on_test_epoch_end()

Log epoch-level test metrics.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Deterministic Uncertainty Estimation (DUE)

DUE Regression

class lightning_uq_box.uq_methods.deterministic_uncertainty_estimation.DUERegression(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, input_size: int, num_targets: int = 1, gp_kernel: str = 'RBF', coeff: float = 0.95, n_power_iterations: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deterministic Uncertainty Estimation (DUE) for Regression.

If you use this model in your research, please cite the following paper:x

• https://arxiv.org/abs/2102.11409

__init__(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, input_size: int, num_targets: int = 1, gp_kernel: str = 'RBF', coeff: float = 0.95, n_power_iterations: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Deterministic Uncertainty Estimation Model.

Initialize a new Deep Kernel Learning Model for Regression.

Parameters:

• feature_extractor – feature extractor model

• n_inducing_points – number of inducing points

• num_targets – number of targets

• gp_kernel – GP kernel choice, supports one of ‘RBF’, ‘Matern12’, ‘Matern32’, ‘Matern52’, ‘RQ’]

• input_size – image input size of data to the model

• coeff – soft normalization only when sigma larger than coeff, should be (0, 1)

• n_power_iterations – number of power iterations for spectral normalization

• freeze_backbone – whether to freeze the feature extractor or not

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

DUE Classification

class lightning_uq_box.uq_methods.deterministic_uncertainty_estimation.DUEClassification(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, input_size: int, num_classes: int, gp_kernel: str = 'RBF', task: str = 'multiclass', coeff: float = 0.95, n_power_iterations: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Deterministic Uncertainty Estimation (DUE) Model.

If you use this model in your research, please cite the following paper:

• https://arxiv.org/abs/2102.11409

__init__(feature_extractor: ~torch.nn.modules.module.Module, n_inducing_points: int, input_size: int, num_classes: int, gp_kernel: str = 'RBF', task: str = 'multiclass', coeff: float = 0.95, n_power_iterations: int = 1, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new Deterministic Uncertainty Estimation Model.

Parameters:

• feature_extractor – feature extractor model

• n_inducing_points – number of inducing points

• input_size – image input size of data to the model

• num_classes – number of classes

• gp_kernel – GP kernel choice, supports one of ‘RBF’, ‘Matern12’, ‘Matern32’, ‘Matern52’, ‘RQ’]

• task – classification task, one of [‘binary’, ‘multiclass’, ‘multilabel’]

• coeff – soft normalization only when sigma larger than coeff should be (0, 1)

• n_power_iterations – number of power iterations for spectral normalization

• freeze_backbone – whether to freeze the feature extractor or not

• optimizer – optimizer used for training

• lr_scheduler – learning rate scheduler

Deep Ensembles

Deep Ensemble Base

class lightning_uq_box.uq_methods.deep_ensemble.DeepEnsemble(ensemble_members: list[dict[str, type[LightningModule] | str]])

Base Class for different Ensemble Models.

If you use this model in your work, please cite:

• https://proceedings.neurips.cc/paper_files/paper/2017/hash/9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html # noqa: E501

__init__(ensemble_members: list[dict[str, type[LightningModule] | str]]) → None

Initialize a new instance of DeepEnsembleModel Wrapper.

Parameters:

• ensemble_members – List of dicts where each element specifies the LightningModule class and a path to a checkpoint

• save_dir – path to directory where to store prediction

• quantiles – quantile values to compute for prediction

setup_task() → None

Set up task.

forward(X: Tensor) → Tensor

Forward step of Deep Ensemble.

Parameters:

X – input tensor of shape [batch_size, input_di]

Returns:

Ensemble member outputs stacked over last dimension for output of [batch_size, num_outputs, num_ensemble_members]

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test step.

generate_ensemble_predictions(X: Tensor) → Tensor

Generate DeepEnsemble Predictions.

Parameters:

X – input tensor of shape [batch_size, input_di]

Returns:

the ensemble predictions

Deep Ensemble Regression

class lightning_uq_box.uq_methods.deep_ensemble.DeepEnsembleRegression(ensemble_members: list[dict[str, type[LightningModule] | str]])

Deep Ensemble Model for regression.

If you use this model in your work, please cite:

• https://proceedings.neurips.cc/paper_files/paper/2017/hash/9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html

pred_file_name = 'preds.csv'

setup_task() → None

Set up task for regression.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → Any

Compute prediction step for a deep ensemble.

Parameters:

• X – input tensor of shape [batch_size, input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

mean and standard deviation of MC predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Deep Ensemble Classification

class lightning_uq_box.uq_methods.deep_ensemble.DeepEnsembleClassification(ensemble_members: list[dict[str, type[LightningModule] | str]], num_classes: int, task: str = 'multiclass')

Deep Ensemble Model for classification.

If you use this model in your work, please cite:

• https://proceedings.neurips.cc/paper_files/paper/2017/hash/9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html

valid_tasks = ['multiclass', 'binary', 'multilabel']

pred_file_name = 'preds.csv'

__init__(ensemble_members: list[dict[str, type[LightningModule] | str]], num_classes: int, task: str = 'multiclass') → None

Initialize a new instance of DeepEnsemble for Classification.

Parameters:

• ensemble_members – List of dicts where each element specifies the LightningModule class and a path to a checkpoint

• num_classes – number of classes

• task – classification task, one of “multiclass”, “binary” or “multilabel”

setup_task() → None

Set up task for classification.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → Any

Compute prediction step for a deep ensemble.

Parameters:

• X – input tensor of shape [batch_size, input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

mean and standard deviation of MC predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Deep Ensemble Segmentation

class lightning_uq_box.uq_methods.deep_ensemble.DeepEnsembleSegmentation(ensemble_members: list[dict[str, type[LightningModule] | str]], num_classes: int, task: str = 'multiclass', save_preds: bool = False)

Deep Ensemble Model for segmentation.

If you use this model in your work, please cite:

• https://proceedings.neurips.cc/paper_files/paper/2017/hash/9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html

pred_dir_name = 'preds'

__init__(ensemble_members: list[dict[str, type[LightningModule] | str]], num_classes: int, task: str = 'multiclass', save_preds: bool = False) → None

Initialize a new instance of DeepEnsemble for Segmentation.

Parameters:

• ensemble_members – List of dicts where each element specifies the LightningModule class and a path to a checkpoint

• num_classes – number of classes

• task – classification task, one of “multiclass”, “binary” or “multilabel”

• save_preds – whether to save predictions

setup_task() → None

Set up task for segmentation.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → Any

Compute prediction step for a deep ensemble.

Parameters:

• X – input tensor of shape [batch_size, input_di]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

mean and standard deviation of MC predictions

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

Deep Ensemble Pixelwise Regression

class lightning_uq_box.uq_methods.deep_ensemble.DeepEnsemblePxRegression(ensemble_members: list[dict[str, type[LightningModule] | str]], save_preds: bool = False)

Deep Ensemble Model for pixelwise regression.

If you use this model in your work, please cite:

• https://proceedings.neurips.cc/paper_files/paper/2017/hash/9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html

Added in version 0.2.0.

__init__(ensemble_members: list[dict[str, type[LightningModule] | str]], save_preds: bool = False) → None

Initialize a new instance of DeepEnsemble for Pixelwise Regression.

Parameters:

• ensemble_members – List of dicts where each element specifies the LightningModule class and a path to a checkpoint

• save_preds – whether to save predictions

pred_dir_name = 'preds'

setup_task() → None

Set up task specific attributes.

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Masked Ensemble

Masked Ensemble Base

class lightning_uq_box.uq_methods.masked_ensemble.MasksemblesBase(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Base class for Masked Ensemble models.

If you use this model in your work, please cite:

• https://arxiv.org/abs/2012.08334

The input from the dataloader will be repeated for each estimator, so consider this when defining the batch size regarding memory usage.

Added in version 0.2.

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Initialize Masked Ensemble model.

Parameters:

• model – PyTorch model to turn into a Masked Ensemble and train it.

• loss_fn – Loss function.

• num_estimators – The number of estimators (masks) to generate.

• scale – The scale factor for mask generation. Muste be a scaler in the interval [1, 6].

• optimizer – Optimizer to use.

• lr_scheduler – Learning rate scheduler.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for the metrics.

Parameters:

out – model output

Returns:

adapted output

convert_to_masked_ensemble(model: Module) → None

Convert model to a Masked Ensemble model.

Parameters:

model – PyTorch model to turn into a Masked Ensemble

forward(x: Tensor)

Forward pass.

Parameters:

x – Input tensor of shape [batch_size, *input_shape]

Returns:

Output tensor of shape [batch_size * num_estimators, *output_shape]

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

on_train_epoch_end()

Log epoch-level training metrics.

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

on_validation_epoch_end()

Log epoch-level validation metrics.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Compute and return the test prediction.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

test prediction with uncertainty

on_test_epoch_end()

Log epoch-level test metrics.

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

Masked Ensemble Regression

class lightning_uq_box.uq_methods.masked_ensemble.MasksemblesRegression(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Masked Ensemble for regression tasks.

If you use this model in your work, please cite:

• https://arxiv.org/abs/2012.08334

The input from the dataloader will be repeated for each estimator, so consider this when defining the batch size regarding memory usage.

Added in version 0.2.

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Initialize Masked Ensemble for regression tasks.

Parameters:

• model – PyTorch model to turn into a Masked Ensemble and train it.

• loss_fn – Loss function to train the model.

• num_estimators – The number of estimators (masks) to generate.

• scale – The scale factor for mask generation. Muste be a scaler in the interval [1, 6].

• optimizer – Optimizer to use.

• lr_scheduler – Learning rate scheduler.

setup_task() → None

Set up task specific attributes.

predict_step(x: Tensor) → Tensor

Predict the output of the model.

Parameters:

x – Input tensor of shape [batch_size, *input_shape]

Returns:

Output tensor of shape [batch_size, *output_shape]

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Masked Ensemble Classification

class lightning_uq_box.uq_methods.masked_ensemble.MasksemblesClassification(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Masked Ensemble for classification tasks.

If you use this model in your work, please cite:

• https://arxiv.org/abs/2012.08334

The input from the dataloader will be repeated for each estimator, so consider this when defining the batch size regarding memory usage.

Added in version 0.2.

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, num_estimators: int, scale: float, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None)

Initialize Masked Ensemble for classification tasks.

Parameters:

• model – PyTorch model to turn into a Masked Ensemble and train it.

• loss_fn – Loss function to train the model.

• num_estimators – The number of estimators (masks) to generate.

• scale – The scale factor for mask generation. Muste be a scaler in the interval [1, 6].

• task – what kind of classification task, choose one of [“binary”, “multiclass”, “multilabel”]

• optimizer – Optimizer to use.

• lr_scheduler – Learning rate scheduler.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for the metrics.

Parameters:

out – model output

Returns:

adapted output

predict_step(x: Tensor) → Tensor

Predict the output of the model.

Parameters:

x – Input tensor of shape [batch_size, *input_shape]

Returns:

Output tensor of shape [batch_size, *output_shape]

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Density Uncertainty Model

Density Uncertainty Base

class lightning_uq_box.uq_methods.density_uncertainty.DensityLayerModelBase(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1.0, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 1, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Density Layer Model.

If you use this module in your work, please cite the following paper:

• https://arxiv.org/abs/2306.12497

pred_file_name = 'preds.csv'

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module, prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1.0, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 1, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a Density Layer Model.

Parameters:

• model – PyTorch model that will be converted to a Density Layer Model.

• loss_fn – Loss function used for the target minimization, could be a custom loss function depending on the regression or classification task.

• prior_std – Standard deviation of the prior.

• posterior_std_init – Initial standard deviation of the posterior.

• kl_beta – KL divergence weight.

• ll_scale – Log likelihood scaling factor.

• pretrain_epochs – Number of pretraining epochs for generative energy model, which can stabilize training, before switching to normal training regime that includes KL divergence.

• num_samples_test – Number of samples to use for test time predictions.

• stochastic_module_names – List of module names that should become density layers.

• freeze_backbone – If True, freeze the backbone.

• optimizer – Optimizer.

• lr_scheduler – Learning rate scheduler.

Raises:

AssertionError – If num_samples_test is less than or equal to 0.

setup_task() → None

Set up task.

compute_kl_divergence() → Tensor

Compute the KL divergence of the model.

gather_loglikelihood() → Tensor

Gather loglikelihood terms from the density layers.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Test step.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dictionary

Density Uncertainty Regression

class lightning_uq_box.uq_methods.density_uncertainty.DensityLayerModelRegression(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 5, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Density Layer Model for Regression Tasks.

If you use this module in your work, please cite the following paper:

• https://arxiv.org/abs/2306.12497

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 5, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a Density Layer Model for Regression Tasks.

Parameters:

• model – PyTorch model that will be converted to a Density Layer Model.

• loss_fn – Loss function used for the target minimization, defaults to MSE Loss.

• prior_std – Standard deviation of the prior.

• posterior_std_init – Initial standard deviation of the posterior.

• kl_beta – KL divergence weight.

• ll_scale – Log likelihood scaling factor.

• pretrain_epochs – Number of pretraining epochs for generative energy model, which can stabilize training, before switching to normal training regime that includes KL divergence.

• num_samples_test – Number of samples to use for test time predictions.

• stochastic_module_names – List of module names that should become density layers.

• freeze_backbone – If True, freeze the backbone.

• optimizer – Optimizer.

• lr_scheduler – Learning rate scheduler

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for the metrics.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – input tensor

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

dictionary of predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Density Uncertainty Classification

class lightning_uq_box.uq_methods.density_uncertainty.DensityLayerModelClassification(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, task: str = 'multiclass', prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 5, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

Density Layer Model for Classification Tasks.

valid_tasks = ['binary', 'multiclass', 'multilable']

__init__(model: ~torch.nn.modules.module.Module, loss_fn: ~torch.nn.modules.module.Module | None = None, task: str = 'multiclass', prior_std: float = 0.1, posterior_std_init: float = 0.001, kl_beta: float = 1, ll_scale: float = 0.01, pretrain_epochs: int = 0, num_samples_test: int = 5, stochastic_module_names: list[int | str] | None = None, freeze_backbone: bool = False, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a Density Layer Model for Classification Tasks.

Parameters:

• model – PyTorch model that will be converted to a Density Layer Model.

• loss_fn – Loss function used for the target minimization, defaults to CrossEntropy Loss.

• task – Classification task type, one of “binary”, “multiclass”, or “multilabel”.

• prior_std – Standard deviation of the prior.

• posterior_std_init – Initial standard deviation of the posterior.

• kl_beta – KL divergence weight.

• ll_scale – Log likelihood scaling factor.

• pretrain_epochs – Number of pretraining epochs for generative energy model, which can stabilize training, before switching to normal training regime that includes KL divergence.

• num_samples_test – Number of samples to use for test time predictions.

• stochastic_module_names – List of module names that should become density layers.

• freeze_backbone – If True, freeze the backbone.

• optimizer – Optimizer.

• lr_scheduler – Learning rate scheduler

Raises:

AssertionError – If task is not one of the valid tasks.

setup_task() → None

Set up task specific attributes.

adapt_output_for_metrics(out: Tensor) → Tensor

Adapt the output for the metrics.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – input tensor

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

dictionary of predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Generative Models

Classification and Regression Diffusion (CARD)

CARD Base

class lightning_uq_box.uq_methods.cards.CARDBase(cond_mean_model: ~torch.nn.modules.module.Module, guidance_model: ~torch.nn.modules.module.Module, n_steps: int = 1000, beta_schedule: str = 'linear', beta_start: float = 1e-05, beta_end: float = 0.01, n_z_samples: int = 100, ema_decay: float = 0.995, ema_update_every: float = 10, ema_update_after_step: int = 0, guidance_optim: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

CARD Model.

Diffusion Model based on CARD paper.

If you use this in your research, please cite the following paper:

• https://arxiv.org/abs/2206.07275

pred_file_name = 'predictions.csv'

__init__(cond_mean_model: ~torch.nn.modules.module.Module, guidance_model: ~torch.nn.modules.module.Module, n_steps: int = 1000, beta_schedule: str = 'linear', beta_start: float = 1e-05, beta_end: float = 0.01, n_z_samples: int = 100, ema_decay: float = 0.995, ema_update_every: float = 10, ema_update_after_step: int = 0, guidance_optim: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of the CARD Model.

Parameters:

• cond_mean_model – conditional mean model, should be pretrained model that estimates $E[y|x]$

• guidance_model – guidance diffusion model

• n_steps – number of diffusion steps

• beta_schedule – what type of noise scheduling to conduct

• beta_start – start value of beta scheduling

• beta_end – end value of beta scheduling

• n_z_samples – number of samples during prediction

• ema_decay – exponential moving average decay

• ema_update_every – How often to update the EMA model, in terms of every n gradient steps.

• ema_update_after_step – after which step to start updating the EMA model

• guidance_optim – optimizer for the guidance model

• lr_scheduler – learning rate scheduler

Changed in version 0.2.0: Added arguments ema_decay, ema_update_every, ema_update_after_step for EMA support.

setup_task() → None

Setup task specific attributes.

diffusion_process(batch: dict[str, Tensor]) → Tensor

Diffusion process during training.

Parameters:

batch – the output of your DataLoader

Returns:

loss from diffusion process

on_after_backward()

Update EMA after each backward pass.

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, ndarray]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

diffusion samples for each time step

p_sample(x: Tensor, y: Tensor, y_0_hat: Tensor, y_T_mean: Tensor, t: int, alphas: Tensor, one_minus_alphas_bar_sqrt: Tensor) → Tensor

Reverse diffusion process sampling, one time step.

This is the process of generating a sample from the model’s prior distribution and then evolving it through the diffusion process. It starts from the final time step and goes backwards to the initial time step. At each time step, a noise variable is sampled and the state is updated according to the reverse diffusion process.

Parameters:

• x – input features

• y – sampled y at time step t, y_t.

• y_0_hat – prediction of pre-trained guidance model.

• y_T_mean – mean of prior distribution at timestep T.

• t – time step

• alphas – noise schedule alpha

• one_minus_alphas_bar_sqrt – noise schedule one minus alpha sqrt

Returns:

reverse process sample

p_sample_t_1to0(x: Tensor, y: Tensor, y_0_hat: Tensor, y_T_mean: Tensor, one_minus_alphas_bar_sqrt: Tensor) → Tensor

Reverse sample function, sample y_0 given y_1.

Parameters:

• x – input

• y – sampled y at time step t, y_t.

• y_0_hat – prediction of pre-trained guidance model.

• y_T_mean – mean of prior distribution at timestep T.

• one_minus_alphas_bar_sqrt – noise schedule one minus alpha bar sqrt

Returns:

y_0 sample

p_sample_loop(x: Tensor, y_0_hat: Tensor, y_T_mean: Tensor, n_steps: int, alphas: Tensor, one_minus_alphas_bar_sqrt: Tensor, only_last_sample: bool = False) → list[Tensor]

P sample loop for the entire chain.

Parameters:

• x – input

• y_0_hat – prediction of pre-trained guidance model.

• y_T_mean – mean of prior distribution at timestep T.

• n_steps – number of diffusion steps

• alphas – noise schedule alpha

• one_minus_alphas_bar_sqrt – noise schedule one minus alpha

• only_last_sample – whether to only return the last sample

Returns:

list of samples for each diffusion time step

q_sample(y: Tensor, y_0_hat: Tensor, alphas_bar_sqrt: Tensor, one_minus_alphas_bar_sqrt: Tensor, t: int, noise: Tensor | None = None) → Tensor

Q sampling process.

This is the process of approximating the posterior distribution of the latent variables given the observed data. It starts from the initial time step and goes forward to the final time step. At each time step, a noise variable is sampled and the state is updated according to the forward diffusion process.

Parameters:

• y – sampled y at time step t, y_t.

• y_0_hat – prediction of pre-trained guidance model.

• alphas_bar_sqrt – noise schedule alpha bar

• one_minus_alphas_bar_sqrt – noise schedule one minus alpha bar

• t – time step

• noise – optional noise tensor

Returns:

q sample at time step t

extract(input: Tensor, t: int, x: Tensor) → Tensor

Extract noise level at time step t from schedule.

Parameters:

• input – noise input

• t – time step

• x – tensor to make noisy version of

Returns:

noisy version of x

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the test loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

test loss

configure_optimizers() → Any

Configure optimizers.

CARD Regression

class lightning_uq_box.uq_methods.cards.CARDRegression(cond_mean_model: ~torch.nn.modules.module.Module, guidance_model: ~torch.nn.modules.module.Module, n_steps: int = 1000, beta_schedule: str = 'linear', beta_start: float = 1e-05, beta_end: float = 0.01, n_z_samples: int = 100, ema_decay: float = 0.995, ema_update_every: float = 10, ema_update_after_step: int = 0, guidance_optim: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

CARD Regression Model.

If you use this in your research, please cite the following paper:

• https://arxiv.org/abs/2206.07275

setup_task() → None

Setup task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

prediction dictionary with uncertainty estimates and samples

on_test_batch_end(outputs: dict[str, ndarray], batch: Any, batch_idx: int, dataloader_idx=0)

Test batch end save predictions.

CARD Classification

class lightning_uq_box.uq_methods.cards.CARDClassification(cond_mean_model: ~torch.nn.modules.module.Module, guidance_model: ~torch.nn.modules.module.Module, n_steps: int = 1000, beta_schedule: str = 'linear', beta_start: float = 1e-05, beta_end: float = 0.01, n_z_samples: int = 100, task: str = 'multiclass', ema_decay: float = 0.995, ema_update_every: float = 10, ema_update_after_step: int = 0, guidance_optim: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None)

CARD Classification Model.

If you use this in your research, please cite the following paper:

• https://arxiv.org/abs/2206.07275

valid_tasks = ['binary', 'multiclass']

__init__(cond_mean_model: ~torch.nn.modules.module.Module, guidance_model: ~torch.nn.modules.module.Module, n_steps: int = 1000, beta_schedule: str = 'linear', beta_start: float = 1e-05, beta_end: float = 0.01, n_z_samples: int = 100, task: str = 'multiclass', ema_decay: float = 0.995, ema_update_every: float = 10, ema_update_after_step: int = 0, guidance_optim: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] = None) → None

Initialize a new instance of the CARD Classification.

Parameters:

• cond_mean_model – conditional mean model, should be pretrained model that estimates $E[y|x]$

• guidance_model – guidance diffusion model

• n_steps – number of diffusion steps

• beta_schedule – what type of noise scheduling to conduct

• beta_start – start value of beta scheduling

• beta_end – end value of beta scheduling

• n_z_samples – number of samples during prediction

• task – classification task, either binary or multiclass

• ema_decay – exponential moving average decay

• ema_update_every – How often to update the EMA model, in terms of every n gradient steps.

• ema_update_after_step – after which step to start updating the EMA model

• guidance_optim – optimizer for the guidance model

• lr_scheduler – learning rate scheduler

Changed in version 0.2.0: Added arguments ema_decay, ema_update_every, ema_update_after_step for EMA support.

setup_task() → None

Setup task specific attributes.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Prediction step.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

predictions

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Noise Scheduler

class lightning_uq_box.uq_methods.cards.NoiseScheduler(schedule: str = 'linear', n_steps: int = 1000, beta_start: float = 1e-05, beta_end: float = 0.01)

Noise Scheduler for Diffusion Training.

valid_schedules = ['linear', 'const', 'quad', 'jsd', 'sigmoid', 'cosine', 'cosine_anneal']

__init__(schedule: str = 'linear', n_steps: int = 1000, beta_start: float = 1e-05, beta_end: float = 0.01) → None

Initialize a new instance of the noise scheduler.

Parameters:

• schedule – type of noise schedule

• n_steps – number of diffusion time steps

• beta_start – beta noise start value

• beta_end – beta noise end value

Raises:

AssertionError if schedule is invalid –

linear_schedule() → Tensor

Linear Schedule.

constant_schedule() → Tensor

Constant Schedule.

quadratic_schedule() → Tensor

Quadratic Schedule.

sigmoid_schedule() → Tensor

Sigmoid Schedule.

cosine_schedule() → Tensor

Cosine Schedule.

cosine_anneal_schedule() → Tensor

Cosine Annealing Schedule.

get_noisy_x_at_t(t, x) → Tensor

Retrieve a noisy representation at time step t.

Parameters:

• input – schedule version

• t – time step

• x – tensor ot make noisy version of

Returns:

A noisy

Probabilistic UNet

class lightning_uq_box.uq_methods.prob_unet.ProbUNet(model: ~torch.nn.modules.module.Module, latent_dim: int = 6, num_filters: list[int] = [32, 64, 128, 192], num_convs_per_block: int = 3, num_convs_fcomb: int = 4, fcomb_filter_size: int = 32, beta: float = 10.0, num_samples: int = 5, task: str = 'multiclass', criterion: ~torch.nn.modules.module.Module | None = None, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None, save_preds: bool = False)

Probabilistic U-Net.

If you use this code, please cite the following paper:

• https://arxiv.org/abs/1806.05034

valid_tasks = ['multiclass', 'binary']

pred_dir_name = 'preds'

__init__(model: ~torch.nn.modules.module.Module, latent_dim: int = 6, num_filters: list[int] = [32, 64, 128, 192], num_convs_per_block: int = 3, num_convs_fcomb: int = 4, fcomb_filter_size: int = 32, beta: float = 10.0, num_samples: int = 5, task: str = 'multiclass', criterion: ~torch.nn.modules.module.Module | None = None, optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None, save_preds: bool = False) → None

Initialize a new instance of ProbUNet.

Parameters:

• model – Unet model

• latent_dim – latent dimension

• num_filters – number of filters per block in AxisAlignedConvGaussian

• num_convs_per_block – num of convs per block in AxisAlignedConvGaussian

• num_convs_fcomb – number of convolutions in fcomb

• fcomb_filter_size – filter size for the fcomb network

• beta – beta parameter

• num_samples – number of latent samples to use during prediction

• task – task type, either “multiclass” or “binary”

• criterion – reconstruction criterion, Cross Entropy Loss by default

• optimizer – optimizer

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

setup_task() → None

Set up the task.

compute_loss(batch: dict[str, Tensor]) → dict[str, Tensor]

Compute the evidence lower bound (ELBO) of the log-likelihood of P(Y|X).

Parameters:

batch – batch of data with input and target key

Returns:

A dictionary containing the total loss, reconstruction loss, KL loss, and the reconstruction

kl_divergence(analytic: bool = True, z_posterior: Tensor | None = None) → Tensor

Compute the KL divergence between the posterior and prior KL(Q||P).

Parameters:

• analytic – calculate KL analytically or via sampling from the posterior

• z_posterior – if we use sampling to approximate KL we can sample here or supply a sample

Returns:

The KL divergence

reconstruct(use_posterior_mean: bool = False, z_posterior: Tensor | None = None) → Tensor

Reconstruct a segmentation from a posterior sample.

Decoding a posterior sample and UNet feature map

Parameters:

• use_posterior_mean – use posterior_mean instead of sampling z_q

• z_posterior – use a provided sample or sample from posterior latent space

Returns:

The reconstructed segmentation

sample(testing: bool = False) → Tensor

Sample a segmentation via reconstructing from a prior sample.

Parameters:

testing – whether to sample from the prior or use the mean

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

on_train_epoch_end()

Log epoch-level test metrics.

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

on_validation_epoch_end()

Log epoch-level test metrics.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Compute and return the test loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

test prediction dict

on_test_start() → None

Create logging directory and initialize metrics.

on_test_epoch_end()

Log epoch-level test metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Compute and return the prediction.

Parameters:

• X – the input image

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dict

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

Hierachical Probabilistic UNet

class lightning_uq_box.uq_methods.hierarchical_prob_unet.HierarchicalProbUNet(latent_dims: tuple[int, ...] = (1, 1, 1, 1), channels_per_block: tuple[int, ...] | None = None, num_in_channels: int = 3, num_classes: int = 2, down_channels_per_block: tuple[int, ...] | None = None, activation_fn: ~collections.abc.Callable[[~torch.Tensor], ~torch.Tensor] = <function relu>, convs_per_block: int = 3, blocks_per_level: int = 3, loss_type: str = 'geco', kappa: int = 10, beta: float | None = 100, top_k_percentage: float | None = None, deterministic_top_k: int | None = None, num_samples: int = 5, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None, save_preds: bool = False)

Hierarchical Probabilistic U-Net.

If you use this model, please cite the following paper:

• https://arxiv.org/pdf/1905.13077.pdf

valid_loss_types = ['elbo', 'geco']

valid_tasks = ['multiclass', 'binary']

pred_dir_name = 'preds'

__init__(latent_dims: tuple[int, ...] = (1, 1, 1, 1), channels_per_block: tuple[int, ...] | None = None, num_in_channels: int = 3, num_classes: int = 2, down_channels_per_block: tuple[int, ...] | None = None, activation_fn: ~collections.abc.Callable[[~torch.Tensor], ~torch.Tensor] = <function relu>, convs_per_block: int = 3, blocks_per_level: int = 3, loss_type: str = 'geco', kappa: int = 10, beta: float | None = 100, top_k_percentage: float | None = None, deterministic_top_k: int | None = None, num_samples: int = 5, task: str = 'multiclass', optimizer: ~typing.Callable[[~typing.Iterable], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, lr_scheduler: ~typing.Callable[[~torch.optim.optimizer.Optimizer], ~torch.optim.lr_scheduler.LRScheduler | ~lightning.pytorch.cli.ReduceLROnPlateau] | None = None, save_preds: bool = False) → None

Initialize a new HierarchicalProbUNET.

Parameters:

• latent_dims – A tuple of integers specifying the dimensions of the latents at each scale. The length of the tuple indicates the number of U-Net decoder scales that have latents

• channels_per_block – A tuple of integers specifying the number of output channels for each block or None. If None, the default values are used

• num_in_channels – the number of input channels

• num_classes – An integer specifying the number of output classes.

• down_channels_per_block – A tuple of integers specifying the number of intermediate channels for each block or None. If None, the intermediate channels are chosen equal to channels_per_block

• activation_fn – A callable activation function

• convs_per_block – An integer specifying the number of convolutional layers

• blocks_per_level – An integer specifying the number of residual blocks per level

• model – the hierarchical probabilistic u-net model

• loss_type – the type of loss to use, either “elbo” or “geco”

• kappa – the kappa parameter for the geco loss

• beta – the beta parameter for the elbo loss

• top_k_percentage – the percentage of top loss_mask pixels to use for the loss

• deterministic_top_k – An optional percentage for the top-k loss_mask.

• num_samples – the number of samples to use during prediction

• task – task type, either “multiclass” or “binary”

• optimizer – optimizer

• lr_scheduler – learning rate scheduler

• save_preds – whether to save predictions

setup_task() → None

Set up the task.

forward(X: Tensor) → Tensor

Forward pass through the model.

Parameters:

X – the input data

Returns:

the output of the model

construct_latent_space(img: Tensor, seg_mask: Tensor) → None

Construct the latent space.

Parameters:

• img – the input image

• seg_mask – the segmentation mask

compute_loss(batch: dict[str, Tensor], loss_mask: Tensor | None = None) → Tensor

Compute the loss from the output of the model.

Parameters:

• batch – the batch of data containing image and segmentation mask

• loss_mask – Optional tensor that masks some pixels from being included in the loss

Returns:

the loss

reconstruct(mean: bool = False) → dict[str, Any]

Reconstruct the input.

Parameters:

mean (bool, optional) – Whether to use the mean. Defaults to False.

Returns:

A dictionary containing encoder and decoder features.

Return type:

dict[str, Any]

kl() → dict[int, Tensor]

Compute the KL divergence.

Returns:

A dictionary containing the KL divergence for each level

sample(img: Tensor, mean: bool = False, z_q: Tensor | None = None) → dict[str, Any]

Sample from the model.

Parameters:

• img – The image tensor.

• mean – Whether to use the mean

• z_q – Latent tensor

Returns:

A dictionary containing encoder and decoder features

training_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the training loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

training loss

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the validation loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

validation loss

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → Tensor

Compute and return the test loss.

Parameters:

• batch – the output of your DataLoader

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

test prediction dict

on_test_start() → None

Create logging directory and initialize metrics.

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → Tensor

Compute and return the prediction.

Parameters:

• X – the input data

• batch_idx – the index of the batch

• dataloader_idx – the index of the dataloader

Returns:

prediction dict

configure_optimizers() → dict[str, Any]

Initialize the optimizer and learning rate scheduler.

Returns:

a “lr dict” according to the pytorch lightning documentation

UQ Calibration Methods

Test Time Augmentation (TTA)

class lightning_uq_box.uq_methods.inference_time_augmentation.TTABase(model: LightningModule | Module, tt_augmentation: list[Callable] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean')

Test Time Augmentation Module.

In addition to a prediction with no test time augmentation, an additional prediction will be made for each element in tt_augmentation.

valid_merge_strategies: list[str] = ['mean', 'median', 'sum', 'max', 'min']

__init__(model: LightningModule | Module, tt_augmentation: list[Callable] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean') → None

Initialize a new instance of TTA module.

Parameters:

• model – LightningModule or nn.Module used for prediction

• tt_augmentation – list of test time augmentation function, assumed to accept an input that is a dictionary with key self.input_key and self.target_key which is input and target, if None, a set of default augmentations will be used

• merge_strategy – strategy to merge the predictions from the different augmentations

compute_predictive_uncertainty(aug_tensor: Tensor) → dict[str, Tensor]

Merge predictions via different strategies to compute predictive uncertainty.

Parameters:

aug_tensor – The tensor containing predictions from different augmentations

Returns:

dict with predictive mean and uncertainty

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step for TTA procedure.

predict_step(X: Tensor, aug: list[Callable] = None, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with TTA applied.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• aug – augmentation function to apply to X

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

logits and conformalized prediction sets

setup_task() → None

Setup task.

validation_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

No validation step in TTA.

on_validation_start() → None

No validation step in TTA.

class lightning_uq_box.uq_methods.inference_time_augmentation.TTAClassification(model: LightningModule | Module, tt_augmentation: list[Callable[[...], Any]] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean', task: Literal['binary', 'multiclass', 'multilable'] = 'multiclass')

Classification Test Time Augmentation Module.

pred_file_name = 'preds.csv'

valid_tasks: list[str] = ['binary', 'multiclass', 'multilable']

__init__(model: LightningModule | Module, tt_augmentation: list[Callable[[...], Any]] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean', task: Literal['binary', 'multiclass', 'multilable'] = 'multiclass') → None

Initialize a new instance of TTA Classification module.

Parameters:

• model – LightningModule or nn.Module used for prediction

• tt_augmentation – list of test time augmentation function, assumed to accept an input that is a dictionary with key self.input_key and self.target_key which is input and target, if None, a set of default augmentations will be used

• merge_strategy – strategy to merge the predictions from the different augmentations

• task – task type, one of “binary”, “multiclass”, “multilabel”

setup_task() → None

Set up task specific attributes.

compute_predictive_uncertainty(aug_tensor: Tensor) → dict[str, Tensor]

Merge predictions according to merge_strategy.

Parameters:

aug_tensor – The tensor containing predictions from different augmentations

Returns:

The tensor after applying the merge strategy

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

class lightning_uq_box.uq_methods.inference_time_augmentation.TTARegression(model: LightningModule | Module, tt_augmentation: list[Callable[[...], Any]] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean')

Regression Test Time Augmentation Module.

pred_file_name = 'preds.csv'

__init__(model: LightningModule | Module, tt_augmentation: list[Callable[[...], Any]] | None = None, merge_strategy: Literal['mean', 'median', 'sum', 'max', 'min'] = 'mean') → None

Initialize a new instance of TTA Regression module.

Parameters:

• model – LightningModule or nn.Module used for prediction

• tt_augmentation – list of test time augmentation function, assumed to accept an input that is a dictionary with key self.input_key and self.target_key which is input and target, if None, a set of default augmentations will be used

• merge_strategy – strategy to merge the predictions from the different augmentations

setup_task() → None

Set up task specific attributes.

compute_predictive_uncertainty(aug_tensor: Tensor) → dict[str, Tensor]

Merge predictions according to merge_strategy.

Parameters:

aug_tensor – The tensor containing predictions from different augmentations

Returns:

The tensor after applying the merge strategy

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Conformal Quantile Regression

class lightning_uq_box.uq_methods.conformal_qr.ConformalQR(model: Module | LightningModule, quantiles: list[float] = [0.1, 0.5, 0.9], alpha: float = 0.1)

Conformalized Quantile Regression.

If you use this model, please cite the following paper:

• https://papers.nips.cc/paper_files/paper/2019/hash/5103c3584b063c431bd1268e9b5e76fb-Abstract.html

pred_file_name = 'preds.csv'

__init__(model: Module | LightningModule, quantiles: list[float] = [0.1, 0.5, 0.9], alpha: float = 0.1) → None

Initialize a new CQR model.

Parameters:

• model – underlying model to be wrapped

• quantiles – quantiles to be used for CQR

• alpha – 1 - alpha is desired error rate for quantile

setup_task() → None

Set up task.

forward(X: Tensor) → dict[str, Tensor]

Forward pass of model.

Parameters:

X – input tensor of shape [batch_size x input_dims]

Returns:

model output tensor of shape [batch_size x num_outputs]

adjust_model_logits(model_output: dict[str, Tensor] | Tensor) → dict[str, Tensor]

Conformalize underlying model output.

Parameters:

model_output – model output tensor of shape [batch_size x num_outputs]

Returns:

conformalized model predictions

on_train_end() → None

Perform CQR computation to obtain q_hat for predictions.

Parameters:

outputs – list of dictionaries containing model outputs and labels

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

Parameters:

• batch – batch of testing data

• batch_idx – batch index

• dataloader_idx – dataloader index

predict_step(X: Tensor) → dict[str, Tensor]

Prediction step that produces conformalized prediction sets.

Parameters:

• X – input tensor of shape [batch_size x input_dims]

• batch_idx – batch index

• dataloader_idx – dataloader index

configure_optimizers() → Optimizer | Sequence[Optimizer] | Tuple[Sequence[Optimizer], Sequence[LRScheduler | ReduceLROnPlateau | LRSchedulerConfig]] | OptimizerLRSchedulerConfig | Sequence[OptimizerLRSchedulerConfig] | None

No optimizer needed for Conformal QR.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Temperature TempScaling

class lightning_uq_box.uq_methods.temp_scaling.TempScaling(model: LightningModule | Module, optim_lr: float = 0.01, max_iter: int = 50, task: str = 'multiclass')

Temperature Scaling.

If you use this method, please cite the following paper:

• https://arxiv.org/abs/1706.04599

pred_file_name = 'preds.csv'

valid_tasks = ['binary', 'multiclass']

__init__(model: LightningModule | Module, optim_lr: float = 0.01, max_iter: int = 50, task: str = 'multiclass') → None

Initialize Temperature Scaling method.

Parameters:

• model – model to be calibrated with Temperature S

• optim_lr – learning rate for optimizer

• max_iter – maximum number of iterations to run optimizer

• task – classification task, one of “multiclass” or “binary”

setup_task() → None

Set up task.

adjust_model_logits(model_logits: Tensor) → Tensor

Adjust model logits by applying temperature scaling.

Parameters:

model_logits – model output logits of shape [batch_size x num_outputs]

Returns:

adjusted model logits of shape [batch_size x num_outputs]

on_train_epoch_end() → None

Perform CQR computation to obtain q_hat for predictions.

Parameters:

outputs – list of dictionaries containing model outputs and labels

predict_step(X: Tensor) → dict[str, Tensor]

Prediction step with applied temperature scaling.

Parameters:

X – input tensor of shape [batch_size x num_features]

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step after running posthoc fitting methodology.

Parameters:

• batch – batch of testing data

• batch_idx – batch index

• dataloader_idx – dataloader index

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Regularized Adaptive Prediction Sets (RAPS)

class lightning_uq_box.uq_methods.raps.RAPS(model: LightningModule | Module, optim_lr: float = 0.01, max_iter: int = 50, alpha: float = 0.1, kreg: int | None = None, lamda_param: float | None = None, randomized: bool = False, allow_zero_sets: bool = False, pct_param_tune: float = 0.3, lamda_criterion: str = 'size', task: str = 'multiclass')

Regularized Adaptive Prediction Sets (RAPS).

Conformal prediction method for classification tasks, as introduced by Angelopoulos et al. (2020).

If you use this method, please cite the following paper:

• https://arxiv.org/abs/2009.14193

valid_tasks = ['binary', 'multiclass', 'multilable']

valid_lambda_criterion = ['size', 'adaptiveness']

pred_file_name = 'preds.csv'

__init__(model: LightningModule | Module, optim_lr: float = 0.01, max_iter: int = 50, alpha: float = 0.1, kreg: int | None = None, lamda_param: float | None = None, randomized: bool = False, allow_zero_sets: bool = False, pct_param_tune: float = 0.3, lamda_criterion: str = 'size', task: str = 'multiclass') → None

Initialize RAPS.

Parameters:

• model – model to be calibrated with RAPS

• optim_lr – learning rate for optimizer

• max_iter – maximum number of iterations to run optimizer

• alpha – 1 - alpha is the desired coverage

• kreg – regularization param (smaller kreg leads to smaller sets)

• lamda_param – regularization param (larger lamda leads to smaller sets) (any value of kreg and lambda will lead to coverage, but will yield different set sizes)

• randomized – whether to use randomized version of conformal prediction

• allow_zero_sets – whether to allow sets of size zero

• pct_param_tune – fraction of calibration data to use for parameter tuning

• lamda_criterion – optimize for ‘size’ or ‘adaptiveness’

• task – task type, one of ‘binary’, ‘multiclass’, or ‘multilabel’

setup_task() → None

Setup task specific attributes.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step after running posthoc fitting methodology.

predict_step(X: Tensor, batch_idx: int = 0, dataloader_idx: int = 0) → dict[str, Tensor]

Predict step with RAPS applied.

Parameters:

• X – prediction batch of shape [batch_size x input_dims]

• batch_idx – the index of this batch

• dataloader_idx – the index of the dataloader

Returns:

logits and conformalized prediction sets

adjust_model_logits(model_logits: Tensor) → Tensor

Adjust model output according to RAPS with fitted Qhat.

Parameters:

model_logits – model output tensor of shape [batch_size x num_outputs]

Returns:

adjusted model logits tensor of shape [batch_size x num_outputs]

on_train_end() → None

Apply RASP conformal method.

on_test_batch_end(outputs: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch_idx – batch index

• dataloader_idx – dataloader index

Image to Image Conformal

class lightning_uq_box.uq_methods.img2img_conformal.Img2ImgConformal(model: LightningModule | Module, alpha: float = 0.1, delta: float = 0.1, min_lambda: float = 0.0, max_lambda: float = 6.0, num_lambdas: int = 1000, save_preds: bool = False)

Image-to-Image Conformal Uncertainty Estimation.

This module is a wrapper around a base model that provides conformal uncertainty estimates for image-to-image tasks, as introduced by Angelopoulos et al. (2022).

This default implementation uses the quantile regression approach, as it demonstrated good results in the original paper.

But other approaches can be implemented as well, by using a different architecture and overwriting adjust_model_logits and rcps_loss_fn methods.

If you use this model in your research, please cite the following paper:

• https://arxiv.org/abs/2202.05265

Added in version 0.2.0.

pred_dir_name = 'preds'

metrics_file_name = 'test_metrics.json'

__init__(model: LightningModule | Module, alpha: float = 0.1, delta: float = 0.1, min_lambda: float = 0.0, max_lambda: float = 6.0, num_lambdas: int = 1000, save_preds: bool = False)

Initialize a new Img2ImgConformal instance.

The user can select a risk level α ∈ (0, 1), and an error level δ ∈ (0, 1), such as α = δ = 0.1. Then, the conformal procedure construct intervals that contain at least 1 − α of the ground-truth pixel values with probability 1 − δ.

Parameters:

• model – the base model to be conformalized

• alpha – the risk level

• delta – the error level

• min_lambda – the minimum lambda value

• max_lambda – the maximum lambda value

• num_lambdas – the number of lambda values to search

• save_preds – whether to save the predictions

setup_task() → None

Set up task specific attributes.

forward(X: Tensor, lam: float | None) → dict[str, Tensor]

Forward pass of model.

Parameters:

• X – input tensor of shape [batch_size x input_dims]

• lam – The lambda parameter. Default is None.

Returns:

model output tensor of shape [batch_size x num_outputs]

adjust_model_logits(output: Tensor, lam: float | None = None) → tuple

Compute the nested sets from the output of the model.

Parameters:

• output – The output tensor.

• lam – The lambda parameter. Default is None.

Returns:

The lower edge, the output, and the upper edge.

rcps_loss_fn(pset: dict[str, Tensor], label: Tensor)

RCPS Loss function, fraction_missed_loss by default.

Parameters:

• pset – The prediction set, output from adjust_model_logits.

• label – The label

Returns:

The RCPS loss.

get_rcps_losses_from_outputs(logit_dataset: Dataset, lam: float)

Compute the RCPS loss from the model outputs.

Parameters:

• logit_dataset – The logit dataset.

• rcps_loss_fn – The RCPS loss function.

• lam – The lambda parameter

on_train_end() → None

Perform Img2Img calibration.

on_test_start() → None

Create logging directory and initialize metrics.

test_step(batch: dict[str, Tensor], batch_idx: int, dataloader_idx: int = 0) → dict[str, Tensor]

Test step.

Parameters:

• batch – batch of testing data

• batch_idx – batch index

• dataloader_idx – dataloader index

Returns:

test step dictionary with predictions

compute_metrics(pred_dict: dict[str, Tensor], labels: Tensor) → None

Compute metrics for a batch and append them to the metrics lists.

Parameters:

• pred_dict – The prediction dictionary from predict_step for a batch

• labels – The corresponding batch labels

on_test_end() → None

Summarize metrics.

predict_step(X: Tensor, lam: float | None = None) → dict[str, Tensor]

Prediction step with applied temperature scaling.

Parameters:

• X – input tensor of shape [batch_size x num_channels x height x width]

• lam – The lambda parameter

Returns:

prediction dictionary

on_test_batch_end(outputs: dict[str, Tensor], batch: Any, batch_idx: int, dataloader_idx: int = 0) → None

Test batch end save predictions.

Parameters:

• outputs – dictionary of model outputs and aux variables

• batch – batch from dataloader

• batch_idx – batch index

• dataloader_idx – dataloader index