lightning_pose.losses

lightning_pose.losses.factory Module

Factory utilities for building and combining losses from a Hydra config.

Three components work together:

get_loss_classes() — returns the registry mapping loss-name strings to classes.
get_loss_factories() — reads cfg.losses and cfg.model.losses_to_use, assembles per-loss parameter dicts, and returns a {'supervised': LossFactory, 'unsupervised': LossFactory} dict ready to be passed to a model constructor.
LossFactory — a LightningModule that holds instantiated loss objects and computes the total weighted loss in its __call__ method.

Adding a new loss:

Define the class in losses/losses.py, inheriting from Loss. For a heatmap-based loss (operating on heatmaps_targ/heatmaps_pred), inherit from HeatmapLoss instead — it already provides NaN handling and the __call__ pipeline, so only compute_loss needs overriding. Set a loss_name: str class attribute (single name) or multiple LOSS_NAME_*: str class attributes when one class serves several config strings (e.g. PCALoss).
Import the class at the top of this file and add one entry per name to the dict returned by get_loss_classes().
Parameter wiring in get_loss_factories(); only needed in two cases:
- Unsupervised loss, param derived from ``cfg.data``/``cfg.model`` (not a literal cfg.losses.<loss_name> key — literal keys are forwarded automatically): add an elif branch, e.g.:
```
elif loss_name == 'my_new_loss':
    loss_params_dict['unsupervised'][loss_name]['my_param'] = cfg.data.some_field
```
- Supervised loss (any supervised_* name): never auto-forwarded — add an explicit block reading each param you need, e.g.:
```
log_weight = cfg.losses.get('my_new_loss', {}).get('log_weight')
if log_weight is not None:
    loss_params_dict['supervised']['my_new_loss'] = {'log_weight': log_weight}
```
Skipping this step when it’s actually required doesn’t raise a KeyError — the loss is simply never added to loss_params_dict, so it silently never runs.

lightning_pose.losses.losses Module

Supervised and unsupervised losses implemented in pytorch.

The lightning pose package defines each loss as its own class; an initialized loss object, in addition to computing the loss, stores hyperparameters related to the loss (weight in the final objective funcion, epsilon-insensitivity parameter, etc.)

A separate LossFactory class (defined in lightning_pose.losses.factory) collects all losses for a given model and orchestrates their execution, logging, etc.

The general flow of each loss class is as follows: - input: predicted and ground truth data - step 0: remove ground truth samples containing nans if desired - step 1: compute loss for each batch element/keypoint/etc - step 2: epsilon-insensitivity: set loss to zero for any batch element with loss < epsilon - step 3: reduce loss (usually mean) - step 4: log values to a dict - step 5: return loss

Classes

`Loss`	Parent class for all losses.
`HeatmapLoss`	Parent class for heatmap-shaped losses (MSE, KL, JS divergence, etc).
`HeatmapMSELoss`	MSE loss between heatmaps.
`HeatmapKLLoss`	Kullback-Leibler loss between heatmaps.
`HeatmapJSLoss`	Jensen-Shannon loss between heatmaps.
`PCALoss`	Penalize predictions that fall outside a low-dimensional subspace.
`TemporalLoss`	Penalize temporal differences for each target.
`TemporalHeatmapLoss`	Penalize temporal differences for each heatmap.
`RegressionMSELoss`	MSE loss between ground truth and predicted coordinates.
`RegressionRMSELoss`	Root MSE loss between ground truth and predicted coordinates.
`PairwiseProjectionsLoss`	Penalize projections from each pair of cameras into 3D world space.
`ReprojectionHeatmapLoss`	Penalize error between predicted 2D->3D->2D->heatmap and ground truth heatmap.