notebooks/04_pretraining_hyperparam.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Hyperparameter search for pretraining phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# allows for import from `src` directory\n",
    "import sys \n",
    "sys.path.append('../')\n",
    "\n",
    "from src import data\n",
    "from src import utils\n",
    "\n",
    "from models.encoder_decoder_dropout import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.1 Set up data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data already downloaded\n",
      "\n",
      "43910 train rows from 2012-10-02 09:00:00 to 2017-10-05 23:00:00\n",
      "4320 valid rows from 2017-10-05 23:00:00 to 2018-04-03 23:00:00\n",
      "4321 test rows from 2018-04-03 23:00:00 to 2018-09-30 23:00:00\n",
      "\n",
      "16625 samples of 48 input steps and 12 output steps in train\n",
      "3534 samples of 48 input steps and 12 output steps in valid\n",
      "4020 samples of 48 input steps and 12 output steps in test\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# fixed data parameters\n",
    "data_params = {\n",
    "    'n_input_steps': 48,\n",
    "    'n_output_steps': 12\n",
    "}\n",
    "\n",
    "# run the data preprocessing pipeline to create dataset\n",
    "_, _, samples = data.pipeline(data_params['n_input_steps'], data_params['n_output_steps'], '../data')\n",
    "datasets = data.get_datasets(samples, data_params['n_input_steps'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 Bayesian optimisation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get device\n",
    "\n",
    "We are using a GPU on a GCE instances as this is a compute heavy job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "utils.get_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define function to optimise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrap the `utils.train` and `utils.evaluate` to define the function to be optimised.\n",
    "\n",
    "This function takes in a dictionary of hyperparameters and returns the loss value on the validation data using the model trained according to the passed in hyperparameters. We want to know which hyperparameters will minimise this validation loss, so we will use `Ax` to do a guided search using Bayesian optimisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_evaluate(params):\n",
    "    device = utils.get_device()\n",
    "    dataloaders = data.get_dataloaders(datasets, params.get('batch_size'))\n",
    "    in_features = dataloaders['train'].dataset.X.shape[-1]\n",
    "    model = VDEncoderDecoder(in_features=in_features,\n",
    "                             input_steps=params.get('n_input_steps', 48),\n",
    "                             output_steps=params.get('n_output_steps', 12),\n",
    "                             p=params.get('variational_dropout_p')\n",
    "                            ).to(device)\n",
    "    model,_ = utils.train(device=device, model=model, dataloader=dataloaders['train'], params=params, use_tqdm=True)\n",
    "    return utils.evaluate(device=device, model=model, valid_loader=dataloaders['valid'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up `Ax` client\n",
    "\n",
    "We mainly follow the tutorial at https://ax.dev/tutorials/gpei_hartmann_service.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 09-13 18:14:12] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 2 decimal points.\n",
      "[INFO 09-13 18:14:12] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 5 trials, GPEI for subsequent trials]). Iterations after 5 will take longer to generate due to  model-fitting.\n"
     ]
    }
   ],
   "source": [
    "from ax.service.ax_client import AxClient\n",
    "\n",
    "ax_client = AxClient(enforce_sequential_optimization=False)\n",
    "\n",
    "# define hyperparameter bounds \n",
    "ax_client.create_experiment(\n",
    "    name='pretraining',\n",
    "    parameters=[\n",
    "        {\"name\": \"num_epochs\", \"type\": \"range\", \"bounds\": [150, 200]},\n",
    "        {\"name\": \"learning_rate\", \"type\": \"range\", \"bounds\": [5e-4, 1e-3], \"log_scale\": True},\n",
    "        {\"name\": \"batch_size\", \"type\": \"range\", \"bounds\": [128, 1024]},\n",
    "        {\"name\": \"variational_dropout_p\", \"type\": \"range\", \"bounds\": [0.2,0.4]}\n",
    "    ],\n",
    "    objective_name='loss',\n",
    "    minimize=True\n",
    ")\n",
    "\n",
    "def evaluate(parameters):\n",
    "    return {\"pretraining\": train_evaluate(parameters)}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Launch the optimisation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 09-13 16:29:04] ax.service.ax_client: Generated new trial 0 with parameters {'num_epochs': 181, 'learning_rate': 0.0, 'batch_size': 381, 'variational_dropout_p': 0.38}.\n",
      "Epoch=180 | [16625|16625]\tloss=1.2937: 100%|██████████| 181/181 [04:31<00:00,  1.50s/it]\n",
      "[INFO 09-13 16:33:38] ax.service.ax_client: Completed trial 0 with data: {'loss': (0.44, None)}.\n",
      "[INFO 09-13 16:33:38] ax.service.ax_client: Generated new trial 1 with parameters {'num_epochs': 168, 'learning_rate': 0.0, 'batch_size': 725, 'variational_dropout_p': 0.37}.\n",
      "Epoch=167 | [16625|16625]\tloss=1.3428: 100%|██████████| 168/168 [03:07<00:00,  1.11s/it]\n",
      "[INFO 09-13 16:36:46] ax.service.ax_client: Completed trial 1 with data: {'loss': (0.92, None)}.\n",
      "[INFO 09-13 16:36:46] ax.service.ax_client: Generated new trial 2 with parameters {'num_epochs': 188, 'learning_rate': 0.0, 'batch_size': 648, 'variational_dropout_p': 0.23}.\n",
      "Epoch=187 | [16625|16625]\tloss=0.7925: 100%|██████████| 188/188 [03:37<00:00,  1.16s/it]\n",
      "[INFO 09-13 16:40:23] ax.service.ax_client: Completed trial 2 with data: {'loss': (0.27, None)}.\n",
      "[INFO 09-13 16:40:23] ax.service.ax_client: Generated new trial 3 with parameters {'num_epochs': 172, 'learning_rate': 0.0, 'batch_size': 714, 'variational_dropout_p': 0.26}.\n",
      "Epoch=171 | [16625|16625]\tloss=1.1527: 100%|██████████| 172/172 [03:12<00:00,  1.12s/it]\n",
      "[INFO 09-13 16:43:37] ax.service.ax_client: Completed trial 3 with data: {'loss': (0.53, None)}.\n",
      "[INFO 09-13 16:43:37] ax.service.ax_client: Generated new trial 4 with parameters {'num_epochs': 173, 'learning_rate': 0.0, 'batch_size': 239, 'variational_dropout_p': 0.35}.\n",
      "Epoch=172 | [16625|16625]\tloss=1.5751: 100%|██████████| 173/173 [05:47<00:00,  2.01s/it]\n",
      "[INFO 09-13 16:49:24] ax.service.ax_client: Completed trial 4 with data: {'loss': (0.41, None)}.\n",
      "[INFO 09-13 16:49:25] ax.service.ax_client: Generated new trial 5 with parameters {'num_epochs': 188, 'learning_rate': 0.0, 'batch_size': 461, 'variational_dropout_p': 0.27}.\n",
      "Epoch=187 | [16625|16625]\tloss=0.6417: 100%|██████████| 188/188 [04:12<00:00,  1.34s/it]\n",
      "[INFO 09-13 16:53:37] ax.service.ax_client: Completed trial 5 with data: {'loss': (0.34, None)}.\n",
      "[INFO 09-13 16:53:38] ax.service.ax_client: Generated new trial 6 with parameters {'num_epochs': 177, 'learning_rate': 0.0, 'batch_size': 516, 'variational_dropout_p': 0.21}.\n",
      "Epoch=176 | [16625|16625]\tloss=0.8551: 100%|██████████| 177/177 [03:44<00:00,  1.27s/it]\n",
      "[INFO 09-13 16:57:23] ax.service.ax_client: Completed trial 6 with data: {'loss': (0.27, None)}.\n",
      "[INFO 09-13 16:57:23] ax.service.ax_client: Generated new trial 7 with parameters {'num_epochs': 191, 'learning_rate': 0.0, 'batch_size': 464, 'variational_dropout_p': 0.2}.\n",
      "Epoch=190 | [16625|16625]\tloss=0.7043: 100%|██████████| 191/191 [04:14<00:00,  1.33s/it]\n",
      "[INFO 09-13 17:01:38] ax.service.ax_client: Completed trial 7 with data: {'loss': (0.28, None)}.\n",
      "[INFO 09-13 17:01:39] ax.service.ax_client: Generated new trial 8 with parameters {'num_epochs': 190, 'learning_rate': 0.0, 'batch_size': 495, 'variational_dropout_p': 0.2}.\n",
      "Epoch=189 | [16625|16625]\tloss=0.6799: 100%|██████████| 190/190 [04:05<00:00,  1.29s/it]\n",
      "[INFO 09-13 17:05:44] ax.service.ax_client: Completed trial 8 with data: {'loss': (0.22, None)}.\n",
      "[INFO 09-13 17:05:45] ax.service.ax_client: Generated new trial 9 with parameters {'num_epochs': 184, 'learning_rate': 0.0, 'batch_size': 682, 'variational_dropout_p': 0.2}.\n",
      "Epoch=183 | [16625|16625]\tloss=0.7492: 100%|██████████| 184/184 [03:27<00:00,  1.13s/it]\n",
      "[INFO 09-13 17:09:13] ax.service.ax_client: Completed trial 9 with data: {'loss': (0.29, None)}.\n",
      "[INFO 09-13 17:09:13] ax.service.ax_client: Generated new trial 10 with parameters {'num_epochs': 186, 'learning_rate': 0.0, 'batch_size': 237, 'variational_dropout_p': 0.2}.\n",
      "Epoch=185 | [16625|16625]\tloss=0.9811: 100%|██████████| 186/186 [06:18<00:00,  2.04s/it]\n",
      "[INFO 09-13 17:15:32] ax.service.ax_client: Completed trial 10 with data: {'loss': (0.18, None)}.\n",
      "[INFO 09-13 17:15:33] ax.service.ax_client: Generated new trial 11 with parameters {'num_epochs': 198, 'learning_rate': 0.0, 'batch_size': 128, 'variational_dropout_p': 0.2}.\n",
      "Epoch=197 | [16625|16625]\tloss=0.6991: 100%|██████████| 198/198 [10:37<00:00,  3.22s/it]\n",
      "[INFO 09-13 17:26:10] ax.service.ax_client: Completed trial 11 with data: {'loss': (0.32, None)}.\n",
      "[INFO 09-13 17:26:11] ax.service.ax_client: Generated new trial 12 with parameters {'num_epochs': 175, 'learning_rate': 0.0, 'batch_size': 128, 'variational_dropout_p': 0.2}.\n",
      "Epoch=174 | [16625|16625]\tloss=0.6862: 100%|██████████| 175/175 [09:23<00:00,  3.22s/it]\n",
      "[INFO 09-13 17:35:35] ax.service.ax_client: Completed trial 12 with data: {'loss': (0.32, None)}.\n",
      "[INFO 09-13 17:35:36] ax.service.ax_client: Generated new trial 13 with parameters {'num_epochs': 186, 'learning_rate': 0.0, 'batch_size': 308, 'variational_dropout_p': 0.2}.\n",
      "Epoch=185 | [16625|16625]\tloss=0.7924: 100%|██████████| 186/186 [05:15<00:00,  1.70s/it]\n",
      "[INFO 09-13 17:40:52] ax.service.ax_client: Completed trial 13 with data: {'loss': (0.29, None)}.\n",
      "[INFO 09-13 17:40:53] ax.service.ax_client: Generated new trial 14 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 810, 'variational_dropout_p': 0.2}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.7443: 100%|██████████| 200/200 [03:47<00:00,  1.14s/it]\n",
      "[INFO 09-13 17:44:40] ax.service.ax_client: Completed trial 14 with data: {'loss': (0.26, None)}.\n",
      "[INFO 09-13 17:44:42] ax.service.ax_client: Generated new trial 15 with parameters {'num_epochs': 184, 'learning_rate': 0.0, 'batch_size': 954, 'variational_dropout_p': 0.2}.\n",
      "Epoch=183 | [16625|16625]\tloss=0.7027: 100%|██████████| 184/184 [03:23<00:00,  1.10s/it]\n",
      "[INFO 09-13 17:48:05] ax.service.ax_client: Completed trial 15 with data: {'loss': (0.26, None)}.\n",
      "[INFO 09-13 17:48:06] ax.service.ax_client: Generated new trial 16 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 921, 'variational_dropout_p': 0.2}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.5512: 100%|██████████| 200/200 [03:40<00:00,  1.10s/it]\n",
      "[INFO 09-13 17:51:47] ax.service.ax_client: Completed trial 16 with data: {'loss': (0.2, None)}.\n",
      "[INFO 09-13 17:51:47] ax.service.ax_client: Generated new trial 17 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 1024, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.8321: 100%|██████████| 200/200 [03:38<00:00,  1.09s/it]\n",
      "[INFO 09-13 17:55:26] ax.service.ax_client: Completed trial 17 with data: {'loss': (0.65, None)}.\n",
      "[INFO 09-13 17:55:27] ax.service.ax_client: Generated new trial 18 with parameters {'num_epochs': 188, 'learning_rate': 0.0, 'batch_size': 364, 'variational_dropout_p': 0.23}.\n",
      "Epoch=187 | [16625|16625]\tloss=0.7038: 100%|██████████| 188/188 [04:47<00:00,  1.53s/it]\n",
      "[INFO 09-13 18:00:15] ax.service.ax_client: Completed trial 18 with data: {'loss': (0.36, None)}.\n",
      "[INFO 09-13 18:00:16] ax.service.ax_client: Generated new trial 19 with parameters {'num_epochs': 192, 'learning_rate': 0.0, 'batch_size': 746, 'variational_dropout_p': 0.2}.\n",
      "Epoch=191 | [16625|16625]\tloss=0.7917: 100%|██████████| 192/192 [03:39<00:00,  1.14s/it]\n",
      "[INFO 09-13 18:03:55] ax.service.ax_client: Completed trial 19 with data: {'loss': (0.32, None)}.\n"
     ]
    }
   ],
   "source": [
    "# run 20 jobs\n",
    "for i in range(20):\n",
    "    parameters, trial_index = ax_client.get_next_trial()\n",
    "    ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters)['pretraining']['loss'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 09-13 18:14:41] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 2 decimal points.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ax.service.ax_client.AxClient at 0x7f1a4c4236d0>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ax_client.load_from_json_file(\n",
    "    '../model_artifacts/pretraining_hparam_search.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 09-13 18:14:45] ax.service.ax_client: Generated new trial 0 with parameters {'num_epochs': 195, 'learning_rate': 0.0, 'batch_size': 560, 'variational_dropout_p': 0.25}.\n",
      "Epoch=194 | [16625|16625]\tloss=0.8907: 100%|██████████| 195/195 [04:05<00:00,  1.26s/it]\n",
      "[INFO 09-13 18:18:54] ax.service.ax_client: Completed trial 0 with data: {'loss': (0.29, None)}.\n",
      "[INFO 09-13 18:18:54] ax.service.ax_client: Generated new trial 1 with parameters {'num_epochs': 188, 'learning_rate': 0.0, 'batch_size': 594, 'variational_dropout_p': 0.35}.\n",
      "Epoch=187 | [16625|16625]\tloss=1.1547: 100%|██████████| 188/188 [04:20<00:00,  1.39s/it]\n",
      "[INFO 09-13 18:23:15] ax.service.ax_client: Completed trial 1 with data: {'loss': (0.45, None)}.\n",
      "[INFO 09-13 18:23:15] ax.service.ax_client: Generated new trial 2 with parameters {'num_epochs': 196, 'learning_rate': 0.0, 'batch_size': 1002, 'variational_dropout_p': 0.29}.\n",
      "Epoch=195 | [16625|16625]\tloss=1.0670: 100%|██████████| 196/196 [03:35<00:00,  1.10s/it]\n",
      "[INFO 09-13 18:26:50] ax.service.ax_client: Completed trial 2 with data: {'loss': (0.52, None)}.\n",
      "[INFO 09-13 18:26:51] ax.service.ax_client: Generated new trial 3 with parameters {'num_epochs': 168, 'learning_rate': 0.0, 'batch_size': 416, 'variational_dropout_p': 0.38}.\n",
      "Epoch=167 | [16625|16625]\tloss=1.1656: 100%|██████████| 168/168 [03:58<00:00,  1.42s/it]\n",
      "[INFO 09-13 18:30:49] ax.service.ax_client: Completed trial 3 with data: {'loss': (0.73, None)}.\n",
      "[INFO 09-13 18:30:49] ax.service.ax_client: Generated new trial 4 with parameters {'num_epochs': 180, 'learning_rate': 0.0, 'batch_size': 392, 'variational_dropout_p': 0.31}.\n",
      "Epoch=179 | [16625|16625]\tloss=0.9689: 100%|██████████| 180/180 [04:23<00:00,  1.46s/it]\n",
      "[INFO 09-13 18:35:13] ax.service.ax_client: Completed trial 4 with data: {'loss': (0.38, None)}.\n",
      "[INFO 09-13 18:35:14] ax.service.ax_client: Generated new trial 5 with parameters {'num_epochs': 194, 'learning_rate': 0.0, 'batch_size': 445, 'variational_dropout_p': 0.24}.\n",
      "Epoch=193 | [16625|16625]\tloss=0.8654: 100%|██████████| 194/194 [06:07<00:00,  1.90s/it]\n",
      "[INFO 09-13 18:41:22] ax.service.ax_client: Completed trial 5 with data: {'loss': (0.27, None)}.\n",
      "[INFO 09-13 18:41:22] ax.service.ax_client: Generated new trial 6 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 360, 'variational_dropout_p': 0.23}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.4981: 100%|██████████| 200/200 [05:11<00:00,  1.56s/it]\n",
      "[INFO 09-13 18:46:34] ax.service.ax_client: Completed trial 6 with data: {'loss': (0.25, None)}.\n",
      "[INFO 09-13 18:46:34] ax.service.ax_client: Generated new trial 7 with parameters {'num_epochs': 192, 'learning_rate': 0.0, 'batch_size': 445, 'variational_dropout_p': 0.2}.\n",
      "Epoch=191 | [16625|16625]\tloss=0.7277: 100%|██████████| 192/192 [04:24<00:00,  1.38s/it]\n",
      "[INFO 09-13 18:50:59] ax.service.ax_client: Completed trial 7 with data: {'loss': (0.33, None)}.\n",
      "[INFO 09-13 18:51:00] ax.service.ax_client: Generated new trial 8 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 332, 'variational_dropout_p': 0.26}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.4613: 100%|██████████| 200/200 [05:25<00:00,  1.63s/it]\n",
      "[INFO 09-13 18:56:26] ax.service.ax_client: Completed trial 8 with data: {'loss': (0.28, None)}.\n",
      "[INFO 09-13 18:56:27] ax.service.ax_client: Generated new trial 9 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 480, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.7757: 100%|██████████| 200/200 [04:19<00:00,  1.30s/it]\n",
      "[INFO 09-13 19:00:46] ax.service.ax_client: Completed trial 9 with data: {'loss': (0.35, None)}.\n",
      "[INFO 09-13 19:00:47] ax.service.ax_client: Generated new trial 10 with parameters {'num_epochs': 192, 'learning_rate': 0.0, 'batch_size': 292, 'variational_dropout_p': 0.25}.\n",
      "Epoch=191 | [16625|16625]\tloss=0.8193: 100%|██████████| 192/192 [05:35<00:00,  1.74s/it]\n",
      "[INFO 09-13 19:06:23] ax.service.ax_client: Completed trial 10 with data: {'loss': (0.26, None)}.\n",
      "[INFO 09-13 19:06:24] ax.service.ax_client: Generated new trial 11 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 280, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.7208: 100%|██████████| 200/200 [06:01<00:00,  1.81s/it]\n",
      "[INFO 09-13 19:12:26] ax.service.ax_client: Completed trial 11 with data: {'loss': (0.26, None)}.\n",
      "[INFO 09-13 19:12:26] ModelBridge: Leaving out out-of-design observations for arms: 11_0\n",
      "[INFO 09-13 19:12:26] ax.service.ax_client: Generated new trial 12 with parameters {'num_epochs': 199, 'learning_rate': 0.0, 'batch_size': 285, 'variational_dropout_p': 0.24}.\n",
      "Epoch=198 | [16625|16625]\tloss=0.9087: 100%|██████████| 199/199 [05:54<00:00,  1.78s/it]\n",
      "[INFO 09-13 19:18:21] ax.service.ax_client: Completed trial 12 with data: {'loss': (0.36, None)}.\n",
      "[INFO 09-13 19:18:21] ModelBridge: Leaving out out-of-design observations for arms: 12_0, 11_0\n",
      "[INFO 09-13 19:18:22] ax.service.ax_client: Generated new trial 13 with parameters {'num_epochs': 199, 'learning_rate': 0.0, 'batch_size': 283, 'variational_dropout_p': 0.24}.\n",
      "Epoch=198 | [16625|16625]\tloss=0.8146: 100%|██████████| 199/199 [05:55<00:00,  1.79s/it]\n",
      "[INFO 09-13 19:24:18] ax.service.ax_client: Completed trial 13 with data: {'loss': (0.28, None)}.\n",
      "[INFO 09-13 19:24:18] ModelBridge: Leaving out out-of-design observations for arms: 12_0, 11_0, 13_0\n",
      "[INFO 09-13 19:24:18] ax.service.ax_client: Generated new trial 14 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 279, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=1.0061: 100%|██████████| 200/200 [06:02<00:00,  1.81s/it]\n",
      "[INFO 09-13 19:30:21] ax.service.ax_client: Completed trial 14 with data: {'loss': (0.27, None)}.\n",
      "[INFO 09-13 19:30:21] ModelBridge: Leaving out out-of-design observations for arms: 12_0, 14_0, 11_0, 13_0\n",
      "[INFO 09-13 19:30:22] ax.service.ax_client: Generated new trial 15 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 282, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.9741: 100%|██████████| 200/200 [05:58<00:00,  1.79s/it]\n",
      "[INFO 09-13 19:36:20] ax.service.ax_client: Completed trial 15 with data: {'loss': (0.51, None)}.\n",
      "[INFO 09-13 19:36:20] ModelBridge: Leaving out out-of-design observations for arms: 15_0, 13_0, 12_0, 14_0, 11_0\n",
      "[INFO 09-13 19:36:21] ax.service.ax_client: Generated new trial 16 with parameters {'num_epochs': 199, 'learning_rate': 0.0, 'batch_size': 293, 'variational_dropout_p': 0.24}.\n",
      "Epoch=198 | [16625|16625]\tloss=0.8712: 100%|██████████| 199/199 [05:47<00:00,  1.75s/it]\n",
      "[INFO 09-13 19:42:09] ax.service.ax_client: Completed trial 16 with data: {'loss': (0.38, None)}.\n",
      "[INFO 09-13 19:42:09] ModelBridge: Leaving out out-of-design observations for arms: 15_0, 13_0, 12_0, 14_0, 11_0, 16_0\n",
      "[INFO 09-13 19:42:09] ax.service.ax_client: Generated new trial 17 with parameters {'num_epochs': 199, 'learning_rate': 0.0, 'batch_size': 283, 'variational_dropout_p': 0.24}.\n",
      "Epoch=198 | [16625|16625]\tloss=0.9531: 100%|██████████| 199/199 [05:53<00:00,  1.78s/it]\n",
      "[INFO 09-13 19:48:04] ax.service.ax_client: Completed trial 17 with data: {'loss': (0.39, None)}.\n",
      "[INFO 09-13 19:48:04] ModelBridge: Leaving out out-of-design observations for arms: 15_0, 13_0, 12_0, 14_0, 11_0, 17_0, 16_0\n",
      "[INFO 09-13 19:48:04] ax.service.ax_client: Generated new trial 18 with parameters {'num_epochs': 200, 'learning_rate': 0.0, 'batch_size': 280, 'variational_dropout_p': 0.24}.\n",
      "Epoch=199 | [16625|16625]\tloss=0.9529: 100%|██████████| 200/200 [06:00<00:00,  1.80s/it]\n",
      "[INFO 09-13 19:54:05] ax.service.ax_client: Completed trial 18 with data: {'loss': (0.55, None)}.\n",
      "[INFO 09-13 19:54:05] ModelBridge: Leaving out out-of-design observations for arms: 18_0, 15_0, 13_0, 12_0, 14_0, 11_0, 17_0, 16_0\n",
      "[INFO 09-13 19:54:05] ax.service.ax_client: Generated new trial 19 with parameters {'num_epochs': 199, 'learning_rate': 0.0, 'batch_size': 291, 'variational_dropout_p': 0.24}.\n",
      "Epoch=198 | [16625|16625]\tloss=0.8494: 100%|██████████| 199/199 [05:51<00:00,  1.76s/it]\n",
      "[INFO 09-13 19:59:57] ax.service.ax_client: Completed trial 19 with data: {'loss': (0.32, None)}.\n"
     ]
    }
   ],
   "source": [
    "# run 20 jobs\n",
    "for i in range(20):\n",
    "    parameters, trial_index = ax_client.get_next_trial()\n",
    "    ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters)['pretraining']['loss'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get best parameters after 40 iterations of search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'num_epochs': 192,\n",
       "  'learning_rate': 0.0005753955111014699,\n",
       "  'batch_size': 292,\n",
       "  'variational_dropout_p': 0.24793713575125884},\n",
       " ({'loss': 0.259330976534118}, {'loss': {'loss': 0.0007888983119634765}}))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ax_client.get_best_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO 09-13 20:03:24] ax.service.ax_client: Saved JSON-serialized state of optimization to `../model_artifacts/pretraining_hparam_search.json`.\n"
     ]
    }
   ],
   "source": [
    "ax_client.save_to_json_file(\n",
    "    '../model_artifacts/pretraining_hparam_search.json')"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}