RegressionMSELoss

class lightning_pose.losses.losses.RegressionMSELoss[source]

Bases: Loss

MSE loss between ground truth and predicted coordinates.

Attributes Summary

loss_name

Methods Summary

__call__(keypoints_targ, keypoints_pred[, stage])

Compute the regression MSE loss for a batch of predicted keypoints.

compute_loss(targets, predictions)

Compute element-wise MSE between target and predicted coordinates.

remove_nans(targets, predictions)

Mask out NaN coordinate entries from targets and predictions.

Attributes Documentation

loss_name: str = 'regression'

Methods Documentation

__call__(keypoints_targ: Float[Tensor, 'batch two_x_num_keypoints'], keypoints_pred: Float[Tensor, 'batch two_x_num_keypoints'], stage: Literal['train', 'val', 'test'] | None = None, **kwargs: Any) → tuple[Float[Tensor, ''], list[dict]][source]

Compute the regression MSE loss for a batch of predicted keypoints.

Parameters:

  • keypoints_targ – ground-truth (x, y) keypoints; NaN entries are ignored.

  • keypoints_pred – predicted (x, y) keypoints.

  • stage – training stage for logging.

  • **kwargs – ignored extra keyword arguments.

Returns:

Tuple of scalar loss and list of logging dicts.

compute_loss(targets: Float[Tensor, 'batch_x_two_x_num_keypoints'], predictions: Float[Tensor, 'batch_x_two_x_num_keypoints']) → Float[Tensor, 'batch_x_two_x_num_keypoints'][source]

Compute element-wise MSE between target and predicted coordinates.

Parameters:

  • targets – ground-truth coordinate values.

  • predictions – predicted coordinate values.

Returns:

Element-wise squared error tensor.

remove_nans(targets: Float[Tensor, 'batch two_x_num_keypoints'], predictions: Float[Tensor, 'batch two_x_num_keypoints']) → tuple[Float[Tensor, 'num_valid_keypoints'], Float[Tensor, 'num_valid_keypoints']][source]

Mask out NaN coordinate entries from targets and predictions.

Parameters:

  • targets – ground-truth (x, y) keypoints; NaN entries indicate unlabeled keypoints.

  • predictions – predicted (x, y) keypoints.

Returns:

Tuple of (clean_targets, clean_predictions) with NaN positions removed.

__init__(data_module: BaseDataModule | UnlabeledDataModule | None = None, epsilon: float = 0.0, log_weight: float = 0.0, **kwargs: Any) → None[source]

Initialize RegressionMSELoss.

Parameters:

  • data_module – data module providing access to datasets; passed to the parent class.

  • epsilon – loss values below this threshold are zeroed out.

  • log_weight – final weight in front of the loss term in the objective function is computed as 1.0 / (2.0 * exp(log_weight)).

__new__(**kwargs)