{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from htools.ml import BaseModel, GRelu, JRelu, variable_lr_optimizer, stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.path.exists('scratch_inheritance_and_mixins.ipynb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HNet(BaseModel):\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dim):\n",
    "        super().__init__(locals())\n",
    "        layers = [nn.Linear(x_dim, hidden_dim),\n",
    "                  nn.LeakyReLU(),\n",
    "                  nn.Linear(hidden_dim, 3),\n",
    "                  nn.Softmax(-1)]\n",
    "        self.layers = nn.Sequential(*layers)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HNet(\n",
       "  (layers): Sequential(\n",
       "    (0): Linear(in_features=4, out_features=6, bias=True)\n",
       "    (1): LeakyReLU(negative_slope=0.01)\n",
       "    (2): Linear(in_features=6, out_features=3, bias=True)\n",
       "    (3): Softmax(dim=-1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet = HNet(4, 6)\n",
    "hnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[7., 5., 8., 2.],\n",
       "        [1., 3., 8., 2.],\n",
       "        [9., 7., 4., 0.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(10, (3, 4), dtype=torch.float)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2730, 0.3260, 0.4010],\n",
       "        [0.1617, 0.3326, 0.5057],\n",
       "        [0.3660, 0.3071, 0.3270]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-0.071, 0.261), (-0.023, 0.379), (0.018, 0.234), (0.197, 0.093)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 99 weights saved to ../data/model_e99_v2.pth.\n"
     ]
    }
   ],
   "source": [
    "hnet.save(99, '../data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 99 weights loaded from ../data/model_e99.pth.\n",
      "Model parameters: {'x_dim': 4, 'hidden_dim': 6}\n",
      "Currently in eval mode.\n"
     ]
    }
   ],
   "source": [
    "hnet2 = HNet.from_path('../data/model_e99.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.6451, 0.2300, 0.1250],\n",
       "        [0.4431, 0.2086, 0.3483],\n",
       "        [0.8782, 0.0980, 0.0237]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-0.005, 0.298), (-0.018, 0.273), (0.013, 0.198), (0.045, 0.127)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet2.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hnet3 = HNet.from_path('../data/model_e99_v2.pth', verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-0.071, 0.261), (-0.023, 0.379), (0.018, 0.234), (0.197, 0.093)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hnet3.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Nested(BaseModel):\n",
    "    \n",
    "    def __init__(self, x_dim, hidden):\n",
    "        super().__init__(locals())\n",
    "        seq1 = nn.Sequential(nn.Linear(x_dim, hidden),\n",
    "                             nn.Dropout(),\n",
    "                             nn.LeakyReLU())\n",
    "        seq2 = nn.Sequential(nn.Linear(hidden, 1),\n",
    "                             nn.Sigmoid())\n",
    "        self.groups = nn.ModuleList([seq1, seq2])\n",
    "        \n",
    "    def forward(self, x):\n",
    "        for group in self.groups:\n",
    "            x = group(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Nested(\n",
       "  (groups): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "      (1): Dropout(p=0.5, inplace=False)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): Linear(in_features=8, out_features=1, bias=True)\n",
       "      (1): Sigmoid()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nested = Nested(4, 8)\n",
    "nested"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Adam (\n",
       "Parameter Group 0\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 1e-08\n",
       "    lr: 0.003\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variable_lr_optimizer(nested)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Adam (\n",
       "Parameter Group 0\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 1e-08\n",
       "    lr: 1\n",
       "    weight_decay: 0\n",
       "\n",
       "Parameter Group 1\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 1e-08\n",
       "    lr: 2\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variable_lr_optimizer(groups=nested.groups, lrs=[1, 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cpu'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epochs, train_dl, val_dl, model, criterion, hooks, lr=3e-3, \n",
    "          metrics=None, verbose=True):\n",
    "    for hook in hooks:\n",
    "        model.register_backward_hook(hook)\n",
    "        \n",
    "    optim = variable_lr_optimizer()\n",
    "    stats = defaultdict(list)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        model.to(device).train()\n",
    "        e_loss = 0.0\n",
    "        e_total = 0\n",
    "#         for i, (x, y) in enumerate(train_dl):\n",
    "#             optim.zero_grad()\n",
    "#             x.to(device)\n",
    "#             y.to(device)\n",
    "#             bs = x.shape[0]\n",
    "            \n",
    "#             # Forward pass\n",
    "#             y_hat = model(x)\n",
    "#             loss = criterion(y_hat, y, reduction='mean')\n",
    "            \n",
    "#             # Backward pass\n",
    "#             loss.backward()\n",
    "#             optim.step()\n",
    "            \n",
    "#             # Update mini batch stats.\n",
    "#             e_total += bs\n",
    "#             loss += loss * bs\n",
    "            \n",
    "#         # Evaluate on validation set.\n",
    "#         val_stats = validation_metrics()\n",
    "        \n",
    "#         # Update epoch stats.\n",
    "#         stats['loss'].append(e_total)\n",
    "#         stats['val_loss'].append()\n",
    "        \n",
    "        # Print epoch stats.\n",
    "            \n",
    "    return stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_stats_hook(model, grad_in, grad_out):\n",
    "    print(stats(grad_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "variable_lr_optimizer() missing 2 required positional arguments: 'groups' and 'lrs'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-93-213e077428ec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhnet\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBCEWithLogitsLoss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mgradient_stats_hook\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-91-1ddb8bc2cf74>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(epochs, train_dl, val_dl, model, criterion, hooks, metrics, verbose)\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mregister_backward_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m     \u001b[0moptim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvariable_lr_optimizer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m     \u001b[0mstats\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdefaultdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: variable_lr_optimizer() missing 2 required positional arguments: 'groups' and 'lrs'"
     ]
    }
   ],
   "source": [
    "train(3, None, None, hnet, nn.BCEWithLogitsLoss, [gradient_stats_hook])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
