optuna-xgboost/optuna-xgboost.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameter optimization with Optuna and Dask\n",
    "\n",
    "[Optuna](https://optuna.org/) is a popular Python library for hyperparameter optimization. This example walks through a workload using Optuna to optimize an [XGBoost](https://xgboost.readthedocs.io/en/latest/) classification model. and then how to scale the same workload using Dask and Coiled.\n",
    "\n",
    "## Optuna in a nutshell\n",
    "\n",
    "Optuna has three primary concepts:\n",
    "\n",
    "- Objective function: This is some function that depends on the hyperparameters in your model that you would like to optimize. For example, it’s common to maximum a classification model’s prediction accuracy (i.e. the objective function would be the accuracy score).\n",
    "\n",
    "- Optimization trial: A trial is a single evaluation of the objective function with a given set of hyperparameters.\n",
    "\n",
    "- Optimization study: A study is a collection of optimization trials where each trial uses hyperparameters sampled from a set of allowed values.\n",
    "\n",
    "The set of hyperparameters for the trial which gives the optimal value for the objective function are chosen as the best set of hyperparameters.\n",
    "\n",
    "\n",
    "## Scaling Optuna with Dask\n",
    "\n",
    "Below is a snippet which uses Optuna to optimize several hyperparameters for an XGBoost classifier trained on the [breast cancer dataset](https://scikit-learn.org/stable/datasets/index.html#breast-cancer-wisconsin-diagnostic-dataset). We also use [Dask-Optuna](https://jrbourbeau.github.io/dask-optuna/) and [Joblib](https://joblib.readthedocs.io/en/latest/) to run Optuna trials in parallel on a Coiled cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sklearn.datasets\n",
    "import sklearn.metrics\n",
    "from sklearn.model_selection import train_test_split\n",
    "import xgboost as xgb\n",
    "\n",
    "def objective(trial):\n",
    "    # Load our dataset\n",
    "    X, y = sklearn.datasets.load_breast_cancer(return_X_y=True)\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25)\n",
    "    dtrain = xgb.DMatrix(X_train, label=y_train)\n",
    "    dtest = xgb.DMatrix(X_test, label=y_test)\n",
    "\n",
    "    # Get set of hyperparameters\n",
    "    param = {\n",
    "        \"silent\": 1,\n",
    "        \"objective\": \"binary:logistic\",\n",
    "        \"booster\": trial.suggest_categorical(\"booster\", [\"gbtree\", \"dart\"]),\n",
    "        \"lambda\": trial.suggest_float(\"lambda\", 1e-8, 1.0, log=True),\n",
    "        \"alpha\": trial.suggest_float(\"alpha\", 1e-8, 1.0, log=True),\n",
    "        \"max_depth\": trial.suggest_int(\"max_depth\", 1, 9),\n",
    "        \"eta\": trial.suggest_float(\"eta\", 1e-8, 1.0, log=True),\n",
    "        \"gamma\": trial.suggest_float(\"gamma\", 1e-8, 1.0, log=True),\n",
    "        \"grow_policy\": trial.suggest_categorical(\"grow_policy\", [\"depthwise\", \"lossguide\"]),\n",
    "    }\n",
    "\n",
    "    # Train XGBoost model\n",
    "    bst = xgb.train(param, dtrain)\n",
    "    preds = bst.predict(dtest)\n",
    "\n",
    "    # Compute and return model accuracy\n",
    "    pred_labels = np.rint(preds)\n",
    "    accuracy = sklearn.metrics.accuracy_score(y_test, pred_labels)\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to Coiled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import coiled\n",
    "from dask.distributed import Client\n",
    "\n",
    "cluster = coiled.Cluster(\n",
    "    n_workers=10, \n",
    "    software=\"examples/optuna-xgboost\"\n",
    ")\n",
    "client = Client(cluster)\n",
    "client.wait_for_workers(10)\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train with Optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import optuna\n",
    "import dask_optuna\n",
    "import joblib\n",
    "\n",
    "# Create Dask-compatible Optuna storage class\n",
    "storage = dask_optuna.DaskStorage()\n",
    "\n",
    "# Run 500 optimizations trial on our cluster\n",
    "study = optuna.create_study(direction=\"maximize\", storage=storage)\n",
    "with joblib.parallel_backend(\"dask\"):\n",
    "    study.optimize(objective, n_trials=500, n_jobs=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And with that, you’re able to run distributed hyperparameter optimizations using Optuna, Dask, and Coiled!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}