Fix SuperPoint torch.export by replacing data-dependent shapes - #48950
Open
shreyescodes wants to merge 3 commits into
Open
shreyescodes wants to merge 3 commits into
shreyescodes wants to merge 3 commits into
Conversation
Contributor
There was a problem hiding this comment.
Copilot review overview
馃煛 Changes recommended
The default configuration does not preserve the existing eager and integration output contract.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (1)
What changed in this PR
This PR updates SuperPoint keypoint extraction for static-shape torch.export compatibility.
Changes:
- Replaces dynamic masking with penalized
topkselection and validity masks. - Batches decoding and enables static export tests with
max_keypoints=64.
| File | Summary |
|---|---|
tests/鈥媘odels/鈥媠uperpoint/鈥媡est_modeling_superpoint.py |
Enables static export coverage with a fixed keypoint limit. |
src/鈥媡ransformers/鈥媘odels/鈥媠uperpoint/鈥媘odeling_superpoint.py |
Implements static keypoint selection and batched decoding. A critical issue remains: the default max_keypoints=-1 changes eager keypoint behavior, ordering, and output shape. |
馃挕 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Member
|
cc @IlyasMoutawwakil since you commented on the original issue! |
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: superpoint |
Contributor
CI recapDashboard: View test results in Grafana |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR makes the
SuperPointmodel compatible withtorch.exportby removing data-dependent shape operations, addressing the SuperPoint portion of #48243.Previously,
SuperPointInterestPointDecoder._extract_keypointsrelied on boolean masking andtorch.nonzero()to filter keypoints, which resulted in dynamic tensor shapes that blocked tracing. I replaced this with a static-shaped approach usingtorch.topkand score penalties (-inf) for threshold and border constraints. The output is now a fixed-size tensor padded tomax_keypoints, with a validity mask.Since the output shape is now static across the batch, I also refactored the main
forwardmethod to process the entire batch directly without needing the Python list comprehension loop.All tests pass locally, including the
test_torch_export_statictests which I've unskipped for this model.