ELBOLoss

class torch_uncertainty.losses.ELBOLoss(model, inner_loss, kl_weight, num_samples, dist_family=None)

The Evidence Lower Bound (ELBO) loss for Bayesian Neural Networks.

ELBO loss for Bayesian Neural Networks. Use this loss function with the objective that you seek to minimize as inner_loss.

Parameters:

• model (nn.Module) – The Bayesian Neural Network to compute the loss for

• inner_loss (nn.Module) – The loss function to use during training

• kl_weight (float) – The weight of the KL divergence term

• num_samples (int) – The number of samples to use for the ELBO loss

• dist_family (str, optional) – The distribution family to use for the output of the model. None means point-wise prediction. Defaults to None.

Note

Set the model to None if you use the ELBOLoss within the ClassificationRoutine. It will get filled automatically.

forward(inputs, targets)

Gather the KL divergence from the Bayesian modules and aggregate the ELBO loss for a given network.

Parameters:

• inputs (Tensor) – The inputs of the Bayesian Neural Network

• targets (Tensor) – The target values

Returns:

The aggregated ELBO loss

Return type:

Tensor