perf(postprocess): read every image's mask target size once per batch

CHANGELOG.md

@@ -18,6 +18,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed
 
+- Segmentation postprocessing now reads every image's mask resize target before the per-image loop, replacing one CUDA device-to-host synchronization per image with one per batch while preserving mask outputs exactly. `COCOEvalCallback._convert_targets` (train/val mAP accumulation) had the same per-target `orig_size.tolist()` pattern and gets the same fix.
+
 - Training now skips PyTorch Lightning's pre-training sanity validation batches by default; set `TrainConfig(num_sanity_val_steps=N)` to restore it. Training loss component metrics are emitted once per epoch, with decoder and encoder auxiliary terms compacted into `train/<term>_aux` — 17 to 9 tracked metric keys per microbatch on the default `RFDETRSmall` config; `train/loss` preserves its `train_log_on_step` behavior, but component-level metrics no longer honor `train_log_on_step` by default — set `TrainConfig(compact_train_metrics=False)` to restore per-layer keys and `train_log_on_step` honoring for them. Learning-rate metrics now emit only for optimizer updates, including a final partial gradient-accumulation window, cutting LR log calls by about 75% at the default accumulation of four. These are call-count reductions in Lightning's metric bookkeeping, not a measured timing or memory improvement. Update dashboard queries that consume layer-specific auxiliary keys or assume one learning-rate point per microbatch.
 
 ### Fixed

src/rfdetr/models/postprocess.py

