Modeling 2 - HPO with Optuna & XGBoost.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0652f264-b2c3-43b2-9310-7c00e50e7cdc",
   "metadata": {},
   "source": [
    "# HPO of XGBoost with Optuna and Dask\n",
    "\n",
    "In the last notebook we trained a single XGBoost model with fixed hyper-parameters.  Those hyper-parameters were likely wrong.  In this notebook we use [Optuna](https://optuna.org/) to perform hyper-parameter-optimization (HPO) over a space of parameters to find the best model.  This involves training the same dataset repeatedly.\n",
    "\n",
    "The primary difference between this notebook and the previous one is the creation of an `objective` function, and the use of Optuna studies.  At the end we look at the progress during HPO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b929987-5089-4c13-be35-c1812d65fbec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from collections.abc import Iterator\n",
    "from datetime import datetime\n",
    "\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import coiled\n",
    "import distributed\n",
    "import joblib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import optuna\n",
    "import xgboost\n",
    "from dask_ml.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7179b003-e075-4be5-bde1-4ce90b7d726d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Location of feature table\n",
    "FILEPATH = \"s3://coiled-datasets/prefect-dask/nyc-uber-lyft/feature_table.parquet\"\n",
    "\n",
    "# Number of converging serial trials to run in each job\n",
    "N_TRIALS = 5\n",
    "# Number of folds in each trial. This also determines the train/test split\n",
    "# (e.g. N_FOLDS=5 -> train=4/5 of the total data, test=1/5)\n",
    "N_FOLDS = 5\n",
    "\n",
    "# Dask worker instance type and number\n",
    "WORKER_INSTANCE_TYPE = \"r6i.large\"\n",
    "N_WORKERS = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ee232d",
   "metadata": {},
   "source": [
    "### Start coiled cluster\n",
    "\n",
    "**Note:** at the moment of writing, the size of the input dataset dictates the amount of RAM that your cluster must mount. With this dataset of ~55 GiB, you need any combination that will result in 640 GiB cluster memory, so:\n",
    "- 80x m6i.large, or\n",
    "- 40x r6i.large, or\n",
    "- 40x m6i.xlarge,\n",
    "- 20x r6i.xlarge,\n",
    "\n",
    "and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5686386-60ea-4277-9116-daabd3ca5138",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = coiled.Cluster(\n",
    "    worker_vm_types=[WORKER_INSTANCE_TYPE],\n",
    "    scheduler_vm_types=[\"m6i.large\"],\n",
    "    package_sync=True,  # align remote packages to local ones\n",
    "    n_workers=N_WORKERS,\n",
    "    backend_options={\n",
    "        \"region\": \"us-east-2\",\n",
    "        \"multizone\": True,\n",
    "    },\n",
    "    scheduler_options={\"idle_timeout\": \"15 minutes\"},\n",
    ")\n",
    "client = distributed.Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77d715d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load feature table generated by Feature Engineering.ipynb\n",
    "ddf = dd.read_parquet(FILEPATH)\n",
    "\n",
    "# Reduce dataset size. Uncomment to speed up the exercise.\n",
    "# ddf = ddf.partitions[:20]\n",
    "\n",
    "# Under the hood, XGBoost converts floats to `float32`.\n",
    "# Let's do it only once here.\n",
    "float_cols = ddf.select_dtypes(include=\"float\").columns.tolist()\n",
    "ddf = ddf.astype({c: np.float32 for c in float_cols})\n",
    "\n",
    "# We need the categories to be known\n",
    "categorical_vars = ddf.select_dtypes(include=\"category\").columns.tolist()\n",
    "\n",
    "# categorize() reads the whole input and then discards it.\n",
    "# Let's read from disk only once.\n",
    "ddf = ddf.persist()\n",
    "ddf = ddf.categorize(columns=categorical_vars)\n",
    "\n",
    "# We will need to access this multiple times. Let's persist it.\n",
    "ddf = ddf.persist()\n",
    "\n",
    "ddf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dd461c3-e551-420a-a3e1-1ca2515bb613",
   "metadata": {},
   "source": [
    "### Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "161a1197",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we subset data for cross-validation\n",
    "def make_cv_splits(\n",
    "    n_folds: int = N_FOLDS,\n",
    ") -> Iterator[tuple[dd.DataFrame, dd.DataFrame]]:\n",
    "    frac = [1 / n_folds] * n_folds\n",
    "    splits = ddf.random_split(frac, shuffle=True)\n",
    "    for i in range(n_folds):\n",
    "        train = [splits[j] for j in range(n_folds) if j != i]\n",
    "        test = splits[i]\n",
    "        yield dd.concat(train), test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132568f4-ac1b-43ff-85fd-e5a5324a38a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_model(**study_params):\n",
    "    scores = []\n",
    "\n",
    "    for i, (train, test) in enumerate(make_cv_splits()):\n",
    "        print(f\"Training/Test split #{i}\")\n",
    "        y_train = train[\"trip_time\"]\n",
    "        X_train = train.drop(columns=[\"trip_time\"])\n",
    "        y_test = test[\"trip_time\"]\n",
    "        X_test = test.drop(columns=[\"trip_time\"])\n",
    "\n",
    "        print(\"Building DMatrix...\")\n",
    "        d_train = xgboost.dask.DaskDMatrix(\n",
    "            None, X_train, y_train, enable_categorical=True\n",
    "        )\n",
    "\n",
    "        print(\"Training model...\")\n",
    "        model = xgboost.dask.train(\n",
    "            None,\n",
    "            {\"tree_method\": \"hist\", **study_params},\n",
    "            d_train,\n",
    "            num_boost_round=4,\n",
    "            evals=[(d_train, \"train\")],\n",
    "        )\n",
    "\n",
    "        print(\"Running model on test data...\")\n",
    "        predictions = xgboost.dask.predict(None, model, X_test)\n",
    "\n",
    "        print(\"Measuring accuracy of model vs. ground truth...\")\n",
    "        score = mean_squared_error(\n",
    "            y_test.to_dask_array(),\n",
    "            predictions.to_dask_array(),\n",
    "            squared=False,\n",
    "            compute=False,\n",
    "        )\n",
    "        # Compute predictions and mean squared error for this iteration\n",
    "        # while we start the next one\n",
    "        scores.append(score.reshape(1).persist())\n",
    "        del d_train, model, predictions, score\n",
    "        print(\"-\" * 80)\n",
    "\n",
    "    scores = da.concatenate(scores).compute()\n",
    "    print(f\"RSME={scores.mean()} +/- {scores.std()}\")\n",
    "    return scores.mean()\n",
    "    print(\"-\" * 80)\n",
    "\n",
    "\n",
    "def objective(trial):\n",
    "    params = {\n",
    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 75, 125),\n",
    "        \"learning_rate\": trial.suggest_float(\"learning_rate\", 0.5, 0.7),\n",
    "        \"colsample_bytree\": trial.suggest_float(\"colsample_bytree\", 0.5, 1),\n",
    "        \"colsample_bynode\": trial.suggest_float(\"colsample_bynode\", 0.5, 1),\n",
    "        \"colsample_bylevel\": trial.suggest_float(\"colsample_bylevel\", 0.5, 1),\n",
    "        \"reg_lambda\": trial.suggest_float(\"reg_lambda\", 0, 1),\n",
    "        \"max_depth\": trial.suggest_int(\"max_depth\", 1, 6),\n",
    "        \"max_leaves\": trial.suggest_int(\"max_leaves\", 0, 2),\n",
    "        \"max_cat_to_onehot\": trial.suggest_int(\"max_cat_to_onehot\", 1, 10),\n",
    "    }\n",
    "    print(f\"Training model (trial #{trial.number})\")\n",
    "    for k, v in params.items():\n",
    "        print(f\"  {k}={v}\")\n",
    "    return train_model(**params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d70a951",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a single study and run some trials\n",
    "start = datetime.now()\n",
    "study = optuna.create_study(study_name=\"nyc-travel-time-model\")\n",
    "study.optimize(objective, n_trials=N_TRIALS)\n",
    "\n",
    "print(f\"Total time:  {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba79e3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ede5a2bf-c88c-4266-9738-73cb6fcb9d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22a6da34-582c-4e2d-9b58-89a5369d7536",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4928342-4749-4e29-bb65-ccabd8190a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4c41f7-afa1-44a7-ab36-d1ed15a4bd48",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save the results of your study to examine later\n",
    "joblib.dump(study, \"study_single_cluster.pickle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71bc6ef-a34a-45b8-bb04-ce583ac29623",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = optuna.visualization.matplotlib.plot_optimization_history(study)\n",
    "fig.legend(loc=\"upper right\")\n",
    "plt.savefig(\"optimization_history_study_1.png\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}