HeatmapMSELoss

class lightning_pose.losses.losses.HeatmapMSELoss[source]

Bases: HeatmapLoss

MSE loss between heatmaps.

Attributes Summary

loss_name

Methods Summary

compute_loss(targets,Β predictions)

Compute pixel-wise MSE between target and predicted heatmaps.

Attributes Documentation

loss_name: str = 'heatmap_mse'

Methods Documentation

compute_loss(targets: Float[Tensor, 'batch_x_num_keypoints heatmap_height heatmap_width'], predictions: Float[Tensor, 'batch_x_num_keypoints heatmap_height heatmap_width']) → Float[Tensor, 'batch_x_num_keypoints heatmap_height heatmap_width'][source]

Compute pixel-wise MSE between target and predicted heatmaps.

Multiplies by the number of heatmap pixels (h * w) to keep the loss magnitude consistent across different heatmap resolutions.

Parameters:

  • targets – ground-truth heatmaps.

  • predictions – model-predicted heatmaps.

Returns:

element-wise MSE scaled by heatmap area.

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

Initialize HeatmapMSELoss.

Parameters:

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

  • 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)