{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install deepr[cpu]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyper Parameter Tuning\n",
    "\n",
    "This notebook builds upon the previous [pipeline example](https://criteo.github.io/deepr/getting_started/pipeline.html). The goal is to perform hyperparameter search by varying the learning rate and batch size.\n",
    "\n",
    "To launch an HP search, the steps are\n",
    "\n",
    "1. Download / retrieve some config and its macros (usually saved as MLFlow artifacts).\n",
    "2. Prepare the config by adding new macro parameters for a macro named \"params\" (for example change `\"learning_rate\": 0.1` to `\"learning_rate\": \"$params:learning_rate\"`).\n",
    "3. Define a sampler.\n",
    "4. Launch a tuning job using the prepared config, macros and sampler.\n",
    "\n",
    "First, some imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout)\n",
    "logging.getLogger(\"tensorflow\").setLevel(logging.CRITICAL)\n",
    "logging.getLogger(\"cluster_pack\").setLevel(logging.CRITICAL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import deepr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare config\n",
    "\n",
    "The easiest way to prepare a config is to \n",
    "\n",
    "1. Download the config and macros from mlflow with the following command (it will download config and macros artifact to `config.json` and `macros.json` paths from some run id): \n",
    "\n",
    "    ```bash\n",
    "    deepr download_config_and_macros_from_mlflow --run_id=841d3f9f2ba4b69921b426a81fd --tracking_uri=$MLFLOW_TRACKING_URI\n",
    "    ```\n",
    "    \n",
    "    Note that you can also use it in python with\n",
    "    \n",
    "    ```python\n",
    "    deepr.cli.download_config_and_macros_from_mlflow(\n",
    "        run_id=run_id, \n",
    "        path_config=\"config.json\", \n",
    "        path_macros=\"macros.json\", \n",
    "        tracking_uri=TRACKING_URI\n",
    "    )\n",
    "    ```\n",
    "    \n",
    "2. Automatically add new macro references (overwrites the files) with \n",
    "\n",
    "    ```bash\n",
    "    deepr add_macro --config=config.json --macros=macros.json --params=learning_rate,batch_size\n",
    "    ```\n",
    "    \n",
    "    Note that you can also use it in python with\n",
    "    \n",
    "    ```python\n",
    "    deepr.cli.add_macro(\n",
    "        config=\"config.json\", \n",
    "        macros=\"macros.json\", \n",
    "        params=param_grid.keys(),\n",
    "    )\n",
    "    ```\n",
    "    \n",
    "3. Create a short script `my_seeded_study_name.py` that looks like\n",
    "\n",
    "```python\n",
    "import numpy as np\n",
    "import deepr\n",
    "    \n",
    "# Create sampler\n",
    "param_grid = {\n",
    "    \"learning_rate\": np.logspace(start=-5, stop=-3, num=10),\n",
    "}\n",
    "sampler = deepr.jobs.ParamsSampler(param_grid, n_iter=50, seed=42)\n",
    "\n",
    "# Read job config and macros\n",
    "job = deepr.io.read_json(\"config.json\")\n",
    "macros = deepr.io.read_json(\"macros.json\")\n",
    "\n",
    "# Create tuner and run it\n",
    "tuner = deepr.jobs.ParamsTuner(job=job, macros=macros, sampler=sampler)\n",
    "tuner.run()\n",
    "```\n",
    "    \n",
    "Our prepared config should look like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"type\": \"deepr.jobs.yarn_launcher.YarnLauncher\",\n",
    "    \"job\": {\n",
    "        \"type\": \"deepr.jobs.Pipeline\",\n",
    "        \"eval\": None,\n",
    "        \"jobs\": [\n",
    "            {\n",
    "                \"type\": \"deepr.example.jobs.BuildDataset\",\n",
    "                \"path_dataset\": \"viewfs://root/user/deepr/dev/example/data.tfrecord\",\n",
    "                \"num_examples\": 1000,\n",
    "            },\n",
    "            {\n",
    "                \"type\": \"deepr.jobs.YarnTrainer\",\n",
    "                \"trainer\": {\n",
    "                    \"type\": \"deepr.jobs.Trainer\",\n",
    "                    \"eval\": \"skip\",\n",
    "                    \"path_model\": \"$paths:path_model\",  # Uses macro \"paths:path_model\"\n",
    "                    \"pred_fn\": {\"type\": \"deepr.example.layers.Multiply\"},\n",
    "                    \"loss_fn\": {\"type\": \"deepr.example.layers.SquaredL2\"},\n",
    "                    \"optimizer_fn\": {\n",
    "                        \"type\": \"deepr.optimizers.TensorflowOptimizer\",\n",
    "                        \"optimizer\": \"Adam\",\n",
    "                        \"learning_rate\": \"$params:learning_rate\",  # Uses macro \"params:learning_rate\"\n",
    "                    },\n",
    "                    \"prepro_fn\": {\n",
    "                        \"type\": \"deepr.example.prepros.DefaultPrepro\",  # Uses macro \"params:batch_size\"\n",
    "                        \"batch_size\": \"$params:batch_size\",\n",
    "                        \"repeat_size\": 10,\n",
    "                    },\n",
    "                    \"train_input_fn\": {\n",
    "                        \"type\": \"deepr.readers.TFRecordReader\",\n",
    "                        \"path\": \"viewfs://root/user/deepr/dev/example/data.tfrecord\",\n",
    "                        \"num_parallel_reads\": 8,\n",
    "                        \"num_parallel_calls\": 8,\n",
    "                        \"shuffle\": True,\n",
    "                    },\n",
    "                    \"eval_input_fn\": {\n",
    "                        \"type\": \"deepr.readers.TFRecordReader\",\n",
    "                        \"path\": \"viewfs://root/user/deepr/dev/example/data.tfrecord\",\n",
    "                        \"num_parallel_reads\": 8,\n",
    "                        \"num_parallel_calls\": 8,\n",
    "                        \"shuffle\": True,\n",
    "                    },\n",
    "                },\n",
    "                \"config\": {\n",
    "                    \"type\": \"deepr.jobs.YarnTrainerConfig\"\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "    \"config\": {\n",
    "        \"type\": \"deepr.jobs.YarnLauncherConfig\",\n",
    "        \"path_pex_prefix\": \"viewfs://root/user/deepr/dev/example/envs\"\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the learning rate value is a reference to a macro parameter, if we define the following macros, we can parse the config and fill the parameter values.\n",
    "\n",
    "We also used a dynamic macro to set the path to the model dynamically (every experiment needs to use a different model path)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import time\n",
    "\n",
    "class Paths(dict):\n",
    "    \"\"\"Macro that generates new path_model based on date.\"\"\"\n",
    "\n",
    "    def __init__(self, path_model: str = None, path_dataset: str = None):\n",
    "        now = datetime.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "        if path_model is None:\n",
    "            path_model = f\"viewfs://root/user/deepr/dev/example/models/{now}\"\n",
    "        if path_dataset is None:\n",
    "            path_dataset = f\"viewfs://root/user/deepr/dev/example/data.tfrecord\"\n",
    "        super().__init__(path_model=path_model, path_dataset=path_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'path_model': 'viewfs://root/user/deepr/dev/example/models/2020-06-03-15-45-11', 'path_dataset': 'viewfs://root/user/deepr/dev/example/data.tfrecord'}\n",
      "{'path_model': 'viewfs://root/user/deepr/dev/example/models/2020-06-03-15-45-13', 'path_dataset': 'viewfs://root/user/deepr/dev/example/data.tfrecord'}\n"
     ]
    }
   ],
   "source": [
    "print(Paths())\n",
    "time.sleep(2)\n",
    "print(Paths())  # Creates new paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:deepr.config.base:- MACRO PARAM NOT USED: macro = 'paths', param = 'path_dataset'\n",
      "viewfs://root/user/deepr/dev/example/models/2020-06-03-15-45-13\n",
      "0.01\n",
      "16\n"
     ]
    }
   ],
   "source": [
    "macros = {\n",
    "    \"paths\": {\n",
    "        \"type\": \"__main__.Paths\"\n",
    "    },\n",
    "    \"params\": {\n",
    "        \"learning_rate\": 0.01,\n",
    "        \"batch_size\": 16\n",
    "    }\n",
    "}\n",
    "parsed = deepr.parse_config(config, macros)\n",
    "print(parsed[\"job\"][\"jobs\"][1][\"trainer\"][\"path_model\"])\n",
    "print(parsed[\"job\"][\"jobs\"][1][\"trainer\"][\"optimizer_fn\"][\"learning_rate\"])\n",
    "print(parsed[\"job\"][\"jobs\"][1][\"trainer\"][\"prepro_fn\"][\"batch_size\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch\n",
    "\n",
    "Now that the config is prepared, we can launch hyper parameter tuning using the `ParamsSampler` job.\n",
    "\n",
    "The only thing that it does is sampling some learning rate values and use them as macros.\n",
    "\n",
    "First, let's define a sampler for the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_grid = {\n",
    "    \"learning_rate\": np.logspace(start=-3, stop=-1, num=10),\n",
    "    \"batch_size\": [8, 16, 32, 64],\n",
    "}\n",
    "sampler = deepr.jobs.ParamsSampler(param_grid, n_iter=5, seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:deepr.jobs.params_tuner:Sampling with no replacement (parameter grid only has lists of values)\n",
      "INFO:deepr.jobs.params_tuner:Sampled 5 parameters from a total of 40\n",
      "{'batch_size': 64, 'learning_rate': 0.05994842503189409}\n",
      "{'batch_size': 32, 'learning_rate': 0.05994842503189409}\n",
      "{'batch_size': 16, 'learning_rate': 0.007742636826811269}\n",
      "{'batch_size': 8, 'learning_rate': 0.03593813663804626}\n",
      "{'batch_size': 32, 'learning_rate': 0.001}\n"
     ]
    }
   ],
   "source": [
    "for params in sampler:\n",
    "    print(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and use the sampler with our config and macros to launch hyper params tuning with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "HAS_HADOOP = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "if HAS_HADOOP:\n",
    "    tuner = deepr.jobs.ParamsTuner(job=config, macros=macros, sampler=sampler)\n",
    "    tuner.run()"
   ]
  }
 ],
 "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.6.8"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}