diff --git a/doc/source/conf.py b/doc/source/conf.py index 43bd34876..cd0871ea7 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -65,6 +65,7 @@ MOCK_MODULES = [ "torch.utils.data", "torch.utils.data.distributed", "wandb", + "xgboost", "zoopt", ] import scipy.stats diff --git a/doc/source/tune/_tutorials/tune-xgboost.rst b/doc/source/tune/_tutorials/tune-xgboost.rst index f22641c27..d8ac98668 100644 --- a/doc/source/tune/_tutorials/tune-xgboost.rst +++ b/doc/source/tune/_tutorials/tune-xgboost.rst @@ -382,12 +382,19 @@ There are more parameters, which are explained in the :ref:`documentation `. Lastly, we have to report the loss metric to Tune. We do this with a ``Callback`` that -XGBoost accepts and calls after each training iteration. We also tell XGBoost which -loss metrics to calculate in the ``eval_metric`` parameter. These are the metrics -available in ``env.evaluation_result_list`` below. +XGBoost accepts and calls after each evaluation round. Ray Tune comes +with :ref:`two XGBoost callbacks ` +we can use for this. The ``TuneReportCallback`` just reports the evaluation +metrics back to Tune. The ``TuneReportCheckpointCallback`` would also save +checkpoints after each evaluation round. We will just use the former in this +example. + +We also tell XGBoost which loss metrics to calculate in the ``eval_metric`` +parameter in the config. These parameters are then reported to Tune +via the callback. .. code-block:: python - :emphasize-lines: 11,12,13,26,42,44,45,46,47,48,49 + :emphasize-lines: 9,26,42,44-49 import numpy as np import sklearn.datasets @@ -397,12 +404,7 @@ available in ``env.evaluation_result_list`` below. import xgboost as xgb from ray import tune - - - def XGBCallback(env): - # After every training iteration, report loss to Tune - tune.report(**dict(env.evaluation_result_list)) - + from ray.tune.integration.xgboost import TuneReportCallback def train_breast_cancer(config): # Load dataset @@ -414,7 +416,12 @@ available in ``env.evaluation_result_list`` below. train_set = xgb.DMatrix(train_x, label=train_y) test_set = xgb.DMatrix(test_x, label=test_y) # Train the classifier - bst = xgb.train(config, train_set, evals=[(test_set, "eval")], verbose_eval=False, callbacks=[XGBCallback]) + bst = xgb.train( + config, + train_set, + evals=[(test_set, "eval")], + verbose_eval=False, + callbacks=[TuneReportCallback()]) # Predict labels for the test set preds = bst.predict(test_set) pred_labels = np.rint(preds) diff --git a/doc/source/tune/api_docs/integration.rst b/doc/source/tune/api_docs/integration.rst index 7ccb62040..61537af49 100644 --- a/doc/source/tune/api_docs/integration.rst +++ b/doc/source/tune/api_docs/integration.rst @@ -43,4 +43,14 @@ Weights and Biases (tune.integration.wandb) .. autoclass:: ray.tune.integration.wandb.WandbLogger -.. autofunction:: ray.tune.integration.wandb.wandb_mixin \ No newline at end of file +.. autofunction:: ray.tune.integration.wandb.wandb_mixin + + +.. _tune-integration-xgboost: + +XGBoost (tune.integration.xgboost) +---------------------------------- + +.. autoclass:: ray.tune.integration.xgboost.TuneReportCallback + +.. autoclass:: ray.tune.integration.xgboost.TuneReportCheckpointCallback diff --git a/python/ray/tune/examples/xgboost_example.py b/python/ray/tune/examples/xgboost_example.py index 10ad88e6a..73303e221 100644 --- a/python/ray/tune/examples/xgboost_example.py +++ b/python/ray/tune/examples/xgboost_example.py @@ -6,11 +6,7 @@ from sklearn.model_selection import train_test_split import xgboost as xgb from ray import tune - - -def XGBCallback(env): - # After every training iteration, report loss to Tune - tune.report(**dict(env.evaluation_result_list)) +from ray.tune.integration.xgboost import TuneReportCallback def train_breast_cancer(config): @@ -28,7 +24,7 @@ def train_breast_cancer(config): train_set, evals=[(test_set, "eval")], verbose_eval=False, - callbacks=[XGBCallback]) + callbacks=[TuneReportCallback()]) # Predict labels for the test set preds = bst.predict(test_set) pred_labels = np.rint(preds) diff --git a/python/ray/tune/integration/xgboost.py b/python/ray/tune/integration/xgboost.py new file mode 100644 index 000000000..730883c25 --- /dev/null +++ b/python/ray/tune/integration/xgboost.py @@ -0,0 +1,144 @@ +from typing import Dict, List, Union +from ray import tune + +import os + + +class TuneCallback: + """Base class for Tune's XGBoost callbacks.""" + pass + + def __call__(self, env): + raise NotImplementedError + + +class TuneReportCallback(TuneCallback): + """XGBoost to Ray Tune reporting callback + + Reports metrics to Ray Tune. + + Args: + metrics (str|list|dict): Metrics to report to Tune. If this is a list, + each item describes the metric key reported to XGBoost, + and it will reported under the same name to Tune. If this is a + dict, each key will be the name reported to Tune and the respective + value will be the metric key reported to XGBoost. If this is None, + all metrics will be reported to Tune under their default names as + obtained from XGBoost. + + Example: + + .. code-block:: python + + import xgboost + from ray.tune.integration.xgboost import TuneReportCallback + + config = { + # ... + "eval_metric": ["auc", "logloss"] + } + + # Report only log loss to Tune after each validation epoch: + bst = xgb.train( + config, + train_set, + evals=[(test_set, "eval")], + verbose_eval=False, + callbacks=[TuneReportCallback({"loss": "eval-logloss"})]) + + """ + + def __init__(self, + metrics: Union[None, str, List[str], Dict[str, str]] = None): + if isinstance(metrics, str): + metrics = [metrics] + self._metrics = metrics + + def __call__(self, env): + result_dict = dict(env.evaluation_result_list) + if not self._metrics: + report_dict = result_dict + else: + report_dict = {} + for key in self._metrics: + if isinstance(self._metrics, dict): + metric = self._metrics[key] + else: + metric = key + report_dict[key] = result_dict[metric] + tune.report(**report_dict) + + +class _TuneCheckpointCallback(TuneCallback): + """XGBoost checkpoint callback + + Saves checkpoints after each validation step. + + Checkpoint are currently not registered if no ``tune.report()`` call + is made afterwards. Consider using ``TuneReportCheckpointCallback`` + instead. + + Args: + filename (str): Filename of the checkpoint within the checkpoint + directory. Defaults to "checkpoint". + + """ + + def __init__(self, filename: str = "checkpoint"): + self._filename = filename + + def __call__(self, env): + with tune.checkpoint_dir(step=env.iteration) as checkpoint_dir: + env.model.save_model(os.path.join(checkpoint_dir, self._filename)) + + +class TuneReportCheckpointCallback(TuneCallback): + """XGBoost report and checkpoint callback + + Saves checkpoints after each validation step. Also reports metrics to Tune, + which is needed for checkpoint registration. + + Args: + metrics (str|list|dict): Metrics to report to Tune. If this is a list, + each item describes the metric key reported to XGBoost, + and it will reported under the same name to Tune. If this is a + dict, each key will be the name reported to Tune and the respective + value will be the metric key reported to XGBoost. + filename (str): Filename of the checkpoint within the checkpoint + directory. Defaults to "checkpoint". If this is None, + all metrics will be reported to Tune under their default names as + obtained from XGBoost. + + Example: + + .. code-block:: python + + import xgboost + from ray.tune.integration.xgboost import TuneReportCheckpointCallback + + config = { + # ... + "eval_metric": ["auc", "logloss"] + } + + # Report only log loss to Tune after each validation epoch. + # Save model as `xgboost.mdl`. + bst = xgb.train( + config, + train_set, + evals=[(test_set, "eval")], + verbose_eval=False, + callbacks=[TuneReportCheckpointCallback( + {"loss": "eval-logloss"}, "xgboost.mdl)]) + + """ + + def __init__(self, + metrics: Union[None, str, List[str], Dict[str, str]] = None, + filename: str = "checkpoint"): + self._checkpoint = _TuneCheckpointCallback(filename) + self._report = TuneReportCallback(metrics) + + def __call__(self, env): + self._checkpoint(env) + self._report(env)