26 lines
656 B
Python
26 lines
656 B
Python
from collections import defaultdict
|
|
from typing import Dict
|
|
|
|
from ray.tune.stopper.stopper import Stopper
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
|
|
@PublicAPI
|
|
class MaximumIterationStopper(Stopper):
|
|
"""Stop trials after reaching a maximum number of iterations
|
|
|
|
Args:
|
|
max_iter: Number of iterations before stopping a trial.
|
|
"""
|
|
|
|
def __init__(self, max_iter: int):
|
|
self._max_iter = max_iter
|
|
self._iter = defaultdict(lambda: 0)
|
|
|
|
def __call__(self, trial_id: str, result: Dict):
|
|
self._iter[trial_id] += 1
|
|
return self._iter[trial_id] >= self._max_iter
|
|
|
|
def stop_all(self):
|
|
return False
|