@@ -211,6 +211,9 @@ def _postprocess_masks(
             ``upsample_masks_to_image_size``.
         """
         results = []
+        # Read every resize target in one device-to-host synchronization. Calling
+        # target_sizes[i].tolist() inside the loop forces a separate CUDA sync per image.
+        target_sizes_list = target_sizes.tolist() if upsample_masks_to_image_size and out_masks.shape[0] else []
         for i in range(out_masks.shape[0]):
             scores_i, labels_i, boxes_i, k_idx = scores[i], labels[i], boxes[i], topk_boxes[i]
             if score_threshold is not None:
@@ -228,7 +231,7 @@ def _postprocess_masks(
                 res_i["masks"] = (masks_i > 0.0).unsqueeze(1)  # [K,1,Hm,Wm] bool, native resolution
                 results.append(res_i)
                 continue
-            h, w = target_sizes[i].tolist()
+            h, w = target_sizes_list[i]
             # Upsample in chunks and threshold *inside* the comprehension so only one float32 chunk
             # is live at a time; the accumulated list holds bool tensors (1 byte/pixel vs 4 for float32).
             # At K=300, 1080p this reduces peak memory from ~5 GB to ~1.5 GB vs a single F.interpolate

src/rfdetr/training/callbacks/coco_eval.py

@@ -1371,8 +1371,11 @@ def _convert_targets(
                 f"preds and targets must be positionally paired 1:1; got {len(preds)} preds vs {len(targets)} targets"
             )
         out = []
+        # Stack every target's orig_size into one device-to-host synchronization instead of
+        # one per target inside the loop (same fix as PostProcess._postprocess_masks).
+        orig_sizes_list = torch.stack([t["orig_size"] for t in targets]).tolist() if targets else []
         for index, t in enumerate(targets):
-            h, w = t["orig_size"].tolist()
+            h, w = orig_sizes_list[index]
             scale = t["boxes"].new_tensor([w, h, w, h])
             boxes = box_cxcywh_to_xyxy(t["boxes"]) * scale
             entry: dict[str, Tensor] = {"boxes": boxes, "labels": t["labels"]}

tests/models/test_postprocess.py

@@ -257,6 +257,78 @@ def test_native_resolution_thresholds_at_zero_same_as_upsampled_path(self):
         expected = torch.tensor([[True, False], [False, True]])
         assert torch.equal(results[0]["masks"].squeeze(1)[0], expected)
 
+    @pytest.mark.parametrize(
+        ("batch", "upsample", "expected_calls"),
+        [
+            pytest.param(4, True, [(4, 2)], id="upsampled-batch"),
+            pytest.param(4, False, [], id="native-resolution"),
+            pytest.param(0, True, [], id="empty-upsampled-batch"),
+        ],
+    )
+    def test_target_sizes_are_read_once_per_upsampled_batch(
+        self, batch: int, upsample: bool, expected_calls: list[tuple[int, ...]]
+    ) -> None:
+        """A non-empty upsampled batch reads all target sizes together; other paths do not read them."""
+        out_masks = torch.randn(batch, 1, 2, 2)
+        scores = torch.ones(batch, 1)
+        labels = torch.zeros(batch, 1, dtype=torch.long)
+        boxes = torch.zeros(batch, 1, 4)
+        topk_boxes = torch.zeros(batch, 1, dtype=torch.long)
+        target_sizes = torch.full((batch, 2), 8, dtype=torch.long)
+        calls: list[tuple[int, ...]] = []
+        original_tolist = torch.Tensor.tolist
+
+        def tracked_tolist(tensor: torch.Tensor) -> object:
+            """Record the tensor shape passed to ``tolist`` before delegating to PyTorch."""
+            calls.append(tuple(tensor.shape))
+            return original_tolist(tensor)
+
+        with patch.object(torch.Tensor, "tolist", tracked_tolist):
+            PostProcess._postprocess_masks(
+                out_masks,
+                scores,
+                labels,
+                boxes,
+                topk_boxes,
+                target_sizes,
+                upsample_masks_to_image_size=upsample,
+            )
+
+        assert calls == expected_calls
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_target_sizes_are_read_once_per_upsampled_batch_cuda(self) -> None:
+        """Same call-count guarantee as ``test_target_sizes_are_read_once_per_upsampled_batch`` above, with every input
+        tensor on CUDA — the actual site of the device-to-host synchronization this PR collapses from one per image to
+        one per batch."""
+        out_masks = torch.randn(4, 1, 2, 2, device="cuda")
+        scores = torch.ones(4, 1, device="cuda")
+        labels = torch.zeros(4, 1, dtype=torch.long, device="cuda")
+        boxes = torch.zeros(4, 1, 4, device="cuda")
+        topk_boxes = torch.zeros(4, 1, dtype=torch.long, device="cuda")
+        target_sizes = torch.full((4, 2), 8, dtype=torch.long, device="cuda")
+        calls: list[tuple[int, ...]] = []
+        original_tolist = torch.Tensor.tolist
+
+        def tracked_tolist(tensor: torch.Tensor) -> object:
+            """Record the tensor shape passed to ``tolist`` before delegating to PyTorch."""
+            calls.append(tuple(tensor.shape))
+            return original_tolist(tensor)
+
+        with patch.object(torch.Tensor, "tolist", tracked_tolist):
+            PostProcess._postprocess_masks(
+                out_masks,
+                scores,
+                labels,
+                boxes,
+                topk_boxes,
+                target_sizes,
+                upsample_masks_to_image_size=True,
+            )
+
+        assert calls == [(4, 2)]
+
     def test_duplicate_query_selection_repeats_the_same_mask_rows(self):
         """Top-k can pick the same query under two classes; each pick must yield that query's exact mask.
 

tests/training/callbacks/test_coco_eval_callback.py

@@ -1958,6 +1958,66 @@ def test_no_masks_no_iscrowd_keys_absent(self) -> None:
         out = cb._convert_targets(targets)
         assert set(out[0].keys()) == {"boxes", "labels"}
 
+    @pytest.mark.parametrize(
+        ("num_targets", "expected_calls"),
+        [
+            pytest.param(3, [(3, 2)], id="multi-target-batch"),
+            pytest.param(0, [], id="empty-targets"),
+        ],
+    )
+    def test_orig_size_is_read_once_per_batch(self, num_targets: int, expected_calls: list[tuple[int, ...]]) -> None:
+        """A non-empty target batch reads every orig_size together in one device-to-host synchronization instead of one
+        ``tolist()`` call per target inside the loop."""
+        cb = COCOEvalCallback()
+        targets = [
+            {
+                "boxes": torch.zeros(1, 4),
+                "labels": torch.tensor([0]),
+                "orig_size": torch.tensor([10, 10]),
+            }
+            for _ in range(num_targets)
+        ]
+        calls: list[tuple[int, ...]] = []
+        original_tolist = torch.Tensor.tolist
+
+        def tracked_tolist(tensor: torch.Tensor) -> object:
+            """Record the tensor shape passed to ``tolist`` before delegating to PyTorch."""
+            calls.append(tuple(tensor.shape))
+            return original_tolist(tensor)
+
+        with patch.object(torch.Tensor, "tolist", tracked_tolist):
+            cb._convert_targets(targets)
+
+        assert calls == expected_calls
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_orig_size_is_read_once_per_batch_cuda(self) -> None:
+        """Same call-count guarantee as ``test_orig_size_is_read_once_per_batch`` above, with every target tensor on
+        CUDA — the actual site of the device-to-host synchronization this PR collapses from one per target to one per
+        batch."""
+        cb = COCOEvalCallback()
+        targets = [
+            {
+                "boxes": torch.zeros(1, 4, device="cuda"),
+                "labels": torch.tensor([0], device="cuda"),
+                "orig_size": torch.tensor([10, 10], device="cuda"),
+            }
+            for _ in range(3)
+        ]
+        calls: list[tuple[int, ...]] = []
+        original_tolist = torch.Tensor.tolist
+
+        def tracked_tolist(tensor: torch.Tensor) -> object:
+            """Record the tensor shape passed to ``tolist`` before delegating to PyTorch."""
+            calls.append(tuple(tensor.shape))
+            return original_tolist(tensor)
+
+        with patch.object(torch.Tensor, "tolist", tracked_tolist):
+            cb._convert_targets(targets)
+
+        assert calls == [(3, 2)]
+
 
 class TestConvertTargetsWithPreds:
     """_convert_targets(targets, preds) resizes each target's masks to its own paired prediction's grid (preds[index]),