xgboost-optuna.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7653595e-196c-4d7f-b2c3-5e6442599404",
   "metadata": {},
   "source": [
    "<img src=\"https://raw.githubusercontent.com/optuna/optuna/master/docs/image/optuna-logo.png\"\n",
    "     align=\"right\"\n",
    "     width=\"30%\"/>\n",
    "\n",
    "# Hyper-Parameter Optimization with Optuna\n",
    "\n",
    "This trains an XGBoost model and does hyperparameter optimization using Optuna to search and scikit-learn for cross validation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8e7cefa-5e11-4cd5-908f-5b3b054ed844",
   "metadata": {},
   "source": [
    "## Launch Cluster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2386ff41-7fc3-435a-bc8c-a38771073174",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import coiled\n",
    "\n",
    "cluster = coiled.Cluster(\n",
    "    n_workers=20,\n",
    "    name=\"hpo\",\n",
    ")\n",
    "\n",
    "client = cluster.get_client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59064aae-8d76-4296-9404-b00619c5036c",
   "metadata": {},
   "source": [
    "## Optuna Study\n",
    "\n",
    "We use the Dask scheduler to track work between the different experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae2c1a8-14aa-4ebd-86a7-4cf58373d83c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optuna\n",
    "from optuna.integration.dask import DaskStorage\n",
    "\n",
    "study = optuna.create_study(\n",
    "    direction=\"maximize\",\n",
    "    storage=DaskStorage(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c72a7cb-dbbe-4630-8a4c-a59ddb4bfd30",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Objective function\n",
    "\n",
    "We ...\n",
    "\n",
    "-   Load data\n",
    "-   Get recommended hyper-parameters from Optuna\n",
    "-   Train\n",
    "-   Report Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41eaa3e7-9b99-442d-8672-32b085909625",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import load_breast_cancer\n",
    "from sklearn.model_selection import cross_val_score, KFold\n",
    "import xgboost as xgb\n",
    "from optuna.samplers import RandomSampler\n",
    "\n",
    "def objective(trial):\n",
    "    X, y = load_breast_cancer(return_X_y=True)\n",
    "    params = {\n",
    "        \"n_estimators\": 10,\n",
    "        \"verbosity\": 0,\n",
    "        \"lambda\": trial.suggest_float(\"lambda\", 1e-8, 100.0, log=True),\n",
    "        \"alpha\": trial.suggest_float(\"alpha\", 1e-8, 100.0, log=True),\n",
    "        \"colsample_bytree\": trial.suggest_float(\"colsample_bytree\", 0.2, 1.0),\n",
    "        \"max_depth\": trial.suggest_int(\"max_depth\", 2, 10, step=1),\n",
    "        \"min_child_weight\": trial.suggest_float(\"min_child_weight\", 1e-8, 100, log=True),\n",
    "        \"learning_rate\": trial.suggest_float(\"learning_rate\", 1e-8, 1.0, log=True),\n",
    "        \"gamma\": trial.suggest_float(\"gamma\", 1e-8, 1.0, log=True),\n",
    "        \"grow_policy\": \"depthwise\",\n",
    "        \"eval_metric\": \"logloss\"\n",
    "    }\n",
    "    clf = xgb.XGBClassifier(**params)\n",
    "    fold = KFold(n_splits=5, shuffle=True, random_state=0)\n",
    "    score = cross_val_score(clf, X, y, cv=fold, scoring='neg_log_loss')\n",
    "    return score.mean()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0519f886-c142-4c2c-954c-9eb2ba87bfcb",
   "metadata": {},
   "source": [
    "## Execute at Scale\n",
    "\n",
    "All of the actual coordination happens within Optuna.  It's Dask's job just to provide a lot of firepower, which we do by submitting the optimize method many times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a977efd-6301-49b8-8abc-2f5177c03a8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import wait\n",
    "\n",
    "futures = [\n",
    "    client.submit(study.optimize, objective, n_trials=1, pure=False)\n",
    "    for _ in range(500)\n",
    "]\n",
    "\n",
    "_ = wait(futures)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba99faa0-6cf5-44ec-b359-0d9007d622cd",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396e798c-8c4c-47ef-b5bf-73af7767589b",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4098f5-ea7c-47a6-81ba-3b27e1078b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ebfcd6-d9e9-48f0-b256-25072f07ab48",
   "metadata": {},
   "outputs": [],
   "source": [
    "from optuna.visualization.matplotlib import plot_optimization_history, plot_param_importances\n",
    "\n",
    "plot_optimization_history(study)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9940ad3a-e757-4703-b641-b5ebefdc8304",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_param_importances(study)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a41ad2a6-cdfe-48fc-bca8-b0b801977380",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "\n",
    "This cost us about $0.08"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3610787e-0710-406e-a3dd-588b42b42ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:coiled]",
   "language": "python",
   "name": "conda-env-coiled-py"
  },
  "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}