ELBOLoss#

class torch_uncertainty.losses.ELBOLoss(model, inner_loss, kl_weight, num_samples, dist_family=None)[source]#

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 (Module | None) – The Bayesian Neural Network to compute the loss for

  • inner_loss (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 | None) – 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)[source]#

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