Repository: autogluon/autogluon
Pull request: Unify scheduler creation in different tasks. Basic fix to FIFOSchedul…
…er, HyperbandScheduler

*Issue #, if available:*

*Description of changes:*


By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.


Problem statement:
For HyperbandScheduler, epoch must start with 1, not with 0
I noted a mistake in the docstrings, and maybe also in the tabular code.

When time_attr = 'epoch' in HyperbandScheduler, the values coming from the reporter are assumed to start with 1 and run up to max_t. Instead, the examples make them start at 0 and run up to max_t-1.

In other words, HyperbandScheduler treats epoch is "number of epochs done" (which starts at 1). Or think about time or resource spent. The value 0 makes no sense here.

Unfortunately, this really matters. The HB scheduler does certain things when rung levels are reached, and they are 1, 3, 9, ... When train_fn does something like:

for e in range(args.epochs):
    # ...
    reporter(epoch=e, ...)

then the first value returned is not used by the scheduler, but the second is associated with 1. Worse, the final one for max_t is expected by the scheduler, but never comes. This is a problem certainly for the GP methods we are about to contribute.

I could fix this, but I am not very familiar with the code in tabular. I can fix the docstrings.

Hints from issue discussion:

Commit message: Addressed jwmueller comments:
- epochs removed in tabular_prediction, changed compile_scheduler_options
- improved comments and tutorial text

Patch to review:
diff --git a/autogluon/scheduler/fifo.py b/autogluon/scheduler/fifo.py
index e3ef40a9efc5..2e956a0e9004 100644
--- a/autogluon/scheduler/fifo.py
+++ b/autogluon/scheduler/fifo.py
@@ -30,40 +30,40 @@ class FIFOScheduler(TaskScheduler):
 
     Parameters
     ----------
-    train_fn: callable
+    train_fn : callable
         A task launch function for training.
-    args: object (optional)
+    args : object (optional)
         Default arguments for launching train_fn.
-    resource: dict
+    resource : dict
         Computation resources. For example, `{'num_cpus':2, 'num_gpus':1}`
-    searcher: str or BaseSearcher
+    searcher : str or BaseSearcher
         Searcher (get_config decisions). If str, this is passed to
         searcher_factory along with search_options.
-    search_options: dict
+    search_options : dict
         If searcher is str, these arguments are passed to searcher_factory.
-    checkpoint: str
+    checkpoint : str
         If filename given here, a checkpoint of scheduler (and searcher) state
         is written to file every time a job finishes.
         Note: May not be fully supported by all searchers.
-    resume: bool
+    resume : bool
         If True, scheduler state is loaded from checkpoint, and experiment
         starts from there.
         Note: May not be fully supported by all searchers.
-    num_trials: int
+    num_trials : int
         Maximum number of jobs run in experiment.
-    time_out: float
+    time_out : float
         If given, jobs are started only until this time_out (wall clock time)
-    reward_attr: str
+    reward_attr : str
         Name of reward (i.e., metric to maximize) attribute in data obtained
         from reporter
-    time_attr: str
+    time_attr : str
         Name of resource (or time) attribute in data obtained from reporter.
         This attribute is optional for FIFO scheduling, but becomes mandatory
         in multi-fidelity scheduling (e.g., Hyperband).
         Note: The type of resource must be int.
-    dist_ip_addrs: list of str
+    dist_ip_addrs : list of str
         IP addresses of remote machines.
-    training_history_callback: callable
+    training_history_callback : callable
         Callback function func called every time a job finishes, if at least
         training_history_callback_delta_secs seconds passed since the last
         recent call. The call has the form:
@@ -71,9 +71,9 @@ class FIFOScheduler(TaskScheduler):
         Here, self._start_time is time stamp for when experiment started.
         Use this callback to serialize self.training_history after regular
         intervals.
-    training_history_callback_delta_secs: float
+    training_history_callback_delta_secs : float
         See training_history_callback.
-    delay_get_config: bool
+    delay_get_config : bool
         If True, the call to searcher.get_config is delayed until a worker
         resource for evaluation is available. Otherwise, get_config is called
         just after a job has been started.
