From 470d6e9a9bb80793a962c3adacf6650501d73bea Mon Sep 17 00:00:00 2001 From: Henrik Vendelbo Date: Thu, 7 May 2026 17:44:49 +0200 Subject: [PATCH] feat: max_eval_orig_size to restrain eval images loaded to a maximum size on the longest side This helps speed up validation epochs --- src/rfdetr/config.py | 1 + src/rfdetr/training/callbacks/coco_eval.py | 42 +++++- src/rfdetr/training/module_model.py | 8 ++ src/rfdetr/training/trainer.py | 1 + tests/models/test_config.py | 10 ++ tests/training/test_coco_eval_callback.py | 149 +++++++++++++++++++++ 6 files changed, 207 insertions(+), 4 deletions(-) diff --git a/src/rfdetr/config.py b/src/rfdetr/config.py index 1ba175d52..a15fba62e 100644 --- a/src/rfdetr/config.py +++ b/src/rfdetr/config.py @@ -503,6 +503,7 @@ class TrainConfig(BaseModel): eval_max_dets: int = 500 eval_interval: int = 1 log_per_class_metrics: bool = True + max_eval_orig_size: Optional[int] = None aug_config: Optional[Dict[str, Any]] = None augmentation_backend: Literal["cpu", "auto", "gpu"] = "cpu" save_dataset_grids: bool = False diff --git a/src/rfdetr/training/callbacks/coco_eval.py b/src/rfdetr/training/callbacks/coco_eval.py index b1aac64ba..581221126 100644 --- a/src/rfdetr/training/callbacks/coco_eval.py +++ b/src/rfdetr/training/callbacks/coco_eval.py @@ -48,6 +48,12 @@ class COCOEvalCallback(Callback): eval_interval: Run validation metrics every N epochs. Test metrics are always computed when ``trainer.test()`` is called. log_per_class_metrics: When ``False``, skip per-class AP logging/table. + max_eval_orig_size: Cap the longer side of each image's original size to + this value (in pixels) before upsampling masks for COCO evaluation. + Reduces mask buffer memory proportionally to ``(cap/orig)²`` with + negligible impact on mAP since the model input is already much + smaller (e.g. SegNano trains at 312 px). ``None`` disables the cap + and evaluates at full original resolution. Defaults to ``None``. """ def __init__( @@ -57,12 +63,14 @@ def __init__( eval_interval: int = 1, log_per_class_metrics: bool = True, in_notebook: bool | None = None, + max_eval_orig_size: int | None = None, ) -> None: super().__init__() self._max_dets = max_dets self._segmentation = segmentation self._eval_interval = max(1, int(eval_interval)) self._log_per_class_metrics = bool(log_per_class_metrics) + self._max_eval_orig_size = max_eval_orig_size self._class_names: list[str] = [] self._cat_id_to_name: dict[int, str] = {} self._f1_local: dict[int, dict[str, Any]] = init_matching_accumulator() @@ -183,6 +191,12 @@ def on_validation_batch_end( ).to(pl_module.device) samples, _ = batch orig_sizes = torch.stack([t["orig_size"] for t in outputs["targets"]]).to(pl_module.device) + if self._max_eval_orig_size is not None: + # Cap each (H, W) so the longer side ≤ max_eval_orig_size. + # Masks are upsampled to orig_size inside postprocess; capping here + # keeps mask buffers small without affecting training. + scale = (self._max_eval_orig_size / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0) + orig_sizes = (orig_sizes.float() * scale).long() ema_underlying = ema_cb._average_model.module.model with torch.no_grad(): ema_underlying.eval() # AveragedModel deepcopy is not managed by PTL @@ -708,8 +722,24 @@ def _convert_preds(self, preds: list[dict[str, torch.Tensor]]) -> list[dict[str, out = [] for p in preds: entry = dict(p) - if "masks" in entry and entry["masks"].ndim == 4 and entry["masks"].shape[1] == 1: - entry["masks"] = entry["masks"].squeeze(1) + if "masks" in entry: + masks = entry["masks"] + if masks.ndim == 4 and masks.shape[1] == 1: + masks = masks.squeeze(1) + if self._max_eval_orig_size is not None: + h, w = masks.shape[-2:] + if max(h, w) > self._max_eval_orig_size: + scale = self._max_eval_orig_size / max(h, w) + new_h = max(1, int(h * scale)) + new_w = max(1, int(w * scale)) + masks = F.interpolate( + masks.float().unsqueeze(1), + size=(new_h, new_w), + mode="nearest", + ).squeeze(1) + if "boxes" in entry: + entry["boxes"] = entry["boxes"] * scale + entry["masks"] = masks out.append(entry) return out @@ -729,13 +759,17 @@ def _convert_targets(self, targets: list[dict[str, torch.Tensor]]) -> list[dict[ out = [] for t in targets: h, w = t["orig_size"].tolist() + if self._max_eval_orig_size is not None and max(h, w) > self._max_eval_orig_size: + cap_scale = self._max_eval_orig_size / max(h, w) + h = max(1, int(h * cap_scale)) + w = max(1, int(w * cap_scale)) scale = t["boxes"].new_tensor([w, h, w, h]) boxes = box_cxcywh_to_xyxy(t["boxes"]) * scale entry: dict[str, torch.Tensor] = {"boxes": boxes, "labels": t["labels"]} if "masks" in t: masks = t["masks"].bool() - # PostProcess resizes predicted masks to orig_size; resize GT - # masks to match so that mask-IoU comparisons are size-consistent. + # PostProcess resizes predicted masks to (h, w); resize GT masks + # to match so that mask-IoU comparisons are size-consistent. if masks.shape[-2:] != (int(h), int(w)): masks = ( F.interpolate( diff --git a/src/rfdetr/training/module_model.py b/src/rfdetr/training/module_model.py index 5d7f90dff..ee94fdf9b 100644 --- a/src/rfdetr/training/module_model.py +++ b/src/rfdetr/training/module_model.py @@ -220,6 +220,10 @@ def validation_step(self, batch: Tuple, batch_idx: int) -> Dict[str, Any]: self.log("val/loss", loss, prog_bar=True, on_epoch=True, sync_dist=True, batch_size=len(targets)) orig_sizes = torch.stack([t["orig_size"] for t in targets]) + cap = getattr(self.train_config, "max_eval_orig_size", None) + if cap is not None: + scale = (cap / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0) + orig_sizes = (orig_sizes.float() * scale).long() results = self.postprocess(outputs, orig_sizes) return {"results": results, "targets": targets} @@ -355,6 +359,10 @@ def test_step(self, batch: Tuple, batch_idx: int) -> Dict[str, Any]: self.log("test/loss", loss, sync_dist=True, batch_size=len(targets)) orig_sizes = torch.stack([t["orig_size"] for t in targets]) + cap = getattr(self.train_config, "max_eval_orig_size", None) + if cap is not None: + scale = (cap / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0) + orig_sizes = (orig_sizes.float() * scale).long() results = self.postprocess(outputs, orig_sizes) return {"results": results, "targets": targets} diff --git a/src/rfdetr/training/trainer.py b/src/rfdetr/training/trainer.py index 82d56ca49..6dc5dfc26 100644 --- a/src/rfdetr/training/trainer.py +++ b/src/rfdetr/training/trainer.py @@ -211,6 +211,7 @@ def _resolve_precision() -> str: segmentation=model_config.segmentation_head, eval_interval=tc.eval_interval, log_per_class_metrics=tc.log_per_class_metrics, + max_eval_orig_size=tc.max_eval_orig_size, ) ) diff --git a/tests/models/test_config.py b/tests/models/test_config.py index a8eedf011..d41cae793 100644 --- a/tests/models/test_config.py +++ b/tests/models/test_config.py @@ -232,6 +232,7 @@ def test_compute_test_loss_default_is_true(self, tmp_path): pytest.param("train_log_sync_dist", True, id="train_log_sync_dist"), pytest.param("train_log_on_step", True, id="train_log_on_step"), pytest.param("log_per_class_metrics", False, id="log_per_class_metrics"), + pytest.param("max_eval_orig_size", 640, id="max_eval_orig_size"), pytest.param("prefetch_factor", 4, id="prefetch_factor"), pytest.param("pin_memory", False, id="pin_memory"), pytest.param("persistent_workers", False, id="persistent_workers"), @@ -261,6 +262,15 @@ def test_interval_and_prefetch_reject_non_positive_values(self, tmp_path, field, with pytest.raises((ValueError, ValidationError)): self._tc(tmp_path, **{field: value}) + def test_max_eval_orig_size_defaults_to_none(self, tmp_path): + """max_eval_orig_size defaults to None (no cap — evaluates at full original resolution).""" + assert self._tc(tmp_path).max_eval_orig_size is None + + def test_max_eval_orig_size_inherited_by_segmentation_train_config(self): + """SegmentationTrainConfig inherits max_eval_orig_size from TrainConfig.""" + tc = SegmentationTrainConfig(dataset_dir="/tmp", max_eval_orig_size=640) + assert tc.max_eval_orig_size == 640 + def test_batch_size_auto_is_accepted(self, tmp_path): """batch_size accepts the special 'auto' value.""" tc = self._tc(tmp_path, batch_size="auto") diff --git a/tests/training/test_coco_eval_callback.py b/tests/training/test_coco_eval_callback.py index 450df0c08..c66f1fefb 100644 --- a/tests/training/test_coco_eval_callback.py +++ b/tests/training/test_coco_eval_callback.py @@ -747,3 +747,152 @@ def test_no_masks_no_iscrowd_keys_absent(self) -> None: ] out = cb._convert_targets(targets) assert set(out[0].keys()) == {"boxes", "labels"} + + +class TestMaxEvalOrigSize: + """max_eval_orig_size caps mask resolution and box coordinates for COCO eval.""" + + # ------------------------------------------------------------------ + # _convert_preds: mask downsampling + # ------------------------------------------------------------------ + + def test_convert_preds_downsamples_masks_when_above_cap(self) -> None: + """Masks larger than cap are downsampled to fit the longer side.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + preds = [ + { + "masks": torch.ones(2, 1, 1080, 1920), + "boxes": torch.zeros(2, 4), + "scores": torch.ones(2), + "labels": torch.zeros(2, dtype=torch.long), + } + ] + out = cb._convert_preds(preds) + h, w = out[0]["masks"].shape[-2:] + assert max(h, w) == 640 + + def test_convert_preds_preserves_aspect_ratio_on_downsample(self) -> None: + """Downsampled masks preserve the original aspect ratio.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + preds = [ + { + "masks": torch.ones(1, 1, 1080, 1920), + "boxes": torch.zeros(1, 4), + "scores": torch.ones(1), + "labels": torch.zeros(1, dtype=torch.long), + } + ] + out = cb._convert_preds(preds) + h, w = out[0]["masks"].shape[-2:] + assert abs(w / h - 1920 / 1080) < 0.02 + + def test_convert_preds_scales_boxes_with_masks(self) -> None: + """Box coordinates are scaled by the same factor as the mask downsample.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + boxes = torch.tensor([[100.0, 200.0, 300.0, 400.0]]) + preds = [ + { + "masks": torch.ones(1, 1, 1080, 1920), + "boxes": boxes, + "scores": torch.ones(1), + "labels": torch.zeros(1, dtype=torch.long), + } + ] + out = cb._convert_preds(preds) + scale = 640 / 1920 + expected = boxes * scale + assert torch.allclose(out[0]["boxes"], expected, atol=1e-4) + + def test_convert_preds_no_op_when_within_cap(self) -> None: + """Masks already within the cap are not resized.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + preds = [ + { + "masks": torch.ones(1, 400, 600), + "boxes": torch.zeros(1, 4), + "scores": torch.ones(1), + "labels": torch.zeros(1, dtype=torch.long), + } + ] + out = cb._convert_preds(preds) + assert out[0]["masks"].shape[-2:] == (400, 600) + + def test_convert_preds_no_op_when_cap_is_none(self) -> None: + """When max_eval_orig_size is None no downsampling occurs regardless of mask size.""" + cb = COCOEvalCallback(max_eval_orig_size=None) + preds = [ + { + "masks": torch.ones(1, 1, 1080, 1920), + "boxes": torch.zeros(1, 4), + "scores": torch.ones(1), + "labels": torch.zeros(1, dtype=torch.long), + } + ] + out = cb._convert_preds(preds) + assert out[0]["masks"].shape[-2:] == (1080, 1920) + + # ------------------------------------------------------------------ + # _convert_targets: orig_size capping + # ------------------------------------------------------------------ + + def test_convert_targets_caps_gt_mask_size(self) -> None: + """GT masks are resized to the capped (h, w), not the original orig_size.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + targets = [ + { + "boxes": torch.tensor([[0.5, 0.5, 0.1, 0.1]]), + "labels": torch.tensor([0]), + "orig_size": torch.tensor([1080, 1920]), + "masks": torch.ones(1, 1080, 1920, dtype=torch.bool), + } + ] + out = cb._convert_targets(targets) + h, w = out[0]["masks"].shape[-2:] + assert max(h, w) == 640 + + def test_convert_targets_scales_boxes_to_capped_size(self) -> None: + """GT boxes are scaled to the capped (h, w) coordinate space.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + # Normalised box centred at image centre, half the image size + targets = [ + { + "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]), + "labels": torch.tensor([0]), + "orig_size": torch.tensor([1080, 1920]), + } + ] + out = cb._convert_targets(targets) + scale = 640 / 1920 + capped_h = int(1080 * scale) + capped_w = 640 + # box_cxcywh_to_xyxy([0.5,0.5,1.0,1.0]) * [w,h,w,h] → [0,0,w,h] + expected = torch.tensor([[0.0, 0.0, float(capped_w), float(capped_h)]]) + assert torch.allclose(out[0]["boxes"], expected, atol=1.0) + + def test_convert_targets_no_op_when_cap_is_none(self) -> None: + """With max_eval_orig_size=None, GT boxes use full orig_size as scale.""" + cb = COCOEvalCallback(max_eval_orig_size=None) + targets = [ + { + "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]), + "labels": torch.tensor([0]), + "orig_size": torch.tensor([1080, 1920]), + } + ] + out = cb._convert_targets(targets) + expected = torch.tensor([[0.0, 0.0, 1920.0, 1080.0]]) + assert torch.allclose(out[0]["boxes"], expected, atol=1e-4) + + def test_convert_targets_no_op_when_within_cap(self) -> None: + """Images already within the cap are not resized.""" + cb = COCOEvalCallback(max_eval_orig_size=640) + targets = [ + { + "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]), + "labels": torch.tensor([0]), + "orig_size": torch.tensor([480, 640]), + "masks": torch.ones(1, 480, 640, dtype=torch.bool), + } + ] + out = cb._convert_targets(targets) + assert out[0]["masks"].shape[-2:] == (480, 640)