SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[superglue] Fixed the way batch mask was applied to the scores before match assignment computation (#39968)

Browse Source 
fix: mask filling to score was wrong
This commit is contained in:
StevenBucaille
2025-08-07 04:49:39 -04:00
committed by GitHub
parent 2593932f10
commit cdeaad96b7
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/transformers/models/superglue/modeling_superglue.py
View File

@@ -676,8 +676,10 @@ class SuperGlueForKeypointMatching(SuperGluePreTrainedModel):
if mask is not None: if mask is not None:
mask = mask.reshape(batch_size, 2, num_keypoints) mask = mask.reshape(batch_size, 2, num_keypoints)
mask0 = mask[:, 0].unsqueeze(-1).expand(-1, -1, num_keypoints) mask0 = mask[:, 0].unsqueeze(2)
scores = scores.masked_fill(mask0 == 0, -1e9) mask1 = mask[:, 1].unsqueeze(1)
mask = torch.logical_and(mask0, mask1)
scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
# Run the optimal transport. # Run the optimal transport.
scores = log_optimal_transport(scores, self.bin_score, iterations=self.config.sinkhorn_iterations) scores = log_optimal_transport(scores, self.bin_score, iterations=self.config.sinkhorn_iterations)

2
tests/models/superglue/test_modeling_superglue.py
View File

@@ -423,3 +423,5 @@ class SuperGlueModelIntegrationTest(unittest.TestCase):
torch.sum(~torch.isclose(predicted_matching_scores_values, expected_matching_scores_values, atol=1e-2)) < 4 torch.sum(~torch.isclose(predicted_matching_scores_values, expected_matching_scores_values, atol=1e-2)) < 4
) )
self.assertTrue(torch.sum(predicted_matches_values != expected_matches_values) < 4) self.assertTrue(torch.sum(predicted_matches_values != expected_matches_values) < 4)
self.assertTrue(torch.all(outputs.matches[0, 1] < torch.sum(outputs.mask[0, 0])))
self.assertTrue(torch.all(outputs.matches[0, 0] < torch.sum(outputs.mask[0, 1])))