{
 "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": [
    "from collections import defaultdict\n",
    "from collections.abc import Iterable\n",
    "from functools import partial\n",
    "import gc\n",
    "import inspect\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from operator import gt, lt, add, sub\n",
    "import os\n",
    "import pandas as pd\n",
    "from tabulate import tabulate\n",
    "import shutil\n",
    "from sklearn.metrics import (accuracy_score, dcg_score, roc_auc_score, \n",
    "                             precision_score, recall_score)\n",
    "from textblob import TextBlob\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "from tqdm.auto import tqdm\n",
    "import warnings\n",
    "\n",
    "from accio.s3tool import S3tool\n",
    "from htools import (load, save, hdir, LoggerMixin, eprint, assert_raises,\n",
    "                    auto_repr, valuecheck, debug, wrapmethods)\n",
    "from ml_htools.torch_utils import (ModelMixin, variable_lr_optimizer,\n",
    "                                   update_optimizer, DEVICE, stats, adam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reproducible testing.\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data(Dataset):\n",
    "    \n",
    "    def __init__(self, n=64, dim=2):\n",
    "        self.x = torch.rand(n, dim).float()\n",
    "        self.y = torch.clamp(\n",
    "            (self.x[:, 0]*.75 + self.x[:, 1]*.25).round(), 0, 1\n",
    "        ).abs().unsqueeze(-1)\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MulticlassData(Dataset):\n",
    "    \n",
    "    def __init__(self, n=64, dim=2):\n",
    "        self.x = torch.rand(n, dim).float()\n",
    "        # Integer labels between 0 and 4, inclusive.\n",
    "        self.y = torch.clamp(\n",
    "            torch.round(\n",
    "                torch.randint(6, (n, 1)).float() \n",
    "                * (self.x[:, 0]*.75 + self.x[:, 1]*.25).unsqueeze(-1)\n",
    "            ), 0, 4).long().flatten()\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0.4963, 0.7682],\n",
       "         [0.0885, 0.1320],\n",
       "         [0.3074, 0.6341],\n",
       "         [0.4901, 0.8964],\n",
       "         [0.4556, 0.6323]]), tensor([0, 1, 1, 0, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ypred_ = torch.rand(5, 2)\n",
    "y_ = torch.tensor([0, 1, 1, 0, 1])\n",
    "ypred_, y_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.int64"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ypred_.softmax(-1).argmax(-1).dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.7157)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.cross_entropy(ypred_, y_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4324, 0.5676],\n",
       "        [0.4891, 0.5109],\n",
       "        [0.4191, 0.5809],\n",
       "        [0.3998, 0.6002],\n",
       "        [0.4559, 0.5441]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ypred_.softmax(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.7157)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.binary_cross_entropy(ypred_.softmax(-1)[:, -1], y_.float())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.7157)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.cross_entropy(ypred_, y_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.7157)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.nll_loss(ypred_.log_softmax(-1), y_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.4913, -0.2041, -0.0885,  0.5239],\n",
       "         [-0.1435, -0.1116, -0.6136,  0.0316],\n",
       "         [-0.4927,  0.2484,  0.4397,  0.1124],\n",
       "         [-0.8411, -2.3160, -0.1023,  0.7924],\n",
       "         [-0.2897,  0.0525,  0.5229,  2.3022]]), tensor([[0, 1, 0, 0],\n",
       "         [1, 1, 1, 1],\n",
       "         [0, 0, 1, 1],\n",
       "         [0, 0, 0, 0],\n",
       "         [1, 0, 0, 0]]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yhat_ml = torch.randn(5, 4)\n",
    "y_ml = torch.tensor([[0, 1, 0, 0], [1, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 0], [1, 0, 0, 0]])\n",
    "yhat_ml, y_ml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 1.],\n",
       "        [0., 0., 0., 1.],\n",
       "        [0., 1., 1., 1.],\n",
       "        [0., 0., 0., 1.],\n",
       "        [0., 1., 1., 1.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(torch.sigmoid(yhat_ml) > 0.5).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.9687, 0.8004, 0.6499, 0.9890],\n",
       "        [0.7675, 0.7505, 1.0463, 0.6775],\n",
       "        [0.4768, 0.8250, 0.4973, 0.6385],\n",
       "        [0.3585, 0.0941, 0.6433, 1.1659],\n",
       "        [0.8484, 0.7197, 0.9884, 2.3976]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.binary_cross_entropy_with_logits(yhat_ml, y_ml.float(), reduction='none')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.8004, 0.0000, 0.0000],\n",
       "        [0.7675, 0.7505, 1.0463, 0.6775],\n",
       "        [0.0000, 0.0000, 0.4973, 0.6385],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.8484, 0.0000, 0.0000, 0.0000]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "-y_ml.float() * torch.log(torch.sigmoid(yhat_ml))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.9687, 0.0000, 0.6499, 0.9890],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4768, 0.8250, 0.0000, 0.0000],\n",
       "        [0.3585, 0.0941, 0.6433, 1.1659],\n",
       "        [0.0000, 0.7197, 0.9884, 2.3976]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "-(1 - y_ml.float()) * torch.log(1 - torch.sigmoid(yhat_ml))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note: F.cross_entropy combines F.log_softmax and F.nll_loss. However, for last_act we still want F.softmax, as this will create interpretable probabilities.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trainer(LoggerMixin):\n",
    "    \n",
    "    @valuecheck\n",
    "    def __init__(self, net, ds_train, ds_val, dl_train, dl_val,\n",
    "                 criterion, mode:('binary', 'multiclass', 'regression'),\n",
    "                 out_dir, bucket=None, optim_type=Adam, eps=1e-3, \n",
    "                 last_act=None, threshold=0.5, metrics=None, callbacks=None,\n",
    "                 device=DEVICE):\n",
    "        \"\"\"An object to handle model training. This makes it easy for us to\n",
    "        save model weights, optimizer state, datasets and dataloaders all\n",
    "        at once.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        net: BaseModel (inherits from nn.Module)\n",
    "            A pytorch model. The BaseModel implementation from this library\n",
    "            should be used, since Trainer relies on its `unfreeze` method.\n",
    "        ds_train: torch.utils.data.Dataset\n",
    "            Training dataset.\n",
    "        ds_val: torch.utils.data.Dataset\n",
    "            Validation dataset.\n",
    "        dl_train: torch.utils.data.DataLoader\n",
    "            Training dataloader. Lazily retrieves items from train dataset.\n",
    "        dl_val: torch.utils.data.DataLoader\n",
    "            Validation dataloader. Lazily retrieves items from val dataset.\n",
    "        criterion: callable\n",
    "            Typically a PyTorch loss function, but you could define your\n",
    "            own as long as it accepts the same arguments in the same order.\n",
    "            This can be a function (e.g. F.cross_entropy) or a callable\n",
    "            object (e.g. nn.CrossEntropyLoss(). Notice that this is the\n",
    "            object, not the class.)\n",
    "        mode: str\n",
    "            Specifies the problem type. Multi-label classification is \n",
    "            considered 'binary' as well since each example receives a binary\n",
    "            prediction for each class.\n",
    "        out_dir: str\n",
    "            The path to an output directory where logs, model weights, and\n",
    "            more will be stored. If it doesn't already exist, it will be \n",
    "            created.\n",
    "        bucket: str\n",
    "            S3 bucket to store data. This can be used with the S3Upload\n",
    "            callback to automatically upload all files to S3 when training\n",
    "            completes.\n",
    "        optim_type: torch.optim callable\n",
    "            Callable optimizer. The default is Adam. Notice that this is the \n",
    "            class, not the object.\n",
    "        eps: float\n",
    "            The value of epsilon that will be passed to our optimizer. \n",
    "            We use a larger value than PyTorch's default, which empirically\n",
    "            can cause exploding gradients.\n",
    "        last_act: callable or None\n",
    "            Last activation function to be applied outside the model. \n",
    "            For example, for a binary classification problem, if we choose\n",
    "            to use binary_cross_entropy_with_logits loss but want to compute\n",
    "            some metric using soft predictions, we would pass in torch.sigmoid\n",
    "            for `last_act`. For a multi-class problem using F.cross_entropy \n",
    "            loss, we would need to pass in F.softmax to compute predicted \n",
    "            probabilities.  Remember this is ONLY necessary if all of the \n",
    "            following conditions are met:\n",
    "            1. It is a classification problem.\n",
    "            2. We have excluded the final activation from our model for \n",
    "            numerical stability reasons. (I.E. the loss function has the \n",
    "            the final activation built into it.)\n",
    "            3. We wish to compute 1 or more metrics based on soft predictions,\n",
    "            such as AUC-ROC.\n",
    "        threshold: float or None\n",
    "            For a classification problem, pass in the decision threshold to\n",
    "            use when converting soft predictions to hard predictions. For a\n",
    "            regression problem, pass in None.\n",
    "        metrics: list\n",
    "            A list of callable metrics. These will be computed on both the\n",
    "            train and validation sets during training. To maintain \n",
    "            compatibility with scikit-learn metrics, they should accept \n",
    "            two arguments: y_true, followed by either y_score (for soft\n",
    "            predictions) or y_pred (for hard predictions). The name and\n",
    "            order of these arguments matters. If other arguments are\n",
    "            required, pass in a partial with those values specified.\n",
    "        callbacks: list[TorchCallback]\n",
    "            List of callbacks. These will be evaluated during model training\n",
    "            and can be used to track stats, adjust learning rates, clip\n",
    "            gradients, etc.\n",
    "        device: torch.device\n",
    "            Trainer will place the model and current batch of data on this\n",
    "            device during training. The default value uses a GPU if one is\n",
    "            available, otherwise falls back to a CPU.\n",
    "            \n",
    "        Reference\n",
    "        ---------\n",
    "        Classification Loss Function (k = number of classes)\n",
    "\n",
    "        Loss                               y shape  yhat shape  dtype\n",
    "        --------------------------------------------------------------        \n",
    "        binary_cross_entropy_with_logits   (bs, 1)  (bs, 1)     float\n",
    "        \"\" (multilabel case)               (bs, k)  (bs, k)     float\n",
    "        cross_entropy                      (bs,)    (bs, k)     long\n",
    "        \"\"\"\n",
    "        self.net = net\n",
    "        self.ds_train, self.ds_val = ds_train, ds_val\n",
    "        self.dl_train, self.dl_val = dl_train, dl_val\n",
    "        # Optim created in fit() method. Must be after net is on the GPU.\n",
    "        self.optim_type = optim_type\n",
    "        self.optim = None\n",
    "        self.eps = eps\n",
    "        self.criterion = criterion\n",
    "        self.mode = mode\n",
    "        self.device = DEVICE\n",
    "        self.last_act = last_act\n",
    "        self.thresh = threshold\n",
    "        self._stop_training = False\n",
    "        # For now, only print logs. During training, a file will be created.\n",
    "        self.logger = self.get_logger()\n",
    "    \n",
    "        # Storage options.\n",
    "        self.out_dir = out_dir\n",
    "        self.bucket = bucket\n",
    "        os.makedirs(out_dir, exist_ok=True)\n",
    "        \n",
    "        # Dict makes it easier to adjust callbacks after creating model.\n",
    "        self.callbacks = {}\n",
    "        self.add_callbacks(*[BasicConfig(), StatsHandler(), MetricPrinter()] \n",
    "                           + (callbacks or []))\n",
    "        self.metrics = [batch_size] + (metrics or [])\n",
    "    \n",
    "    def save(self, fname):\n",
    "        \"\"\"Save a Trainer for later use. This includes the model, optimizer,\n",
    "        datasets, and data loaders.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        fname: str\n",
    "            File name to save to (not a full path - the trainer already has\n",
    "            an `out_dir` attribute which will be used). The extension must\n",
    "            be .pkl or .zip, and will determine whether the trainer is \n",
    "            compressed.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        save(self, os.path.join(self.out_dir, fname))\n",
    "        \n",
    "    def load(self, fname):\n",
    "        \"\"\"This lets a trainer load a previously saved state. This is NOT\n",
    "        an in-place operation: the new trainer is simply returned.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        fname: str\n",
    "            Name of file where Trainer object is stored. Must end in either\n",
    "            .zip or .pkl. Do not include the full path. This automatically\n",
    "            checks the output directory.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        Trainer\n",
    "        \n",
    "        Examples\n",
    "        --------\n",
    "        trainer = Trainer(...)\n",
    "        trainer.fit(...)\n",
    "        trainer.save('v1')\n",
    "        trainer = trainer.load('v1')\n",
    "        \"\"\"\n",
    "        return load(self, os.path.join(self.out_dir, fname))\n",
    "    \n",
    "    @staticmethod\n",
    "    def from_file(path):\n",
    "        \"\"\"Ths lets us load a previously saved Trainer. Unlike load(), this\n",
    "        does not require us to have a Trainer instance first. The intent is\n",
    "        that load() works well when within a single Jupyter notebook session,\n",
    "        but when returning to work on a different day, we may not have a live\n",
    "        instance of Trainer and a staticmethod lets us load without\n",
    "        remembering all the arguments used at initialization time.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str\n",
    "            Full path to saved file. This differs from load() because here, we\n",
    "            don't have an instance with an out_dir attribute to check.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        Trainer\n",
    "        \"\"\"\n",
    "        return load(path)\n",
    "    \n",
    "    def add_callbacks(self, *callbacks):\n",
    "        \"\"\"Attach additional callbacks to Trainer. Note that callback order \n",
    "        will be determined by their `priority` attribute, not insertion\n",
    "        order.\n",
    "            \n",
    "        Parameters\n",
    "        ----------\n",
    "        callbacks: TorchCallback\n",
    "            One or more callbacks to add. \n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        self.callbacks.update({type(cb).__name__: cb for cb in callbacks})\n",
    "        self.callbacks = dict(sorted(self.callbacks.items(),\n",
    "                                     key=lambda x: x[1].priority))\n",
    "    \n",
    "    def add_metrics(self, *metrics):\n",
    "        \"\"\"Add additional metrics to track. See the `metrics` parameter in \n",
    "        the __init__ docstring for more details.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        metrics: callable\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        self.metrics.extend(metrics)\n",
    "    \n",
    "    def fit(self, epochs, lrs=3e-3, lr_mult=1.0, **kwargs): \n",
    "        \"\"\"Train the model.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        epochs: int\n",
    "            Number of epochs to train for.\n",
    "        lrs: float or Iterable(float)\n",
    "            Pass in one or more learning rates. If lr_mult < 1, these\n",
    "            will be the max LR(s). If the number of values matches the number\n",
    "            of layer groups in the model, they will be matched accordingly,\n",
    "            with the first layer is assigned the first LR. If 1 LR is passed\n",
    "            in and lr_mult < 1, the multiplier will be used to create an\n",
    "            appropriate number of LRs. Example: for a network with 3 groups,\n",
    "            lrs=3e-3 and lr_mult=0.1 will produce LRs of [3e-5, 3e-4, 3e-3].\n",
    "        lr_mult: float\n",
    "            Multiplier used to compute additional learning rates if needed.\n",
    "            See `update_optimizer()` for details.\n",
    "        kwargs: any\n",
    "            Pass in clean=True to remove existing files in out_dir.\n",
    "        \"\"\"\n",
    "        stats = defaultdict(list)\n",
    "        sum_i = 0\n",
    "        _ = self.decide_stop('on_train_begin', epochs, lrs, lr_mult, **kwargs)\n",
    "        for e in range(epochs):\n",
    "            _ = self.decide_stop('on_epoch_begin', e, stats, None)\n",
    "            for i, batch in enumerate(tqdm(self.dl_train)):\n",
    "                sum_i += 1\n",
    "                *xb, yb = map(lambda x: x.to(self.device), batch)\n",
    "                self.optim.zero_grad()\n",
    "                _ = self.decide_stop('on_batch_begin', i, sum_i, stats)\n",
    "                \n",
    "                # Forward and backward passes.\n",
    "                y_score = self.net(*xb)\n",
    "                loss = self.criterion(y_score, yb)\n",
    "                loss.backward()\n",
    "                self.optim.step()\n",
    "                \n",
    "                # Separate because callbacks are only applied during training.\n",
    "                self._update_stats(stats, loss, yb, y_score.detach())\n",
    "                if self.decide_stop('on_batch_end', i, sum_i, stats): break\n",
    "            \n",
    "            # If on_batch_end callback halts training, else block is skipped.  \n",
    "            else: \n",
    "                val_stats = self.validate()\n",
    "                if self.decide_stop('on_epoch_end', e, stats, val_stats): break\n",
    "                continue\n",
    "            break  \n",
    "        _ = self.decide_stop('on_train_end', e, stats, val_stats)\n",
    "    \n",
    "    def validate(self, dl_val=None):\n",
    "        \"\"\"Evaluate the model on a validation set.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        dl_val: torch.utils.data.DataLoader\n",
    "            Accepting an optional dataloader allows the user to pass in \n",
    "            different loaders after training for evaluation. If None is\n",
    "            passed in, self.dl_val is used.\n",
    "        \"\"\"\n",
    "        dl_val = self.dl_val or dl_val\n",
    "        val_stats = defaultdict(list)\n",
    "        self.net.eval()\n",
    "        with torch.no_grad():\n",
    "            for batch in dl_val:\n",
    "                *xb, yb = map(lambda x: x.to(self.device), batch)\n",
    "                y_score = self.net(*xb)\n",
    "                loss = self.criterion(y_score, yb)\n",
    "                self._update_stats(val_stats, loss, yb, y_score)\n",
    "        return val_stats\n",
    "        \n",
    "    def _update_stats(self, stats, loss, yb, y_score):\n",
    "        \"\"\"Update stats in place.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        stats: defaultdict[str, list]\n",
    "        loss: torch.Tensor\n",
    "            Tensor containing single value (mini-batch loss).\n",
    "        yb: torch.Tensor\n",
    "            Mini-batch of labels.\n",
    "        y_score: torch.Tensor\n",
    "            Mini-batch of raw predictions. In the case of \n",
    "            classification, these may still need to be passed\n",
    "            through a sigmoid or softmax.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        # Final activation often excluded from network architecture.\n",
    "        try:\n",
    "            y_score = self.last_act(y_score)\n",
    "        except TypeError:\n",
    "            pass\n",
    "        \n",
    "        # Convert soft predictions to hard predictions.\n",
    "        if self.mode == 'binary':\n",
    "            # In multi-label case, this will have shape (bs, k).\n",
    "            y_pred = (y_score > self.thresh).float()\n",
    "        elif self.mode == 'multiclass':\n",
    "            y_pred = y_score.argmax(-1)\n",
    "        elif self.mode == 'regression':\n",
    "            y_pred = y_score\n",
    "            \n",
    "        stats['loss'].append(loss.detach().cpu().numpy().item())\n",
    "        for m in self.metrics:\n",
    "            yhat = y_pred if hasarg(m, 'y_pred') else y_score\n",
    "            stats[m.__name__.replace('_score', '')].append(m(yb, yhat))\n",
    "        \n",
    "    def decide_stop(self, attr, *args, **kwargs):\n",
    "        \"\"\"Evaluates each of the trainer's callbacks. If any callback\n",
    "        encounters a condition that signals that training should halt,\n",
    "        it will set the attribute trainer._stop_training to True.\n",
    "        This method returns that value. By design, all callbacks will\n",
    "        be called before stopping training.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        attr: str\n",
    "            Determines which method to call for each callback.\n",
    "            One of ('on_train_begin', 'on_train_end', 'on_batch_begin',\n",
    "            'on_batch_end', 'on_epoch_begin', 'on_epoch_end').\n",
    "        args, kwargs: any\n",
    "            Additional arguments to pass to the callbacks.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        bool: If True, halt training.\n",
    "        \"\"\"\n",
    "        self._stop_training = False\n",
    "        # Pass model object as first argument to callbacks.\n",
    "        for cb in self.callbacks.values():\n",
    "            getattr(cb, attr)(self, *args, **kwargs)\n",
    "        return self._stop_training\n",
    "    \n",
    "    def unfreeze(self, n_layers=None, n_groups=None, msg_pre=''):\n",
    "        \"\"\"Pass in either the number of layers or number of groups to \n",
    "        unfreeze. Unfreezing always starts at the end of the network and moves\n",
    "        backward (e.g. n_layers=1 will unfreeze the last 1 layer, or n_groups=2 \n",
    "        will unfreeze the last 2 groups.) Remember than weights and biases are \n",
    "        treated as separate layers.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        n_layers: int or None\n",
    "            Number of layers to unfreeze.\n",
    "        n_groups: int or None\n",
    "            Number of layer groups to unfreeze. For this to work, the model\n",
    "            must define an attribute `groups` containing the layer groups.\n",
    "            Each group can be a layer, a nn.Sequential object, or \n",
    "            nn.Module.\n",
    "        msg_pre: str\n",
    "            Optional: add a prefix to the logged message. For example,\n",
    "            this can be used to record the epoch that unfreezing occurred\n",
    "            during.\n",
    "        \"\"\"\n",
    "        mode = 'layers' if n_layers is not None else 'groups'\n",
    "        msg_pre += f'Unfreezing last {n_layers or n_groups} {mode}.'\n",
    "        self.logger.info(msg_pre)\n",
    "        self.net.unfreeze(n_layers, n_groups)\n",
    "            \n",
    "    def freeze(self):\n",
    "        \"\"\"Freeze whole network. Mostly used for testing.\"\"\"\n",
    "        self.logger.info('Freezing whole network.')\n",
    "        self.net.unfreeze(n_layers=0)\n",
    "    \n",
    "    def cleanup(self, sentinel=None, confirmed=False):\n",
    "        \"\"\"Delete output directory. An empty directory with the same name\n",
    "        will be created in its place.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        sentinel: None\n",
    "            Placeholder to force user to pass confirmed as keyword arg.\n",
    "        confirmed: bool\n",
    "            Placeholder variable. This is just intended to force the user\n",
    "            to confirm their desire to delete files before doing it. If\n",
    "            True, the directory will be deleted. (Technically, any truthy\n",
    "            value will work.)\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        if not confirmed: \n",
    "            self.logger.info('Missing confirmation, cleanup skipped.')\n",
    "            return\n",
    "        self.logger.info('Removing files from output directory.')\n",
    "        shutil.rmtree(self.out_dir)\n",
    "        os.makedirs(self.out_dir)\n",
    "        \n",
    "    def __repr__(self):\n",
    "        r = (f'Trainer(criterion={repr(self.criterion.__name__)}, '\n",
    "             f'out_dir={repr(self.out_dir)}, bucket={repr(self.bucket)})'\n",
    "             f'\\n\\nDatasets: {len(self.ds_train)} train rows, '\n",
    "             f'{len(self.ds_val)} val rows'\n",
    "             f'\\n\\nOptimizer: {repr(self.optim)}'\n",
    "             f'\\n\\n{repr(self.net)})')\n",
    "        return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseModel(nn.Module):\n",
    "    \n",
    "    def unfreeze(self, n_layers=None, n_groups=None):\n",
    "        \"\"\"Pass in either the number of layers or number of groups to \n",
    "        unfreeze. Unfreezing always starts at the end of the network and moves\n",
    "        backward (e.g. n_layers=1 will unfreeze the last 1 layer, or n_groups=2 \n",
    "        will unfreeze the last 2 groups.) Remember than weights and biases are \n",
    "        treated as separate layers.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        n_layers: int or None\n",
    "            Number of layers to unfreeze.\n",
    "        n_groups: int or None\n",
    "            Number of layer groups to unfreeze. For this to work, the model\n",
    "            must define an attribute `groups` containing the layer groups.\n",
    "            Each group can be a layer, a nn.Sequential object, or \n",
    "            nn.Module.\n",
    "        \"\"\"\n",
    "        if n_groups is not None: \n",
    "            self._unfreeze_by_group(n_groups)\n",
    "            return\n",
    "\n",
    "        length = len(self)\n",
    "        for i, p in enumerate(self.parameters()):\n",
    "            p.requires_grad = i >= length - n_layers\n",
    "            \n",
    "    def freeze(self):\n",
    "        \"\"\"Freeze whole network. Mostly used for testing.\"\"\"\n",
    "        self.unfreeze(n_layers=0)\n",
    "        \n",
    "    def _unfreeze_by_group(self, n_groups):\n",
    "        \"\"\"Helper for unfreeze() method.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        n_groups: int\n",
    "            Number of groups to unfreeze, starting at the end of the network.\n",
    "        \"\"\"\n",
    "        length = len(self.groups)\n",
    "        for i, group in enumerate(self.groups):\n",
    "            setting = i >= length - n_groups\n",
    "            for p in group.parameters():\n",
    "                p.requires_grad = setting\n",
    "                \n",
    "    def __len__(self):\n",
    "        \"\"\"Number of parameter matrices in model (basically number of layers, \n",
    "        except that biases are counted separately).\n",
    "        \"\"\"\n",
    "        return sum(1 for p in self.parameters())\n",
    "    \n",
    "    def dims(self):\n",
    "        \"\"\"Get shape of each layer's weights.\"\"\"\n",
    "        return [tuple(p.shape) for p in self.parameters()]\n",
    "\n",
    "    def trainable(self):\n",
    "        \"\"\"Check which layers are trainable.\"\"\"\n",
    "        return [(tuple(p.shape), p.requires_grad) for p in self.parameters()]\n",
    "\n",
    "    def weight_stats(self):\n",
    "        \"\"\"Check mean and standard deviation of each layer's weights.\"\"\"\n",
    "        return [stats(p.data, 3) for p in self.parameters()]\n",
    "\n",
    "    def plot_weights(self):\n",
    "        \"\"\"Plot histograms of each layer's weights.\"\"\"\n",
    "        n_layers = len(self.dims())\n",
    "        fig, ax = plt.subplots(n_layers, figsize=(8, n_layers * 1.25))\n",
    "        if not isinstance(ax, Iterable): ax = [ax]\n",
    "        for i, p in enumerate(self.parameters()):\n",
    "            ax[i].hist(p.data.flatten())\n",
    "            ax[i].set_title(f'Shape: {tuple(p.shape)} Stats: {stats(p.data)}')\n",
    "        plt.tight_layout()\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleModel(BaseModel):\n",
    "    \n",
    "    def __init__(self, dim):\n",
    "        super().__init__()  \n",
    "        self.fc1 = nn.Linear(dim, 2)\n",
    "        self.fc2 = nn.Linear(2, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = F.leaky_relu(self.fc1(x))\n",
    "        return self.fc2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GroupedModel(BaseModel):\n",
    "    \n",
    "    def __init__(self, dim):\n",
    "        super().__init__()  \n",
    "        g1 = nn.Sequential(\n",
    "            nn.Linear(dim, 8),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(8, 4),\n",
    "            nn.LeakyReLU()\n",
    "        )\n",
    "        g2 = nn.Linear(4, 1)\n",
    "        self.groups = nn.ModuleList([g1, g2])\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": [],
   "source": [
    "DIM = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SimpleModel(\n",
       "  (fc1): Linear(in_features=2, out_features=2, bias=True)\n",
       "  (fc2): Linear(in_features=2, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "snet = SimpleModel(DIM)\n",
    "snet"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.002\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(snet, 2e-3)\n",
    "optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got ValueError(Received more learning rates than layer groups.).\n"
     ]
    }
   ],
   "source": [
    "with assert_raises(ValueError) as ar:\n",
    "    optim = variable_lr_optimizer(snet, [3e-3, 1e-1])\n",
    "    optim"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.001\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "update_optimizer(optim, 1e-3, 0.5)\n",
    "optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [((2, 2), False), ((2,), False), ((1, 2), False), ((1,), False)]\n",
      "1 [((2, 2), False), ((2,), False), ((1, 2), False), ((1,), True)]\n",
      "2 [((2, 2), False), ((2,), False), ((1, 2), True), ((1,), True)]\n",
      "3 [((2, 2), False), ((2,), True), ((1, 2), True), ((1,), True)]\n",
      "4 [((2, 2), True), ((2,), True), ((1, 2), True), ((1,), True)]\n"
     ]
    }
   ],
   "source": [
    "snet.freeze()\n",
    "for n in range(5):\n",
    "    snet.unfreeze(n_layers=n)\n",
    "    print(n, snet.trainable())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got AttributeError('SimpleModel' object has no attribute 'groups').\n"
     ]
    }
   ],
   "source": [
    "snet.freeze()\n",
    "with assert_raises(AttributeError) as ar:\n",
    "    for n in range(3):\n",
    "        snet.unfreeze(n_groups=n)\n",
    "        print(n, snet.trainable())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GroupedModel(\n",
       "  (groups): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Linear(in_features=2, out_features=8, bias=True)\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): Linear(in_features=8, out_features=4, bias=True)\n",
       "      (3): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Linear(in_features=4, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gnet = GroupedModel(DIM)\n",
    "gnet"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.001\n",
       "    weight_decay: 0\n",
       "\n",
       "Parameter Group 1\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 0.001\n",
       "    lr: 0.003\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(gnet, [1e-3, 3e-3])\n",
    "optim"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.0006666666666666666\n",
       "    weight_decay: 0\n",
       "\n",
       "Parameter Group 1\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 0.001\n",
       "    lr: 0.002\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "update_optimizer(optim, 2e-3, lr_mult=1/3)\n",
    "optim"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.001\n",
       "    weight_decay: 0\n",
       "\n",
       "Parameter Group 1\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 0.001\n",
       "    lr: 0.001\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variable_lr_optimizer(gnet, 1e-3)"
   ]
  },
  {
   "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: 0.001\n",
       "    lr: 0.0005\n",
       "    weight_decay: 0\n",
       "\n",
       "Parameter Group 1\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 0.001\n",
       "    lr: 0.001\n",
       "    weight_decay: 0\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variable_lr_optimizer(gnet, 1e-3, lr_mult=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), False), ((1, 4), False), ((1,), False)]\n",
      "1 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), False), ((1, 4), True), ((1,), True)]\n",
      "2 [((8, 2), True), ((8,), True), ((4, 8), True), ((4,), True), ((1, 4), True), ((1,), True)]\n"
     ]
    }
   ],
   "source": [
    "gnet.freeze()\n",
    "for n in range(3):\n",
    "    gnet.unfreeze(n_groups=n)\n",
    "    print(n, gnet.trainable())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), False), ((1, 4), False), ((1,), False)]\n",
      "1 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), False), ((1, 4), False), ((1,), True)]\n",
      "2 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), False), ((1, 4), True), ((1,), True)]\n",
      "3 [((8, 2), False), ((8,), False), ((4, 8), False), ((4,), True), ((1, 4), True), ((1,), True)]\n",
      "4 [((8, 2), False), ((8,), False), ((4, 8), True), ((4,), True), ((1, 4), True), ((1,), True)]\n",
      "5 [((8, 2), False), ((8,), True), ((4, 8), True), ((4,), True), ((1, 4), True), ((1,), True)]\n",
      "6 [((8, 2), True), ((8,), True), ((4, 8), True), ((4,), True), ((1, 4), True), ((1,), True)]\n"
     ]
    }
   ],
   "source": [
    "gnet.freeze()\n",
    "for n in range(7):\n",
    "    gnet.unfreeze(n_layers=n)\n",
    "    print(n, gnet.trainable())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Callbacks"
   ]
  },
  {
   "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 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": [
    "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": [
    "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": [
    "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": [
    "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": [
    "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": [
    "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": [
    "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": [
    "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": [
    "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{10: {'n_groups': 1}, 20: {'n_groups': 2}, 30: {'n_groups': 3}}\n",
      "{100: {'n_layers': 5}, 200: {'n_layers': 10}, 350: {'n_layers': 13}}\n"
     ]
    }
   ],
   "source": [
    "mu = ModelUnfreezer({10:1, 20:2, 30:3}, 'groups')\n",
    "print(mu.i2kwargs)\n",
    "\n",
    "mu = ModelUnfreezer({100: 5, 200: 10, 350:13}, 'layers', 'batch')\n",
    "print(mu.i2kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class LRScheduler(TorchCallback):\n",
    "    \n",
    "#     @valuecheck\n",
    "#     def __init__(self, sched, update_on:('batch', 'epoch'), priority=10):\n",
    "#         \"\"\"\n",
    "#         sched: callable\n",
    "#             Accepts two parameters (the max learning rate the number of epochs) and\n",
    "#             returns a sequence of learning rates (if differential learning rates are\n",
    "#             used, these will be the maximum LRs). This will often be a partial\n",
    "#             or an instance of a class with __call__ implemented (a scheduler \n",
    "#             might have a parameter that needs to be specified ahead\n",
    "#             of time, like warmup_percent).\n",
    "#         \"\"\"\n",
    "#         self.sched = sched\n",
    "#         self.update_on = update_on\n",
    "#         # Initialized in `on_train_begin()`.\n",
    "#         self.max_lrs = None\n",
    "#         self.lr_mult = None\n",
    "#         self.priority = priority\n",
    "\n",
    "        \n",
    "#     def on_train_begin(self, trainer, epochs, lrs, lr_mult, clean):\n",
    "#         if not isinstance(lrs, Iterable): lrs = [lrs]\n",
    "#         self.max_lrs = self.sched(lrs[-1], epochs)\n",
    "#         self.lr_mult = lr_mult\n",
    "        \n",
    "#     def on_batch_begin(self, trainer, i, sum_i, stats):\n",
    "#         if self.update_on != 'batch': \n",
    "#             return\n",
    "#         self._update_lr(trainer, sum_i)\n",
    "        \n",
    "#     def on_epoch_begin(self, trainer, epoch, stats, val_stats):\n",
    "#         if self.update_on != 'epoch': \n",
    "#             return\n",
    "#         self._update_lr(trainer, epoch)\n",
    "        \n",
    "#     def _update_lr(self, trainer, n):\n",
    "#         try:\n",
    "#             lr = self.max_lrs[n]\n",
    "#         except IndexError:\n",
    "            \n",
    "#         update_optimizer(trainer.optim, lr, lr_mult=self.lr_mult)\n",
    "#         if self.track: self.history.append(lr)\n",
    "#         trainer.logger.info(f'Set learning rate to {lr}.')\n",
    "\n",
    "\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": [
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def back_translate(text, to, from_lang='en'):\n",
    "    return TextBlob(text)\\\n",
    "        .translate(to=to)\\\n",
    "        .translate(from_lang=to, to=from_lang)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "Visit ESPN to get up-to-the-minute sports news coverage, scores, highlights \n",
    "and commentary for NFL, MLB, NBA, College Football, NCAA Basketball and more.\n",
    "\"\"\"\n",
    "# back_translate(text, 'es')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Metrics\n",
    "\n",
    "Keep sklearn pattern with y_true as first argument.\n",
    "\n",
    "For classification problems, round probabilities once instead of in every metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hasarg(func, arg):\n",
    "    \"\"\"Checks if a function has a given argument. \n",
    "    Works with args and kwargs as well if you exclude the\n",
    "    stars. See example below.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    func: function\n",
    "    arg: str\n",
    "        Name of argument to look for.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    bool\n",
    "    \n",
    "    Example\n",
    "    -------\n",
    "    def foo(a, b=6, *args):\n",
    "        return\n",
    "        \n",
    "    >>> hasarg(foo, 'b')\n",
    "    True\n",
    "    \n",
    "    >>> hasarg(foo, 'args')\n",
    "    True\n",
    "    \n",
    "    >>> hasarg(foo, 'c')\n",
    "    False\n",
    "    \"\"\"\n",
    "    return arg in inspect.signature(func).parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def percent_positive(y_true, y_pred):\n",
    "    \"\"\"Compute the percent of predictions that are positive. This\n",
    "    can help us identify when a model is predicting all ones or zeros.\n",
    "    \"\"\"\n",
    "    return (y_pred == 1).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_soft_prediction(y_true, y_score):\n",
    "    \"\"\"Compute the mean predicted probability.\"\"\"\n",
    "    return y_score.mean() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_size(y_true, y_pred):\n",
    "    \"\"\"Count the number of items in the current batch.\"\"\"\n",
    "    return y_true.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, False]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[hasarg(roc_auc_score, val) for val in ('y_score', 'y_pred')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[False, True]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[hasarg(precision_score, val) for val in ('y_score', 'y_pred')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mish(x):\n",
    "    \"\"\"Mish: A Self Regularized Non-Monotonic Neural Activation Function\n",
    "    https://arxiv.org/pdf/1908.08681v1.pdf\n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor[float]\n",
    "        Input tensor.\n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor[float]: Tensor of same shape as input x.\n",
    "    \"\"\"\n",
    "    return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mish(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=4, out_features=3, bias=True)\n",
       "  (1): Mish()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fc = nn.Linear(4, 3)\n",
    "seq = nn.Sequential(fc,\n",
    "                    Mish())\n",
    "seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1407,  0.2945, -0.0009],\n",
       "        [ 0.2134,  0.3165, -0.0377],\n",
       "        [ 0.0347,  0.2241,  0.1147],\n",
       "        [ 0.5246,  0.4626, -0.2096],\n",
       "        [ 0.2227,  0.4221, -0.0366]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand(5, 4)\n",
    "seq(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1407,  0.2945, -0.0009],\n",
       "        [ 0.2134,  0.3165, -0.0377],\n",
       "        [ 0.0347,  0.2241,  0.1147],\n",
       "        [ 0.5246,  0.4626, -0.2096],\n",
       "        [ 0.2227,  0.4221, -0.0366]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_ = fc(x)\n",
    "mish(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_activations(z, a, mode='scatter', **kwargs):\n",
    "    \"\"\"Plot an input tensor and its corresponding activations.  Both tensors\n",
    "    will be flattened for plotting.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    z: tf.Tensor\n",
    "        Tensor containing values to plot on the x axis (we can often think of\n",
    "        this as the output of a linear layer, where z=f(x) and a=mish(z)).\n",
    "    a: tf.Tensor\n",
    "        Tensor containing values to plot on y axis.\n",
    "    mode: str\n",
    "        'scatter' for scatter plot or 'plot' for line plot.\n",
    "    kwargs: Values to be passed to the matplotlib plotting function, such as \n",
    "        's' when in 'scatter' mode or 'lw' in 'plot' mode.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "    \"\"\"\n",
    "    plt_func = getattr(plt, mode)\n",
    "    kwargs = kwargs or {}\n",
    "    if mode == 'scatter' and not kwargs:\n",
    "        kwargs = {'s': .75}\n",
    "    plt_func(z.numpy().flatten(), a.numpy().flatten(), **kwargs)\n",
    "    plt.axvline(0, lw=.5, alpha=.5)\n",
    "    plt.axhline(0, lw=.5, alpha=.5)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(-5, 5, .05)\n",
    "a = mish(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10384 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_activations(x, a, 'plot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training\n",
    "## Binary Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIM = 2\n",
    "metrics = [accuracy_score, \n",
    "           precision_score, \n",
    "           recall_score, \n",
    "           percent_positive,\n",
    "           mean_soft_prediction\n",
    "          ]\n",
    "\n",
    "# Model starts out unfrozen and freezes last group startingn with epoch 3.\n",
    "# This is not useful but is done here for testing purposes.\n",
    "callbacks = [EarlyStopper('max', 'accuracy', patience=3),\n",
    "             PerformanceThreshold('recall', 'max', 0.25, skip_epochs=5),\n",
    "             MetricHistory(),\n",
    "             ModelUnfreezer({3:1}, 'groups'),\n",
    "             ModelCheckpoint(),\n",
    "             MetricHistory(),\n",
    "             CosineLRScheduler()\n",
    "             S3Uploader('mytorch_test')\n",
    "            ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = Data(n=300, dim=DIM)\n",
    "val = Data(n=30, dim=DIM)\n",
    "\n",
    "dl_train = DataLoader(train, batch_size=8, shuffle=True)\n",
    "dl_val = DataLoader(val, batch_size=8, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Trainer(criterion='binary_cross_entropy_with_logits', out_dir='../data/v1', bucket='datascience-delphi-dev')\n",
       "\n",
       "Datasets: 300 train rows, 30 val rows\n",
       "\n",
       "Optimizer: None\n",
       "\n",
       "GroupedModel(\n",
       "  (groups): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Linear(in_features=2, out_features=8, bias=True)\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): Linear(in_features=8, out_features=4, bias=True)\n",
       "      (3): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Linear(in_features=4, out_features=1, bias=True)\n",
       "  )\n",
       "))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gnet = GroupedModel(DIM)\n",
    "t = Trainer(gnet, train, val, dl_train, dl_val, F.binary_cross_entropy_with_logits, \n",
    "            'binary', '../data/v1', 'datascience-delphi-dev', torch.optim.RMSprop, \n",
    "            last_act=torch.sigmoid, metrics=metrics, callbacks=callbacks)\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-02-14 17:06:46,155 [INFO]: Removing files from output directory.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52057a5cfffb47d18e4151392b50f503",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:46,308\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 1\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.4308 |       0.2395 |\n",
      "| accuracy             |  0.7933 |       0.8667 |\n",
      "| precision            |  0.7367 |       1.0000 |\n",
      "| recall               |  0.8413 |       0.7867 |\n",
      "| percent_positive     |  0.5767 |       0.3667 |\n",
      "| mean_soft_prediction |  0.5086 |       0.3790 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:46,309\n",
      " Saving model. Loss improved from inf to 0.2395.\n",
      "Data written to ../data/v1/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3aa2b694afd4b879a00514ecd7b209d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:46,463\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 2\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.3729 |       0.4348 |\n",
      "| accuracy             |  0.8467 |       0.8333 |\n",
      "| precision            |  0.8354 |       0.7489 |\n",
      "| recall               |  0.8556 |       1.0000 |\n",
      "| percent_positive     |  0.5033 |       0.6667 |\n",
      "| mean_soft_prediction |  0.4888 |       0.7262 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:46,464\n",
      " Epoch 3: Unfreezing last 1 groups.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "339cf1075f1449d79c1f6b312aaaaab9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:46,613\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 3\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1912 |       0.1883 |\n",
      "| accuracy             |  0.9300 |       0.9000 |\n",
      "| precision            |  0.8893 |       1.0000 |\n",
      "| recall               |  0.9867 |       0.8400 |\n",
      "| percent_positive     |  0.5600 |       0.4000 |\n",
      "| mean_soft_prediction |  0.5536 |       0.4145 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:46,614\n",
      " Saving model. Loss improved from 0.2395 to 0.1883.\n",
      "Data written to ../data/v1/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c919047561f48fbbcb1c13ddd72a0d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:46,765\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 4\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1331 |       0.1263 |\n",
      "| accuracy             |  0.9367 |       0.9667 |\n",
      "| precision            |  0.9342 |       0.9333 |\n",
      "| recall               |  0.9620 |       1.0000 |\n",
      "| percent_positive     |  0.5133 |       0.5333 |\n",
      "| mean_soft_prediction |  0.5100 |       0.4703 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:46,766\n",
      " Saving model. Loss improved from 0.1883 to 0.1263.\n",
      "Data written to ../data/v1/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b39e6e9a24ff45a4bb3b506a8e239004",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:46,915\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 5\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1165 |       0.1096 |\n",
      "| accuracy             |  0.9367 |       0.9667 |\n",
      "| precision            |  0.9249 |       0.9333 |\n",
      "| recall               |  0.9387 |       1.0000 |\n",
      "| percent_positive     |  0.5133 |       0.5333 |\n",
      "| mean_soft_prediction |  0.5057 |       0.4992 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:46,916\n",
      " Saving model. Loss improved from 0.1263 to 0.1096.\n",
      "Data written to ../data/v1/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "23adefd14096431cbe0c9963ea1caf0c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:47,065\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 6\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1154 |       0.1125 |\n",
      "| accuracy             |  0.9300 |       0.9667 |\n",
      "| precision            |  0.9324 |       0.9333 |\n",
      "| recall               |  0.9409 |       1.0000 |\n",
      "| percent_positive     |  0.5133 |       0.5333 |\n",
      "| mean_soft_prediction |  0.5086 |       0.4795 |\n",
      "\n",
      "=====\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d08f3503702d42b591be3f50b7abc49a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:47,222\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 7\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1110 |       0.1186 |\n",
      "| accuracy             |  0.9467 |       0.9667 |\n",
      "| precision            |  0.9364 |       1.0000 |\n",
      "| recall               |  0.9660 |       0.9467 |\n",
      "| percent_positive     |  0.5100 |       0.4667 |\n",
      "| mean_soft_prediction |  0.5090 |       0.4657 |\n",
      "\n",
      "=====\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38a5c026066e46df918268b633743f7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=38), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 17:06:47,375\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 8\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  0.1141 |       0.1103 |\n",
      "| accuracy             |  0.9367 |       0.9667 |\n",
      "| precision            |  0.9373 |       0.9333 |\n",
      "| recall               |  0.9397 |       1.0000 |\n",
      "| percent_positive     |  0.5067 |       0.5333 |\n",
      "| mean_soft_prediction |  0.5029 |       0.4783 |\n",
      "\n",
      "=====\n",
      "\n",
      "2020-02-14 17:06:47,376\n",
      " EarlyStopper halting training: validation accuracy has not improved enough in 4 epochs.\n",
      "\n",
      "2020-02-14 17:06:47,377\n",
      " Training complete. Model in eval mode.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 120116 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t.fit(10, 3e-1, 0.5, clean=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data written to ../data/v1/trainer.zip.\n"
     ]
    }
   ],
   "source": [
    "t.save('trainer.zip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Object loaded from ../data/v1/trainer.zip.\n",
      "Trainer(criterion='binary_cross_entropy_with_logits', out_dir='../data/v1', bucket='datascience-delphi-dev')\n",
      "\n",
      "Datasets: 300 train rows, 30 val rows\n",
      "\n",
      "Optimizer: RMSprop (\n",
      "Parameter Group 0\n",
      "    alpha: 0.99\n",
      "    centered: False\n",
      "    eps: 0.001\n",
      "    lr: 0.09752016304205122\n",
      "    momentum: 0\n",
      "    weight_decay: 0\n",
      "\n",
      "Parameter Group 1\n",
      "    alpha: 0.99\n",
      "    centered: False\n",
      "    eps: 0.001\n",
      "    lr: 0.19504032608410243\n",
      "    momentum: 0\n",
      "    weight_decay: 0\n",
      ")\n",
      "\n",
      "GroupedModel(\n",
      "  (groups): ModuleList(\n",
      "    (0): Sequential(\n",
      "      (0): Linear(in_features=2, out_features=8, bias=True)\n",
      "      (1): LeakyReLU(negative_slope=0.01)\n",
      "      (2): Linear(in_features=8, out_features=4, bias=True)\n",
      "      (3): LeakyReLU(negative_slope=0.01)\n",
      "    )\n",
      "    (1): Linear(in_features=4, out_features=1, bias=True)\n",
      "  )\n",
      "))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "18311"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t2 = Trainer.from_file(os.path.join('..', 'data', 'v1', 'trainer.zip'))\n",
    "print(t2)\n",
    "del t2; gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiclass classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleMulticlassModel(BaseModel):\n",
    "    \n",
    "    def __init__(self, dim, classes):\n",
    "        super().__init__()  \n",
    "        self.fc1 = nn.Linear(dim, 10)\n",
    "        self.fc2 = nn.Linear(10, classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = F.leaky_relu(self.fc1(x))\n",
    "        return self.fc2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIM = 2\n",
    "metrics = [accuracy_score,\n",
    "           mean_soft_prediction\n",
    "          ]\n",
    "\n",
    "# Model starts out unfrozen and freezes last group startingn with epoch 3.\n",
    "# This is not useful but is done here for testing purposes.\n",
    "callbacks = [EarlyStopper('max', 'accuracy', patience=3),\n",
    "             PerformanceThreshold('loss', 'min', 2, skip_epochs=5),\n",
    "             MetricHistory(),\n",
    "             ModelUnfreezer({3:2}, 'layers'),\n",
    "             ModelCheckpoint(),\n",
    "             MetricHistory(),\n",
    "             CosineLRScheduler()\n",
    "             S3Uploader('mytorch_test')\n",
    "            ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = MulticlassData(n=34, dim=DIM)\n",
    "val = MulticlassData(n=30, dim=DIM)\n",
    "\n",
    "dl_train = DataLoader(train, batch_size=8, shuffle=True)\n",
    "dl_val = DataLoader(val, batch_size=8, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SimpleMulticlassModel(\n",
       "  (fc1): Linear(in_features=2, out_features=10, bias=True)\n",
       "  (fc2): Linear(in_features=10, out_features=5, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smnet = SimpleMulticlassModel(DIM, 5)\n",
    "smnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09378d6c7f1d42238ae697ee1af27d58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,398\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 1\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.5517 |       1.6642 |\n",
      "| accuracy             |  0.3235 |       0.2667 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n",
      "\n",
      "2020-02-14 16:58:00,399\n",
      " Saving model. Loss improved from inf to 1.6642.\n",
      "Data written to ../data/v2/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a2d4cbef0f64ceea4093814ff7a35e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,430\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 2\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.4351 |       1.8731 |\n",
      "| accuracy             |  0.2647 |       0.2667 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n",
      "\n",
      "2020-02-14 16:58:00,431\n",
      " Unfreezing last 2 layers.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4aafd1ba2f0548eebf92cbb918800262",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,457\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 3\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.3237 |       1.8345 |\n",
      "| accuracy             |  0.3529 |       0.3000 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ddbe30d3d2bb4129a7dc9a6db1773a69",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,485\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 4\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2768 |       1.7265 |\n",
      "| accuracy             |  0.4118 |       0.3667 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8be20e1726e24171a8e662846a1d9f10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,515\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 5\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2725 |       1.6096 |\n",
      "| accuracy             |  0.3824 |       0.4000 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n",
      "\n",
      "2020-02-14 16:58:00,515\n",
      " Saving model. Loss improved from 1.6642 to 1.6096.\n",
      "Data written to ../data/v2/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "371ce1a26b8d48f59971c1fad078cf9c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,546\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 6\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2853 |       1.6049 |\n",
      "| accuracy             |  0.3235 |       0.3333 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n",
      "\n",
      "2020-02-14 16:58:00,547\n",
      " Saving model. Loss improved from 1.6096 to 1.6049.\n",
      "Data written to ../data/v2/trainer.pkl.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8fa75f119c684b3c916c162bb4e19716",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,574\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 7\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2391 |       1.6560 |\n",
      "| accuracy             |  0.4118 |       0.3333 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb298f23b3674626860f85d8764214ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,599\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 8\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2090 |       1.6696 |\n",
      "| accuracy             |  0.4412 |       0.3000 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a6288979d024c69b40875d1297b297e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=5), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "2020-02-14 16:58:00,624\n",
      " \n",
      "=====\n",
      "\n",
      "Epoch 9\n",
      "\n",
      "| Metric               |   Train |   Validation |\n",
      "|----------------------|---------|--------------|\n",
      "| loss                 |  1.2167 |       1.6823 |\n",
      "| accuracy             |  0.4118 |       0.3000 |\n",
      "| mean_soft_prediction |  0.2000 |       0.2000 |\n",
      "\n",
      "\n",
      "2020-02-14 16:58:00,625\n",
      " EarlyStopper halting training: validation accuracy has not improved enough in 4 epochs.\n",
      "\n",
      "2020-02-14 16:58:00,626\n",
      " Training complete. Model in eval mode.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51980 characters]\n",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t.fit(10, 3e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([tensor([[0.5219, 0.1604],\n",
       "          [0.1135, 0.4161],\n",
       "          [0.3907, 0.3769],\n",
       "          [0.9464, 0.3559],\n",
       "          [0.0547, 0.3759],\n",
       "          [0.0930, 0.3026],\n",
       "          [0.0189, 0.5636],\n",
       "          [0.3094, 0.2666]])], tensor([2, 0, 0, 2, 0, 0, 0, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "*x, y = next(iter(dl_train))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0467, -0.0801,  0.4771, -0.1857,  0.2001],\n",
       "        [ 0.0752, -0.0900,  0.3997, -0.3018,  0.1799],\n",
       "        [ 0.0133, -0.1059,  0.4362, -0.2662,  0.2148],\n",
       "        [-0.1016, -0.1393,  0.5033, -0.2052,  0.2878],\n",
       "        [ 0.0808, -0.0785,  0.3978, -0.2938,  0.1649],\n",
       "        [ 0.0617, -0.0693,  0.4109, -0.2671,  0.1590],\n",
       "        [ 0.1163, -0.1104,  0.3881, -0.3379,  0.1813],\n",
       "        [ 0.0125, -0.0808,  0.4401, -0.2375,  0.1854]],\n",
       "       grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yhat = smnet(*x)\n",
    "yhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8]), torch.Size([8, 5]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.shape, yhat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.1722, 0.1666, 0.2908, 0.1499, 0.2205],\n",
       "        [0.1989, 0.1686, 0.2752, 0.1364, 0.2209],\n",
       "        [0.1854, 0.1646, 0.2830, 0.1402, 0.2268],\n",
       "        [0.1620, 0.1560, 0.2967, 0.1461, 0.2392],\n",
       "        [0.2000, 0.1705, 0.2745, 0.1375, 0.2175],\n",
       "        [0.1954, 0.1714, 0.2771, 0.1407, 0.2154],\n",
       "        [0.2078, 0.1657, 0.2727, 0.1320, 0.2218],\n",
       "        [0.1848, 0.1683, 0.2834, 0.1439, 0.2196]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.softmax(yhat, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8, 5]), torch.Size([8]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yhat.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.5432, grad_fn=<NllLossBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.cross_entropy(yhat, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Trainer(criterion='cross_entropy', out_dir='../data/v2', bucket='datascience-delphi-dev')\n",
       "\n",
       "Datasets: 34 train rows, 30 val rows\n",
       "\n",
       "Optimizer: None\n",
       "\n",
       "SimpleMulticlassModel(\n",
       "  (fc1): Linear(in_features=2, out_features=10, bias=True)\n",
       "  (fc2): Linear(in_features=10, out_features=5, bias=True)\n",
       "))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = Trainer(smnet, train, val, dl_train, dl_val, F.cross_entropy, 'multiclass',\n",
    "        '../data/v2', 'datascience-delphi-dev', last_act=partial(F.softmax, dim=-1),\n",
    "            metrics=metrics, callbacks=callbacks)\n",
    "t"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
