Bayesian KL Loss

torchbnn.functional.bayesian_kl_loss(model, reduction='mean', last_layer_only=False)[source]

An method for calculating KL divergence of whole layers in the model.

  • model (nn.Module) – a model to be calculated for KL-divergence.
  • reduction (string, optional) – Specifies the reduction to apply to the output: 'mean': the sum of the output will be divided by the number of elements of the output. 'sum': the output will be summed.
  • last_layer_only (Bool) – True for return only the last layer’s KL divergence.