@@ -175,16 +175,6 @@ def __init__(self, train_fn, args=None, resource=None,
         # str(task.task_id)
         self.training_history = OrderedDict()
         self.config_history = OrderedDict()
-        # Resume experiment from checkpoint?
-        if resume:
-            assert checkpoint is not None, \
-                "Need checkpoint to be set if resume = True"
-            if os.path.isfile(checkpoint):
-                self.load_state_dict(load(checkpoint))
-            else:
-                msg = f'checkpoint path {checkpoint} is not available for resume.'
-                logger.exception(msg)
-                raise FileExistsError(msg)
         # Needed for training_history callback mechanism, which is used to
         # serialize training_history after each
         # training_history_call_delta_secs seconds
@@ -195,6 +185,16 @@ def __init__(self, train_fn, args=None, resource=None,
         self.training_history_callback_delta_secs = \
             training_history_callback_delta_secs
         self._delay_get_config = delay_get_config
+        # Resume experiment from checkpoint?
+        if resume:
+            assert checkpoint is not None, \
+                "Need checkpoint to be set if resume = True"
+            if os.path.isfile(checkpoint):
+                self.load_state_dict(load(checkpoint))
+            else:
+                msg = f'checkpoint path {checkpoint} is not available for resume.'
+                logger.exception(msg)
+                raise FileExistsError(msg)
 
     def run(self, **kwargs):
         """Run multiple number of trials
@@ -361,9 +361,6 @@ def _run_reporter(self, task, task_job, reporter):
         last_result = None
         while not task_job.done():
             reported_result = reporter.fetch()
-            # Time since start of experiment
-            elapsed_time = self._elapsed_time()
-            reported_result['time_since_start'] = elapsed_time
             if 'traceback' in reported_result:
                 # Evaluation has failed
                 logger.exception(reported_result['traceback'])
@@ -374,6 +371,14 @@ def _run_reporter(self, task, task_job, reporter):
             if reported_result.get('done', False):
                 reporter.move_on()
                 break
+            if len(reported_result) == 0:
+                # An empty dict should just be skipped
+                logger.warning("Skipping empty dict received from reporter")
+                continue
+            # Time since start of experiment
+            elapsed_time = self._elapsed_time()
+            reported_result['time_since_start'] = elapsed_time
+
             # Extra information from searcher (optional)
             dataset_size = self.searcher.dataset_size()
             if dataset_size > 0:
diff --git a/autogluon/scheduler/hyperband.py b/autogluon/scheduler/hyperband.py
index cf6e96a81bca..3501d5df02da 100644
--- a/autogluon/scheduler/hyperband.py
+++ b/autogluon/scheduler/hyperband.py
@@ -3,11 +3,14 @@
 import threading
 import numpy as np
 import multiprocessing as mp
+import os
 
 from .fifo import FIFOScheduler
 from .hyperband_stopping import HyperbandStopping_Manager
 from .hyperband_promotion import HyperbandPromotion_Manager
 from .reporter import DistStatusReporter
+from ..utils import load
+
 
 __all__ = ['HyperbandScheduler',
            'HyperbandStopping_Manager',
@@ -186,9 +189,11 @@ def __init__(self, train_fn, args=None, resource=None,
         _search_options['resource_attribute'] = time_attr
         _search_options['min_epochs'] = grace_period
         _search_options['max_epochs'] = max_t
+        # Pass resume=False here. Resume needs members of this object to be
+        # created
         super().__init__(
             train_fn=train_fn, args=args, resource=resource, searcher=searcher,
-            search_options=_search_options, checkpoint=checkpoint, resume=resume,
+            search_options=_search_options, checkpoint=checkpoint, resume=False,
             num_trials=num_trials, time_out=time_out, max_reward=max_reward,
             reward_attr=reward_attr, time_attr=time_attr,
             visualizer=visualizer, dist_ip_addrs=dist_ip_addrs,
@@ -236,6 +241,16 @@ def __init__(self, train_fn, args=None, resource=None,
         # This lock protects both _running_tasks and terminator, the latter
         # does not define its own lock
         self._hyperband_lock = mp.Lock()
+        if resume:
+            assert checkpoint is not None, \
+                "Need checkpoint to be set if resume = True"
+            if os.path.isfile(checkpoint):
+                self.load_state_dict(load(checkpoint))
+            else:
+                msg = f'checkpoint path {checkpoint} is not available for resume.'
+                logger.exception(msg)
+                raise FileExistsError(msg)
+
 
     def add_job(self, task, **kwargs):
         """Adding a training task to the scheduler.
@@ -330,12 +345,9 @@ def _update_searcher(self, task, result):
     def _run_reporter(self, task, task_job, reporter):
         last_result = None
         last_updated = None
+        task_key = str(task.task_id)
         while not task_job.done():
             reported_result = reporter.fetch()
-            # Time since start of experiment
-            elapsed_time = self._elapsed_time()
-            reported_result['time_since_start'] = elapsed_time
-            task_key = str(task.task_id)
             if 'traceback' in reported_result:
                 # Evaluation has failed
                 logger.exception(reported_result['traceback'])
@@ -353,6 +365,13 @@ def _run_reporter(self, task, task_job, reporter):
                     # Cleanup
                     del self._running_tasks[task_key]
                 break
+            if len(reported_result) == 0:
+                # An empty dict should just be skipped
+                logger.warning("Skipping empty dict received from reporter")
+                continue
+            # Time since start of experiment
+            elapsed_time = self._elapsed_time()
+            reported_result['time_since_start'] = elapsed_time
 
             # Call before _add_training_results, since we may be able to report
             # extra information from the bracket
@@ -431,14 +450,6 @@ def _run_reporter(self, task, task_job, reporter):
                 # variant. It means that the *task* terminates, while the
                 # evaluation of the config is just paused
                 last_result['terminated'] = True
-                resource = int(reported_result[self._time_attr])
-                if self.type == 'stopping' or resource >= self.max_t:
-                    act_str = 'terminating'
-                else:
-                    act_str = 'pausing'
-                logger.debug(
-                    'Stopping task ({} evaluation, resource = {}):\n{}'.format(
-                        act_str, resource, task))
                 with self._hyperband_lock:
                     self.terminator.on_task_remove(task)
                     # Cleanup
diff --git a/autogluon/task/base/base_task.py b/autogluon/task/base/base_task.py
index 43bfd572dc86..6ed10eaae2b2 100644
--- a/autogluon/task/base/base_task.py
+++ b/autogluon/task/base/base_task.py
@@ -8,7 +8,11 @@
 from ...scheduler import *
 from ...utils import in_ipynb
 
-__all__ = ['BaseDataset', 'BaseTask']
+__all__ = [
+    'BaseDataset',
+    'BaseTask',
+    'compile_scheduler_options',
+    'create_scheduler']
 
 Results = collections.namedtuple('Results', 'model reward config time metadata')
 
@@ -17,8 +21,18 @@
     'random': FIFOScheduler,
     'skopt': FIFOScheduler,
     'hyperband': HyperbandScheduler,
-    'rl': RLScheduler,
-}
+    'rl': RLScheduler}
+
+
+def create_scheduler(train_fn, search_strategy, scheduler_options):
+    if isinstance(search_strategy, str):
+        scheduler_cls = schedulers[search_strategy.lower()]
+    else:
+        assert callable(search_strategy)
+        scheduler_cls = search_strategy
+        scheduler_options = copy.copy(scheduler_options)
+        scheduler_options['searcher'] = 'random'
+    return scheduler_cls(train_fn, **scheduler_options)
 
 
 class BaseDataset(mx.gluon.data.Dataset):
