{
 "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",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import warnings\n",
    "\n",
    "from accio.s3tool import S3tool\n",
    "from htools import auto_repr, valuecheck\n",
    "from incendio.core import DEVICE"
   ]
  },
  {
   "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",
    "        trainer.optim = variable_lr_optimizer(\n",
    "            trainer.net, lrs, lr_mult, trainer.optim_type, trainer.eps\n",
    "        )\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",
    "        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, goal:('max', 'min'), metric='loss', min_improvement=0.0, \n",
    "                 patience=3, priority=15):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\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",
    "        metric: str\n",
    "            Quantity to monitor. This will always be computed on the \n",
    "            validation set.\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.model_dir = None\n",
    "        \n",
    "    def on_train_begin(self, trainer, *args, **kwargs):\n",
    "        self.best_metric = self.init_metric\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",
    "            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, *args, **kwargs):\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.to_csv(os.path.join(trainer.out_dir, self.fname), \n",
    "                       index=False)\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, prefix, priority=95):\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, trainer.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
}
