{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-05 22:18:20,653\tINFO worker.py:1774 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://172.29.235.222:8265 \u001b[39m\u001b[22m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Using device cuda:0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = (\n",
    "    \"expandable_segments:True\"  # avoid fragmented CUDA memory\n",
    ")\n",
    "\n",
    "from config import LOGS_PATH, RUNS_PATH, TRAIN_DATA, TEST_DATA\n",
    "from utils import set_up_logging\n",
    "\n",
    "set_up_logging(LOGS_PATH)\n",
    "\n",
    "from utils import get_device\n",
    "from training import train_with_ray_factory\n",
    "from ray import tune\n",
    "import ray\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "from ray.air import RunConfig\n",
    "\n",
    "\n",
    "TRIAL_COUNT = 50\n",
    "EPOCH_COUNT = 4\n",
    "CHUNK_COUNT = EPOCH_COUNT * 40\n",
    "\n",
    "ray.init(include_dashboard=True, dashboard_host=\"0.0.0.0\")\n",
    "\n",
    "device = get_device()\n",
    "f\"Using device {device}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"tuneStatus\">\n",
       "  <div style=\"display: flex;flex-direction: row\">\n",
       "    <div style=\"display: flex;flex-direction: column;\">\n",
       "      <h3>Tune Status</h3>\n",
       "      <table>\n",
       "<tbody>\n",
       "<tr><td>Current time:</td><td>2024-09-05 22:23:47</td></tr>\n",
       "<tr><td>Running for: </td><td>00:05:24.00        </td></tr>\n",
       "<tr><td>Memory:      </td><td>22.7/54.9 GiB      </td></tr>\n",
       "</tbody>\n",
       "</table>\n",
       "    </div>\n",
       "    <div class=\"vDivider\"></div>\n",
       "    <div class=\"systemInfo\">\n",
       "      <h3>System Info</h3>\n",
       "      Using AsyncHyperBand: num_stopped=0<br>Bracket: Iter 128.000: None | Iter 32.000: None | Iter 8.000: None | Iter 2.000: None<br>Logical resource usage: 32.0/32 CPUs, 1.0/1 GPUs (0.0/1.0 accelerator_type:G)\n",
       "    </div>\n",
       "    <div class=\"vDivider\"></div>\n",
       "<div class=\"messages\">\n",
       "  <h3>Messages</h3>\n",
       "  \n",
       "  ... 30 more trials not shown (30 PENDING)\n",
       "  \n",
       "</div>\n",
       "<style>\n",
       ".messages {\n",
       "  color: var(--jp-ui-font-color1);\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "  padding-left: 1em;\n",
       "  overflow-y: auto;\n",
       "}\n",
       ".messages h3 {\n",
       "  font-weight: bold;\n",
       "}\n",
       ".vDivider {\n",
       "  border-left-width: var(--jp-border-width);\n",
       "  border-left-color: var(--jp-border-color0);\n",
       "  border-left-style: solid;\n",
       "  margin: 0.5em 1em 0.5em 1em;\n",
       "}\n",
       "</style>\n",
       "\n",
       "  </div>\n",
       "  <div class=\"hDivider\"></div>\n",
       "  <div class=\"trialStatus\">\n",
       "    <h3>Trial Status</h3>\n",
       "    <table>\n",
       "<thead>\n",
       "<tr><th>Trial name                </th><th>status  </th><th>loc                   </th><th style=\"text-align: right;\">  dropout_prob</th><th>features            </th><th style=\"text-align: right;\">  kernel_size</th><th style=\"text-align: right;\">  leaky_relu_alpha</th><th style=\"text-align: right;\">  leaky_relu_slope</th><th style=\"text-align: right;\">  learning_rate</th><th style=\"text-align: right;\">  scheduler_gamma</th></tr>\n",
       "</thead>\n",
       "<tbody>\n",
       "<tr><td>train_with_ray_61b01_00000</td><td>RUNNING </td><td>172.29.235.222:3274755</td><td style=\"text-align: right;\">    0.0240677 </td><td>[16, 32, 64, 12_25c0</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0147726 </td><td style=\"text-align: right;\">       0.0164835  </td><td style=\"text-align: right;\">    0.00409419 </td><td style=\"text-align: right;\">         0.977701</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00001</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0252525 </td><td>[16, 32, 64, 12_9940</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0607274 </td><td style=\"text-align: right;\">       0.00769014 </td><td style=\"text-align: right;\">    0.00175452 </td><td style=\"text-align: right;\">         0.971626</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00002</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0492015 </td><td>[16, 32, 64, 12_0940</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.0522751 </td><td style=\"text-align: right;\">       0.00532775 </td><td style=\"text-align: right;\">    0.000102449</td><td style=\"text-align: right;\">         0.985839</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00003</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0091843 </td><td>[16, 32, 64, 12_8e80</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.0215584 </td><td style=\"text-align: right;\">       0.00986109 </td><td style=\"text-align: right;\">    0.0046811  </td><td style=\"text-align: right;\">         0.969517</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00004</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0200707 </td><td>[16, 32, 64, 12_e200</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.00839403</td><td style=\"text-align: right;\">       0.0236732  </td><td style=\"text-align: right;\">    0.00456458 </td><td style=\"text-align: right;\">         0.949592</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00005</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.071518  </td><td>[16, 32, 64, 12_4200</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.00955745</td><td style=\"text-align: right;\">       0.000338297</td><td style=\"text-align: right;\">    0.0012211  </td><td style=\"text-align: right;\">         0.959114</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00006</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.00712943</td><td>[16, 32, 64, 12_f8c0</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0128474 </td><td style=\"text-align: right;\">       0.0214861  </td><td style=\"text-align: right;\">    0.000319318</td><td style=\"text-align: right;\">         0.954027</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00007</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0554306 </td><td>[16, 32, 64, 12_1900</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0686442 </td><td style=\"text-align: right;\">       0.0103842  </td><td style=\"text-align: right;\">    0.00856337 </td><td style=\"text-align: right;\">         0.966035</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00008</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0637642 </td><td>[16, 32, 64, 12_8f80</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.0125659 </td><td style=\"text-align: right;\">       0.00386933 </td><td style=\"text-align: right;\">    0.00135947 </td><td style=\"text-align: right;\">         0.955525</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00009</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0526107 </td><td>[16, 32, 64, 12_16c0</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.00576232</td><td style=\"text-align: right;\">       0.0120612  </td><td style=\"text-align: right;\">    0.000143394</td><td style=\"text-align: right;\">         0.983664</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00010</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0321631 </td><td>[16, 32, 64, 12_fc00</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0100027 </td><td style=\"text-align: right;\">       0.016943   </td><td style=\"text-align: right;\">    0.000461478</td><td style=\"text-align: right;\">         0.981847</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00011</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.015057  </td><td>[16, 32, 64, 12_f080</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0599078 </td><td style=\"text-align: right;\">       0.0293855  </td><td style=\"text-align: right;\">    0.00182771 </td><td style=\"text-align: right;\">         0.965469</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00012</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0995986 </td><td>[16, 32, 64, 12_d240</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0349584 </td><td style=\"text-align: right;\">       0.000219269</td><td style=\"text-align: right;\">    0.000288736</td><td style=\"text-align: right;\">         0.966034</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00013</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0227281 </td><td>[16, 32, 64, 12_c600</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.0324662 </td><td style=\"text-align: right;\">       0.0125853  </td><td style=\"text-align: right;\">    0.000173035</td><td style=\"text-align: right;\">         0.948412</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00014</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0307992 </td><td>[16, 32, 64, 12_d040</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.062771  </td><td style=\"text-align: right;\">       0.0198744  </td><td style=\"text-align: right;\">    0.000639558</td><td style=\"text-align: right;\">         0.94965 </td></tr>\n",
       "<tr><td>train_with_ray_61b01_00015</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0896525 </td><td>[16, 32, 64, 12_2e80</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0150036 </td><td style=\"text-align: right;\">       0.0171698  </td><td style=\"text-align: right;\">    0.00206888 </td><td style=\"text-align: right;\">         0.99602 </td></tr>\n",
       "<tr><td>train_with_ray_61b01_00016</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0932697 </td><td>[16, 32, 64, 12_6540</td><td style=\"text-align: right;\">            5</td><td style=\"text-align: right;\">        0.0699308 </td><td style=\"text-align: right;\">       0.0154819  </td><td style=\"text-align: right;\">    0.00384655 </td><td style=\"text-align: right;\">         0.954913</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00017</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.00240956</td><td>[16, 32, 64, 12_8640</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0318032 </td><td style=\"text-align: right;\">       0.00959876 </td><td style=\"text-align: right;\">    0.000135257</td><td style=\"text-align: right;\">         0.952462</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00018</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0883309 </td><td>[16, 32, 64, 12_b800</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0455664 </td><td style=\"text-align: right;\">       0.00565    </td><td style=\"text-align: right;\">    0.00066994 </td><td style=\"text-align: right;\">         0.997372</td></tr>\n",
       "<tr><td>train_with_ray_61b01_00019</td><td>PENDING </td><td>                      </td><td style=\"text-align: right;\">    0.0500287 </td><td>[16, 32, 64, 12_8440</td><td style=\"text-align: right;\">            3</td><td style=\"text-align: right;\">        0.0674193 </td><td style=\"text-align: right;\">       0.0220701  </td><td style=\"text-align: right;\">    0.000987128</td><td style=\"text-align: right;\">         0.94602 </td></tr>\n",
       "</tbody>\n",
       "</table>\n",
       "  </div>\n",
       "</div>\n",
       "<style>\n",
       ".tuneStatus {\n",
       "  color: var(--jp-ui-font-color1);\n",
       "}\n",
       ".tuneStatus .systemInfo {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "}\n",
       ".tuneStatus td {\n",
       "  white-space: nowrap;\n",
       "}\n",
       ".tuneStatus .trialStatus {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "}\n",
       ".tuneStatus h3 {\n",
       "  font-weight: bold;\n",
       "}\n",
       ".tuneStatus .hDivider {\n",
       "  border-bottom-width: var(--jp-border-width);\n",
       "  border-bottom-color: var(--jp-border-color0);\n",
       "  border-bottom-style: solid;\n",
       "}\n",
       ".tuneStatus .vDivider {\n",
       "  border-left-width: var(--jp-border-width);\n",
       "  border-left-color: var(--jp-border-color0);\n",
       "  border-left-style: solid;\n",
       "  margin: 0.5em 1em 0.5em 1em;\n",
       "}\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33m(raylet)\u001b[0m Warning: The actor ImplicitFunc is very large (26 MiB). Check that its definition is not implicitly capturing a large array or other object in scope. Tip: use ray.put() to put large objects in the Ray object store.\n"
     ]
    }
   ],
   "source": [
    "config = {\n",
    "    \"batch_size\": 48,\n",
    "    \"edit_count\": EPOCH_COUNT,\n",
    "    \"bin_count\": 32,\n",
    "    \"learning_rate\": tune.loguniform(1e-4, 1e-2),\n",
    "    \"scheduler_gamma\": tune.uniform(0.94, 0.9999),\n",
    "    # \"elu_alpha\": tune.uniform(0.5, 2),\n",
    "    \"leaky_relu_slope\": tune.uniform(0, 0.03),\n",
    "    \"dropout_prob\": tune.uniform(0, 0.1),\n",
    "    \"chunk_count\": CHUNK_COUNT,\n",
    "    \"features\": tune.choice(\n",
    "        [\n",
    "            [16, 32, 64],\n",
    "            [16, 32, 64, 128],\n",
    "            [16, 32, 64, 128, 256],\n",
    "            [16, 32, 32, 32, 64],\n",
    "            [32, 64],\n",
    "            [32, 128],\n",
    "            [32, 64, 128],\n",
    "            [32, 64, 128, 256],\n",
    "            [16, 16, 16, 16, 16],\n",
    "            [16, 16, 16, 16, 16, 16, 16, 16],\n",
    "            [16, 16, 16, 16, 16, 16, 16, 16, 16, 16],\n",
    "            [32, 32, 32],\n",
    "            [32, 32, 32, 32],\n",
    "            [64, 64, 64],\n",
    "            [64, 64, 64, 64],\n",
    "            [64, 64, 64, 64, 64],\n",
    "            [256, 64, 256],\n",
    "        ]\n",
    "    ),\n",
    "    \"use_residual\": True,\n",
    "    \"kernel_size\": tune.choice([3, 5]),\n",
    "    \"model_type\": \"HistogramNet\",\n",
    "    \"use_instance_norm\": True,\n",
    "    \"use_elu\": False,\n",
    "    \"leaky_relu_alpha\": tune.uniform(0, 0.07),\n",
    "}\n",
    "\n",
    "tuner = tune.Tuner(\n",
    "    tune.with_resources(\n",
    "        tune.with_parameters(\n",
    "            train_with_ray_factory(\n",
    "                train_data_paths=TRAIN_DATA,\n",
    "                test_data_paths=TEST_DATA,\n",
    "                device=device,\n",
    "                log_dir=RUNS_PATH / \"custom\",\n",
    "            )\n",
    "        ),\n",
    "        resources={\"cpu\": 32, \"gpu\": 1},\n",
    "    ),\n",
    "    run_config=RunConfig(storage_path=RUNS_PATH, name=\"tune\"),\n",
    "    tune_config=tune.TuneConfig(\n",
    "        metric=\"chunk_test_loss\",\n",
    "        mode=\"min\",\n",
    "        scheduler=ASHAScheduler(\n",
    "            max_t=config[\"chunk_count\"], grace_period=2, time_attr=\"chunk_id\"\n",
    "        ),\n",
    "        num_samples=TRIAL_COUNT,\n",
    "    ),\n",
    "    param_space=config,\n",
    ")\n",
    "results = tuner.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best trial config: {'batch_size': 32, 'edit_count': 3, 'bin_count': 32, 'learning_rate': 0.0029226033808016005, 'scheduler_gamma': 0.8167516482513361, 'elu_alpha': 1.3510723758569865, 'leaky_relu_slope': 0.0066452562138349025, 'dropout_prob': 0.0039589213934103865, 'chunk_count': 4, 'features': [8, 16, 32], 'use_residual': True, 'kernel_size': 5, 'model_type': 'HistogramNet', 'use_instance_norm': True, 'use_elu': True, 'leaky_relu_alpha': 0.04487148648446764}\n",
      "Best trial final validation loss: 5.2214884757995605\n"
     ]
    }
   ],
   "source": [
    "best_result = results.get_best_result(\"chunk_test_loss\", \"min\")\n",
    "\n",
    "print(\"Best trial config: {}\".format(best_result.config))\n",
    "print(\n",
    "    \"Best trial final validation loss: {}\".format(\n",
    "        best_result.metrics[\"chunk_test_loss\"]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hparams = {\n",
    "#     \"batch_size\": 64,\n",
    "#     \"edit_count\": 12,\n",
    "#     \"bin_count\": 16,\n",
    "#     \"learning_rate\": 0.0006126108207352808,\n",
    "#     \"scheduler_gamma\": 0.9382286228762693,\n",
    "#     \"num_epochs\": 10,\n",
    "#     \"elu_alpha\": 1.3092260477215776,\n",
    "#     \"leaky_relu_slope\": 0.029438156325552762,\n",
    "#     \"dropout_prob\": 0.06261255195786307,\n",
    "#     \"features\": [8, 16, 32],\n",
    "#     \"use_residual\": True,\n",
    "#     \"kernel_size\": 5,\n",
    "#     \"model_type\": \"HistogramNet\",\n",
    "#     \"use_instance_norm\": True,\n",
    "#     \"use_elu\": False,\n",
    "#     \"leaky_relu_alpha\": 0.03745605986732464,\n",
    "# }\n",
    "\n",
    "# train(\n",
    "#     hparams,\n",
    "#     train_data_paths=TRAIN_DATA,\n",
    "#     test_data_paths=TEST_DATA,\n",
    "#     log_dir=RUNS_PATH,\n",
    "#     max_duration=None,\n",
    "#     use_tqdm=True,\n",
    "#     device=device,\n",
    "#     **hparams\n",
    "# )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bipolaroid",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
