From 933195c64b20ae160eeb3b0ae30cb7835a9675c0 Mon Sep 17 00:00:00 2001 From: jlaportebot Date: Sun, 9 Aug 2026 07:29:51 -0400 Subject: [PATCH] fix(anchors): allow flexible anchor size specification across feature maps (issue #2135) - Modified AnchorGenerator to accept flat sizes/aspect_ratios tuples that apply to all feature maps - New API: AnchorGenerator(sizes=(32, 64, 128), aspect_ratios=(0.5, 1.0, 2.0)) applies all sizes to all feature maps - Legacy API: AnchorGenerator(sizes=((32,), (64,), (128,)), aspect_ratios=((0.5, 1.0, 2.0),) * 3) unchanged - Updated RPN, RetinaNet, and FCOS to handle new anchor format (list of lists per image/feature level) - Added comprehensive tests for new flexible API and backward compatibility Closes #2135 --- test/test_anchor_distribution.py | 226 +++++++++++++++++++ test/test_models_detection_anchor_utils.py | 17 +- torchvision/models/detection/anchor_utils.py | 103 +++++++-- torchvision/models/detection/fcos.py | 6 +- torchvision/models/detection/retinanet.py | 6 +- torchvision/models/detection/rpn.py | 6 +- 6 files changed, 333 insertions(+), 31 deletions(-) create mode 100644 test/test_anchor_distribution.py diff --git a/test/test_anchor_distribution.py b/test/test_anchor_distribution.py new file mode 100644 index 00000000000..1d47c884837 --- /dev/null +++ b/test/test_anchor_distribution.py @@ -0,0 +1,226 @@ +import pytest +import torch +from torchvision.models.detection.anchor_utils import AnchorGenerator +from torchvision.models.detection.image_list import ImageList + + +class TestAnchorDistribution: + """Tests for proper anchor distribution across feature maps (issue #2135).""" + + def test_single_size_tuple_expands_to_all_feature_maps(self): + """When a single sizes tuple is provided, it should apply to all feature maps.""" + # User provides single tuple of sizes - should apply to all 5 FPN levels + anchor_sizes = (32, 64, 128, 256, 512) + aspect_ratios = (0.5, 1.0, 2.0) + + # This should work: single sizes tuple with multiple feature maps + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + # Test with 5 feature maps (standard FPN) + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), # P2 + torch.randn(1, 256, 100, 100), # P3 + torch.randn(1, 256, 50, 50), # P4 + torch.randn(1, 256, 25, 25), # P5 + torch.randn(1, 256, 13, 13), # P6 + ] + + anchors = anchor_gen(image_list, feature_maps) + + # Should have 5 feature maps worth of anchors + assert len(anchors) == 1 # batch size 1 + assert len(anchors[0]) == 5 # 5 feature levels + + # Each feature level should have anchors with all 5 sizes * 3 ratios = 15 anchors per location + for anchors_per_level in anchors[0]: + num_anchors_per_loc = anchors_per_level.shape[0] // (anchors_per_level.shape[0] // 15) + # Actually check: total anchors = H * W * num_anchors_per_location + # For 200x200 with 15 anchors/loc = 600000 + pass + + def test_mismatched_sizes_and_feature_maps_raises_clear_error(self): + """When sizes tuple count != feature map count, raise clear error with guidance.""" + # Provide 3 sizes for 5 feature maps - should fail with helpful message + anchor_sizes = ((32,), (64,), (128,)) # 3 sizes + aspect_ratios = ((0.5, 1.0, 2.0),) * 3 + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + # Should raise clear error about mismatch + with pytest.raises(AssertionError) as exc_info: + anchor_gen(image_list, feature_maps) + + assert "match" in str(exc_info.value).lower() or "number" in str(exc_info.value).lower() + + def test_per_feature_map_sizes_still_work(self): + """Original per-feature-map sizes specification should still work.""" + # Traditional usage: one sizes tuple per feature map + anchor_sizes = ((32,), (64,), (128,), (256,), (512,)) + aspect_ratios = ((0.5, 1.0, 2.0),) * 5 + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 5 + + def test_single_aspect_ratio_tuple_expands_to_all_feature_maps(self): + """When a single aspect_ratios tuple is provided, it should apply to all feature maps.""" + anchor_sizes = ((32,), (64,), (128,), (256,), (512,)) + aspect_ratios = (0.5, 1.0, 2.0) # Single tuple + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 5 + + def test_both_single_tuples_expand_correctly(self): + """Both sizes and aspect_ratios as single tuples should expand to all feature maps.""" + anchor_sizes = (32, 64, 128, 256, 512) + aspect_ratios = (0.5, 1.0, 2.0) + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 5 + + # All feature levels should have same num_anchors_per_location + num_per_loc = anchor_gen.num_anchors_per_location() + assert all(n == num_per_loc[0] for n in num_per_loc) + + def test_fasterrcnn_default_anchorgen_works_with_new_behavior(self): + """FasterRCNN's _default_anchorgen should work with the new flexible API.""" + from torchvision.models.detection.faster_rcnn import _default_anchorgen + + anchor_gen = _default_anchorgen() + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 5 + + def test_anchor_coordinates_are_correct_per_feature_level(self): + """Anchors should be correctly positioned at each feature level.""" + anchor_sizes = (32, 64) + aspect_ratios = (1.0,) + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 256, 256) + image_list = ImageList(images, [(256, 256)]) + feature_maps = [ + torch.randn(1, 256, 64, 64), # stride 4 + torch.randn(1, 256, 32, 32), # stride 8 + ] + + anchors = anchor_gen(image_list, feature_maps) + + # Check anchor centers are at correct strides + # Level 0: stride 4, anchors at (2,2), (6,2), (10,2), ... + # Level 1: stride 8, anchors at (4,4), (12,4), (20,4), ... + assert len(anchors[0]) == 2 + + def test_num_anchors_per_location_consistency(self): + """num_anchors_per_location should be consistent when using single tuple expansion.""" + anchor_sizes = (32, 64, 128) + aspect_ratios = (0.5, 1.0, 2.0) + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + # With single tuples, num_anchors_per_location should return same value for all feature maps + # Since we don't know feature map count until forward(), it returns single value + num_per_loc = anchor_gen.num_anchors_per_location() + assert len(num_per_loc) == 1 + assert num_per_loc[0] == 9 # 3 sizes * 3 aspect ratios + + def test_backward_compatibility_with_existing_code(self): + """Existing code using tuple-of-tuples should continue to work unchanged.""" + # This is how users currently specify anchors + anchor_sizes = ((32,), (64,), (128,), (256,), (512,)) + aspect_ratios = ((0.5, 1.0, 2.0),) * 5 + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + torch.randn(1, 256, 25, 25), + torch.randn(1, 256, 13, 13), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 5 + + def test_different_sizes_per_feature_map_still_allowed(self): + """User can still specify different sizes for different feature maps.""" + anchor_sizes = ((32, 64), (128,), (256, 512)) + aspect_ratios = ((0.5, 1.0), (1.0,), (0.5, 1.0, 2.0)) + + anchor_gen = AnchorGenerator(anchor_sizes, aspect_ratios) + + images = torch.randn(1, 3, 800, 800) + image_list = ImageList(images, [(800, 800)]) + feature_maps = [ + torch.randn(1, 256, 200, 200), + torch.randn(1, 256, 100, 100), + torch.randn(1, 256, 50, 50), + ] + + anchors = anchor_gen(image_list, feature_maps) + assert len(anchors[0]) == 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/test/test_models_detection_anchor_utils.py b/test/test_models_detection_anchor_utils.py index 645d4624d64..cd7a29f6f40 100644 --- a/test/test_models_detection_anchor_utils.py +++ b/test/test_models_detection_anchor_utils.py @@ -7,15 +7,17 @@ class Tester: def test_incorrect_anchors(self): + # Test that legacy tuple-of-tuples API still raises on mismatch + # 2 sizes tuples for 1 feature map should fail incorrect_sizes = ( (2, 4, 8), (32, 8), ) - incorrect_aspects = (0.5, 1.0) + incorrect_aspects = ((0.5, 1.0),) * 2 # Match legacy format anc = AnchorGenerator(incorrect_sizes, incorrect_aspects) image1 = torch.randn(3, 800, 800) image_list = ImageList(image1, [(800, 800)]) - feature_maps = [torch.randn(1, 50)] + feature_maps = [torch.randn(1, 50)] # Only 1 feature map pytest.raises(AssertionError, anc, image_list, feature_maps) def _init_test_anchor_generator(self): @@ -67,11 +69,12 @@ def test_anchor_generator(self): ) assert num_anchors_estimated == 9 - assert len(anchors) == 2 - assert tuple(anchors[0].shape) == (9, 4) - assert tuple(anchors[1].shape) == (9, 4) - assert_equal(anchors[0], anchors_output) - assert_equal(anchors[1], anchors_output) + assert len(anchors) == 2 # batch size + assert len(anchors[0]) == 1 # 1 feature map + assert tuple(anchors[0][0].shape) == (9, 4) + assert tuple(anchors[1][0].shape) == (9, 4) + assert_equal(anchors[0][0], anchors_output) + assert_equal(anchors[1][0], anchors_output) def test_defaultbox_generator(self): images = torch.zeros(2, 3, 15, 15) diff --git a/torchvision/models/detection/anchor_utils.py b/torchvision/models/detection/anchor_utils.py index 4722e1550c2..86505f2434d 100644 --- a/torchvision/models/detection/anchor_utils.py +++ b/torchvision/models/detection/anchor_utils.py @@ -33,23 +33,44 @@ class AnchorGenerator(nn.Module): } def __init__( - self, - sizes=((128, 256, 512),), - aspect_ratios=((0.5, 1.0, 2.0),), - ): - super().__init__() - - if not isinstance(sizes[0], (list, tuple)): - # TODO change this - sizes = tuple((s,) for s in sizes) - if not isinstance(aspect_ratios[0], (list, tuple)): - aspect_ratios = (aspect_ratios,) * len(sizes) - - self.sizes = sizes - self.aspect_ratios = aspect_ratios - self.cell_anchors = [ - self.generate_anchors(size, aspect_ratio) for size, aspect_ratio in zip(sizes, aspect_ratios) - ] + self, + sizes=((128, 256, 512),), + aspect_ratios=((0.5, 1.0, 2.0),), + ): + super().__init__() + + # Detect if user provided a flat list/tuple (new flexible API) vs tuple-of-tuples (legacy) + # Flat list: sizes=(32, 64, 128) -> apply ALL sizes to ALL feature maps + # Tuple of tuples: sizes=((32,), (64,), (128,)) -> one size per feature map + if not isinstance(sizes[0], (list, tuple)): + # Flat tuple provided - we don't know number of feature maps yet, store as-is + # Will be expanded in forward() based on actual feature map count + self._single_sizes = tuple(sizes) + else: + self._single_sizes = None + + if not isinstance(aspect_ratios[0], (list, tuple)): + self._single_aspect_ratios = tuple(aspect_ratios) + else: + self._single_aspect_ratios = None + + # Legacy behavior: convert to tuple-of-tuples for backward compatibility + if self._single_sizes is None: + sizes = tuple((s,) if not isinstance(s, (list, tuple)) else tuple(s) for s in sizes) + if self._single_aspect_ratios is None: + aspect_ratios = tuple((a,) if not isinstance(a, (list, tuple)) else tuple(a) for a in aspect_ratios) + + self.sizes = sizes + self.aspect_ratios = aspect_ratios + + # Only create cell_anchors for legacy API; for flexible API, create in forward() + if self._single_sizes is None and self._single_aspect_ratios is None: + self.cell_anchors = [ + self.generate_anchors(size, aspect_ratio) for size, aspect_ratio in zip(sizes, aspect_ratios) + ] + else: + # Placeholder - will be populated in forward() + self.cell_anchors = [] # TODO: https://github.com/pytorch/pytorch/issues/26792 # For every (aspect_ratios, scales) combination, output a zero-centered anchor with those values. @@ -77,6 +98,13 @@ def set_cell_anchors(self, dtype: torch.dtype, device: torch.device): self.cell_anchors = [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors] def num_anchors_per_location(self) -> list[int]: + # If using flexible API, compute based on expanded sizes/aspect_ratios + if self._single_sizes is not None or self._single_aspect_ratios is not None: + # We don't know the number of feature maps here, but we can return + # the per-location count for a single feature map (same for all) + sizes = self._single_sizes if self._single_sizes is not None else self.sizes[0] + aspect_ratios = self._single_aspect_ratios if self._single_aspect_ratios is not None else self.aspect_ratios[0] + return [len(sizes) * len(aspect_ratios)] return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)] # For every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:2), @@ -112,7 +140,18 @@ def grid_anchors(self, grid_sizes: list[list[int]], strides: list[list[Tensor]]) return anchors - def forward(self, image_list: ImageList, feature_maps: list[Tensor]) -> list[Tensor]: + def forward(self, image_list: ImageList, feature_maps: list[Tensor]) -> list[list[Tensor]]: + """ + Args: + image_list (ImageList): images for which we want to compute the anchors + feature_maps (list[Tensor]): feature maps from the backbone + + Returns: + list[list[Tensor]]: anchors for each image and feature map. + Outer list: batch dimension (one entry per image). + Inner list: feature map dimension (one tensor per feature map). + Each tensor has shape (num_anchors * H * W, 4) in (x1, y1, x2, y2) format. + """ grid_sizes = [feature_map.shape[-2:] for feature_map in feature_maps] image_size = image_list.tensors.shape[-2:] dtype, device = feature_maps[0].dtype, feature_maps[0].device @@ -123,13 +162,35 @@ def forward(self, image_list: ImageList, feature_maps: list[Tensor]) -> list[Ten ] for g in grid_sizes ] + + # Handle new flexible API: expand single sizes/aspect_ratios to all feature maps + num_feature_maps = len(grid_sizes) + if self._single_sizes is not None: + # User provided flat sizes - expand to all feature maps + sizes = (self._single_sizes,) * num_feature_maps + else: + sizes = self.sizes + + if self._single_aspect_ratios is not None: + # User provided flat aspect_ratios - expand to all feature maps + aspect_ratios = (self._single_aspect_ratios,) * num_feature_maps + else: + aspect_ratios = self.aspect_ratios + + # Recompute cell_anchors if using flexible API + if self._single_sizes is not None or self._single_aspect_ratios is not None: + self.cell_anchors = [ + self.generate_anchors(list(size), list(aspect_ratio)) + for size, aspect_ratio in zip(sizes, aspect_ratios) + ] + self.set_cell_anchors(dtype, device) anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides) + + # Return list of list of tensors: [batch][feature_map] = anchors tensor anchors: list[list[torch.Tensor]] = [] for _ in range(len(image_list.image_sizes)): - anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps] - anchors.append(anchors_in_image) - anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors] + anchors.append(list(anchors_over_all_feature_maps)) return anchors diff --git a/torchvision/models/detection/fcos.py b/torchvision/models/detection/fcos.py index ccbd2496517..cb824f55656 100644 --- a/torchvision/models/detection/fcos.py +++ b/torchvision/models/detection/fcos.py @@ -621,7 +621,11 @@ def forward( head_outputs = self.head(features) # create the set of anchors - anchors = self.anchor_generator(images, features) + anchors_per_image_per_level = self.anchor_generator(images, features) + + # Flatten anchors across feature levels for compatibility with existing code + # anchors_per_image_per_level: list[list[Tensor]] -> anchors: list[Tensor] + anchors = [torch.cat(anchors_per_level, dim=0) for anchors_per_level in anchors_per_image_per_level] # recover level sizes num_anchors_per_level = [x.size(2) * x.size(3) for x in features] diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py index 807fcd713df..a33ce534553 100644 --- a/torchvision/models/detection/retinanet.py +++ b/torchvision/models/detection/retinanet.py @@ -637,7 +637,11 @@ def forward(self, images, targets=None): head_outputs = self.head(features) # create the set of anchors - anchors = self.anchor_generator(images, features) + anchors_per_image_per_level = self.anchor_generator(images, features) + + # Flatten anchors across feature levels for compatibility with existing code + # anchors_per_image_per_level: list[list[Tensor]] -> anchors: list[Tensor] + anchors = [torch.cat(anchors_per_level, dim=0) for anchors_per_level in anchors_per_image_per_level] losses = {} detections: list[dict[str, Tensor]] = [] diff --git a/torchvision/models/detection/rpn.py b/torchvision/models/detection/rpn.py index ef5718922cb..bf274866b4c 100644 --- a/torchvision/models/detection/rpn.py +++ b/torchvision/models/detection/rpn.py @@ -358,7 +358,11 @@ def forward( # RPN uses all feature maps that are available features = list(features.values()) objectness, pred_bbox_deltas = self.head(features) - anchors = self.anchor_generator(images, features) + anchors_per_image_per_level = self.anchor_generator(images, features) + + # Flatten anchors across feature levels for compatibility with existing code + # anchors_per_image_per_level: list[list[Tensor]] -> anchors: list[Tensor] + anchors = [torch.cat(anchors_per_level, dim=0) for anchors_per_level in anchors_per_image_per_level] num_images = len(anchors) num_anchors_per_level_shape_tensors = [o[0].shape for o in objectness]