Merge pull request optuna#1430 from bigbird555/fix-pruner-bug

Experimental commit to fix pruner bug.

Author: HideakiImamura

Diff for: optuna/pruners/_percentile.py

@@ -55,8 +55,8 @@ def _is_first_in_interval_step(step, intermediate_steps, n_warmup_steps, interva
     # type: (int, KeysView[int], int, int) -> bool
 
     nearest_lower_pruning_step = (
-        (step - n_warmup_steps - 1) // interval_steps * interval_steps + n_warmup_steps + 1
-    )
+        step - n_warmup_steps
+    ) // interval_steps * interval_steps + n_warmup_steps
     assert nearest_lower_pruning_step >= 0
 
     # `intermediate_steps` may not be sorted so we must go through all elements.
@@ -167,7 +167,7 @@ def prune(self, study, trial):
             return False
 
         n_warmup_steps = self._n_warmup_steps
-        if step <= n_warmup_steps:
+        if step < n_warmup_steps:
             return False
 
         if not _is_first_in_interval_step(

Diff for: optuna/pruners/_threshold.py

@@ -118,7 +118,7 @@ def prune(self, study: "optuna.study.Study", trial: "optuna.trial.FrozenTrial")
             return False
 
         n_warmup_steps = self._n_warmup_steps
-        if step <= n_warmup_steps:
+        if step < n_warmup_steps:
             return False
 
         if not _is_first_in_interval_step(

Diff for: tests/pruners_tests/test_median.py

@@ -95,28 +95,28 @@ def test_median_pruner_n_warmup_steps():
     study = optuna.study.create_study()
 
     trial = optuna.trial.Trial(study, study._storage.create_new_trial(study._study_id))
+    trial.report(1, 0)
     trial.report(1, 1)
-    trial.report(1, 2)
     study._storage.set_trial_state(trial._trial_id, TrialState.COMPLETE)
 
     trial = optuna.trial.Trial(study, study._storage.create_new_trial(study._study_id))
-    trial.report(2, 1)
+    trial.report(2, 0)
     # A pruner is not activated during warm-up steps.
     assert not pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(2, 2)
+    trial.report(2, 1)
     # A pruner is activated after warm-up steps.
     assert pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
 
 @pytest.mark.parametrize(
     "n_warmup_steps,interval_steps,report_steps,expected_prune_steps",
     [
-        (1, 2, 1, [2, 4]),
-        (0, 3, 10, list(range(1, 30))),
-        (2, 3, 10, list(range(11, 30))),
-        (0, 10, 3, [1, 2, 3, 13, 14, 15, 22, 23, 24]),
-        (2, 10, 3, [4, 5, 6, 13, 14, 15, 25, 26, 27]),
+        (1, 2, 1, [1, 3]),
+        (0, 3, 10, list(range(29))),
+        (2, 3, 10, list(range(10, 29))),
+        (0, 10, 3, [0, 1, 2, 12, 13, 14, 21, 22, 23]),
+        (2, 10, 3, [3, 4, 5, 12, 13, 14, 24, 25, 26]),
     ],
 )
 def test_median_pruner_interval_steps(
@@ -129,7 +129,7 @@ def test_median_pruner_interval_steps(
 
     trial = optuna.trial.Trial(study, study._storage.create_new_trial(study._study_id))
     n_steps = max(expected_prune_steps)
-    base_index = 1
+    base_index = 0
     for i in range(base_index, base_index + n_steps):
         trial.report(base_index, i)
     study._storage.set_trial_state(trial._trial_id, TrialState.COMPLETE)
@@ -139,5 +139,5 @@ def test_median_pruner_interval_steps(
         if (i - base_index) % report_steps == 0:
             trial.report(2, i)
         assert pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id)) == (
-            i > n_warmup_steps and i in expected_prune_steps
+            i >= n_warmup_steps and i in expected_prune_steps
         )

Diff for: tests/pruners_tests/test_threshold.py

@@ -82,10 +82,10 @@ def test_threshold_pruner_n_warmup_steps() -> None:
     trial = optuna.trial.Trial(study, study._storage.create_new_trial(study._study_id))
     pruner = optuna.pruners.ThresholdPruner(lower=0.0, upper=1.0, n_warmup_steps=2)
 
-    trial.report(-10.0, 1)
+    trial.report(-10.0, 0)
     assert not pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(100.0, 2)
+    trial.report(100.0, 1)
     assert not pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
     trial.report(-100.0, 3)
@@ -103,17 +103,17 @@ def test_threshold_pruner_interval_steps() -> None:
     trial = optuna.trial.Trial(study, study._storage.create_new_trial(study._study_id))
     pruner = optuna.pruners.ThresholdPruner(lower=0.0, upper=1.0, interval_steps=2)
 
-    trial.report(-10.0, 1)
+    trial.report(-10.0, 0)
     assert pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(100.0, 2)
+    trial.report(100.0, 1)
     assert not pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(-100.0, 3)
+    trial.report(-100.0, 2)
     assert pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(10.0, 4)
+    trial.report(10.0, 3)
     assert not pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))
 
-    trial.report(1000.0, 5)
+    trial.report(1000.0, 4)
     assert pruner.prune(study=study, trial=study._storage.get_trial(trial._trial_id))