{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "25705204",
   "metadata": {},
   "source": [
    "# Running Tune experiments with Optuna\n",
    "\n",
    "This example demonstrates the usage of Optuna with Ray Tune via `OptunaSearch`, including conditional search spaces and the multi-objective use case.\n",
    "\n",
    "While you may use a scheduler with `OptunaSearch`, e.g. `AsyncHyperBandScheduler`, please note that schedulers may not work correctly with multi-objective optimization, since they often expect a scalar score.\n",
    "\n",
    "Click below to see all the imports we need for this example.\n",
    "You can also launch directly into a Binder instance to run this notebook yourself.\n",
    "Just click on the rocket symbol at the top of the navigation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2fbc3e",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from typing import Dict, Optional, Any\n",
    "\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.tune.suggest import ConcurrencyLimiter\n",
    "from ray.tune.schedulers import AsyncHyperBandScheduler\n",
    "from ray.tune.suggest.optuna import OptunaSearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56404ca0",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "ray.init(configure_logging=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac0a86b6",
   "metadata": {},
   "source": [
    "Let's start by defining a simple evaluation function.\n",
    "An explicit math formula is queried here for demonstration, yet in practice this is typically a black-box function-- e.g. the performance results after training an ML model.\n",
    "We artificially sleep for a bit (`0.1` seconds) to simulate a long-running ML experiment.\n",
    "This setup assumes that we're running multiple `step`s of an experiment while tuning three hyperparameters,\n",
    "namely `width`, `height`, and `activation`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa1a9a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(step, width, height, activation):\n",
    "    time.sleep(0.1)\n",
    "    activation_boost = 10 if activation==\"relu\" else 0\n",
    "    return (0.1 + width * step / 100) ** (-1) + height * 0.1 + activation_boost"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f066da2",
   "metadata": {},
   "source": [
    "Next, our ``objective`` function to be optimized takes a Tune ``config``, evaluates the `score` of your experiment in a training loop,\n",
    "and uses `tune.report` to report the `score` back to Tune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd6c61a1",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "def objective(config):\n",
    "    for step in range(config[\"steps\"]):\n",
    "        score = evaluate(step, config[\"width\"], config[\"height\"], config[\"activation\"])\n",
    "        tune.report(iterations=step, mean_loss=score)\n",
    "              "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fcf8aef",
   "metadata": {},
   "source": [
    "Next we define a search space. The critical assumption is that the optimal hyperparamters live within this space. Yet, if the space is very large, then those hyperparamters may be difficult to find in a short amount of time.\n",
    "\n",
    "The simplest case is a search space with independent dimensions. In this case, a config dictionary will suffice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ba9fbfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_space = {\n",
    "    \"steps\": 100,\n",
    "    \"width\": tune.uniform(0, 20),\n",
    "    \"height\": tune.uniform(-100, 100),\n",
    "    \"activation\": tune.choice([\"relu\", \"tanh\"]),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93afec28",
   "metadata": {},
   "source": [
    "While defining the search algorithm, we may choose to provide an initial set of hyperparameters that we believe are especially promising or informative, and\n",
    "pass this information as a helpful starting point for the `OptunaSearch` object.\n",
    "\n",
    "Here we define the Optuna search algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8ded912",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = OptunaSearch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a001c1be",
   "metadata": {},
   "source": [
    "We also constrain the the number of concurrent trials to `4` with a `ConcurrencyLimiter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ada67a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "algo = ConcurrencyLimiter(searcher, max_concurrent=4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ce35bf5",
   "metadata": {},
   "source": [
    "Lastly, we set the number of samples for this Tune run to `1000`\n",
    "(you can decrease this if it takes too long on your machine)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6c39b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8855163",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# We override here for our smoke tests.\n",
    "num_samples = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "410b7464",
   "metadata": {},
   "source": [
    "Furthermore, we define a `scheduler` to go along with our algorithm. This is optional, and only to demonstrate that we don't need to compromise other great features of Ray Tune while using Optuna."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b5dbea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = AsyncHyperBandScheduler()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24792770",
   "metadata": {},
   "source": [
    "Now all that's left is running the experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "677eeb36",
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis = tune.run(\n",
    "    objective,\n",
    "    search_alg=algo,\n",
    "    scheduler=scheduler,\n",
    "    metric=\"mean_loss\",\n",
    "    mode=\"min\",\n",
    "    num_samples=num_samples,\n",
    "    config=search_space\n",
    ")\n",
    "\n",
    "print(\"Best hyperparameters found were: \", analysis.best_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "574cca0b",
   "metadata": {},
   "source": [
    "While defining the search algorithm, we may choose to provide an initial set of hyperparameters that we believe are especially promising or informative, and\n",
    "pass this information as a helpful starting point for the `OptunaSearch` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2e38018",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_params = [\n",
    "    {\"width\": 1, \"height\": 2, \"activation\": \"relu\"},\n",
    "    {\"width\": 4, \"height\": 2, \"activation\": \"relu\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b25729",
   "metadata": {},
   "source": [
    "Now the `search_alg` built using `OptunaSearch` takes `points_to_evaluate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d85142ae",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "searcher = OptunaSearch(points_to_evaluate=initial_params)\n",
    "algo = ConcurrencyLimiter(searcher, max_concurrent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7c8c0f7",
   "metadata": {},
   "source": [
    "And run the experiment with initial hyperparameter evaluations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26815194",
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis = tune.run(\n",
    "    objective,\n",
    "    search_alg=algo,\n",
    "    metric=\"mean_loss\",\n",
    "    mode=\"min\",\n",
    "    num_samples=num_samples,\n",
    "    config=search_space\n",
    ")\n",
    "\n",
    "print(\"Best hyperparameters found were: \", analysis.best_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8b9278c",
   "metadata": {},
   "source": [
    "Sometimes we may want to build a more complicated search space that has conditional dependencies on other hyperparameters. In this case, we pass a define-by-run function to the `search_alg` argument in `ray.tune()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ea7416",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "def define_by_run_func(trial) -> Optional[Dict[str, Any]]:\n",
    "    \"\"\"Define-by-run function to create the search space.\n",
    "\n",
    "    Ensure no actual computation takes place here. That should go into\n",
    "    the trainable passed to ``tune.run`` (in this example, that's\n",
    "    ``objective``).\n",
    "\n",
    "    For more information, see https://optuna.readthedocs.io/en/stable\\\n",
    "    /tutorial/10_key_features/002_configurations.html\n",
    "\n",
    "    This function should either return None or a dict with constant values.\n",
    "    \"\"\"\n",
    "\n",
    "    activation = trial.suggest_categorical(\"activation\", [\"relu\", \"tanh\"])\n",
    "\n",
    "    # Define-by-run allows for conditional search spaces.\n",
    "    if activation == \"relu\":\n",
    "        trial.suggest_float(\"width\", 0, 20)\n",
    "        trial.suggest_float(\"height\", -100, 100)\n",
    "    else:\n",
    "        trial.suggest_float(\"width\", -1, 21)\n",
    "        trial.suggest_float(\"height\", -101, 101)\n",
    "        \n",
    "    # Return all constants in a dictionary.\n",
    "    return {\"steps\": 100}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "283aec6a",
   "metadata": {},
   "source": [
    "As before, we create the `search_alg` from `OptunaSearch` and `ConcurrencyLimiter`, this time we define the scope of search via the `space` argument and provide no initialization. We also must specific metric and mode when using `space`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4818fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = OptunaSearch(space=define_by_run_func, metric=\"mean_loss\", mode=\"min\")\n",
    "algo = ConcurrencyLimiter(searcher, max_concurrent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "020f4e6e",
   "metadata": {},
   "source": [
    "Running the experiment with a define-by-run search space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f478ec68",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "analysis = tune.run(\n",
    "    objective,\n",
    "    search_alg=algo,\n",
    "    num_samples=num_samples\n",
    ")\n",
    "\n",
    "print(\"Best hyperparameters for loss found were: \", analysis.get_best_config(\"mean_loss\", \"min\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc1ac46f",
   "metadata": {},
   "source": [
    "Finally, let's take a look at the multi-objective case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b08eefa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multi_objective(config):\n",
    "    # Hyperparameters\n",
    "    width, height = config[\"width\"], config[\"height\"]\n",
    "\n",
    "    for step in range(config[\"steps\"]):\n",
    "        # Iterative training function - can be any arbitrary training procedure\n",
    "        intermediate_score = evaluate(step, config[\"width\"], config[\"height\"], config[\"activation\"])\n",
    "        # Feed the score back back to Tune.\n",
    "        tune.report(\n",
    "           iterations=step, loss=intermediate_score, gain=intermediate_score * width\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73943695",
   "metadata": {},
   "source": [
    "We define the `OptunaSearch` object this time with metric and mode as list arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab840970",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = OptunaSearch(metric=[\"loss\", \"gain\"], mode=[\"min\", \"max\"])\n",
    "algo = ConcurrencyLimiter(searcher, max_concurrent=4)\n",
    "\n",
    "analysis = tune.run(\n",
    "    multi_objective,\n",
    "    search_alg=algo,\n",
    "    num_samples=num_samples,\n",
    "    config=search_space\n",
    ")\n",
    "\n",
    "print(\"Best hyperparameters for loss found were: \", analysis.get_best_config(\"loss\", \"min\"))\n",
    "print(\"Best hyperparameters for gain found were: \", analysis.get_best_config(\"gain\", \"max\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c06be1e",
   "metadata": {},
   "source": [
    "We can mix-and-match the use of initial hyperparameter evaluations, conditional search spaces via define-by-run functions, and multi-objective tasks. This is also true of scheduler usage, with the exception of multi-objective optimization-- schedulers typically rely on a single scalar score, rather than the two scores we use here: loss, gain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e662ef66",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}