{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Core\n",
    "\n",
    "> The basics for building and training models are contained in this module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from collections import defaultdict\n",
    "from collections.abc import Iterable\n",
    "from functools import partial, wraps\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import shutil\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from tqdm.auto import tqdm\n",
    "import warnings\n",
    "\n",
    "from htools import load, save, LoggerMixin, valuecheck, hasarg\n",
    "from incendio.callbacks import BasicConfig, StatsHandler, MetricPrinter\n",
    "from incendio.metrics import batch_size\n",
    "from incendio.optimizers import variable_lr_optimizer, update_optimizer\n",
    "from incendio.utils import quick_stats, DEVICE, identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Used in notebook but not needed in package.\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from htools import assert_raises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "`BaseModel` allows models to freeze/unfreeze layers and provides several methods for weight diagnostics. It should not be instantiated directly, but used as a parent class for a model. Like all PyTorch models, its children will still need to call `super().__init__()` and implement a `forward()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "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 numel(self):\n",
    "        return sum(p.numel() 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 [quick_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(\n",
    "                f'Shape: {tuple(p.shape)} Stats: {quick_stats(p.data)}'\n",
    "            )\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "        \n",
    "    def predict(self, *xb):\n",
    "        \"\"\"Predict on one batch of data. This is almost identical to \n",
    "        self.__call__: the only differences are that it first puts the model\n",
    "        in eval mode and it doesn't compute gradients.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        xb: torch.tensors\n",
    "            One or more tensors comprising the inputs of a single mini batch.\n",
    "        \"\"\"\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            return self(*xb)"
   ]
  },
  {
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unfrozen [False, False, False, False]\n",
      "Unfrozen [False, False, False, True]\n",
      "Unfrozen [False, False, True, True]\n",
      "Unfrozen [False, True, True, True]\n",
      "Unfrozen [True, True, True, True]\n"
     ]
    }
   ],
   "source": [
    "snet = SimpleModel(2)\n",
    "snet.freeze()\n",
    "for n in range(5):\n",
    "    snet.unfreeze(n_layers=n)\n",
    "    unfrozen = [x[1] for x in snet.trainable()]\n",
    "    print('Unfrozen', unfrozen)\n",
    "    assert sum(unfrozen) == n\n",
    "    assert not any(unfrozen[:-n])"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unfrozen [False, False, False, False, False, False]\n",
      "Unfrozen [False, False, False, False, True, True]\n",
      "Unfrozen [True, True, True, True, True, True]\n"
     ]
    }
   ],
   "source": [
    "gnet = GroupedModel(2)\n",
    "gnet.freeze()\n",
    "n_unfrozen = [0, 2, 6]\n",
    "for n, nu in zip(range(3), n_unfrozen):\n",
    "    gnet.unfreeze(n_groups=n)\n",
    "    unfrozen = [x[1] for x in gnet.trainable()]\n",
    "    print('Unfrozen', unfrozen)\n",
    "    assert sum(unfrozen) == nu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unfrozen [False, False, False, False, False, False]\n",
      "Unfrozen [False, False, False, False, False, True]\n",
      "Unfrozen [False, False, False, False, True, True]\n",
      "Unfrozen [False, False, False, True, True, True]\n",
      "Unfrozen [False, False, True, True, True, True]\n",
      "Unfrozen [False, True, True, True, True, True]\n",
      "Unfrozen [True, True, True, True, True, True]\n"
     ]
    }
   ],
   "source": [
    "gnet.freeze()\n",
    "for n in range(7):\n",
    "    gnet.unfreeze(n_layers=n)\n",
    "    unfrozen = [x[1] for x in gnet.trainable()]\n",
    "    print('Unfrozen', unfrozen)\n",
    "    assert sum(unfrozen) == n\n",
    "    assert not any(unfrozen[:-n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      ")\n",
      "As expected, got ValueError(Received more learning rates than layer groups.).\n"
     ]
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(snet, 2e-3)\n",
    "print(optim)\n",
    "\n",
    "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": [],
   "source": [
    "update_optimizer(optim, 1e-3, 0.5)\n",
    "assert len(optim.param_groups) == 1\n",
    "assert optim.param_groups[0]['lr'] == 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      ")\n",
      "[0.0006666666666666666, 0.002]\n"
     ]
    }
   ],
   "source": [
    "lrs = [1e-3, 3e-3]\n",
    "optim = variable_lr_optimizer(gnet, lrs)\n",
    "print(optim)\n",
    "assert [group['lr'] for group in optim.param_groups] == lrs\n",
    "\n",
    "update_optimizer(optim, 2e-3, lr_mult=1/3)\n",
    "print([group['lr'] for group in optim.param_groups])\n",
    "assert np.isclose(optim.param_groups[1]['lr'], optim.param_groups[0]['lr'] * 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0005, 0.001]\n"
     ]
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(gnet, 1e-3, lr_mult=0.5)\n",
    "print([group['lr'] for group in optim.param_groups])\n",
    "assert np.isclose(optim.param_groups[1]['lr'], optim.param_groups[0]['lr'] * 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      ")\n",
      "As expected, got ValueError(Received more learning rates than layer groups.).\n"
     ]
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(snet, 2e-3)\n",
    "print(optim)\n",
    "\n",
    "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": [],
   "source": [
    "update_optimizer(optim, 1e-3, 0.5)\n",
    "assert len(optim.param_groups) == 1\n",
    "assert optim.param_groups[0]['lr'] == 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      ")\n",
      "[0.0006666666666666666, 0.002]\n"
     ]
    }
   ],
   "source": [
    "lrs = [1e-3, 3e-3]\n",
    "optim = variable_lr_optimizer(gnet, lrs)\n",
    "print(optim)\n",
    "assert [group['lr'] for group in optim.param_groups] == lrs\n",
    "\n",
    "update_optimizer(optim, 2e-3, lr_mult=1/3)\n",
    "print([group['lr'] for group in optim.param_groups])\n",
    "assert np.isclose(optim.param_groups[1]['lr'], optim.param_groups[0]['lr'] * 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0005, 0.001]\n"
     ]
    }
   ],
   "source": [
    "optim = variable_lr_optimizer(gnet, 1e-3, lr_mult=0.5)\n",
    "print([group['lr'] for group in optim.param_groups])\n",
    "assert np.isclose(optim.param_groups[1]['lr'], optim.param_groups[0]['lr'] * 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "adam = partial(torch.optim.Adam, eps=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def handle_interrupt(meth):\n",
    "    \"\"\"Decorator for Trainer.fit() method that allows the user to\n",
    "    interrupt training with Ctrl+c while still running the\n",
    "    `on_train_end` method for each of its callbacks. Without this,\n",
    "    cutting training short would lose that functionality, so we\n",
    "    couldn't do things like uploading to S3.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    meth: callable\n",
    "        The method to decorate.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    callable: The wrapped method.\n",
    "    \"\"\"\n",
    "    @wraps(meth)\n",
    "    def wrapper(*args, **kwargs):\n",
    "        instance = args[0]\n",
    "        res = None\n",
    "        try:\n",
    "            res = meth(*args, **kwargs)\n",
    "        except KeyboardInterrupt:\n",
    "            instance.logger.info(f'Stop training due to KeyboardInterrupt.')\n",
    "            instance._stop_training = True\n",
    "            # Dummy values used for epoch and stats to indicate that\n",
    "            # training was interrupted. `fit()` method returns None.\n",
    "            _ = instance.decide_stop('on_train_end', -1, {})\n",
    "        return res\n",
    "    return wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "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, optim=None, optim_type=Adam, eps=1e-3, last_act=None,\n",
    "                 threshold=0.5, metrics=None, callbacks=None, device=DEVICE):\n",
    "        \"\"\"An object to handle model training. This makes it easy for us to\n",
    "        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",
    "        optim: torch.optim\n",
    "            Optional: an optimizer object\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",
    "        if last_act is None and mode != 'regression':\n",
    "            warnings.warn(\n",
    "                'Last activation is None for a classification problem. This '\n",
    "                'means your network must include a sigmoid or softmax at the '\n",
    "                'end if you wish to compute any metrics using soft '\n",
    "                'predictions.'\n",
    "            )\n",
    "\n",
    "        if optim:\n",
    "            optim_type = type(optim)\n",
    "            warnings.warn('Inferring optim_type from optim argument.')\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 = optim\n",
    "        self.eps = eps\n",
    "        self.criterion = criterion\n",
    "        self.mode = mode\n",
    "        self.device = DEVICE\n",
    "        self.last_act = last_act or identity\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",
    "        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 model and optimizer state dicts for later use. This\n",
    "        includes the model, optimizer. Datasets and data loaders are\n",
    "        excluded since:\n",
    "        a. It seems that they can't be pickled in some cases (e.g. on Ubuntu).\n",
    "        b. They have no learnable parameters to track during training.\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",
    "        data = {'model': self.net.state_dict()}\n",
    "        try:\n",
    "            data['optim'] = self.optim.state_dict()\n",
    "        except AttributeError:\n",
    "            self.logger.warning('No optimizer. Only saving model state dict.')\n",
    "        torch.save(data, os.path.join(self.out_dir, fname))\n",
    "\n",
    "    def load(self, fname=None, old_path=None):\n",
    "        \"\"\"This lets a trainer load previously saved model and optimizer\n",
    "        weights. This is an in-place operation, so nothing is 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",
    "        old_path: str\n",
    "            Full path to file where a previous Trainer object is stored.\n",
    "            This allows us to load model and optimizer weights from a\n",
    "            different round of training and store the results in a new\n",
    "            directory, potentially using different hyperparameters or\n",
    "            datasets.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "\n",
    "        Examples\n",
    "        --------\n",
    "        trainer = Trainer(...)\n",
    "        trainer.fit(...)\n",
    "        trainer.save('v1')\n",
    "        trainer = trainer.load('v1')\n",
    "        \"\"\"\n",
    "        path = old_path or os.path.join(self.out_dir, fname)\n",
    "        self.logger.info(f'Loading weights from {path}.')\n",
    "        data = torch.load(path, map_location=self.device)\n",
    "        self.net.load_state_dict(data['model'])\n",
    "\n",
    "        # Create optimizer to load state dict. LR will be updated later.\n",
    "        if not self.optim:\n",
    "            self.net.to(self.device)\n",
    "            self.optim = variable_lr_optimizer(self.net,\n",
    "                                               optimizer=self.optim_type)\n",
    "        try:\n",
    "            self.optim.load_state_dict(data['optim'])\n",
    "        except (AttributeError, KeyError) as e:\n",
    "            self.logger.warning('Could not load optimizer. '\n",
    "                                ' Loading model weights only.\\n' + repr(e))\n",
    "\n",
    "    def add_callbacks(self, *callbacks):\n",
    "        \"\"\"Attach additional callbacks to Trainer. Note that callback order\n",
    "        will be determined by their `order` 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].order))\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 set_callback_attr(self, cb_name, attr, val):\n",
    "        \"\"\"Convenience method to change an attribute of an existing\n",
    "        callback.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        cb_name: str\n",
    "            Name of callback to update.\n",
    "        attr: str\n",
    "            Name of attribute to update.\n",
    "        val: any\n",
    "            Value of attribute to set.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \n",
    "        Examples\n",
    "        --------\n",
    "        # This reduces the frequency with with we update the batch stats\n",
    "        # in the progress bar.\n",
    "        trainer.set_callback_attr('MetricPrinter', 'batch_freq', 10)\n",
    "        \"\"\"\n",
    "        setattr(self.callbacks[cb_name], attr, val)\n",
    "\n",
    "    @handle_interrupt\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",
    "        self.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, None)\n",
    "            for i, batch in enumerate(self.pbar):\n",
    "                _ = self.decide_stop('on_batch_begin', i, sum_i)\n",
    "                sum_i += 1\n",
    "                *xb, yb = map(lambda x: x.to(self.device), batch)\n",
    "                self.optim.zero_grad()\n",
    "                _ = self.decide_stop('after_zero_grad', i, sum_i, xb, yb)\n",
    "\n",
    "                # Forward and backward passes.\n",
    "                y_score = self.net(*xb)\n",
    "                if self.decide_stop('after_forward', i, sum_i): break\n",
    "                loss = self.criterion(y_score, yb)\n",
    "                if self.decide_stop('after_loss', i, sum_i): break\n",
    "                loss.backward()\n",
    "                if self.decide_stop('after_backward', i, sum_i): break\n",
    "                self.optim.step()\n",
    "                if self.decide_stop('after_step', i, sum_i): break\n",
    "\n",
    "                # Separate because callbacks are only applied during training.\n",
    "                self._update_stats(self.stats, loss, yb, y_score)\n",
    "                if self.decide_stop('on_batch_end', i, sum_i): 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, val_stats): break\n",
    "                continue\n",
    "            break\n",
    "        _ = self.decide_stop('on_train_end', e, 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 = dl_val or self.dl_val\n",
    "        val_stats = defaultdict(list)\n",
    "        self.net.eval()\n",
    "        with torch.no_grad():\n",
    "            for batch in tqdm(dl_val, leave=False):\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 predict(self, *xb, logits=True):\n",
    "        \"\"\"Make predictions on a batch of data. This automatically does things\n",
    "        like putting the data and model on the same device, putting the model\n",
    "        in eval mode, and ensuring that gradients are not computed (reduces\n",
    "        time and memory usage).\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        xb: torch.tensors\n",
    "            Inputs to the model. This will often just be one x tensor, but\n",
    "            sometimes other inputs are required as well (e.g. attention masks).\n",
    "        logits: bool\n",
    "            If True, output logits. If False, the last activation function is \n",
    "            applied.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        torch.tensor: Model predictions.\n",
    "        \"\"\"\n",
    "        xb = map(lambda x: x.to(self.device), xb)\n",
    "        self.net.to(self.device)\n",
    "        res = self.net.predict(*xb)\n",
    "        if not logits: res = self.last_act(res)\n",
    "        return res\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",
    "        yb, y_score = yb.detach().cpu(), y_score.detach().cpu()\n",
    "        # Final activation often excluded from network architecture.\n",
    "        y_score = self.last_act(y_score)\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)})'\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"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
