Distributions

Both losses compare (B, N, H, W) probability heatmaps. Use kl_div_loss_2d() when target is the reference distribution, and js_div_loss_2d() for a symmetric loss that stays finite where one heatmap is zero.

kornia.losses.js_div_loss_2d(pred, target, reduction='mean')

Compute the Jensen-Shannon divergence between heatmaps.

loss = js_div_loss_2d(pred, target) is the symmetric divergence of each of the B N pairs of (H, W) probability heatmaps.

\[\text{JSD}(P, Q) = \frac{1}{2} \mathrm{KL}(P \,\|\, M) + \frac{1}{2} \mathrm{KL}(Q \,\|\, M), \quad M = \frac{P + Q}{2}\]

Convention:

• The divergence is in nats and at most \(\ln 2\). The inputs and the reductions are those of kl_div_loss_2d(); see its Convention block.

• A cell that is zero in both inputs contributes zero, including to the gradients. Both inputs must have the same shape; unsupported reductions raise NotImplementedError.

• The mixture is evaluated in log space so halving a tiny probability cannot round it to zero. Float16 and bfloat16 inputs are evaluated in float32. The returned loss retains the promoted input dtype.

Parameters:

• pred (Tensor) – predicted heatmaps with shape (B, N, H, W).

• target (Tensor) – target heatmaps with the same shape.

• reduction (str, optional) – 'mean' (default) averages the B N divergences, 'sum' adds them, 'none' returns them with shape (B, N). Default: "mean"

Return type:

Tensor

Returns:

the loss, a scalar or, for reduction='none', with shape (B, N).

See also

• kl_div_loss_2d(): not symmetric, and inf where pred is zero and target is positive.

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')

Compute the Kullback-Leibler divergence KL(target || pred) between heatmaps.

loss = kl_div_loss_2d(pred, target) is the divergence of each of the B N pairs of (H, W) probability heatmaps, in nats.

\[\mathrm{KL}(\text{target} \,\|\, \text{pred}) = \sum_{H \times W} \text{target} \left(\log \text{target} - \log \text{pred}\right)\]

Convention:

• The order is that of torch’s F.kl_div(pred.log(), target). Losses and metrics maps it onto torch, scipy and torchmetrics.

• The inputs are probabilities, each \((b, n)\) slice a distribution over \(H \times W\). The log is taken inside, so log-probabilities give NaN, and nothing is normalised. A zero in pred where target is positive gives inf.

• The default 'mean' averages over the \(B N\) distributions, not over every element as torch’s does.

• Both inputs must have the same shape. Unsupported reductions, including torch’s 'batchmean', raise NotImplementedError. A zero in target with a finite nonnegative pred contributes zero to the value and gradients, including when both entries are zero.

• Under CPU autocast, the logarithm of a float16/bfloat16 pred is taken in float32, as CUDA and MPS autocast already do; F.kl_div alone would promote to float32 only after that logarithm had been rounded.

Parameters:

• pred (Tensor) – predicted heatmaps with shape (B, N, H, W).

• target (Tensor) – target heatmaps with the same shape.

• reduction (str, optional) – 'mean' (default) averages the B N divergences, 'sum' adds them, 'none' returns them with shape (B, N). Default: "mean"

Return type:

Tensor

Returns:

the loss, a scalar or, for reduction='none', with shape (B, N).

See also

• js_div_loss_2d(): symmetric and at most \(\ln 2\), finite where pred is zero.

Examples

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