Investigate training support for Sapiens2ForPoseEstimation when labels are provided
Author: Sainava
Pull request: no
Labels: name: Feature request

### Feature request

While exploring the recently added `Sapiens2ForPoseEstimation` implementation, I noticed that the model's `forward()` method accepts `labels` (documented as heatmap ground truth with shape `(batch_size, num_keypoints, height, width)`), but currently raises `NotImplementedError("Training is not yet supported")` when labels are provided.

I investigated the original Meta Sapiens2 pose-estimation implementation and found that pose training is configured through OpenMMLab's `KeypointMSELoss` with `use_target_weight=True`.

Relevant references:

* Hugging Face implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/sapiens2/modular_sapiens2.py#L1713

* Meta Sapiens2 pose configuration: https://github.com/facebookresearch/sapiens2/blob/main/sapiens/pose/configs/keypoints308/shutterstock_goliath_3po/sapiens2_1b_keypoints308_shutterstock_goliath_3po-1024x768.py#L270-297

* OpenMMLab `KeypointMSELoss` implementation: https://github.com/open-mmlab/mmpose/blob/main/mmpose/models/losses/heatmap_loss.py

* Sapiens2 paper: https://arxiv.org/abs/2604.21681

I would like to better understand whether the lack of training support is intentional or if support for supervised pose-estimation training would be considered within scope for the Transformers implementation.


### Motivation

While exploring `Sapiens2ForPoseEstimation`, I noticed that the model accepts `labels` and documents them as heatmap ground truth targets for loss computation, but currently raises `NotImplementedError("Training is not yet supported")` whenever labels are provided.

To understand whether this was specific to Sapiens2 or part of a broader design decision, I investigated the original Meta Sapiens2 implementation as well as the OpenMMLab components used during training.

My findings were:

* The Hugging Face implementation already supports the complete inference pipeline, including feature extraction, heatmap decoding, and returning pose heatmaps.
* The forward signature explicitly accepts `labels` and documents them as heatmap supervision targets.
* The original Meta Sapiens2 pose-estimation configuration uses OpenMMLab's `KeypointMSELoss` with `use_target_weight=True` and `loss_weight=10.0`.
* The corresponding OpenMMLab implementation computes a heatmap-based MSE loss and supports target weighting through a masking mechanism used for keypoint visibility.
* I also checked `VitPoseForPoseEstimation` in Transformers and found the same pattern: labels are accepted but training currently raises `NotImplementedError`.

Because both Sapiens2 and ViTPose expose pose-estimation outputs and accept supervision targets but currently stop before loss computation, I wanted to understand whether pose-estimation training support is intentionally omitted from Transformers or whether support for supervised training would be considered within scope.

### Your contribution

I would be happy to contribute a PR if maintainers consider pose-estimation training support to be within scope for Transformers.

Based on the investigation above, a potential direction would be to implement the supervised pose-estimation loss path when `labels` are provided, using the original Sapiens2 training setup and OpenMMLab's `KeypointMSELoss` behavior as references while remaining consistent with existing Transformers conventions.

If this aligns with maintainer expectations, I would be interested in contributing the implementation, accompanying tests, and documentation.



Comments:
Hi @guarin , gentle ping on this issue.

While investigating the original Sapiens2 training setup, I found that pose training relies on OpenMMLab's `KeypointMSELoss`, whereas the current Transformers implementation accepts `labels` but raises `NotImplementedError` when they are provided.

Before spending time working on a PR, I wanted to check whether supervised pose-estimation training support is something maintainers would consider in scope for `Sapiens2ForPoseEstimation`.

I'd be happy to help with an implementation if this is a direction the team would be interested in pursuing.


Ah sorry missed this issue, thanks for the ping! A PR with the loss would be really appreciated and a great addition to the library. I didn't include it in the initial PR because the focus was on getting inference working first. 

Thanks for the clarification!

That makes sense. I'll take a closer look at the original Sapiens2/OpenMMLab training setup and start working on a PR for the pose-estimation loss path. I'll open a draft PR once I have something working so it's easier to discuss implementation details.