Implementation:Kornia Kornia Divergence Loss

Knowledge Sources	Kornia
Domains	Vision, Loss_Functions
Last Updated	2026-02-09 15:00 GMT

Overview

Divergence Loss provides Kullback-Leibler (KL) and Jensen-Shannon (JS) divergence losses between 2D probability distribution heatmaps.

Description

This module implements two statistical divergence measures for comparing probability distributions represented as 2D heatmaps:

KL Divergence (Kullback-Leibler): Measures how one probability distribution diverges from a second expected probability distribution. It is asymmetric, meaning KL(P||Q) is not equal to KL(Q||P).

JS Divergence (Jensen-Shannon): A symmetrized and smoothed version of KL divergence, computed as:

$\text{JSD}(P\Vert Q)=\frac{1}{2}{D}_{KL}(P\Vert M)+\frac{1}{2}{D}_{KL}(Q\Vert M)$

where $M=\frac{1}{2}(P+Q)$.

Both losses operate on 4D tensors shaped as (B, N, H, W) representing batches of 2D heatmaps (e.g., keypoint heatmaps in pose estimation). The heatmaps are reshaped to compute divergences per channel.

Usage

Import these losses when comparing predicted heatmaps with ground-truth probability distributions, such as in human pose estimation, keypoint detection, or any task that involves predicting spatial probability maps.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/losses/divergence.py
• Lines: 1-90

Signature

def js_div_loss_2d(
    pred: torch.Tensor,
    target: torch.Tensor,
    reduction: str = "mean",
) -> torch.Tensor: ...

def kl_div_loss_2d(
    pred: torch.Tensor,
    target: torch.Tensor,
    reduction: str = "mean",
) -> torch.Tensor: ...

Import

from kornia.losses import js_div_loss_2d
from kornia.losses import kl_div_loss_2d

I/O Contract

Inputs

Name	Type	Required	Description
pred	torch.Tensor	Yes	Input heatmap tensor with shape (B, N, H, W)
target	torch.Tensor	Yes	Target heatmap tensor with shape (B, N, H, W)
reduction	str	No	Reduction mode: 'none', 'mean' (default), or 'sum'

Outputs

Name	Type	Description
loss	torch.Tensor	Computed divergence loss; shape (B, N) for 'none', scalar for 'mean'/'sum'

Usage Examples

import torch
from kornia.losses import js_div_loss_2d, kl_div_loss_2d

# Create uniform probability heatmaps
pred = torch.full((1, 1, 2, 4), 0.125)
target = torch.full((1, 1, 2, 4), 0.125)

# JS divergence (symmetric)
js_loss = js_div_loss_2d(pred, target)
# Returns 0.0 for identical distributions

# KL divergence (asymmetric)
kl_loss = kl_div_loss_2d(pred, target)
# Returns 0.0 for identical distributions

# With different distributions
pred_diff = torch.rand(2, 3, 8, 8)
pred_diff = pred_diff / pred_diff.sum(dim=(-1, -2), keepdim=True)
target_diff = torch.rand(2, 3, 8, 8)
target_diff = target_diff / target_diff.sum(dim=(-1, -2), keepdim=True)
loss = js_div_loss_2d(pred_diff, target_diff, reduction="mean")

Related Pages

• Principle:Kornia_Kornia_Distribution_Divergence