Fix SuperAnimal / pretrained load for RTMPose: implement convert_weights on RTMCCHead - #3270
Conversation
JiwaniZakir
left a comment
JiwaniZakir
left a comment
There was a problem hiding this comment.
The convert_weights method in rtmcc_head.py checks if fl_w in state_dict before applying the conversion, but then unconditionally accesses state_dict[fl_b] without a corresponding guard. If a checkpoint uses a bias-free final_layer, this will raise a KeyError rather than silently skipping — worth adding a separate if fl_b in state_dict check or bundling both under a single check that verifies both keys exist.
The fallback torch.rand((), ...) when old_vals is empty introduces non-determinism into weight loading, which is particularly problematic for reproducibility of fine-tuning runs. Since this branch handles positions where no old keypoint pair contributes (an edge case whose practical frequency is unclear), torch.zeros or w_old.mean() would be safer and deterministic defaults.
It's also worth documenting the assumption baked into old_idx = int(conv[i] - conv[j]) + k_old - 1: this treats differences in original keypoint indices as relative positional distances in the GAU weight table. That coupling between semantic keypoint ordering and positional encoding is non-obvious and could silently produce incorrect weight remappings if the pretrained model was trained with a different keypoint ordering convention.
|
@JiwaniZakir excellent feedback, thanks! Based on your suggestions, I incorporated the following changes:
Note: for convenience, I've set the default behavior Thanks again for your review! |
…nsure deterministic init.
4a9d990 to
1d1b59c
Compare
There was a problem hiding this comment.
Pull request overview
This PR fixes SuperAnimal pretrained checkpoint loading for RTMPose models by adding weight-conversion support to RTMCCHead, aligning it with other heads that can subset/reorder keypoints during transfer learning.
Changes:
- Make
RTMCCHeadsubclassWeightConversionMixinto participate in the SuperAnimal conversion workflow. - Implement
RTMCCHead.convert_weights(...)to remapfinal_layeroutput channels and adaptgau.wto new keypoint counts.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Description
Training or fine-tuning with SuperAnimal weight initialization and with_decoder=True loads a checkpoint and calls
head.convert_weights(...)on each head so bodyparts can be subset and reordered. HeatmapHead already provides this via WeightConversionMixin, but RTMCCHead did not, which caused:This change makes RTMCCHead consistent with other heads that support SuperAnimal transfer.
What changed
RTMCCHead now subclasses
WeightConversionMixinalongsideBaseHead, matching the pattern used byHeatmapHead.convert_weightsimplementationfinal_layer: Reorders/subsets the Conv2d output channels using the same convention as
DeconvModule.convert_weightsfor per–bodypart outputs (weight and bias indexed by the conversion tensor).gau.w: TheGatedAttentionUnitrelative bias table changes size when the number of keypoints changes; it cannot be remapped by a simple channel permutation. The PR rebuilds w to shape2 * K_new - 1using a mapping from new token pairs to pretrained relative indices (with averaging when multiple pairs collapse to the same new relative offset), so the checkpoint can be loaded without shape errors. (If reviewers prefer a more conservative transfer, this part can be swapped for re-initialization ofgau.wonly.)