diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dd3e75ca..aa400ce21 100644 --- a/CHANGELOG.md +++ b/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. + - `SetCriterion.loss_masks` now samples matched ground-truth mask labels via direct tensor indexing instead of `point_sample`'s grid-sample machinery on CPU, guarded to the case where the matched masks are large enough (at least 1048576 total elements per image group, at least 2 matches per image group, and at least 16x the number of points being sampled in aggregate), contiguous with a shared spatial shape, and paired with float32 point coordinates plus 1-D int64 CPU-side matcher indices whose values are in range; CUDA and anything else outside that contract keeps using the previous `point_sample` path unchanged, and the sampled labels are bit-identical either way. The per-group match-count floor exists because a single large mask can clear the element floor on its own while carrying only one match — the fixed per-image loop overhead (slicing, a device transfer, a gather) then dominates regardless of resolution. These routing floors are conservative bounds, not calibrated crossover points. Measured on an i9-13900HX (single thread, 8 ground-truth masks x 13 repeated match groups, 312x312 masks, 380 sampled points, median of 100 calls, 7 repeated trials): the sampling step itself is 50.8x-58.4x faster and `loss_masks` end-to-end is 6.7x-7.1x faster, with 92.034% less CPU allocation traffic recorded by the profiler for the sampling step. - 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/_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. @@ -25,7 +27,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Fixed - Fixed TFLite export failure when `onnx2tf` could not resolve the installed `onnxsim` console script from a non-activated virtual environment. `onnx2tf` invokes the bare `onnxsim` name; if that lookup raises `FileNotFoundError` it logs `Failed to optimize the onnx file` (the same warning also appears in working runs). On a stock `RFDETRSmall()` export this then failed with `RuntimeError: onnx2tf conversion failed: Output tensors of a Functional model must be the output of a TensorFlow Layer`. RF-DETR now temporarily adds the running interpreter's script directory to `PATH` during conversion. ([#1365](https://github.com/roboflow/rf-detr/issues/1365)) + - The default (torchvision-native) training pipeline no longer silently corrupts keypoint annotations when `keypoint_flip_pairs` is empty on a schema that has genuine left/right pairs. `RandomHorizontalFlip` on this backend always mirrored keypoint x-coordinates when a flip was drawn, but only relabeled left/right joints `if self.keypoint_flip_pairs:` — with an empty list (the pydantic default, and one possible outcome when automatic flip-pair inference from dataset metadata doesn't match an asymmetric schema), affected training samples got their keypoints mirrored in position while keeping their original left/right label, with no warning. `_build_torchvision_pipeline` now drops the flip entirely for an empty-but-not-`None` `keypoint_flip_pairs`, logging the same warning the Albumentations backend already emits via `filter_keypoint_hflip_augmentations` (worded for this backend's lack of an editable `aug_config`), matching the annotation-safety behavior that backend has had since #1122. An empty list can also legitimately mean the schema has no left/right pairs at all (e.g. a single midline keypoint) — the unpatched flip was already harmless there since nothing needed relabeling; this fix disables it there too, for consistency with the Albumentations backend's existing contract, at the cost of a now-unavailable-by-default augmentation for that narrower case. Detection-only pipelines (`keypoint_flip_pairs=None`) and keypoint pipelines with real pairs are unaffected. + - `BestModelCallback` no longer treats PyTorch Lightning's pre-training sanity-check validation pass as a real epoch's result. Its EMA-checkpoint tracking and the `smooth_alpha` smoothing accumulator are custom bookkeeping that sit outside `ModelCheckpoint`'s own `trainer.sanity_checking` guard (the guard the regular-checkpoint path already inherits), so a positive sanity-check score — common when starting a new training run initialized with `pretrain_weights` from a checkpoint pretrained on a different dataset — could get written out as the permanent "best" `checkpoint_best_ema.pth` before a single real epoch ran, and real training could then never surpass it. Note this is distinct from PTL's own `resume`/`ckpt_path` restart, which PTL itself skips the sanity check for (`not val_loop.restarting`). ([#1348](https://github.com/roboflow/rf-detr/issues/1348)) ## [1.9.3] — 2026-08-17 diff --git a/src/rfdetr/models/postprocess.py b/src/rfdetr/models/postprocess.py index c100508a7..71506ad7c 100644 --- a/src/rfdetr/models/postprocess.py +++ b/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 diff --git a/src/rfdetr/training/callbacks/coco_eval.py b/src/rfdetr/training/callbacks/coco_eval.py index 1ae5b4483..a2b7f01cf 100644 --- a/src/rfdetr/training/callbacks/coco_eval.py +++ b/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"]} diff --git a/tests/models/test_postprocess.py b/tests/models/test_postprocess.py index 1163b3e76..2370916cf 100644 --- a/tests/models/test_postprocess.py +++ b/tests/models/test_postprocess.py @@ -257,6 +257,113 @@ 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 + + def test_upsampled_batch_pairs_non_square_target_sizes_with_their_mask_rows(self) -> None: + """Each image keeps its own non-square target size and mask geometry after batched size conversion. + + This prevents a regression where ``target_sizes.tolist()`` is batched but the resulting rows are swapped or + height/width are transposed before per-image interpolation. + """ + out_masks = torch.tensor( + [ + [[[5.0, -5.0, -5.0], [5.0, -5.0, -5.0]]], + [[[-5.0, -5.0, 5.0], [-5.0, -5.0, 5.0]]], + ] + ) + scores = torch.tensor([[0.9], [0.8]]) + labels = torch.tensor([[1], [2]]) + boxes = torch.zeros(2, 1, 4) + topk_boxes = torch.zeros(2, 1, dtype=torch.long) + target_sizes = torch.tensor([[4, 9], [9, 4]]) # (H, W): landscape, then portrait + + results = PostProcess._postprocess_masks(out_masks, scores, labels, boxes, topk_boxes, target_sizes) + + expected_masks = [ + torch.nn.functional.interpolate( + out_masks[0, 0][None, None], size=(4, 9), mode="bilinear", align_corners=False + ) + > 0.0, + torch.nn.functional.interpolate( + out_masks[1, 0][None, None], size=(9, 4), mode="bilinear", align_corners=False + ) + > 0.0, + ] + assert results[0]["masks"].shape == (1, 1, 4, 9) + assert results[1]["masks"].shape == (1, 1, 9, 4) + assert torch.equal(results[0]["masks"], expected_masks[0]) + assert torch.equal(results[1]["masks"], expected_masks[1]) + + @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. diff --git a/tests/training/callbacks/test_coco_eval_callback.py b/tests/training/callbacks/test_coco_eval_callback.py index 4610cadbe..407416a2d 100644 --- a/tests/training/callbacks/test_coco_eval_callback.py +++ b/tests/training/callbacks/test_coco_eval_callback.py @@ -1958,6 +1958,95 @@ 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 + + def test_batch_pairs_non_square_orig_sizes_with_their_target_rows(self) -> None: + """Each target's box uses its own non-square orig_size after batched size conversion. + + This prevents swapped ``orig_size.tolist()`` rows and height/width transposition from assigning a sibling + image's scale to the wrong target. + """ + cb = COCOEvalCallback() + targets = [ + { + "boxes": torch.tensor([[0.25, 0.5, 0.2, 0.4]]), + "labels": torch.tensor([3]), + "orig_size": torch.tensor([100, 300]), # H=100, W=300 + }, + { + "boxes": torch.tensor([[0.5, 0.25, 0.4, 0.2]]), + "labels": torch.tensor([7]), + "orig_size": torch.tensor([240, 80]), # H=240, W=80 + }, + ] + + out = cb._convert_targets(targets) + + assert out[0]["boxes"].shape == (1, 4) + assert out[1]["boxes"].shape == (1, 4) + assert out[0]["labels"].item() == 3 + assert out[1]["labels"].item() == 7 + torch.testing.assert_close(out[0]["boxes"], torch.tensor([[45.0, 30.0, 105.0, 70.0]])) + torch.testing.assert_close(out[1]["boxes"], torch.tensor([[24.0, 36.0, 56.0, 84.0]])) + + @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]),