{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from collections.abc import Iterable\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from operator import lt, gt, add, sub\n",
    "import os\n",
    "import pandas as pd\n",
    "import requests\n",
    "from tabulate import tabulate\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import warnings\n",
    "\n",
    "from accio.s3tool import S3tool\n",
    "from htools import auto_repr, valuecheck, save\n",
    "from incendio.utils import DEVICE\n",
    "from incendio.optimizers import variable_lr_optimizer, update_optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@auto_repr\n",
    "class TorchCallback:\n",
    "\n",
    "    def on_train_begin(self, trainer, epochs, lrs, lr_mult, **kwargs):\n",
    "        pass\n",
    "\n",
    "    def on_train_end(self, trainer, epoch, stats, val_stats):\n",
    "        pass\n",
    "\n",
    "    def on_epoch_begin(self, trainer, epoch, stats, val_stats):\n",
    "        pass\n",
    "\n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        pass\n",
    "\n",
    "    def on_batch_begin(self, trainer, i, sum_i, stats):\n",
    "        pass\n",
    "\n",
    "    def on_batch_end(self, trainer, i, sum_i, stats):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SchedulerMixin(TorchCallback):\n",
    "\n",
    "    verbose = False\n",
    "\n",
    "    def on_train_end(self, trainer, *args, **kwargs):\n",
    "        self.plot_lrs(os.path.join(trainer.out_dir, 'lrs.png'))\n",
    "\n",
    "    def update_lr(self, trainer, n):\n",
    "        try:\n",
    "            lr = self.lrs[n]\n",
    "        except IndexError as e:\n",
    "            return\n",
    "\n",
    "        update_optimizer(trainer.optim, lr, lr_mult=self.lr_mult)\n",
    "        if self.verbose:\n",
    "            trainer.logger.info(f'Set learning rate to {lr:.4f}.')\n",
    "\n",
    "    def plot_lrs(self, path=None):\n",
    "        \"\"\"Display learning rate by iteration.\n",
    "\n",
    "        Note: If the plot is not as smooth as expected, this likely\n",
    "        means that there are very few iterations per epoch\n",
    "        (i.e. the batch size is very large, at least in relative terms).\n",
    "        \"\"\"\n",
    "        plt.plot(self.lrs)\n",
    "        plt.xlabel('Iteration')\n",
    "        plt.ylabel('Learning Rate')\n",
    "        plt.title('Learning Rate Schedule')\n",
    "        if path:\n",
    "            plt.savefig(path)\n",
    "            plt.close()\n",
    "        else:\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SawtoothScheduler(SchedulerMixin):\n",
    "\n",
    "    def __init__(self, add=1e-5, scale=0.6, patience=5, priority=10):\n",
    "        self.add = add\n",
    "        self.scale = scale\n",
    "        self.patience = patience\n",
    "        self.priority = priority\n",
    "        \n",
    "        # These are reset in `on_train_begin`, but types remain the same.\n",
    "        self.lrs = []\n",
    "        self.since_improve = 0\n",
    "        self.recent_best = float('inf')\n",
    "        self.lr_mult = 1.0\n",
    "\n",
    "    def on_train_begin(self, trainer, epochs, lrs, lr_mult, **kwargs):\n",
    "        \"\"\"Wrapper to schedule learning rates depending on chosen method.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        restarts: bool\n",
    "            If True, use schedule with restarts. If False, use regular\n",
    "            cosine annealing that spans whole duration of training.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: LR for each iteration (i.e. output[i] is the LR to use\n",
    "            at iteration i).\n",
    "        \"\"\"\n",
    "        self.lrs.clear()\n",
    "        self.since_improve = 0\n",
    "        self.recent_best = float('inf')\n",
    "        self.lr_mult = lr_mult\n",
    "\n",
    "    def on_batch_begin(self, trainer, i, sum_i, stats):\n",
    "        loss = stats.get('loss')\n",
    "        if loss is None: return\n",
    "        \n",
    "        lr = max(p['lr'] for p in trainer.optim.param_groups)\n",
    "        if loss < self.recent_best:\n",
    "            self.recent_best = loss\n",
    "            self.since_improve = 0\n",
    "            lr += self.add\n",
    "        elif loss >= self.recent_best and self.since_improve < self.patience:\n",
    "            self.since_improve += 1\n",
    "            lr += self.add / (self.since_improve+1)\n",
    "        else:\n",
    "            self.since_improve += 1\n",
    "            lr *= self.scale\n",
    "        update_optimizer(trainer.optim, lr, lr_mult=self.lr_mult)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.003"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Net(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_size, out_size):\n",
    "        super().__init__()\n",
    "        fc1 = nn.Linear(in_size, out_size)\n",
    "        fc2 = nn.Linear(out_size, 1)\n",
    "        self.groups = nn.ModuleList([fc1, fc2])\n",
    "        \n",
    "    def forward(self, x):\n",
    "        for g in self.groups:\n",
    "            x = torch.sigmoid(g(x))\n",
    "        return x\n",
    "\n",
    "net = Net(10, 4)\n",
    "opt = variable_lr_optimizer(net, [2e-3, 3e-3])\n",
    "max(p['lr'] for p in opt.param_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CosineLRScheduler(SchedulerMixin):\n",
    "    \"\"\"Learning rate scheduler that makes updates each batch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, warm=0.3, restarts=False, cycle_len=5, cycle_decay=0.0,\n",
    "                 min_lr=None, verbose=False, priority=10):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        warm: float\n",
    "            Percent of training run (or cycle length) devoted to the increasing\n",
    "            portion of the schedule. Default 0.3.\n",
    "        restarts: bool\n",
    "            Specifies whether to use restarts, i.e. use a cyclical LR.\n",
    "            True: Version of cosine annealing with restarts. In one\n",
    "                  cycle, LR starts high and gradually decreases.\n",
    "                  At the start of the next cycle, it is\n",
    "                  immediately increased again.\n",
    "            False: Version of cosine annealing where LR increases\n",
    "                   for first 30% of training, then decreases for\n",
    "                   remaining 70%.\n",
    "        cycle_len: int\n",
    "            Number of epochs contained in a single cycle. Only used\n",
    "            when scheduler uses restarts.\n",
    "        cycle_decay: float\n",
    "            Scalar to decay the learning rate at the end of each cycle.\n",
    "            This is only used with restarts, since the regular cosine\n",
    "            annealing already decays the LR over time.\n",
    "            E.g. 1.0 will use no decay.\n",
    "            0.9 means that cycle 2 LRs = cycle 1 LRs * 0.9,\n",
    "            cycle 3 LRs = cycle 1 LRs * .81,\n",
    "            etc.\n",
    "        min_lr: float\n",
    "            Minimum learning rate. If None is specified, it will be set\n",
    "            to max_lr / 10.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.warm = warm\n",
    "        self.cycle_len = cycle_len\n",
    "        self.cycle_decay = cycle_decay\n",
    "        self.restarts = restarts\n",
    "        self.verbose = verbose\n",
    "        self.min_lr = min_lr\n",
    "        self.priority = priority\n",
    "\n",
    "        # Set in `on_train_begin()`.\n",
    "        self.lrs = None             # Iterable[float]\n",
    "        self.batches_per_e = None   # int\n",
    "        self.batches = None         # int\n",
    "        self.max_lr = None          # float\n",
    "        self.lr_mult = None         # float\n",
    "\n",
    "    def on_train_begin(self, trainer, epochs, lrs, lr_mult, **kwargs):\n",
    "        \"\"\"Wrapper to schedule learning rates depending on chosen method.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        restarts: bool\n",
    "            If True, use schedule with restarts. If False, use regular\n",
    "            cosine annealing that spans whole duration of training.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: LR for each iteration (i.e. output[i] is the LR to use\n",
    "            at iteration i).\n",
    "        \"\"\"\n",
    "        self.batches_per_e = len(trainer.dl_train)\n",
    "        self.batches = epochs * self.batches_per_e\n",
    "        self.max_lr = max(lrs) if isinstance(lrs, Iterable) else lrs\n",
    "        self.lr_mult = lr_mult\n",
    "        if not self.min_lr: self.min_lr = self.max_lr / 10\n",
    "\n",
    "        if self.restarts and self.batches < self.cycle_len:\n",
    "            warnings.warn('Training will be less than 1 full cycle.')\n",
    "\n",
    "        if self.restarts:\n",
    "            self.lrs = self._cosine_restarts_schedule()\n",
    "        else:\n",
    "            self.lrs = self._cosine_schedule()\n",
    "\n",
    "    def on_batch_begin(self, trainer, i, sum_i, stats):\n",
    "        self.update_lr(trainer, sum_i)\n",
    "\n",
    "    @staticmethod\n",
    "    def _cosine_anneal(batches, lr1, lr2):\n",
    "        \"\"\"Helper function for _cosine_schedule().\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batches: int\n",
    "            Number of batches in segment.\n",
    "        lr1: float\n",
    "            Learning rate at start of segment.\n",
    "        lr2: float\n",
    "            Learning rate at end of segment.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array\n",
    "        \"\"\"\n",
    "        i = np.arange(batches)\n",
    "        return lr2 + (lr1 - lr2)*(1 + np.cos(np.pi * i/batches))/2\n",
    "\n",
    "    def _cosine_schedule(self):\n",
    "        \"\"\"Cosine annealing scheduler. Computes learning rates for each\n",
    "        iteration.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array\n",
    "        \"\"\"\n",
    "        seg1 = self._cosine_anneal(int(self.warm * self.batches),\n",
    "                                   self.min_lr, self.max_lr)\n",
    "        seg2 = self._cosine_anneal(int(np.ceil((1 - self.warm) * self.batches)),\n",
    "                                   self.max_lr, self.min_lr)\n",
    "        return np.concatenate((seg1, seg2))\n",
    "\n",
    "    def _cosine_restarts_schedule(self):\n",
    "        \"\"\"Cosine annealing with restarts.\"\"\"\n",
    "        cycles = int(np.ceil(self.batches / (self.cycle_len * self.batches_per_e)))\n",
    "        cycle_batches = self.cycle_len * self.batches_per_e\n",
    "        lrs = [self._cosine_anneal(cycle_batches, self.max_lr, self.min_lr)\n",
    "               / (1 + self.cycle_decay * i) for i in range(cycles)]\n",
    "        return np.concatenate(lrs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
