{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": false
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "[[chapter_accel_sgd]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variants of SGD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you know all about how the architectures are put together, it's time to start exploring the training process.\n",
    "\n",
    "We explained earlier the basis of Stochastic Gradient Descent: pass a minibatch in the model, compare it to our target with the loss function then compute the gradients of this loss function with regards to each weight before updating the weights with the formula:\n",
    "\n",
    "```python\n",
    "new_weight = weight - lr * weight.grad\n",
    "```\n",
    "\n",
    "We implemented this from scratch in a training loop, and also saw that Pytorch provides a simple `nn.SGD` class that does this calculation for each parameter for us. Let's now build some faster optimizers, using a flexible foundation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's start with SGD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll create a baseline, using plain SGD, and compare it to fastai's default optimizer. We'll start by grabbing Imagenette with the same `get_data` we used in <<chapter_resnet>>:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide_input\n",
    "def get_data(url, presize, resize):\n",
    "    path = untar_data(url)\n",
    "    return DataBlock(\n",
    "        blocks=(ImageBlock, CategoryBlock), get_items=get_image_files, \n",
    "        splitter=GrandparentSplitter(valid_name='val'),\n",
    "        get_y=parent_label, item_tfms=Resize(presize),\n",
    "        batch_tfms=[*aug_transforms(min_scale=0.5, size=resize),\n",
    "                    Normalize.from_stats(*imagenet_stats)],\n",
    "    ).dataloaders(path, bs=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dls = get_data(URLs.IMAGENETTE_160, 160, 128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create a ResNet34 without pretraining, and pass along any arguments received:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_learner(**kwargs):\n",
    "    return cnn_learner(dls, resnet34, pretrained=False,\n",
    "                    metrics=accuracy, **kwargs).to_fp16()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the default fastai optimizer, with the usual 3e-3 learning rate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.571932</td>\n",
       "      <td>2.685040</td>\n",
       "      <td>0.322548</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.904674</td>\n",
       "      <td>1.852589</td>\n",
       "      <td>0.437452</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.586909</td>\n",
       "      <td>1.374908</td>\n",
       "      <td>0.594904</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn = get_learner()\n",
    "learn.fit_one_cycle(3, 0.003)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try plain SGD. We can pass `opt_func` (optimization function) to `cnn_learner` to get fastai to use any optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = get_learner(opt_func=SGD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing to look at is `lr_find`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(0.017378008365631102, 3.019951861915615e-07)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12620 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn.lr_find()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like we'll need to use a higher learning rate than we normally use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.969412</td>\n",
       "      <td>2.214596</td>\n",
       "      <td>0.242038</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.442730</td>\n",
       "      <td>1.845950</td>\n",
       "      <td>0.362548</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>2.157159</td>\n",
       "      <td>1.741143</td>\n",
       "      <td>0.408917</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn.fit_one_cycle(3, 0.03, moms=(0,0,0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Because accelerated SGD using momentum with is such a good idea, fastai uses it by default in `fit_one_cycle`, so we turn it off with `moms=(0,0,0)`; we'll be learning about momentum shortly.)\n",
    "\n",
    "Clearly, plain SGD isn't training as fast as we'd like. So let's learn the tricks to get accelerated training!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A generic optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to build up our accelerated SGD tricks, we'll need to start with a nice flexible optimizer foundation. No library prior to fastai provided such a foundation, but during fastai's development we realized that all optimizer improvements we'd seen in the academic literature could be handled using *optimizer callbacks*. These are small pieces of code that an optimizer can add to the optimizer `step`. They are called by fastai's `Optimizer` class. This is a small class (less than a screen of code); these are the definitions in `Optimizer` of the two key methods that we've been using in this book:\n",
    "\n",
    "```python\n",
    "def zero_grad(self):\n",
    "    for p,*_ in self.all_params():\n",
    "        p.grad.detach_()\n",
    "        p.grad.zero_()\n",
    "\n",
    "def step(self):\n",
    "    for p,pg,state,hyper in self.all_params():\n",
    "        for cb in self.cbs:\n",
    "            state = _update(state, cb(p, **{**state, **hyper}))\n",
    "        self.state[p] = state\n",
    "```\n",
    "\n",
    "As we saw when training an MNIST model from scratch, `zero_grad` just loops through the parameters of the model and sets the gradients to zero. It also calls `detach_`, which removes any history of gradient computation, since it won't be needed after `zero_grad`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The more interesting method is `step`, which loops through the callbacks (`cbs`) and calls them to update the parameters (the `_update` function just calls `state.update` if there's anything returned by `cb(...)`). As you can see, `Optimizer` doesn't actually do any SGD steps itself. Let's see how we can add SGD to `Optimizer`.\n",
    "\n",
    "Here's an optimizer callback that does a single SGD step, by multiplying `-lr` by the gradients, and adding that to the parameter (when `Tensor.add_` in PyTorch is passed two parameters, they are multiplied together before the addition): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sgd_cb(p, lr, **kwargs): p.data.add_(-lr, p.grad.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can pass this to `Optimizer` using the `cbs` parameter; we'll need to use `partial` since `Learner` will call this function to create our optimizer later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_func = partial(Optimizer, cbs=[sgd_step])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if this trains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.730918</td>\n",
       "      <td>2.009971</td>\n",
       "      <td>0.332739</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.204893</td>\n",
       "      <td>1.747202</td>\n",
       "      <td>0.441529</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.875621</td>\n",
       "      <td>1.684515</td>\n",
       "      <td>0.445350</td>\n",
       "      <td>00:09</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn = get_learner(opt_func=opt_func)\n",
    "learn.fit(3, 0.03)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's working! So that's how we create SGD from scratch in fastai."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Momentum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SGD is the idea of taking a step in the direction of the steepest slope at each point of time. But what if we have a ball rolling down the mountain? It won't, at each given point, exactly follow the direction of the gradient, as it will have *momentum*. A ball with more momentum (for instance, a heavier ball) will skip over little bumps and holes, and be more likely to get to the bottom of a bumpy mountain. A ping pong ball, on the other hand, will get stuck in every little crevice.\n",
    "\n",
    "So how could we bring this idea over to SGD? We can use a moving average, instead of only the current gradient, to make our step:\n",
    "\n",
    "```python\n",
    "weight.avg = beta * weight.avg + (1-beta) * weight.grad\n",
    "new_weight = weight - lr * weight.avg\n",
    "```\n",
    "\n",
    "Here `beta` is some number we choose which defines how much momentum to use. If `beta` is zero, then the first equation above becomes  `weight.avg = weight.grad`, so we end up with plain SGD. But if it's a number close to one, then the main direction chosen is an average of previous steps. (If you have done a bit of statistics, you may recognize in the first equation an *exponentially weighted moving average*, which is very often used to denoise data and get the underlying tendency.)\n",
    "\n",
    "Note that we are writing `weight.avg` to highlight the fact we need to store thoe moving averages for each parameter of the model (and they all their own independent moving averages).\n",
    "\n",
    "Here is an example of noisy data for a single parameter, with the momentum curve plotted in red, and the gradients of the parameter plotted in blue. The gradients increase, and then decrease, and the momentum does a good job of following the general trend, without getting too influenced by noise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17588 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(-4, 4, 100)\n",
    "y = 1 - (x/3) ** 2\n",
    "x1 = x + np.random.randn(100) * 0.1\n",
    "y1 = y + np.random.randn(100) * 0.1\n",
    "plt.scatter(x1,y1)\n",
    "idx = x1.argsort()\n",
    "beta,avg,res = 0.7,0,[]\n",
    "for i in idx:\n",
    "    avg = beta * avg + (1-beta) * y1[i]\n",
    "    res.append(avg/(1-beta**(i+1)))\n",
    "plt.plot(x1[idx],np.array(res), color='red');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works particularly well if the loss function has narrow canyons we need to navigate: vanilla SGD would send us from one side to the other while SGD with momentum will average those to roll down inside. The parameter `beta` determines the strength of that momentum we are using: with a small beta we stay closer to the actual gradient values whereas with a high beta, we will mostly go in the direction of the average of the gradients and it will take a while before any change in the gradients makes that trend move.\n",
    "\n",
    "With a large beta, we might miss that the gradients have changed directions and roll over a small local minima which is a desired side-effect: intuitively, when we show a new picture/text/data to our model, it will look like something in the training set but won't be exactly like it. That means it will correspond to a point in the loss function that is closest to the minimum we ended up with at the end of training, but not exactly *at* that minimum. We then would rather end up training in a wide minimum, where nearby points have approximately the same loss (or if you prefer, a point where the loss is as flat as possible). Here's how the above chart varies as we change beta:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57736 characters]\n",
      "text/plain": [
       "<Figure size 864x576 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(-4, 4, 100)\n",
    "y = 1 - (x/3) ** 2\n",
    "x1 = x + np.random.randn(100) * 0.1\n",
    "y1 = y + np.random.randn(100) * 0.1\n",
    "_,axs = plt.subplots(2,2, figsize=(12,8))\n",
    "betas = [0.5,0.7,0.9,0.99]\n",
    "idx = x1.argsort()\n",
    "for beta,ax in zip(betas, axs.flatten()):\n",
    "    ax.scatter(x1,y1)\n",
    "    avg,res = 0,[]\n",
    "    for i in idx:\n",
    "        avg = beta * avg + (1-beta) * y1[i]\n",
    "        res.append(avg)#/(1-beta**(i+1)))\n",
    "    ax.plot(x1[idx],np.array(res), color='red');\n",
    "    ax.set_title(f'beta={beta}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see in these examples that a beta that's too high results in the overall changes in gradient getting ignored. In SGD with momentum, a value of `beta` that is often used is 0.9.\n",
    "\n",
    "`fit_one_cycle` by default starts with a beta of 0.95, gradually adjusts it to 0.85, and then gradually moves it back to 0.95 at the end of training. Let's see how our training goes with momentum added to plain SGD:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to add momentum to our optimizer, we'll first need to keep track of the moving average gradient, which we can do with another callback. When an optimizer callback returns a dict, it is used to update the state of the optimizer, and is passed back to the optimizer on the next step. So this callback will keep track of the gradient averages in a parameter called `grad_avg`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_grad(p, mom, grad_avg=None, **kwargs):\n",
    "    if grad_avg is None: grad_avg = torch.zeros_like(p.grad.data)\n",
    "    return {'grad_avg': grad_avg*mom + p.grad.data}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use it, we just have to replace `p.grad.data` with `grad_avg` in our step function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def momentum_step(p, lr, grad_avg, **kwargs): p.data.add_(-lr, grad_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_func = partial(Optimizer, cbs=[average_grad,momentum_step], mom=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Learner` will automatically schedule `mom` and `lr`, so fit_one_cycle will even work with our custom Optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.856000</td>\n",
       "      <td>2.493429</td>\n",
       "      <td>0.246115</td>\n",
       "      <td>00:10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.504205</td>\n",
       "      <td>2.463813</td>\n",
       "      <td>0.348280</td>\n",
       "      <td>00:10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>2.187387</td>\n",
       "      <td>1.755670</td>\n",
       "      <td>0.418853</td>\n",
       "      <td>00:10</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn = get_learner(opt_func=opt_func)\n",
    "learn.fit_one_cycle(3, 0.03)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28636 characters]\n",
      "text/plain": [
       "<Figure size 864x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn.recorder.plot_sched()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're still not getting great results, so let's see what else we can do."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RMSProp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RMSProp is another variant of SGD introduced by Geoffrey Hinton in [Lecture 6e of his Coursera class](http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf). The main difference with SGD is that it uses an adaptive learning rate: instead of using the same learning rate for every parameter, each parameter gets it's own specific learning rate controlled by a global learning rate. That way we can speed up training by giving a high learning rate to the weights that needs to change a lot while the ones that are good enough get a lower learning rate.\n",
    "\n",
    "How do we decide which parameter should have a high learning rate and which should not? We can look at the gradients to get an idea. Not just the one we computed, but all of them: if they have been close to 0 for a while, it means this parameter will need a higher learning rate because the loss is very flat. On the opposite, if they are all over the place, we should probably be careful and pick a low learning rate to avoid divergence. We can't just average the gradients to see if they're changing a lot, since the average of a large positive and a large negative number is close to zero. So we can use the usual trick of either taking the absolute value, or the squared values (and then taking the square root after the mean).\n",
    "\n",
    "Once again, to pick the general tendency behind the noise, we will use a moving average, specifically the moving average of the gradients squared. Then, we will update the corresponding weight by using the current gradient (for the direction) divided by the square root of this moving average (that way if it's low, the effective learning rate will be higher, and if it's big, the effective learning rate will be lower).\n",
    "\n",
    "```python\n",
    "w.square_avg = alpha * w.square_avg + (1-alpha) * (w.grad ** 2)\n",
    "new_w = w - lr * w.grad / math.sqrt(w.square_avg + eps)\n",
    "```\n",
    "\n",
    "The `eps` (*epsilon*) is added for numerical stability (usually set at 1e-8) and the default value for `alpha` is usually 0.99."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can add this to `Optimizer` by doing much the same thing we did for `avg_grad`, but with an extra `**2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_sqr_grad(p, sqr_mom, sqr_avg=None, **kwargs):\n",
    "    if sqr_avg is None: sqr_avg = torch.zeros_like(p.grad.data)\n",
    "    return {'sqr_avg': sqr_avg*sqr_mom + p.grad.data**2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can define our step function and optimizer as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rms_prop_step(p, lr, sqr_avg, eps, grad_avg=None, **kwargs):\n",
    "    denom = sqr_avg.sqrt().add_(eps)\n",
    "    p.data.addcdiv_(-lr, p.grad, denom)\n",
    "\n",
    "opt_func = partial(Optimizer, cbs=[average_sqr_grad,rms_prop_step],\n",
    "                   sqr_mom=0.99, eps=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.766912</td>\n",
       "      <td>1.845900</td>\n",
       "      <td>0.402548</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.194586</td>\n",
       "      <td>1.510269</td>\n",
       "      <td>0.504459</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.869099</td>\n",
       "      <td>1.447939</td>\n",
       "      <td>0.544968</td>\n",
       "      <td>00:11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn = get_learner(opt_func=opt_func)\n",
    "learn.fit_one_cycle(3, 0.003)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much better! Now we just have to bring these ideas together, and we have Adam, fastai's default optimizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adam mixes the ideas of SGD with momentum and RMSProp together: it uses the moving average of the gradients as a direction and divides by the square root of the moving average of the gradients squared to give an adaptive learning rate to each parameter.\n",
    "\n",
    "There is one other difference with how Adam calculates moving averages, is that it takes the *unbiased* moving average which is:\n",
    "\n",
    "``` python\n",
    "w.avg = beta * w.avg + (1-beta) * w.grad\n",
    "unbias_avg = w.avg / (1 - (beta**(i+1)))\n",
    "```\n",
    "\n",
    "if we are the `i`-th iteration (starting at 0 like python does). This divisor of `1 - (beta**(i+1))` makes sure the unbiased average looks more like the gradients at the beginning (since `beta < 1` the denominator is very quickly very close to 1).\n",
    "\n",
    "Putting everything together, our update step looks like:\n",
    "``` python\n",
    "w.avg = beta1 * w.avg + (1-beta1) * w.grad\n",
    "unbias_avg = w.avg / (1 - (beta1**(i+1)))\n",
    "w.sqr_avg = beta2 * w.sqr_avg + (1-beta2) * (w.grad ** 2)\n",
    "new_w = w - lr * unbias_avg / sqrt(w.sqr_avg + eps)\n",
    "```\n",
    "\n",
    "Like for RMSProp, `eps` is usually set to 1e-8, and the default for `(beta1,beta2)` suggested by the literature `(0.9,0.999)`. \n",
    "\n",
    "In fastai, Adam is the default optimizer we use since it allows faster training, but we found that `beta2=0.99` is better suited for the type of schedule we are using. `beta1` is the momentum parameter, which we specify with the argument `moms` in our call to `fit_one_cycle`. As for `eps`, fastai uses a default of 1e-5. `eps` is not just useful for numerical stability. A higher `eps` limits the maximum value of the adjusted learning rate. To take an extreme example, if `eps` is 1, then the adjusted learning will never be higher than the base learning rate. \n",
    "\n",
    "Rather than show all the code for this in the book, we'll let you look at the optimizer notebook in fastai's GitHub repository--you'll see all the code we've seen so far, along with Adam and other optimizers, and lots of examples and tests."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoupled weight_decay"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've discussed weight decay before, which is equivalent to (in the case of vanilla SGD) updating the parameters\n",
    "with:\n",
    "\n",
    "``` python\n",
    "new_weight = weight - lr*weight.grad - lr*wd*weight\n",
    "```\n",
    "\n",
    "This last formula explains why the name of this technique is weight decay, as each weight is decayed by a factor `lr * wd`. \n",
    "\n",
    "However, this only works correctly  for standard SGD, because we have seen that with momentum, RMSProp or in Adam, the update has some additional formulas around the gradient. In those cases, the formula that comes from L2 regularization:\n",
    "\n",
    "``` python\n",
    "weight.grad += wd*weight\n",
    "```\n",
    "\n",
    "is different than weight decay:\n",
    "\n",
    "``` python\n",
    "new_weight = weight - lr*weight.grad - lr*wd*weight\n",
    "```\n",
    "\n",
    "Most libraries use the first formulation, but it was pointed out in [Decoupled Weight Regularization](https://arxiv.org/pdf/1711.05101.pdf) by Ilya Loshchilov and Frank Hutter, second one is the only correct approach with the Adam optimizer or momentum, which is why fastai makes it its default.\n",
    "\n",
    "Now you know everything that is hidden behind the line `learn.fit_one_cycle`!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Questionnaire"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. What is the equation for a step of SGD, in math or code (as you prefer)?\n",
    "1. What do we pass to `cnn_learner` to use a non-default optimizer?\n",
    "1. What are optimizer callbacks?\n",
    "1. What does `zero_grad` do in an optimizer?\n",
    "1. What does `step` do in an optimizer? How is it implemented in the general optimizer?\n",
    "1. Rewrite `sgd_cb` to use the `+=` operator, instead of `add_`.\n",
    "1. What is momentum? Write out the equation.\n",
    "1. What's a physical analogy for momentum? How does it apply in our model training settings?\n",
    "1. What does a bigger value for momentum do to the gradients?\n",
    "1. What are the default values of momentum for 1cycle training?\n",
    "1. What is RMSProp? Write out the equation.\n",
    "1. What do the squared values of the gradients indicate?\n",
    "1. How does Adam differ from momentum and RMSProp?\n",
    "1. Write out the equation for Adam.\n",
    "1. Calculate the value of `unbias_avg` and `w.avg` for a few batches of dummy values.\n",
    "1. What's the impact of having a high eps in Adam?\n",
    "1. Read through the optimizer notebook in fastai's repo, and execute it.\n",
    "1. In what situations do dynamic learning rate methods like Adam change the behaviour of weight decay?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Further research"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Look up the \"rectified Adam\" paper and implement it using the general optimizer framework, and try it out. Search for other recent optimizers that work well in practice, and pick one to implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "split_at_heading": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