@@ -32,18 +46,12 @@ class BaseTask(object):
     Dataset = BaseDataset
 
     @classmethod
-    def run_fit(cls, train_fn, search_strategy, scheduler_options):
+    def run_fit(cls, train_fn, search_strategy, scheduler_options,
+                plot_results=False):
         start_time = time.time()
         # create scheduler and schedule tasks
-        if isinstance(search_strategy, str):
-            scheduler_cls = schedulers[search_strategy.lower()]
-        else:
-            assert callable(search_strategy)
-            scheduler_cls = search_strategy
-            scheduler_options['searcher'] = 'random'
-        plot_results = scheduler_options.pop('plot_results') \
-            if 'plot_results' in scheduler_options else False
-        scheduler = scheduler_cls(train_fn, **scheduler_options)
+        scheduler = create_scheduler(
+            train_fn, search_strategy, scheduler_options)
         print('scheduler:', scheduler)
         scheduler.run()
         scheduler.join_jobs()
@@ -74,3 +82,77 @@ def run_fit(cls, train_fn, search_strategy, scheduler_options):
     @abstractmethod
     def fit(cls, *args, **kwargs):
         pass
+
+
+# These search_strategies use HyperbandScheduler, along with certain
+# searchers.
+searcher_for_hyperband_strategy = {
+    'hyperband': 'random'}
+
+
+def compile_scheduler_options(
+        scheduler_options, search_strategy, search_options, nthreads_per_trial,
+        ngpus_per_trial, checkpoint, num_trials, time_out, resume, visualizer,
+        time_attr, reward_attr, dist_ip_addrs, epochs=None):
+    """
+    Updates a copy of scheduler_options (scheduler-specific options, can be
+    empty) with general options. The result can be passed to __init__ of the
+    scheduler.
+
+    Special role of epochs for HyperbandScheduler: If the search_strategy
+    involves HyperbandScheduler and epochs is given, then this value is
+    copied to scheduler_options['max_t']. Pass epochs for applications
+    where the time_attr is epoch, and epochs is the maximum number of
+    epochs.
+
+    :param scheduler_options:
+    :param search_strategy:
+    :param search_options:
+    :param nthreads_per_trial:
+    :param ngpus_per_trial:
+    :param checkpoint:
+    :param num_trials:
+    :param time_out:
+    :param resume:
+    :param visualizer:
+    :param time_attr:
+    :param reward_attr:
+    :param dist_ip_addrs:
+    :param kwargs:
+    :param epochs: See above. Optional
+    :return: Copy of scheduler_options with updates
+
+    """
+    if scheduler_options is None:
+        scheduler_options = dict()
+    else:
+        assert isinstance(scheduler_options, dict)
+    assert isinstance(search_strategy, str)
+    if search_options is None:
+        search_options = dict()
+    if visualizer is None:
+        visualizer = 'none'
+    if time_attr is None:
+        time_attr = 'epoch'
+    if reward_attr is None:
+        reward_attr = 'accuracy'
+    scheduler_options = copy.copy(scheduler_options)
+    scheduler_options.update({
+        'resource': {
+            'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
+        'searcher': search_strategy,
+        'search_options': search_options,
+        'checkpoint': checkpoint,
+        'resume': resume,
+        'num_trials': num_trials,
+        'time_out': time_out,
+        'reward_attr': reward_attr,
+        'time_attr': time_attr,
+        'visualizer': visualizer,
+        'dist_ip_addrs': dist_ip_addrs})
+    searcher = searcher_for_hyperband_strategy.get(search_strategy)
+    if searcher is not None:
+        scheduler_options['searcher'] = searcher
+        if epochs is not None:
+            scheduler_options['max_t'] = epochs
+    return scheduler_options
diff --git a/autogluon/task/image_classification/image_classification.py b/autogluon/task/image_classification/image_classification.py
index 259e3f950c42..616b1b9b7ffd 100644
--- a/autogluon/task/image_classification/image_classification.py
+++ b/autogluon/task/image_classification/image_classification.py
@@ -8,8 +8,7 @@
 from .nets import *
 from .pipeline import train_image_classification
 from .utils import *
-from ..base import BaseTask
-from ..base.base_task import schedulers
+from ..base import BaseTask, compile_scheduler_options, create_scheduler
 from ...core import *
 from ...core.loss import *
 from ...core.optimizer import *
@@ -83,14 +82,14 @@ def fit(dataset,
             search_strategy='random',
             plot_results=False,
             verbose=False,
-            search_options={},
+            search_options=None,
             time_limits=None,
             resume=False,
             output_directory='checkpoint/',
             visualizer='none',
             num_trials=2,
-            dist_ip_addrs=[],
-            grace_period=None,
+            dist_ip_addrs=None,
+            scheduler_options=None,
             auto_search=True,
             lr_config=Dict(
                 lr_mode='cosine',
@@ -153,24 +152,27 @@ def fit(dataset,
         ngpus_per_trial : int
             How many GPUs to use in each trial (ie. single training run of a model).
         output_directory : str
-            Dir to save search results.
+            Checkpoints of the search state are written to
+            os.path.join(output_directory, 'exp1.ag')
         search_strategy : str
             Which hyperparameter search algorithm to use.
             Options include: 'random' (random search), 'skopt' (SKopt Bayesian optimization), 'grid' (grid search), 'hyperband' (Hyperband), 'rl' (reinforcement learner)
         search_options : dict
             Auxiliary keyword arguments to pass to the searcher that performs hyperparameter optimization.
         resume : bool
-            If a model checkpoint file exists, model training will resume from there when specified.
+            If True, the hyperparameter search is started from state loaded from
+            os.path.join(output_directory, 'exp1.ag')
         dist_ip_addrs : list
             List of IP addresses corresponding to remote workers, in order to leverage distributed computation.
+        scheduler_options : dict
+            Extra arguments passed to __init__ of scheduler, to configure the
+            orchestration of training jobs during hyperparameter-tuning.
         verbose : bool
             Whether or not to print out intermediate information during training.
         plot_results : bool
             Whether or not to generate plots summarizing training process.
         visualizer : str
             Describes method to visualize training progress during `fit()`. Options: ['mxboard', 'tensorboard', 'none'].
-        grace_period : int
-            The grace period in early stopping when using Hyperband to tune hyperparameters. If None, this is set automatically.
         auto_search : bool
             If True, enables automatic suggestion of network types and hyper-parameter ranges adaptively based on provided dataset.
 
@@ -268,30 +270,15 @@ def fit(dataset,
             lr_config=lr_config
         )
 
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'checkpoint': checkpoint,
-            'num_trials': num_trials,
-            'time_out': time_limits,
-            'resume': resume,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',
-            'reward_attr': 'classification_reward',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-            'plot_results': plot_results
-        }
-        if search_strategy == 'hyperband':
-            scheduler_options.update(
-                {
-                    'searcher': 'random',
-                    'max_t': epochs,
-                    'grace_period': grace_period if grace_period else epochs // 4
-                }
-            )
-
-        results = BaseTask.run_fit(train_image_classification, search_strategy, scheduler_options)
+        scheduler_options = compile_scheduler_options(
+            scheduler_options, search_strategy, search_options,
+            nthreads_per_trial, ngpus_per_trial, checkpoint, num_trials,
+            time_out=time_limits, resume=resume, visualizer=visualizer,
+            time_attr='epoch', reward_attr='classification_reward',
+            dist_ip_addrs=dist_ip_addrs, epochs=epochs)
+        results = BaseTask.run_fit(
+            train_image_classification, search_strategy, scheduler_options,
+            plot_results=plot_results)
         args = sample_config(train_image_classification.args, results['best_config'])
 
         kwargs = {'num_classes': results['num_classes'], 'ctx': mx.cpu(0)}
@@ -301,14 +288,8 @@ def fit(dataset,
 
         if ensemble > 1:
             models = [model]
-            if isinstance(search_strategy, str):
-                scheduler_cls = schedulers[search_strategy.lower()]
-            else:
-                assert callable(search_strategy)
-                scheduler_cls = search_strategy
-                scheduler_options['searcher'] = 'random'
-            scheduler = scheduler_cls(
-                train_image_classification, **scheduler_options)
+            scheduler = create_scheduler(
+                train_image_classification, search_strategy, scheduler_options)
             for i in range(1, ensemble):
                 resultsi = scheduler.run_with_config(results['best_config'])
                 kwargs = {
diff --git a/autogluon/task/object_detection/object_detection.py b/autogluon/task/object_detection/object_detection.py
index 93b4cbe493f9..b5e6d3b2623f 100644
--- a/autogluon/task/object_detection/object_detection.py
+++ b/autogluon/task/object_detection/object_detection.py
@@ -7,7 +7,7 @@
 from .detector import Detector
 from .pipeline import train_object_detection
 from .utils import get_network
-from ..base import BaseTask
+from ..base import BaseTask, compile_scheduler_options
 from ...core.decorator import sample_config
 from ...core.space import Categorical
 from ...scheduler.resource import get_cpu_count, get_gpu_count
@@ -43,15 +43,16 @@ def fit(dataset='voc',
             ngpus_per_trial=1,
             hybridize=True,
             search_strategy='random',
-            search_options={},
+            search_options=None,
             time_limits=None,
             verbose=False,
             transfer='coco',
-            resume='',
+            resume_model='',
             checkpoint='checkpoint/exp1.ag',
+            resume=False,
             visualizer='none',
-            dist_ip_addrs=[],
-            grace_period=None,
+            dist_ip_addrs=None,
+            scheduler_options=None,
             auto_search=True,
             seed=223,
             data_shape=416,
@@ -101,7 +102,7 @@ def fit(dataset='voc',
             The final returned model may be fit to all of the data (after hyperparameters have been selected).
         batch_size : int
             How many images to group in each mini-batch during gradient computations in training.
-        epochs: int
+        epochs : int
             How many epochs to train the neural networks for at most.
         num_trials : int
             Maximal number of hyperparameter configurations to try out.
@@ -123,16 +124,20 @@ def fit(dataset='voc',
             `fit()` will stop training new models after this amount of time has elapsed (but models which have already started training will continue to completion). 
         verbose : bool
             Whether or not to print out intermediate information during training.
-        checkpoint: str
-            The path to local directory where trained models will be saved.
-        resume : str
+        resume_model : str
             Path to checkpoint file of existing model, from which model training should resume.
+        checkpoint : str or None
+            State of hyperparameter search is stored to this local file
+        resume : bool
+            If True, the hyperparameter search is started from state loaded
+            from checkpoint
         visualizer : str
             Describes method to visualize training progress during `fit()`. Options: ['mxboard', 'tensorboard', 'none']. 
         dist_ip_addrs : list
             List of IP addresses corresponding to remote workers, in order to leverage distributed computation.
-        grace_period : int
-            The grace period in early stopping when using Hyperband to tune hyperparameters. If None, this is set automatically.
+        scheduler_options : dict
+            Extra arguments passed to __init__ of scheduler, to configure the
+            orchestration of training jobs during hyperparameter-tuning.
         auto_search : bool
             If True, enables automatic suggestion of network types and hyper-parameter ranges adaptively based on provided dataset.
         seed : int
@@ -242,28 +247,16 @@ def fit(dataset='voc',
             mixup=mixup,
             no_mixup_epochs=no_mixup_epochs,
             label_smooth=label_smooth,
-            resume=resume,
+            resume=resume_model,
             syncbn=syncbn,
             reuse_pred_weights=reuse_pred_weights)
 
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'checkpoint': checkpoint,
-            'num_trials': num_trials,
-            'time_out': time_limits,
-            'resume': resume,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',
-            'reward_attr': 'map_reward',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-        }
-        if search_strategy == 'hyperband':
-            scheduler_options.update({
-                'searcher': 'random',
-                'max_t': epochs,
-                'grace_period': grace_period if grace_period else epochs // 4})
+        scheduler_options = compile_scheduler_options(
+            scheduler_options, search_strategy, search_options,
+            nthreads_per_trial, ngpus_per_trial, checkpoint, num_trials,
+            time_out=time_limits, resume=resume, visualizer=visualizer,
+            time_attr='epoch', reward_attr='map_reward',
+            dist_ip_addrs=dist_ip_addrs, epochs=epochs)
         results = BaseTask.run_fit(
             train_object_detection, search_strategy, scheduler_options)
         logger.info(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> finish model fitting")
diff --git a/autogluon/task/tabular_prediction/tabular_prediction.py b/autogluon/task/tabular_prediction/tabular_prediction.py
index 774eb51145c9..3448855e0662 100644
--- a/autogluon/task/tabular_prediction/tabular_prediction.py
+++ b/autogluon/task/tabular_prediction/tabular_prediction.py
@@ -8,7 +8,7 @@
 from .hyperparameter_configs import get_hyperparameter_config
 from .predictor import TabularPredictor
 from .presets_configs import set_presets, unpack
-from ..base import BaseTask
+from ..base import BaseTask, compile_scheduler_options
 from ..base.base_task import schedulers
 from ...utils import verbosity2loglevel
 from ...utils.tabular.features.auto_ml_feature_generator import AutoMLFeatureGenerator
@@ -59,13 +59,33 @@ def load(output_directory, verbosity=2):
 
     @staticmethod
     @unpack(set_presets)
-    def fit(train_data, label, tuning_data=None, time_limits=None, output_directory=None, presets=None,
-            problem_type=None, eval_metric=None, stopping_metric=None,
-            auto_stack=False, hyperparameter_tune=False, feature_prune=False,
-            holdout_frac=None, num_bagging_folds=0, num_bagging_sets=None, stack_ensemble_levels=0,
-            hyperparameters=None, num_trials=None, search_strategy='random', search_options=None,
-            nthreads_per_trial=None, ngpus_per_trial=None, dist_ip_addrs=None, visualizer='none',
-            verbosity=2, **kwargs):
+    def fit(train_data,
+            label,
+            tuning_data=None,
+            time_limits=None,
+            output_directory=None,
+            presets=None,
+            problem_type=None,
+            eval_metric=None,
+            stopping_metric=None,
+            auto_stack=False,
+            hyperparameter_tune=False,
+            feature_prune=False,
+            holdout_frac=None,
+            num_bagging_folds=0,
+            num_bagging_sets=None,
+            stack_ensemble_levels=0,
+            hyperparameters=None,
+            num_trials=None,
+            search_strategy='random',
+            search_options=None,
+            nthreads_per_trial=None,
+            ngpus_per_trial=None,
+            dist_ip_addrs=None,
+            visualizer='none',
+            verbosity=2,
+            scheduler_options=None,
+            **kwargs):
         """
         Fit models to predict a column of data table based on the other columns.
 
@@ -264,6 +284,10 @@ def fit(train_data, label, tuning_data=None, time_limits=None, output_directory=
             Higher levels correspond to more detailed print statements (you can set verbosity = 0 to suppress warnings).
             If using logging, you can alternatively control amount of information printed via `logger.setLevel(L)`,
             where `L` ranges from 0 to 50 (Note: higher values of `L` correspond to fewer print statements, opposite of verbosity levels)
+        scheduler_options : dict
+            Extra arguments passed to __init__ of scheduler, to configure the
+            orchestration of training jobs during hyperparameter-tuning. This
+            is ignored if hyperparameter_tune=False.
 
         Kwargs can include additional arguments for advanced users:
             feature_generator_type : `FeatureGenerator` class, default=`AutoMLFeatureGenerator`
@@ -497,31 +521,14 @@ def fit(train_data, label, tuning_data=None, time_limits=None, output_directory=
         stopping_metric = get_metric(stopping_metric, problem_type, 'stopping_metric')
 
         # All models use the same scheduler:
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'num_trials': num_trials,
-            'time_out': time_limits_hpo,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',  # For tree ensembles, each new tree (ie. boosting round) is considered one epoch
-            'reward_attr': 'validation_performance',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-        }
-        if isinstance(search_strategy, str):
-            scheduler_cls = schedulers[search_strategy.lower()]
-            # This is a fix for now. But we need to separate between scheduler
-            # (mainly FIFO and Hyperband) and searcher. Currently, most searchers
-            # only work with FIFO, and Hyperband works only with random searcher,
-            # but this will be different in the future.
-            if search_strategy == 'hyperband':
-                # Currently, HyperbandScheduler only supports random searcher
-                scheduler_options['searcher'] = 'random'
-        else:
-            # TODO: Check that search_strategy is a subclass of TaskScheduler
-            assert callable(search_strategy)
-            scheduler_cls = search_strategy
-            scheduler_options['searcher'] = 'random'
+        scheduler_options = compile_scheduler_options(
+            scheduler_options, search_strategy, search_options,
+            nthreads_per_trial, ngpus_per_trial, checkpoint=None,
+            num_trials=num_trials, time_out=time_limits_hpo, resume=False,
+            visualizer=visualizer, time_attr='epoch',
+            reward_attr='validation_performance',
+            dist_ip_addrs=dist_ip_addrs)
+        scheduler_cls = schedulers[search_strategy.lower()]
         scheduler_options = (scheduler_cls, scheduler_options)  # wrap into tuple
         learner = Learner(path_context=output_directory, label=label, problem_type=problem_type, objective_func=eval_metric, stopping_metric=stopping_metric,
                           id_columns=id_columns, feature_generator=feature_generator, trainer_type=trainer_type,
diff --git a/autogluon/task/text_classification/text_classification.py b/autogluon/task/text_classification/text_classification.py
index 505b4a2385d3..5de60bec195e 100644
--- a/autogluon/task/text_classification/text_classification.py
+++ b/autogluon/task/text_classification/text_classification.py
@@ -5,7 +5,7 @@
 
 from ...core import *
 from ...scheduler.resource import get_cpu_count, get_gpu_count
-from ..base import BaseTask
+from ..base import BaseTask, compile_scheduler_options
 from ...utils import update_params
 
 from .network import get_network
@@ -49,14 +49,14 @@ def fit(dataset='SST',
             ngpus_per_trial=1,
             hybridize=True,
             search_strategy='random',
-            search_options={},
+            search_options=None,
             time_limits=None,
             resume=False,
             checkpoint='checkpoint/exp1.ag',
             visualizer='none',
             num_trials=2,
-            dist_ip_addrs=[],
-            grace_period=None,
+            dist_ip_addrs=None,
+            scheduler_options=None,
             auto_search=True,
             verbose=False,
             **kwargs):
@@ -115,16 +115,18 @@ def fit(dataset='SST',
             `fit()` will stop training new models after this amount of time has elapsed (but models which have already started training will continue to completion). 
         verbose : bool
             Whether or not to print out intermediate information during training.
-        checkpoint: str
-            The path to local directory where trained models will be saved.
-        resume : str
-            Path to checkpoint file of existing model, from which model training should resume.
+        checkpoint : str or None
+            State of hyperparameter search is stored to this local file
+        resume : bool
+            If True, the hyperparameter search is started from state loaded
+            from checkpoint
         visualizer : str
             Describes method to visualize training progress during `fit()`. Options: ['mxboard', 'tensorboard', 'none']. 
         dist_ip_addrs : list
             List of IP addresses corresponding to remote workers, in order to leverage distributed computation.
-        grace_period : int
-            The grace period in early stopping when using Hyperband to tune hyperparameters. If None, this is set automatically.
+        scheduler_options : dict
+            Extra arguments passed to __init__ of scheduler, to configure the
+            orchestration of training jobs during hyperparameter-tuning.
         auto_search : bool
             If True, enables automatic suggestion of network types and hyper-parameter ranges adaptively based on provided dataset.
         
@@ -176,24 +178,12 @@ def fit(dataset='SST',
             final_fit=False,
             **kwargs)
 
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'checkpoint': checkpoint,
-            'num_trials': num_trials,
-            'time_out': time_limits,
-            'resume': resume,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',
-            'reward_attr': 'accuracy',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-        }
-        if search_strategy == 'hyperband':
-            scheduler_options.update({
-                'searcher': 'random',
-                'max_t': epochs,
-                'grace_period': grace_period if grace_period else epochs//4})
+        scheduler_options = compile_scheduler_options(
+            scheduler_options, search_strategy, search_options,
+            nthreads_per_trial, ngpus_per_trial, checkpoint, num_trials,
+            time_out=time_limits, resume=resume, visualizer=visualizer,
+            time_attr='epoch', reward_attr='accuracy',
+            dist_ip_addrs=dist_ip_addrs, epochs=epochs)
         results = BaseTask.run_fit(
             train_text_classification, search_strategy, scheduler_options)
         args = sample_config(train_text_classification.args, results['best_config'])
diff --git a/docs/tutorials/image_classification/hpo.md b/docs/tutorials/image_classification/hpo.md
index 6e899c416521..cabf0587b253 100644
--- a/docs/tutorials/image_classification/hpo.md
+++ b/docs/tutorials/image_classification/hpo.md
@@ -118,23 +118,32 @@ print('Top-1 test acc: %.3f' % test_acc)
 
 ### Hyperband Early Stopping
 
-AutoGluon currently supports scheduling trials in serial order and with early stopping
-(e.g., if the performance of the model early within training already looks bad, the trial may be terminated early to free up resources).
-Here is an example of using an early stopping scheduler :class:`autogluon.scheduler.HyperbandScheduler`:
+AutoGluon currently supports scheduling trials in serial order and with early
+stopping (e.g., if the performance of the model early within training already
+looks bad, the trial may be terminated early to free up resources).
+Here is an example of using an early stopping scheduler
+:class:`autogluon.scheduler.HyperbandScheduler`. `scheduler_options` is used
+to configure the scheduler. In this example, we run Hyperband with a single
+bracket, and stop/go decisions are made after 1 and 2 epochs (`grace_period`,
+`grace_period * reduction_factor`):
 
 ```{.python .input}
 search_strategy = 'hyperband'
+scheduler_options = {
+    'grace_period': 1,
+    'reduction_factor': 2,
+    'brackets': 1}
 
 classifier = task.fit(dataset,
                       net=net,
                       optimizer=optimizer,
                       search_strategy=search_strategy,
-                      epochs=epochs,
+                      epochs=4,
                       num_trials=2,
                       verbose=False,
                       plot_results=True,
                       ngpus_per_trial=1,
-                      grace_period=1)
+                      scheduler_options=scheduler_options)
 
 print('Top-1 val acc: %.3f' % classifier.results[classifier.results['reward_attr']])
 ```


Code review comments:
diff_hunk: @@ -59,13 +59,33 @@ def load(output_directory, verbosity=2):
 
     @staticmethod
     @unpack(set_presets)
-    def fit(train_data, label, tuning_data=None, time_limits=None, output_directory=None, presets=None,
-            problem_type=None, eval_metric=None, stopping_metric=None,
-            auto_stack=False, hyperparameter_tune=False, feature_prune=False,
-            holdout_frac=None, num_bagging_folds=0, num_bagging_sets=None, stack_ensemble_levels=0,
-            hyperparameters=None, num_trials=None, search_strategy='random', search_options=None,
-            nthreads_per_trial=None, ngpus_per_trial=None, dist_ip_addrs=None, visualizer='none',
-            verbosity=2, **kwargs):
+    def fit(train_data,
+            label,
+            tuning_data=None,
+            time_limits=None,
+            output_directory=None,
+            presets=None,
+            problem_type=None,
+            eval_metric=None,
+            stopping_metric=None,
+            auto_stack=False,
+            hyperparameter_tune=False,
+            feature_prune=False,
+            holdout_frac=None,
+            num_bagging_folds=0,
+            num_bagging_sets=None,
+            stack_ensemble_levels=0,
+            hyperparameters=None,
+            num_trials=None,
+            search_strategy='random',
+            search_options=None,
+            nthreads_per_trial=None,
+            ngpus_per_trial=None,
+            dist_ip_addrs=None,
+            visualizer='none',
+            verbosity=2,
+            scheduler_options=None,
path: autogluon/task/tabular_prediction/tabular_prediction.py
text: @user1:
prefer to move scheduler_options to be directly before search_strategy so it is near similar arguments.

@author:
OK

diff_hunk: @@ -268,30 +270,15 @@ def fit(dataset,
             lr_config=lr_config
         )
 
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'checkpoint': checkpoint,
-            'num_trials': num_trials,
-            'time_out': time_limits,
-            'resume': resume,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',
-            'reward_attr': 'classification_reward',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-            'plot_results': plot_results
-        }
-        if search_strategy == 'hyperband':
-            scheduler_options.update(
-                {
-                    'searcher': 'random',
-                    'max_t': epochs,
-                    'grace_period': grace_period if grace_period else epochs // 4
-                }
-            )
-
-        results = BaseTask.run_fit(train_image_classification, search_strategy, scheduler_options)
+        scheduler_options = compile_scheduler_options(
+            scheduler_options, search_strategy, search_options,
+            nthreads_per_trial, ngpus_per_trial, checkpoint, num_trials,
+            time_out=time_limits, resume=resume, visualizer=visualizer,
+            time_attr='epoch', reward_attr='classification_reward',
+            dist_ip_addrs=dist_ip_addrs, epochs=epochs)
+        results = BaseTask.run_fit(
path: autogluon/task/image_classification/image_classification.py
text: @user1:
super().run_fit is cleaner

@author:
OK

@author:
Sorry this does not work, these are class and static methods. Reverting to how it was before.

diff_hunk: @@ -268,30 +270,15 @@ def fit(dataset,
             lr_config=lr_config
         )
 
-        scheduler_options = {
-            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
-            'checkpoint': checkpoint,
-            'num_trials': num_trials,
-            'time_out': time_limits,
-            'resume': resume,
-            'visualizer': visualizer,
-            'time_attr': 'epoch',
-            'reward_attr': 'classification_reward',
-            'dist_ip_addrs': dist_ip_addrs,
-            'searcher': search_strategy,
-            'search_options': search_options,
-            'plot_results': plot_results
-        }
-        if search_strategy == 'hyperband':
-            scheduler_options.update(
-                {
-                    'searcher': 'random',
-                    'max_t': epochs,
-                    'grace_period': grace_period if grace_period else epochs // 4
-                }
-            )
-
-        results = BaseTask.run_fit(train_image_classification, search_strategy, scheduler_options)
+        scheduler_options = compile_scheduler_options(
path: autogluon/task/image_classification/image_classification.py
text: @user1:
keyword arguments should be used with this many inputs, it is very hard to notice if an arg is in the wrong position currently.

@author:
OK

diff_hunk: @@ -43,15 +43,16 @@ def fit(dataset='voc',
             ngpus_per_trial=1,
             hybridize=True,
             search_strategy='random',
-            search_options={},
+            search_options=None,
             time_limits=None,
             verbose=False,
             transfer='coco',
-            resume='',
+            resume_model='',
path: autogluon/task/object_detection/object_detection.py
text: @user1:
any reason to change the name? This will cause backward incompatibility even though we are in early stages

@author:
Yes, but it was a bug before. You passed resume (a str) to the scheduler, which needs resume to be a bool. This fixes that bug.


@author:
The schedulers currently can only resume from the file whose name is passed in checkpoint. This is a little funny, because that same file would then be overwritten by subsequent checkpoints. But it is simple and works, and I'd not change it for now.

@user2:
The defect/bug is unrelated to the name change, we should keep it as resume, to be consistent with image_classification and avoid backwards breaking change.

@user2:
Why not have the boolean resume you require specified in scheduler_options instead?

@author:
Well, the thing is for all other task's, resume is bool and is passed to the scheduler.

My guess is this code would have never worked, certainly not as intended. In all other task's, you checkpoint the searcher state and resume from there. Here, somehow, you resume from an NN model.

I don't really understand this code. Maybe the best would be to pass resume=False to the scheduler? I can also pass (len(resume) > 0) to the resume (bool) argument of the scheduler, maybe that was intended here? Let me try that.

But I think whoever wrote this task, needs to say what is intended here, because the behaviour differs from the other ones, and that is confusing to the user.


@author:
OK, I am now passing len(resume) > 0 to resume (bool) of the scheduler. Note what this means: If resume is a non-empty file name, the train_fn will somehow load model pars from that file, and then the scheduler/searcher will resume from the checkpoint loaded from the file name `checkpoint`.

@user2:
@user3 Could you weigh in here?