{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp optimizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizers\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": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from collections.abc import Iterable\n",
    "from functools import partial\n",
    "from torch.optim import Adam\n",
    "\n",
    "from incendio.metrics import batch_size\n",
    "from incendio.utils import quick_stats, DEVICE"
   ]
  },
  {
   "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": [
    "## Optimizers\n",
    "\n",
    "Optimizers like Adam or RMSProp can contain multiple \"parameter groups\", each with a different learning rate. (Other hyperparameters can vary as well, but we ignore that for now.) The functions below allow us to get a new optimizer or update an existing one. It allows us to easily use differential learning rate, but that is not required: it can also use the same LR for each parameter group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def variable_lr_optimizer(model, lr=3e-3, lr_mult=1.0, optimizer=Adam, \n",
    "                          eps=1e-3, **kwargs):\n",
    "    \"\"\"Get an optimizer that uses different learning rates for different layer\n",
    "    groups. Additional keyword arguments can be used to alter momentum and/or\n",
    "    weight decay, for example, but for the sake of simplicity these values\n",
    "    will be the same across layer groups.\n",
    "\n",
    "    Parameters\n",
    "    -----------\n",
    "    model: nn.Module\n",
    "        A model object. If you intend to use differential learning rates,\n",
    "        the model must have an attribute `groups` containing a ModuleList of\n",
    "        layer groups in the form of Sequential objects. The number of layer\n",
    "        groups must match the number of learning rates passed in.\n",
    "    lr: float, Iterable[float]\n",
    "        A number of list of numbers containing the learning rates to use for\n",
    "        each layer group. There should generally be one LR for each layer group\n",
    "        in the model. If fewer LR's are provided, lr_mult will be used to\n",
    "        compute additional LRs. See `update_optimizer` for details.\n",
    "    lr_mult: float\n",
    "        If you pass in fewer LRs than layer groups, `lr_mult` will be used to\n",
    "        compute additional learning rates from the one that was passed in.\n",
    "    optimizer: torch optimizer\n",
    "        The Torch optimizer to be created (Adam by default).\n",
    "    eps: float\n",
    "        Hyperparameter used by optimizer. The default of 1e-8 can lead to\n",
    "        exploding gradients, so we typically override this.\n",
    "\n",
    "    Examples\n",
    "    ---------\n",
    "    optim = variable_lr_optimizer(model, lrs=[3e-3, 3e-2, 1e-1])\n",
    "    \"\"\"\n",
    "    groups = getattr(model, 'groups', [model])\n",
    "    # Placeholder LR used. We update this afterwards.\n",
    "    data = [{'params': group.parameters(), 'lr': 0} for group in groups]\n",
    "    optim = optimizer(data, eps=eps, **kwargs)\n",
    "    update_optimizer(optim, lr, lr_mult)\n",
    "    return optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def update_optimizer(optim, lrs, lr_mult=1.0):\n",
    "    \"\"\"Pass in 1 or more learning rates, 1 for each layer group, and update the\n",
    "    optimizer accordingly. The optimizer is updated in place so nothing is\n",
    "    returned.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    optim: torch.optim\n",
    "        Optimizer object.\n",
    "    lrs: float, Iterable[float]\n",
    "        One or more learning rates. If using multiple values, usually the\n",
    "        earlier values will be smaller and later values will be larger. This\n",
    "        can be achieved by passing in a list of LRs that is the same length as\n",
    "        the number of layer groups in the optimizer, or by passing in a single\n",
    "        LR and a value for lr_mult.\n",
    "    lr_mult: float\n",
    "        If you pass in fewer LRs than layer groups, `lr_mult` will be used to\n",
    "        compute additional learning rates from the one that was passed in.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    If optim has 3 layer groups, this will result in LRs of [3e-5, 3e-4, 3e-3]\n",
    "    in that order:\n",
    "    update_optimizer(optim, lrs=3e-3, lr_mult=0.1)\n",
    "\n",
    "    Again, optim has 3 layer groups. We leave the default lr_mult of 1.0 so\n",
    "    each LR will be 3e-3.\n",
    "    update_optimizer(optim, lrs=3e-3)\n",
    "\n",
    "    Again, optim has 3 layer groups. 3 LRs are passed in so lr_mult is unused.\n",
    "    update_optimizer(optim, lrs=[1e-3, 1e-3, 3e-3])\n",
    "    \"\"\"\n",
    "    if not isinstance(lrs, Iterable): lrs = [lrs]\n",
    "    n_missing = len(optim.param_groups) - len(lrs)\n",
    "\n",
    "    if n_missing < 0:\n",
    "        raise ValueError('Received more learning rates than layer groups.')\n",
    "    while n_missing > 0:\n",
    "        lrs.insert(0, lrs[0] * lr_mult)\n",
    "        n_missing -= 1\n",
    "\n",
    "    for group, lr in zip(optim.param_groups, lrs):\n",
    "        group['lr'] = lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "adam = partial(Adam, eps=1e-3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
