{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Search Space Basics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will walk you through the process of defining a search space, and how to incorporate them into ablator for ablation study on various hyperparameters."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a search space with `ablator.config.hpo.SearchSpace`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Ablator, [`ablator.config.hpo.SearchSpace`](../config.train.parallel.experiment.rst#configurations-for-parallel-models-experiments) is used to define the search space for a hyperparameter, based on which ablator creates many trials of different values for that hyperparameter. It allows you to specify the range of values for different types of data, and the type of data the hyperparameter is.\n",
    "\n",
    "Import `SearchSpace`:\n",
    "\n",
    "```python\n",
    "from ablator.config.hpo import SearchSpace\n",
    "```\n",
    "\n",
    "The `SearchSpace` class (one object created for one hyperparameter) takes the following arguments:\n",
    "\n",
    "- **`value_range`**: defines a continuous range for a continuous hyperparameter. It is specified in the format of `[<lower_bound>, <upper_bound>]`. In each trial, the hyperparameter will be sampled with a value taken from this range. \n",
    "\n",
    "- **`categorical_values`**: defines a discrete set for a discrete hyperparameter. In each trial, the hyperparameter will be sampled with a value taken from this set. \n",
    "\n",
    "- **`value_type`**: specifies the hyperparameter's data type. Ablator supports `\"int\"` for integer values and `\"float\"` for decimal or floating-point values. This argument is required for hyperparameters that take values from a `value_range`.\n",
    "\n",
    "Note that categorical values do not require `value_type`.\n",
    "\n",
    "In the example below, we create a search space with a continuous float range `[0.05, 0.1]` and a search space with a discrete set `[32, 64, 128]`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from ablator.config.hpo import SearchSpace\n",
    "\n",
    "SearchSpace(value_range=[0.05, 0.1], value_type=\"float\")\n",
    "SearchSpace(categorical_values=[32, 64, 128])\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating search space for hyperparameters"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall from [Configuration Basics](./Configuration-Basics.ipynb) tutorial, `ParallelConfig` has an argument for search space, which is `search_space`. This argument is defined as a dictionary of `SearchSpace` objects, which captures all search spaces for all hyperparameters that we want to run ablation study on."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`SearchSpace` can be created for hyperparameters that are ablator-predefined configuration attributes, or custom configuration attributes:\n",
    "\n",
    "- **Predefined Configurations**: Ablator offers predefined configurations for optimizers, schedulers, batch size, epochs, and more. These configurations are readily available for users to use in their experiments.\n",
    "\n",
    "- **Custom Configurations** (added by users): Users can define custom configurations for hyperparameters specific to their models. For example, the number of hidden layer in a neural network, activation functions, and other relevant hyperparameters. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `SearchSpace` for predefined configurations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For optimizers\n",
    "\n",
    "Ablator supports three predefined optimizers: `SGD`, `Adam`, and `AdamW`. For an optimizer chosen for the training process, you can create a search space for any of its parameters. For example, to create search space for AdamW optimizer's (parameters are learning rate, epsilon, weight decay, etc.), you can do the following:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "my_search_space = {\n",
    "    \"train_config.optimizer_config.arguments.lr\": SearchSpace(\n",
    "        value_range  = [0.01, 0.05],\n",
    "        value_type = \"float\"\n",
    "    ), \n",
    "    \"train_config.optimizer_config.arguments.eps\": SearchSpace(\n",
    "        value_range  = [1e-9, 1e-7],\n",
    "        value_type = \"float\"\n",
    "    ), \n",
    "    \"train_config.optimizer_config.arguments.weight_decay\": SearchSpace(\n",
    "        value_range  = [1e-4, 1e-3],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The syntax for creating search space for optimizers in ablator is:\n",
    "\n",
    "```python\n",
    "search_space = {\n",
    "    \"train_config.optimizer_config.arguments.<parameter 1>\": search_space_1,\n",
    "    \"train_config.optimizer_config.arguments.<parameter 2>\": search_space_2,\n",
    "    ...\n",
    "}\n",
    "```\n",
    "where `<parameter 1>` and `<parameter 2>` are the parameters for the corresponding optimizer. You can find parameters for all optimizers in the [Configuration Basics](./Configuration-Basics.ipynb) tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For schedulers\n",
    "\n",
    "Ablator supports three predefined schedulers: `StepLR`, `OneCycleLR`, and `ReduceLROnPlateau`. For a scheduler chosen for the training process, you can create a search space for any of its parameters. For example, to create search space for `ReduceOnPlateau` scheduler (parameters are min learning rate, patience, factor, threshold, etc.), you can do the following:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "my_search_space = {\n",
    "    \"train_config.scheduler_config.arguments.min_lr\": SearchSpace(\n",
    "        value_range  = [1e-6, 1e-4],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "    \"train_config.scheduler_config.arguments.threshold\": SearchSpace(\n",
    "        value_range  = [1e-5, 1e-3],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The syntax for creating search space for schedulers in ablator is:\n",
    "\n",
    "```python\n",
    "search_space = {\n",
    "    \"train_config.scheduler_config.arguments.<parameter 1>\": search_space_1,\n",
    "    \"train_config.scheduler_config.arguments.<parameter 2>\": search_space_2,\n",
    "    ...\n",
    "}\n",
    "```\n",
    "where `<parameter 1>` and `<parameter 2>` are the parameters for the corresponding scheduler. You can find parameters for all schedulers in the [Configuration Basics](./Configuration-Basics.ipynb) tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For other parameters\n",
    " \n",
    "We can also provide `SearchSpace` to other parameters like epochs, batch_size, etc. from `TrainConfig`.\n",
    "\n",
    "The syntax will be:\n",
    "\n",
    "```python\n",
    "search_space = {\n",
    "    \"train_config.<parameter 1>\": search_space_1,\n",
    "    \"train_config.<parameter 2>\": search_space_2,\n",
    "    ...\n",
    "}\n",
    "```\n",
    "where `<parameter 1>` and `<parameter 2>` are the attributes of `TrainConfig`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, trying different batch_size or epochs can be easily done with the following snippet:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "my_search_space = {\n",
    "    \"train_config.batch_size\": SearchSpace(\n",
    "        categorical_values = [32, 64, 128]\n",
    "    ),\n",
    "    \"train_config.epochs\": SearchSpace(\n",
    "        value_range  = [10, 20],\n",
    "        value_type = \"int\"\n",
    "    ),\n",
    "}\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `SearchSpace` for custom configurqations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous tutorials, we have shown that you can run ablation experiment to study different components of a model. For example, we want to study the impact of the hyperparameters `hidden_size` and `activation` on the performance of a model. So we're first creating a custom model configuration with these hyperparameters and using this configuration to build the model:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "class CustomModelConfig(ModelConfig):   # hyperparameters to be studied\n",
    "    hidden_size :int \n",
    "    activation: str\n",
    "\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self, config: CustomModelConfig) -> None:\n",
    "        activation_list = {\"relu\" : nn.ReLU(), \"elu\": nn.ELU()}\n",
    "        \n",
    "        input_size = 100\n",
    "        self.fc1 = nn.Linear(input_size, config.hidden_size)\n",
    "        self.act1 = activation_list[config.activation]\n",
    "\n",
    "model_config = CustomModelConfig(\n",
    "    hidden_size = 256,\n",
    "    activation = \"relu\"\n",
    ")\n",
    "```\n",
    "\n",
    "Note that we still need to create a model config object with initial values for the hyperparameters (and pass that to the running configuration), even though later we will create multiple trials with different values for them, taken from the search space. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now create `SearchSpace` for `hidden_size` (an integer range) and `activation` (a discrete set):\n",
    "\n",
    "The syntax for search space for model hyperparameters is:\n",
    "```python\n",
    "search_space = {\n",
    "    \"model_config.<parameter 1>\": search_space_1,\n",
    "    \"model_config.<parameter 2>\": search_space_2,\n",
    "    ...\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "my_search_space = {\n",
    "    \"model_config.hidden_size\": SearchSpace(\n",
    "        value_range=[250, 500], value_type=\"int\"\n",
    "    ),\n",
    "    \"model_config.activation\": SearchSpace(\n",
    "        categorical_values = [\"relu\",\"elu\"]\n",
    "    ), \n",
    "}\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting everything into one:\n",
    "\n",
    "```python\n",
    "my_search_space = {\n",
    "    \"train_config.optimizer_config.arguments.lr\": SearchSpace(\n",
    "        value_range  = [0.01, 0.05],\n",
    "        value_type = \"float\"\n",
    "    ), \n",
    "    \"train_config.optimizer_config.arguments.eps\": SearchSpace(\n",
    "        value_range  = [1e-9, 1e-7],\n",
    "        value_type = \"float\"\n",
    "    ), \n",
    "    \"train_config.optimizer_config.arguments.weight_decay\": SearchSpace(\n",
    "        value_range  = [1e-4, 1e-3],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "    \"train_config.scheduler_config.arguments.min_lr\": SearchSpace(\n",
    "        value_range  = [1e-6, 1e-4],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "    \"train_config.scheduler_config.arguments.threshold\": SearchSpace(\n",
    "        value_range  = [1e-5, 1e-3],\n",
    "        value_type = \"float\"\n",
    "    ),\n",
    "    \"train_config.batch_size\": SearchSpace(\n",
    "        categorical_values = [32, 64, 128]\n",
    "    ),\n",
    "    \"train_config.epochs\": SearchSpace(\n",
    "        value_range  = [10, 20],\n",
    "        value_type = \"int\"\n",
    "    ),\n",
    "    \"model_config.hidden_size\": SearchSpace(\n",
    "        value_range=[250, 500], value_type=\"int\"\n",
    "    ),\n",
    "    \"model_config.activation\": SearchSpace(\n",
    "        categorical_values = [\"relu\",\"elu\"]\n",
    "    )\n",
    "}\n",
    "```\n",
    "\n",
    "Finally, `my_search_space` dictionary is passed to the `ParallelConfig`. This will be explored in more detail in [Hyperparameter Optimization](./HPO-tutorial.ipynb) tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SearchSpace in YAML files\n",
    "\n",
    "If you are using a YAML file to define configurations, we can specify a search space as follows:\n",
    "\n",
    "```yaml\n",
    "# other configurations ...\n",
    "search_space:\n",
    "  model_config.hidden_size:\n",
    "    value_range:\n",
    "    - '250'\n",
    "    - '500'\n",
    "    categorical_values: null\n",
    "    value_type: int\n",
    "  train_config.optimizer_config.arguments.lr:\n",
    "    value_range:\n",
    "    - '0.001'\n",
    "    - '0.01'\n",
    "    categorical_values: null\n",
    "    value_type: float\n",
    "  model_config.activation:\n",
    "    value_range: null\n",
    "    categorical_values:\n",
    "    - relu\n",
    "    - leakyRelu\n",
    "    - elu\n",
    "    value_type: float\n",
    "# other configurations ...\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial, we have demonstrated how to create search space objects and how to utilize them to define search space for various hyperparameters. In the subsequent tutorial, we will explain how to use `search_space` with `ParallelConfig` to launch a parallel ablation experiment."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}