{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Callbacks\n",
    "\n",
    "> API details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "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": [
    "# export\n",
    "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",
    "from tabulate import tabulate\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": [
    "# export\n",
    "@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": [
    "# export\n",
    "class BasicConfig(TorchCallback):\n",
    "    \"\"\"Handles basic model tasks like putting the model on the GPU\n",
    "    and switching between train and eval modes.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, priority=0):\n",
    "        self.priority = priority\n",
    "    \n",
    "    def on_train_begin(self, trainer, epochs, lrs, lr_mult, **kwargs):\n",
    "        trainer.net.to(DEVICE)\n",
    "        if not trainer.optim:\n",
    "            trainer.optim = variable_lr_optimizer(\n",
    "                trainer.net, lrs, lr_mult, trainer.optim_type, trainer.eps\n",
    "            )\n",
    "        else:\n",
    "            update_optimizer(trainer.optim, lrs, lr_mult=lr_mult)\n",
    "        trainer.logger.info(trainer.optim)\n",
    "        if kwargs.get('clean') is True: trainer.cleanup(confirmed=True)\n",
    "\n",
    "    def on_epoch_begin(self, trainer, *args, **kwargs):\n",
    "        trainer.net.train()\n",
    "        \n",
    "    def on_train_end(self, trainer, *args, **kwargs):\n",
    "        trainer.logger.info('Training complete. Model in eval mode.')\n",
    "        trainer.net.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class StatsHandler(TorchCallback):\n",
    "    \"\"\"This updates metrics at the end of each epoch to account for\n",
    "    potentially varying batch sizes.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, priority=5):\n",
    "        self.priority = priority\n",
    "        \n",
    "    def on_epoch_begin(self, trainer, epoch, stats, val_stats):\n",
    "        \"\"\"Resets stats at the start of each epoch.\"\"\"\n",
    "        stats.clear()\n",
    "        \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        \"\"\"Computes (possibly weighted) averages of mini-batch stats\n",
    "        at the end of each epoch.\n",
    "        \"\"\"\n",
    "        for group in (stats, val_stats):\n",
    "            for k, v in group.items():\n",
    "                if k == 'batch_size': continue\n",
    "                group[k] = np.average(v, weights=group['batch_size'])\n",
    "            group.pop('batch_size')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class MetricPrinter(TorchCallback):\n",
    "    \"\"\"Prints metrics at the end of each epoch. This is one of the \n",
    "    default callbacks provided in BaseModel - it does not need to\n",
    "    be passed in explicitly.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, priority=10):\n",
    "        \"\"\"Priority must be higher than StatsHandler, otherwise\n",
    "        metrics will be printed before they're aggregated.\n",
    "        \"\"\"\n",
    "        self.priority = priority\n",
    "        \n",
    "    def on_train_begin(self, trainer, *args, **kwargs):\n",
    "        trainer.logger = trainer.get_logger(\n",
    "            os.path.join(trainer.out_dir, 'train.log'),\n",
    "            fmt='\\n%(asctime)s\\n %(message)s'\n",
    "        )\n",
    "    \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        data = [[k, v, val_stats[k]] for k, v in stats.items()]\n",
    "        table = tabulate(data, headers=['Metric', 'Train', 'Validation'], \n",
    "                         tablefmt='github', floatfmt='.4f')\n",
    "        trainer.logger.info(\n",
    "            f'\\n{\"=\"*5}\\n\\nEpoch {epoch}\\n\\n{table}\\n\\n{\"=\"*5}'\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class EarlyStopper(TorchCallback):\n",
    "    \n",
    "    @valuecheck\n",
    "    def __init__(self, metric, goal:('max', 'min'), min_improvement=0.0, \n",
    "                 patience=3, priority=15):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        metric: str\n",
    "            Quantity to monitor. This will always be computed on the \n",
    "            validation set.\n",
    "        goal: str\n",
    "            Indicates what we want to do to the metric in question.\n",
    "            Either 'min' or 'max'. E.g. metric 'loss' should have goal 'min'\n",
    "            while metric 'precision' should have goal 'max'.\n",
    "        min_improvement: float\n",
    "            Amount of change needed to qualify as improvement. For example,\n",
    "            min_improvement of 0.0 means any improvement is sufficient. With\n",
    "            a min_improvent of 0.2, we will stop training even if the\n",
    "            quantity improves by, for example, 0.1.\n",
    "        patience: int\n",
    "            Number of acceptable epochs without improvement. E.g. patience=0 \n",
    "            means the metric must improve every epoch for training to continue.            \n",
    "        \"\"\"\n",
    "        # Will use op like: self.op(new_val, current_best)\n",
    "        if goal == 'min':\n",
    "            self.init_metric = self.best_metric = float('inf')\n",
    "            self.op = lt\n",
    "            self.op_best = sub\n",
    "        elif goal == 'max':\n",
    "            self.init_metric = self.best_metric = float('-inf')\n",
    "            self.op = gt\n",
    "            self.op_best = add\n",
    "           \n",
    "        self.priority = priority\n",
    "        self.metric = metric\n",
    "        self.min_improvement = min_improvement\n",
    "        self.patience = patience\n",
    "        self.since_improvement = 0\n",
    "        \n",
    "    def on_train_begin(self, trainer, *args, **kwargs):\n",
    "        \"\"\"Resets tracked variables at start of training.\"\"\"\n",
    "        self.best_metric = self.init_metric\n",
    "        self.since_improvement = 0\n",
    "    \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        # Error handling.\n",
    "        new_val = val_stats.get(self.metric)\n",
    "        if new_val is None:\n",
    "            trainer.logger.info(f'EarlyStopper could not find {self.metric}. '\n",
    "                                f'Callback behavior may not be enforced.')\n",
    "            return\n",
    "        \n",
    "        # Expected behavior.\n",
    "        if self.op(new_val, self.op_best(self.best_metric, self.min_improvement)):\n",
    "            self.best_metric = new_val\n",
    "            self.since_improvement = 0\n",
    "        else:\n",
    "            self.since_improvement += 1\n",
    "            if self.since_improvement > self.patience:\n",
    "                trainer.logger.info(\n",
    "                    f'EarlyStopper halting training: validation {self.metric} '\n",
    "                    f'has not improved enough in {self.since_improvement} epochs.'\n",
    "                )\n",
    "                trainer._stop_training = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class PerformanceThreshold(TorchCallback):\n",
    "    \n",
    "    @valuecheck\n",
    "    def __init__(self, metric, goal:('min', 'max'), threshold, skip_epochs=0, \n",
    "                 split:('train', 'val')='val', priority=15):\n",
    "        self.priority = priority\n",
    "        self.metric = metric\n",
    "        self.threshold = threshold\n",
    "        self.skip_epochs = skip_epochs\n",
    "        self.split = split\n",
    "        self.op = gt if goal == 'min' else lt\n",
    "        \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        if epoch < self.skip_epochs:\n",
    "            return\n",
    "        \n",
    "        # Error handling.\n",
    "        data = val_stats if self.split == 'val' else stats\n",
    "        new_val = data.get(self.metric)\n",
    "        if new_val is None:\n",
    "            trainer.logger.info(f'{self.metric.title()} not found in metrics. '\n",
    "                                 'PerformanceThreshold may not be enforced.')\n",
    "            return\n",
    "        \n",
    "        # Expected behavior.\n",
    "        if self.op(new_val, self.threshold):\n",
    "            trainer.logger.info(\n",
    "                f'PerformanceThreshold halting training: {self.metric} '\n",
    "                f'of {new_val:.4f} did not meet threshold.'\n",
    "            )\n",
    "            trainer._stop_training = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ModelCheckpoint(TorchCallback):\n",
    "    \n",
    "    @valuecheck\n",
    "    def __init__(self, metric='loss', goal:('max', 'min')='min', priority=25):\n",
    "        # Will use op like: self.op(new_val, current_best)\n",
    "        if goal == 'min':\n",
    "            self.init_metric = self.best_metric = float('inf')\n",
    "            self.op = lt\n",
    "            self.op_best = sub\n",
    "        elif goal == 'max':\n",
    "            self.init_metric = self.best_metric = float('-inf')\n",
    "            self.op = gt\n",
    "            self.op_best = add\n",
    "\n",
    "        self.priority = priority\n",
    "        self.metric = metric\n",
    "        self.metric_path = None\n",
    "        \n",
    "    def on_train_begin(self, trainer, *args, **kwargs):\n",
    "        self.best_metric = self.init_metric\n",
    "        self.metric_path = os.path.join(trainer.out_dir, \n",
    "                                        'best_val_metrics.json')\n",
    "        \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        new_val = val_stats.get(self.metric)\n",
    "        # Error handling.\n",
    "        if new_val is None:\n",
    "            trainer.logger.info(f'{self.metric} not found in metrics.'\n",
    "                                 'ModelCheckpoint may not save models.')\n",
    "            return\n",
    "            \n",
    "        # Expected behavior.\n",
    "        if self.op(new_val, self.best_metric):\n",
    "            trainer.logger.info(\n",
    "                f'Saving model. {self.metric.title()} improved from '\n",
    "                f'{self.best_metric:.4f} to {new_val:.4f}.'\n",
    "            )\n",
    "            trainer.save(f'trainer.pkl')\n",
    "            save({k: round(v, 5) for k, v in val_stats.items()},\n",
    "                 self.metric_path)\n",
    "            self.best_metric = new_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class MetricHistory(TorchCallback):\n",
    "    \"\"\"Separate from StatsHandler in case we don't want to log outputs.\"\"\"\n",
    "    \n",
    "    def __init__(self, fname='history.csv', plot_fname='history.png',\n",
    "                 priority=90):\n",
    "        self.train_hist = []\n",
    "        self.val_hist = []\n",
    "        self.fname = fname\n",
    "        self.plot_fname = plot_fname\n",
    "        self.priority = priority\n",
    "        \n",
    "    def on_train_begin(self, trainer, *args, **kwargs):\n",
    "        self.train_hist.clear()\n",
    "        self.val_hist.clear()\n",
    "        \n",
    "    def on_epoch_end(self, trainer, epoch, stats, val_stats):\n",
    "        self.train_hist.append(stats.copy())\n",
    "        self.val_hist.append(val_stats.copy())\n",
    "        \n",
    "    def on_train_end(self, trainer, epoch, stats, val_stats):\n",
    "        self.df = pd.concat([\n",
    "            pd.DataFrame(self.train_hist),\n",
    "            pd.DataFrame(self.val_hist)\\\n",
    "              .rename(lambda x: f'val_{x}', axis='columns')\n",
    "        ], axis=1)\n",
    "        self.df.round(5).to_csv(\n",
    "            os.path.join(trainer.out_dir, self.fname), index=False\n",
    "        )\n",
    "        self.plot(os.path.join(trainer.out_dir, self.plot_fname))\n",
    "        \n",
    "    def plot(self, path=None):\n",
    "        cols = self.df.shape[1]\n",
    "        fig, ax = plt.subplots(cols//4, 2, figsize=(12, cols))\n",
    "        for i, axi in zip(range(cols // 2), ax.flatten()):\n",
    "            col = self.df.columns[i]\n",
    "            axi.plot(self.df[col], label='train')\n",
    "            axi.plot(self.df[f'val_{col}'], label='val')\n",
    "            axi.set_title(col.title())\n",
    "            axi.set_xlabel('Epoch')\n",
    "            axi.set_ylabel('Score')\n",
    "            axi.legend()\n",
    "        plt.tight_layout()\n",
    "        if path: \n",
    "            plt.savefig(path)\n",
    "        else:\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class S3Uploader(TorchCallback):\n",
    "    \"\"\"Upload model and logs to S3 when training finishes.\"\"\"\n",
    "    \n",
    "    def __init__(self, bucket, prefix, priority=95):\n",
    "        self.bucket = bucket\n",
    "        self.prefix = prefix\n",
    "        self.priority = priority\n",
    "    \n",
    "    def on_train_end(self, trainer, *args, **kwargs):\n",
    "        paths = [f.path for f in os.scandir(trainer.out_dir)\n",
    "                 if f.is_file() and not f.name.startswith('.')]\n",
    "        s3 = S3tool()\n",
    "        try:\n",
    "            s3.upload_files(paths, self.bucket, self.prefix)\n",
    "        except Exception as e:\n",
    "            trainer.logger.error(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class EC2Closer(TorchCallback):\n",
    "    \n",
    "    def __init__(self, timeout=5, priority=100):\n",
    "        self.timeout = timeout\n",
    "        self.priority = priority\n",
    "        \n",
    "    def on_train_end(self, trainer, *args, **kwargs):\n",
    "        try:\n",
    "            r = requests.get(url, timeout=self.timeout).json()\n",
    "        except requests.ReadTimeout as e:\n",
    "            trainer.logger.info('Request timed out. Failed to '\n",
    "                                'shutdown instance.')\n",
    "            return\n",
    "        \n",
    "        id_, region = r['instanceId'], r['region']\n",
    "        ec2 = boto3.client('ec2', region_name=region)\n",
    "        ec2.stop_instances(InstanceIds=[id_], DryRun=debug)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ModelUnfreezer(TorchCallback):\n",
    "    \"\"\"Gradually unfreeze a model during training.\n",
    "    \"\"\"\n",
    "    \n",
    "    @valuecheck\n",
    "    def __init__(self, i2n, unfreeze_type:('groups', 'layers')='groups', \n",
    "                 mode:('batch', 'epoch')='epoch', priority=25):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        i2n: dict\n",
    "            Maps index of batch/epoch to the number of layers or groups \n",
    "            to unfreeze at that point in time. Batches and epochs are \n",
    "            both zero-indexed. Note that batch refers to the global \n",
    "            batch number (e.g. if there are 100 batches per epoch, the\n",
    "            first batch of the second epoch is batch #101.)\n",
    "        unfreeze_type: str\n",
    "            Specifies whether to unfreeze groups or layers.\n",
    "        mode: str\n",
    "            Specifies whether the indices in `i2n` refer to batches or\n",
    "            epochs.\n",
    "        priority: int\n",
    "            Determine place in the callback queue. Smaller numbers are\n",
    "            executed earlier.\n",
    "            \n",
    "        Examples\n",
    "        --------\n",
    "        This will create a callback that unfreezes the last 2 layer\n",
    "        groups at epoch 2, the last 3 groups at epoch 10, and the\n",
    "        last 4 groups at epoch 25.\n",
    "        \n",
    "        ModelUnfreezer(\n",
    "            i2n={2: 2, 10: 3, 25: 4},\n",
    "            unfreeze_type='groups',\n",
    "            mode='epoch'\n",
    "        )\n",
    "        \"\"\"        \n",
    "        self.priority = priority\n",
    "        self.i2kwargs = {i: {f'n_{unfreeze_type}': n} \n",
    "                         for i, n in i2n.items()}\n",
    "        self.mode = mode\n",
    "        \n",
    "    def on_batch_begin(self, trainer, i, sum_i, stats):\n",
    "        if self.mode != 'batch': return\n",
    "        \n",
    "        kwargs = self.i2kwargs.get(sum_i, None)\n",
    "        if kwargs: trainer.unfreeze(**kwargs, \n",
    "                                    msg_pre=f'Global batch {sum_i}: ')\n",
    "    \n",
    "    def on_epoch_begin(self, trainer, epoch, stats, val_stats):\n",
    "        if self.mode != 'epoch': return\n",
    "        \n",
    "        kwargs = self.i2kwargs.get(epoch, None)\n",
    "        if kwargs: trainer.unfreeze(**kwargs,\n",
    "                                    msg_pre=f'Epoch {epoch}: ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "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",
    "            lr = self.lrs[-1]\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": [
    "# export\n",
    "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
}
