Distributions#

kornia.losses.js_div_loss_2d(pred, target, reduction='mean')[source]#

Calculate the Jensen-Shannon divergence loss between heatmaps.

Parameters:

pred (Tensor) – the input torch.Tensor with shape \((B, N, H, W)\).
target (Tensor) – the target torch.Tensor with shape \((B, N, H, W)\).
reduction (str, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Default: "mean"

Return type:

Tensor

Examples

>>> pred = torch.full((1, 1, 2, 4), 0.125)
>>> loss = js_div_loss_2d(pred, pred)
>>> loss.item()
0.0

kornia.losses.kl_div_loss_2d(pred, target, reduction='mean')[source]#

Calculate the Kullback-Leibler divergence loss between heatmaps.

Parameters:

pred (Tensor) – the input torch.Tensor with shape \((B, N, H, W)\).
target (Tensor) – the target torch.Tensor with shape \((B, N, H, W)\).
reduction (str, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Default: "mean"

Return type:

Tensor

Examples

>>> pred = torch.full((1, 1, 2, 4), 0.125)
>>> loss = kl_div_loss_2d(pred, pred)
>>> loss.item()
0.0