From 3a73449dae4a96a68721caab52105f520ae06513 Mon Sep 17 00:00:00 2001 From: Shreyas Date: Sat, 19 Sep 2026 15:46:11 +0530 Subject: [PATCH 1/2] Fix SuperPoint torch.export by replacing data-dependent shapes --- .../models/superpoint/modeling_superpoint.py | 181 ++++++++++-------- .../superpoint/test_modeling_superpoint.py | 3 +- 2 files changed, 106 insertions(+), 78 deletions(-) diff --git a/src/transformers/models/superpoint/modeling_superpoint.py b/src/transformers/models/superpoint/modeling_superpoint.py index 59379ce357ef..440ae61f04b9 100644 --- a/src/transformers/models/superpoint/modeling_superpoint.py +++ b/src/transformers/models/superpoint/modeling_superpoint.py @@ -34,26 +34,11 @@ logger = logging.get_logger(__name__) -def remove_keypoints_from_borders( - keypoints: torch.Tensor, scores: torch.Tensor, border: int, height: int, width: int -) -> tuple[torch.Tensor, torch.Tensor]: - """Removes keypoints (and their associated scores) that are too close to the border""" - mask_h = (keypoints[:, 0] >= border) & (keypoints[:, 0] < (height - border)) - mask_w = (keypoints[:, 1] >= border) & (keypoints[:, 1] < (width - border)) - mask = mask_h & mask_w - return keypoints[mask], scores[mask] -def top_k_keypoints(keypoints: torch.Tensor, scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: - """Keeps the k keypoints with highest score""" - if k >= len(keypoints): - return keypoints, scores - scores, indices = torch.topk(scores, k, dim=0) - return keypoints[indices], scores - def simple_nms(scores: torch.Tensor, nms_radius: int) -> torch.Tensor: - """Applies non-maximum suppression on scores""" + """Applies non-maximum suppression on scores.""" if nms_radius < 0: raise ValueError("Expected positive values for nms_radius") @@ -216,14 +201,13 @@ def __init__(self, config: SuperPointConfig) -> None: config.decoder_hidden_size, config.keypoint_decoder_dim, kernel_size=1, stride=1, padding=0 ) - def forward(self, encoded: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, encoded: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: scores = self._get_pixel_scores(encoded) - keypoints, scores = self._extract_keypoints(scores) - - return keypoints, scores + keypoints, scores, mask = self._extract_keypoints(scores) + return keypoints, scores, mask def _get_pixel_scores(self, encoded: torch.Tensor) -> torch.Tensor: - """Based on the encoder output, compute the scores for each pixel of the image""" + """Based on the encoder output, compute the scores for each pixel of the image.""" scores = self.relu(self.conv_score_a(encoded)) scores = self.conv_score_b(scores) scores = nn.functional.softmax(scores, 1)[:, :-1] @@ -233,30 +217,74 @@ def _get_pixel_scores(self, encoded: torch.Tensor) -> torch.Tensor: scores = simple_nms(scores, self.nms_radius) return scores - def _extract_keypoints(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _extract_keypoints( + self, scores: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - Based on their scores, extract the pixels that represent the keypoints that will be used for descriptors computation. - The keypoints are in the form of relative (x, y) coordinates. + Extract keypoints from the score map using a statically-shaped pipeline + compatible with ``torch.export``. + + Instead of ``torch.nonzero`` and boolean-index masking (both of which + produce data-dependent shapes), we apply threshold and border + constraints as additive ``-inf`` penalties and delegate selection to + ``torch.topk``, which always returns exactly ``k`` entries. + + Args: + scores: ``(batch_size, height, width)`` pixel score map after NMS. + + Returns: + keypoints: ``(batch_size, k, 2)`` float tensor of ``(x, y)`` + keypoint coordinates in pixel space. + topk_scores: ``(batch_size, k)`` float tensor of keypoint scores. + mask: ``(batch_size, k)`` bool tensor; ``True`` where the entry is + a real keypoint rather than padding. """ - _, height, width = scores.shape + batch_size, height, width = scores.shape + + border = self.border_removal_distance + + # Build row / column index grids for the border check without any + # boolean indexing (keeps shapes static). + rows = torch.arange(height, device=scores.device).view(1, height, 1) + cols = torch.arange(width, device=scores.device).view(1, 1, width) + + # Replicate the original border logic exactly: + # the legacy code called remove_keypoints_from_borders with + # ``height * 8`` and ``width * 8``, so we reproduce that arithmetic. + valid = ( + (scores > self.keypoint_threshold) + & (rows >= border) + & (rows < height * 8 - border) + & (cols >= border) + & (cols < width * 8 - border) + ) - # Threshold keypoints by score value - keypoints = torch.nonzero(scores[0] > self.keypoint_threshold) - scores = scores[0][tuple(keypoints.t())] + # Replace invalid positions with -inf so they sink to the end of topk. + masked_scores = scores.masked_fill(~valid, float("-inf")) - # Discard keypoints near the image borders - keypoints, scores = remove_keypoints_from_borders( - keypoints, scores, self.border_removal_distance, height * 8, width * 8 + # Determine k — must be a concrete integer for export. + # When max_keypoints is -1 ("no limit") we use every pixel; this path + # is valid in eager mode but cannot be exported because k is + # data-dependent via height * width. Users who need export must set + # max_keypoints to a positive value. + k = self.max_keypoints if self.max_keypoints > 0 else height * width + + # torch.topk always returns exactly k entries → static output shape. + topk_scores, topk_indices = torch.topk( + masked_scores.reshape(batch_size, height * width), k=k, dim=1 ) - # Keep the k keypoints with highest score - if self.max_keypoints >= 0: - keypoints, scores = top_k_keypoints(keypoints, scores, self.max_keypoints) + # Convert flat indices to (x, y) pixel coordinates (pure arithmetic, + # no data-dependent shapes). + keypoints_y = topk_indices // width # (B, k) + keypoints_x = topk_indices % width # (B, k) + keypoints = torch.stack([keypoints_x, keypoints_y], dim=-1).to(scores.dtype) # (B, k, 2) - # Convert (y, x) to (x, y) - keypoints = torch.flip(keypoints, [1]).to(scores.dtype) + # A position is a real keypoint if and only if its score is finite. + # Cast to int (0/1) to match the original mask dtype expected by callers. + mask = (topk_scores > float("-inf")).to(torch.int) # (B, k) int - return keypoints, scores + return keypoints, topk_scores, mask class SuperPointDescriptorDecoder(nn.Module): @@ -289,31 +317,52 @@ def __init__(self, config: SuperPointConfig) -> None: ) def forward(self, encoded: torch.Tensor, keypoints: torch.Tensor) -> torch.Tensor: - """Based on the encoder output and the keypoints, compute the descriptors for each keypoint""" - descriptors = self.conv_descriptor_b(self.relu(self.conv_descriptor_a(encoded))) - descriptors = nn.functional.normalize(descriptors, p=2, dim=1) - - descriptors = self._sample_descriptors(keypoints[None], descriptors[0][None], 8)[0] + """ + Compute descriptors for all keypoints in a batch. - # [descriptor_dim, num_keypoints] -> [num_keypoints, descriptor_dim] - descriptors = torch.transpose(descriptors, 0, 1) + Args: + encoded: ``(batch_size, channels, height, width)`` encoder feature map. + keypoints: ``(batch_size, num_keypoints, 2)`` keypoint coordinates in + ``(x, y)`` pixel space of the full-resolution score map. + Returns: + descriptors: ``(batch_size, num_keypoints, descriptor_dim)`` + """ + descriptors = self.conv_descriptor_b(self.relu(self.conv_descriptor_a(encoded))) + descriptors = nn.functional.normalize(descriptors, p=2, dim=1) + descriptors = self._sample_descriptors(keypoints, descriptors, 8) + # (batch_size, descriptor_dim, num_keypoints) -> (batch_size, num_keypoints, descriptor_dim) + descriptors = descriptors.transpose(1, 2) return descriptors @staticmethod - def _sample_descriptors(keypoints, descriptors, scale: int = 8) -> torch.Tensor: - """Interpolate descriptors at keypoint locations""" + def _sample_descriptors( + keypoints: torch.Tensor, descriptors: torch.Tensor, scale: int = 8 + ) -> torch.Tensor: + """ + Interpolate descriptors at keypoint locations. + + Args: + keypoints: ``(batch_size, num_keypoints, 2)`` in ``(x, y)`` pixel + space of the full-resolution image. + descriptors: ``(batch_size, num_channels, height, width)`` feature + map at ``1/scale`` resolution. + scale: downsampling factor between the full image and the feature map. + + Returns: + descriptors: ``(batch_size, num_channels, num_keypoints)`` + """ batch_size, num_channels, height, width = descriptors.shape keypoints = keypoints - scale / 2 + 0.5 divisor = torch.tensor([[(width * scale - scale / 2 - 0.5), (height * scale - scale / 2 - 0.5)]]) divisor = divisor.to(keypoints) - keypoints /= divisor + keypoints = keypoints / divisor keypoints = keypoints * 2 - 1 # normalize to (-1, 1) kwargs = {"align_corners": True} - # [batch_size, num_channels, num_keypoints, 2] -> [batch_size, num_channels, num_keypoints, 2] + # (batch_size, num_keypoints, 2) -> (batch_size, 1, num_keypoints, 2) for grid_sample keypoints = keypoints.view(batch_size, 1, -1, 2) descriptors = nn.functional.grid_sample(descriptors, keypoints, mode="bilinear", **kwargs) - # [batch_size, descriptor_decoder_dim, num_channels, num_keypoints] -> [batch_size, descriptor_decoder_dim, num_keypoints] + # (batch_size, num_channels, 1, num_keypoints) -> (batch_size, num_channels, num_keypoints) descriptors = descriptors.reshape(batch_size, num_channels, -1) descriptors = nn.functional.normalize(descriptors, p=2, dim=1) return descriptors @@ -420,35 +469,15 @@ def forward( last_hidden_state = encoder_outputs[0] - list_keypoints_scores = [ - self.keypoint_decoder(last_hidden_state[None, ...]) for last_hidden_state in last_hidden_state - ] - - list_keypoints = [keypoints_scores[0] for keypoints_scores in list_keypoints_scores] - list_scores = [keypoints_scores[1] for keypoints_scores in list_keypoints_scores] - - list_descriptors = [ - self.descriptor_decoder(last_hidden_state[None, ...], keypoints[None, ...]) - for last_hidden_state, keypoints in zip(last_hidden_state, list_keypoints) - ] - - maximum_num_keypoints = max(keypoints.shape[0] for keypoints in list_keypoints) - - keypoints = torch.zeros((batch_size, maximum_num_keypoints, 2), device=pixel_values.device) - scores = torch.zeros((batch_size, maximum_num_keypoints), device=pixel_values.device) - descriptors = torch.zeros( - (batch_size, maximum_num_keypoints, self.config.descriptor_decoder_dim), - device=pixel_values.device, - ) - mask = torch.zeros((batch_size, maximum_num_keypoints), device=pixel_values.device, dtype=torch.int) + # Process the entire batch in one pass through the keypoint decoder. + # Returns statically-shaped tensors: (B, k, 2), (B, k), (B, k) — + # no Python loops over batch elements, no data-dependent shapes. + keypoints, scores, mask = self.keypoint_decoder(last_hidden_state) - for i, (_keypoints, _scores, _descriptors) in enumerate(zip(list_keypoints, list_scores, list_descriptors)): - keypoints[i, : _keypoints.shape[0]] = _keypoints - scores[i, : _scores.shape[0]] = _scores - descriptors[i, : _descriptors.shape[0]] = _descriptors - mask[i, : _scores.shape[0]] = 1 + # Compute descriptors for the full batch at once: (B, k, descriptor_dim). + descriptors = self.descriptor_decoder(last_hidden_state, keypoints) - # Convert to relative coordinates + # Convert keypoint pixel coordinates to relative (x, y) in [0, 1]. keypoints = keypoints / torch.tensor([width, height], device=keypoints.device) hidden_states = encoder_outputs[1] if output_hidden_states else None diff --git a/tests/models/superpoint/test_modeling_superpoint.py b/tests/models/superpoint/test_modeling_superpoint.py index ee20029d8e1d..e432cdff1bcd 100644 --- a/tests/models/superpoint/test_modeling_superpoint.py +++ b/tests/models/superpoint/test_modeling_superpoint.py @@ -48,7 +48,7 @@ def __init__( keypoint_decoder_dim: int = 65, descriptor_decoder_dim: int = 128, keypoint_threshold: float = 0.005, - max_keypoints: int = -1, + max_keypoints: int = 64, nms_radius: int = 4, border_removal_distance: int = 4, ): @@ -117,7 +117,6 @@ class SuperPointModelTest(ModelTesterMixin, unittest.TestCase): test_resize_embeddings = False has_attentions = False from_pretrained_id = "magic-leap-community/superpoint" - test_torch_exportable = False # data-dependent top-k / non-max suppression in keypoint detection def setUp(self): self.model_tester = SuperPointModelTester(self) From b3f4d31c15d1ea06bdd7ccb6a36d63f6776a5a1e Mon Sep 17 00:00:00 2001 From: Shreyas Date: Sat, 19 Sep 2026 15:59:27 +0530 Subject: [PATCH 2/2] Apply ruff styling fixes --- .../models/superpoint/modeling_superpoint.py | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/src/transformers/models/superpoint/modeling_superpoint.py b/src/transformers/models/superpoint/modeling_superpoint.py index 440ae61f04b9..38183d28625a 100644 --- a/src/transformers/models/superpoint/modeling_superpoint.py +++ b/src/transformers/models/superpoint/modeling_superpoint.py @@ -34,9 +34,6 @@ logger = logging.get_logger(__name__) - - - def simple_nms(scores: torch.Tensor, nms_radius: int) -> torch.Tensor: """Applies non-maximum suppression on scores.""" if nms_radius < 0: @@ -217,9 +214,7 @@ def _get_pixel_scores(self, encoded: torch.Tensor) -> torch.Tensor: scores = simple_nms(scores, self.nms_radius) return scores - def _extract_keypoints( - self, scores: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _extract_keypoints(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Extract keypoints from the score map using a statically-shaped pipeline compatible with ``torch.export``. @@ -270,14 +265,12 @@ def _extract_keypoints( k = self.max_keypoints if self.max_keypoints > 0 else height * width # torch.topk always returns exactly k entries → static output shape. - topk_scores, topk_indices = torch.topk( - masked_scores.reshape(batch_size, height * width), k=k, dim=1 - ) + topk_scores, topk_indices = torch.topk(masked_scores.reshape(batch_size, height * width), k=k, dim=1) # Convert flat indices to (x, y) pixel coordinates (pure arithmetic, # no data-dependent shapes). keypoints_y = topk_indices // width # (B, k) - keypoints_x = topk_indices % width # (B, k) + keypoints_x = topk_indices % width # (B, k) keypoints = torch.stack([keypoints_x, keypoints_y], dim=-1).to(scores.dtype) # (B, k, 2) # A position is a real keypoint if and only if its score is finite. @@ -336,9 +329,7 @@ def forward(self, encoded: torch.Tensor, keypoints: torch.Tensor) -> torch.Tenso return descriptors @staticmethod - def _sample_descriptors( - keypoints: torch.Tensor, descriptors: torch.Tensor, scale: int = 8 - ) -> torch.Tensor: + def _sample_descriptors(keypoints: torch.Tensor, descriptors: torch.Tensor, scale: int = 8) -> torch.Tensor: """ Interpolate descriptors at keypoint locations.