{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "from utils import set_up_logging\n",
    "from training import train, random_hparam_search\n",
    "from config import LOGS_PATH, RUNS_PATH, TRAIN_DATA, TEST_DATA, MODELS_PATH\n",
    "\n",
    "set_up_logging(LOGS_PATH)\n",
    "\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "f\"Using device {device}\""
   ]
  },
  {
   "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.001,\n",
    "#     \"scheduler_gamma\": 0.9,\n",
    "#     \"num_epochs\": 12,\n",
    "#     \"model_type\": \"SimpleCNN\",\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",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import loguniform, uniform, randint\n",
    "from models import MODELS, test_models\n",
    "\n",
    "\n",
    "hyperparameters = [\n",
    "    {\n",
    "        \"batch_size\": [32, 64, 128],\n",
    "        \"edit_count\": [12],\n",
    "        \"bin_count\": [16],\n",
    "        \"learning_rate\": loguniform(5e-4, 5e-3),\n",
    "        \"scheduler_gamma\": uniform(loc=0.8, scale=0.15),\n",
    "        \"num_epochs\": [12],\n",
    "        \"elu_alpha\": uniform(0.5, 1.5),\n",
    "        \"leaky_relu_slope\": uniform(0, 0.03),\n",
    "        \"dropout_prob\": uniform(0, 0.1),\n",
    "        \"features\": [\n",
    "            [16, 32],\n",
    "            [16, 32, 64],\n",
    "            [16, 32, 64, 128],\n",
    "            [32, 64],\n",
    "            [32, 64, 128],\n",
    "            [8, 16, 32],\n",
    "            [8, 8, 8],\n",
    "            [8, 8, 8, 8, 8],\n",
    "            [8, 8, 8, 8, 8, 8, 8],\n",
    "            [16, 16, 16, 16, 16],\n",
    "            [16, 16, 16],\n",
    "            [32, 32],    \n",
    "            [32, 32, 32],\n",
    "            [32, 32, 32, 32],\n",
    "            [64, 64],\n",
    "            [64, 64, 64]\n",
    "        ],\n",
    "        \"use_residual\": [True, False],\n",
    "        \"kernel_size\": [3, 5],\n",
    "        \"model_type\": [\"HistogramNet\"],\n",
    "        \"use_instance_norm\": [True, False],\n",
    "        \"use_elu\": [True, False],\n",
    "        \"leaky_relu_alpha\": uniform(0, 0.05),\n",
    "    }\n",
    "]\n",
    "\n",
    "test_models()\n",
    "\n",
    "random_hparam_search(\n",
    "    hyperparameters=hyperparameters,\n",
    "    train_data_paths=TRAIN_DATA,\n",
    "    test_data_paths=TEST_DATA,\n",
    "    models_path=MODELS_PATH,\n",
    "    tensorboard_path=RUNS_PATH,\n",
    "    timeout_hours=4,\n",
    "    device=device,\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.1.-1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}