chore: import upstream snapshot with attribution

tests/training/callbacks/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/training/callbacks/test_drop_path_callback.py

@@ -0,0 +1,293 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for :class:`rfdetr.training.callbacks.drop_schedule.DropPathCallback`."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from rfdetr.training.callbacks.drop_schedule import DropPathCallback
+from rfdetr.training.drop_schedule import drop_scheduler
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_trainer(global_step: int = 0, estimated_stepping_batches: int = 50) -> MagicMock:
+    """Create a minimal mock Trainer with controllable step metadata."""
+    trainer = MagicMock()
+    trainer.global_step = global_step
+    trainer.estimated_stepping_batches = estimated_stepping_batches
+    return trainer
+
+
+def _make_mock_pl_module(epochs: int = 5) -> MagicMock:
+    """Create a minimal mock RFDETRModule with ``train_config.epochs``."""
+    pl_module = MagicMock()
+    pl_module.train_config.epochs = epochs
+    return pl_module
+
+
+# ---------------------------------------------------------------------------
+# TestDropPathCallbackInit
+# ---------------------------------------------------------------------------
+
+
+class TestDropPathCallbackInit:
+    """Verify constructor defaults."""
+
+    def test_default_args(self) -> None:
+        """Default rates are zero and vit_encoder_num_layers is 12."""
+        cb = DropPathCallback()
+        assert cb._drop_path == 0.0
+        assert cb._dropout == 0.0
+        assert cb._vit_encoder_num_layers == 12
+        assert cb._dp_schedule is None
+        assert cb._do_schedule is None
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainStart:
+    """Verify schedule arrays built in ``on_train_start``."""
+
+    def test_dp_schedule_matches_drop_scheduler_standard(self) -> None:
+        """drop_path schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(drop_path=0.3)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10)
+        assert cb._dp_schedule is not None
+        np.testing.assert_array_equal(cb._dp_schedule, expected)
+
+    def test_do_schedule_matches_drop_scheduler_standard(self) -> None:
+        """Dropout schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.1, 5, 10)
+        assert cb._do_schedule is not None
+        np.testing.assert_array_equal(cb._do_schedule, expected)
+
+    def test_no_dp_schedule_when_rate_zero(self) -> None:
+        """drop_path=0.0 leaves ``_dp_schedule`` as None."""
+        cb = DropPathCallback(drop_path=0.0)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        assert cb._dp_schedule is None
+
+    def test_dp_schedule_early_mode(self) -> None:
+        """Early mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="early")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "early")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+    def test_dp_schedule_late_mode(self) -> None:
+        """Late mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="late")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "late")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainBatchStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainBatchStart:
+    """Verify model update calls in ``on_train_batch_start``."""
+
+    def test_update_drop_path_called_with_correct_rate(self) -> None:
+        """``update_drop_path`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[0], 6)
+
+    def test_update_dropout_called_with_correct_rate(self) -> None:
+        """``update_dropout`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._do_schedule is not None
+        pl_module.model.update_dropout.assert_called_once_with(cb._do_schedule[0])
+
+    def test_no_update_when_step_out_of_bounds(self) -> None:
+        """No model updates when ``global_step`` exceeds schedule length."""
+        cb = DropPathCallback(drop_path=0.3, dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 9999
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        pl_module.model.update_drop_path.assert_not_called()
+        pl_module.model.update_dropout.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "step",
+        [
+            pytest.param(0, id="first_step"),
+            pytest.param(5, id="mid_step"),
+            pytest.param(9, id="last_of_first_epoch"),
+        ],
+    )
+    def test_drop_rates_at_multiple_steps_match_schedule(self, step: int) -> None:
+        """Each step uses the correct value from the pre-computed schedule."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = step
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[step], 6)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerValidation
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerValidation:
+    """Verify drop_scheduler raises for invalid inputs."""
+
+    @pytest.mark.parametrize(
+        "cutoff_epoch",
+        [
+            pytest.param(6, id="above_epochs"),
+            pytest.param(-1, id="negative"),
+        ],
+    )
+    def test_raises_for_invalid_cutoff_epoch(self, cutoff_epoch: int) -> None:
+        """drop_scheduler raises ValueError when cutoff_epoch is outside [0, epochs]."""
+        with pytest.raises(ValueError, match="cutoff_epoch must be in"):
+            drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode="early")
+
+    @pytest.mark.parametrize(
+        ("epochs", "niter_per_ep", "match"),
+        [
+            pytest.param(0, 10, "epochs must be >= 1", id="epochs_zero"),
+            pytest.param(5, 0, "niter_per_ep must be >= 1", id="niter_per_ep_zero"),
+        ],
+    )
+    def test_raises_for_invalid_epoch_counts(self, epochs: int, niter_per_ep: int, match: str) -> None:
+        """drop_scheduler raises ValueError when epochs or niter_per_ep is less than 1."""
+        with pytest.raises(ValueError, match=match):
+            drop_scheduler(0.3, epochs, niter_per_ep)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerBoundary
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerBoundary:
+    """Verify drop_scheduler with cutoff_epoch at the inclusive boundaries 0 and epochs."""
+
+    @pytest.mark.parametrize(
+        ("cutoff_epoch", "mode", "expected_first", "expected_last"),
+        [
+            pytest.param(0, "early", 0.0, 0.0, id="early_cutoff_zero_all_zeros"),
+            pytest.param(5, "early", 0.3, 0.3, id="early_cutoff_full_all_rate"),
+            pytest.param(0, "late", 0.3, 0.3, id="late_cutoff_zero_all_rate"),
+            pytest.param(5, "late", 0.0, 0.0, id="late_cutoff_full_all_zeros"),
+        ],
+    )
+    def test_boundary_cutoff_epoch(
+        self,
+        cutoff_epoch: int,
+        mode: str,
+        expected_first: float,
+        expected_last: float,
+    ) -> None:
+        """Boundary cutoff_epoch values (0 and epochs) produce correct first and last rates."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode=mode)
+        assert schedule[0] == expected_first
+        assert schedule[-1] == expected_last
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerLinear
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerLinear:
+    """Verify drop_scheduler with schedule='linear' in early mode."""
+
+    def test_linear_early_starts_at_drop_rate(self) -> None:
+        """Linear early schedule first value equals drop_rate."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[0] == pytest.approx(0.3, abs=1e-9)
+
+    def test_linear_early_ends_early_phase_at_zero(self) -> None:
+        """Linear early schedule last value of the early phase equals 0."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[19] == pytest.approx(0.0, abs=1e-9)
+
+    def test_linear_early_late_phase_is_zero(self) -> None:
+        """Linear early schedule: all values after cutoff_epoch are zero."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        np.testing.assert_array_equal(schedule[20:], 0.0)
+
+    def test_linear_early_decreases_monotonically(self) -> None:
+        """Linear early schedule values decrease monotonically during the early phase."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert np.all(np.diff(schedule[:20]) <= 0)
+
+    def test_linear_same_shape_as_constant(self) -> None:
+        """Schedule='linear' output has the same length as schedule='constant'."""
+        linear = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        constant = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="constant")
+        assert linear.shape == constant.shape

tests/training/callbacks/test_ema_callback.py

@@ -0,0 +1,260 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit and parity tests for RFDETREMACallback."""
+
+from __future__ import annotations
+
+import math
+import warnings
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+from torch.optim.swa_utils import AveragedModel
+
+from rfdetr.training.callbacks.ema import RFDETREMACallback
+from rfdetr.training.model_ema import ModelEma
+
+
+class _EMAContainerModule(nn.Module):
+    """Minimal module with `.model` to mirror RFDETRModelModule shape."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = nn.Linear(4, 2)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+
+class TestAvgFnDecayFormula:
+    """Verify the tau / no-tau decay formula matches ModelEma."""
+
+    @pytest.mark.parametrize(
+        "num_averaged",
+        [
+            pytest.param(0, id="step-0"),
+            pytest.param(5, id="step-5"),
+            pytest.param(99, id="step-99"),
+        ],
+    )
+    def test_tau_zero_uses_fixed_decay(self, num_averaged: int) -> None:
+        """With tau=0 the effective decay equals the base decay at every step."""
+        decay = 0.99
+        cb = RFDETREMACallback(decay=decay, tau=0)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged)
+
+        expected = ema_val * decay + model_val * (1.0 - decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+    def test_tau_warmup_at_step_1(self) -> None:
+        """At the first call (num_averaged=0) with tau>0 the effective decay uses updates=1 matching ModelEma's
+        1-indexed counter."""
+        decay = 0.993
+        tau = 100
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged=0)
+
+        updates = 1  # num_averaged + 1
+        effective_decay = decay * (1 - math.exp(-updates / tau))
+        expected = ema_val * effective_decay + model_val * (1.0 - effective_decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+
+class TestModelEmaParity:
+    """Ensure N-step EMA weights match ModelEma exactly."""
+
+    def test_avg_fn_matches_modelema_weight_parity(self) -> None:
+        """Simulate 500 update steps and compare final EMA weights with ModelEma.module to confirm numerical parity."""
+        torch.manual_seed(42)
+        n_steps = 500
+        decay = 0.993
+        tau = 100
+
+        model = nn.Linear(4, 4)
+        model_ema = ModelEma(model, decay=decay, tau=tau)
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+
+        # Initialise manual EMA state from model (same as ModelEma deepcopy)
+        ema_weights: dict[str, torch.Tensor] = {name: p.clone() for name, p in model.named_parameters()}
+
+        for step in range(n_steps):
+            # Perturb model parameters
+            with torch.no_grad():
+                for p in model.parameters():
+                    p.add_(torch.randn_like(p) * 0.01)
+
+            # Update legacy ModelEma
+            model_ema.update(model)
+
+            # Replicate update via callback avg_fn
+            model_weights = {name: p.clone() for name, p in model.named_parameters()}
+            for name in ema_weights:
+                ema_weights[name] = cb._avg_fn(ema_weights[name], model_weights[name], step)
+
+        # Compare
+        legacy_state = dict(model_ema.module.named_parameters())
+        for name, cb_val in ema_weights.items():
+            assert torch.allclose(cb_val, legacy_state[name], atol=1e-5), (
+                f"Parity failed for {name}: max diff = {(cb_val - legacy_state[name]).abs().max().item()}"
+            )
+
+
+class TestShouldUpdate:
+    """Verify should_update triggers on steps and epochs."""
+
+    def test_should_update_on_step(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(step_idx=42) is True
+
+    def test_should_update_on_epoch(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(epoch_idx=3) is True
+
+    def test_should_update_neither(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update() is False
+
+
+class TestInit:
+    """Construction and EMA-state access behavior."""
+
+    def test_init_emits_no_user_warning(self) -> None:
+        """Instantiation should not emit runtime UserWarnings."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            RFDETREMACallback()
+        user_warns = [w for w in caught if issubclass(w.category, UserWarning)]
+        assert not user_warns
+
+    def test_get_ema_model_state_dict_none_before_setup(self) -> None:
+        """EMA state accessor returns None before averaged model is created."""
+        cb = RFDETREMACallback()
+        assert cb.get_ema_model_state_dict() is None
+
+    def test_get_ema_model_state_dict_returns_model_weights(self) -> None:
+        """EMA state accessor returns the wrapped `.model` state dict."""
+
+        class _Container(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.model = nn.Linear(4, 2)
+
+        cb = RFDETREMACallback()
+        container = _Container()
+        cb._average_model = AveragedModel(container, avg_fn=cb._avg_fn)
+
+        state = cb.get_ema_model_state_dict()
+
+        assert state is not None
+        assert "weight" in state
+        assert "bias" in state
+
+
+class TestUpdateInterval:
+    """Verify update_interval_steps throttles EMA updates on step hooks."""
+
+    def test_updates_only_on_interval_steps(self) -> None:
+        """update_interval_steps=2 updates on steps 2, 4, ...
+
+        only.
+        """
+        cb = RFDETREMACallback(update_interval_steps=2)
+        cb._average_model = MagicMock()
+
+        trainer = MagicMock()
+        pl_module = MagicMock()
+
+        for step in (1, 2, 3, 4):
+            trainer.global_step = step
+            cb.on_train_batch_end(trainer, pl_module, outputs=None, batch=None, batch_idx=step - 1)
+
+        assert cb._average_model.update_parameters.call_count == 2
+
+
+class TestLegacyEMAResume:
+    """Legacy checkpoint EMA payload is consumed by the callback setup path."""
+
+    def test_setup_loads_pending_legacy_ema_state_into_average_model(self) -> None:
+        """`_pending_legacy_ema_state` must initialize EMA weights at fit setup."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        trainer = MagicMock()
+
+        legacy_ema_state = {k: torch.full_like(v, 2.0) for k, v in pl_module.model.state_dict().items()}
+        pl_module._pending_legacy_ema_state = legacy_ema_state
+
+        cb.setup(trainer, pl_module, stage="fit")
+
+        assert cb._average_model is not None
+        restored = cb._average_model.module.model.state_dict()
+        for key, expected in legacy_ema_state.items():
+            assert torch.allclose(restored[key], expected)
+        assert not hasattr(pl_module, "_pending_legacy_ema_state")
+
+
+class TestSuppressTestSwap:
+    """suppress_test_swap must disable the test-time EMA weight swap while leaving defaults unchanged."""
+
+    @staticmethod
+    def _make_swap_scenario() -> tuple[RFDETREMACallback, _EMAContainerModule]:
+        """Build a module at weight 7.0 with an EMA average model captured at weight 5.0."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(5.0)
+        cb._average_model = AveragedModel(model=pl_module, use_buffers=True, avg_fn=cb._avg_fn)
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(7.0)
+        return cb, pl_module
+
+    def test_default_flag_is_false(self) -> None:
+        """The suppression flag defaults to False so standalone trainer.test() keeps EMA evaluation."""
+        cb = RFDETREMACallback()
+        assert cb.suppress_test_swap is False
+
+    def test_on_test_epoch_start_swaps_by_default(self) -> None:
+        """Without suppression, the test hooks swap live weights (7.0) for EMA weights (5.0)."""
+        cb, pl_module = self._make_swap_scenario()
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 5.0))
+
+    def test_on_test_epoch_start_suppressed_keeps_live_weights(self) -> None:
+        """With suppress_test_swap=True the live weights (7.0) must stay in place during test."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))
+
+    def test_on_test_epoch_end_suppressed_does_not_swap(self) -> None:
+        """With suppression active, on_test_epoch_end must not swap EMA weights in unpaired."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+        cb.on_test_epoch_end(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))