# ------------------------------------------------------------------------ # 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