deeplabcut.pose_estimation_pytorch.models.criterions.kl_discrete

SimCC Discrete KL Divergence loss with Gaussian Label Smoothing.

Can be used for SimCC-type heads. Modified from the mmpose implementation. For more details, see https://github.com/open-mmlab/mmpose.

Classes:

Name	Description
KLDiscreteLoss	KLDiscrete loss.

KLDiscreteLoss

Bases: BaseCriterion

KLDiscrete loss.

Parameters:

Name	Type	Description	Default
beta	float	Temperature for the softmax.	1.0
label_softmax	bool	Use softmax on the labels.	False
label_beta	float	Temperature for the softmax on the labels.	10.0
use_target_weight	bool	Allows the use a weighted loss for different joints.	True
mask	list[int] | None	Indices of masked keypoints.	None
mask_weight	float	Weight for masked keypoints.	1.0

Source code in deeplabcut/pose_estimation_pytorch/models/criterions/kl_discrete.py

@CRITERIONS.register_module
class KLDiscreteLoss(BaseCriterion):
    """KLDiscrete loss.

    Args:
        beta: Temperature for the softmax.
        label_softmax: Use softmax on the labels.
        label_beta: Temperature for the softmax on the labels.
        use_target_weight: Allows the use a weighted loss for different joints.
        mask: Indices of masked keypoints.
        mask_weight: Weight for masked keypoints.
    """

    def __init__(
        self,
        beta: float = 1.0,
        label_softmax: bool = False,
        label_beta: float = 10.0,
        use_target_weight: bool = True,
        mask: list[int] | None = None,
        mask_weight: float = 1.0,
    ):
        super().__init__()
        self.beta = beta
        self.label_softmax = label_softmax
        self.label_beta = label_beta
        self.use_target_weight = use_target_weight
        self.mask = mask
        self.mask_weight = mask_weight

        self.log_softmax = nn.LogSoftmax(dim=1)
        self.kl_loss = nn.KLDivLoss(reduction="none")

    def forward(
        self,
        output: torch.Tensor,
        target: torch.Tensor,
        weights: torch.Tensor | float = 1.0,
        **kwargs,
    ) -> torch.Tensor:
        n, k, _ = output.shape
        if self.use_target_weight and isinstance(weights, torch.Tensor):
            weight = weights.reshape(-1)
        else:
            weight = 1.0

        pred = output.reshape(-1, output.size(-1))
        target = target.reshape(-1, target.size(-1))
        loss = self.criterion(pred, target).mul(weight)
        if self.mask is not None:
            loss = loss.reshape(n, k)
            loss[:, self.mask] = loss[:, self.mask] * self.mask_weight

        return loss.sum() / k

    def criterion(self, dec_outs, labels):
        log_pt = self.log_softmax(dec_outs * self.beta)
        if self.label_softmax:
            labels = F.softmax(labels * self.label_beta, dim=1)
        loss = torch.mean(self.kl_loss(log_pt, labels), dim=1)
        return loss