06_Optimizers_workset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "executionInfo": {
     "elapsed": 5680,
     "status": "ok",
     "timestamp": 1658778337919,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "PPIg8ZrxJ2Eh"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 6,
     "status": "ok",
     "timestamp": 1658778337920,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "kybq1eS5J7u6"
   },
   "outputs": [],
   "source": [
    "N_SAMPLES = 1000\n",
    "noise = torch.randn(N_SAMPLES) * 3\n",
    "# Pick the values of the parameters we want to learn\n",
    "param_1_actual = 8.765\n",
    "param_2_actual = -4.321\n",
    "\n",
    "def f(xs, param1=param_1_actual, param2=param_2_actual, add_noise=True):\n",
    "    \"\"\"\n",
    "    This function generates Ys from xs. xs should be a torch.tensor with two columns.\n",
    "    The default for param1 and param2 are given, but we may want to change them later.\n",
    "    By default, we'll add noise but we'll see some other visualizations are \n",
    "    cleaner when we remove it.\n",
    "    \"\"\"\n",
    "    params = torch.tensor([param1, param2]).float()\n",
    "    out = xs@params\n",
    "    if add_noise:\n",
    "        out.add_(torch.randn(xs.shape[0]) * 3)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1658778337920,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "IP052fbmLPIM"
   },
   "outputs": [],
   "source": [
    "# Create the xs\n",
    "xs = torch.randn(N_SAMPLES, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "executionInfo": {
     "elapsed": 248,
     "status": "ok",
     "timestamp": 1658778338163,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "xwhr3GM3Lgph"
   },
   "outputs": [],
   "source": [
    "# Create the ys\n",
    "ys = f(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CAXUC5IaEQ-3"
   },
   "source": [
    "# Exercise 6.1: Make a dataset\n",
    "\n",
    "Use what you know from previous lessions to complete the class `MyDataset`.\n",
    "Then, use the `xs` and `ys` we created to instantiate your dataset and a `DataLoader` with `batch_size=32` and `shuffle=True`.\n",
    "<!-- startquestion -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1658778338164,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "NJARjluAEgrK"
   },
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    ...\n",
    "\n",
    "# Create the datset\n",
    "ds = ...\n",
    "# Create the dataloader\n",
    "dl = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "executionInfo": {
     "elapsed": 4,
     "status": "ok",
     "timestamp": 1658778338164,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "yhnRCzC1E85p"
   },
   "outputs": [],
   "source": [
    "assert isinstance(ds, Dataset)\n",
    "assert isinstance(dl, DataLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 683
    },
    "executionInfo": {
     "elapsed": 697,
     "status": "ok",
     "timestamp": 1658778338858,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "Z-IYXLU-LxHA",
    "outputId": "cf191152-8c9c-47f5-be65-2cf5e1d00413"
   },
   "outputs": [],
   "source": [
    "# Visualize a sample of our data\n",
    "_df = pd.DataFrame(xs.numpy(), columns=['x1', 'x2'])\n",
    "_df['y'] = ys.numpy()\n",
    "data_fig = plt.figure(figsize=(12,12))\n",
    "ax = data_fig.add_subplot(projection='3d')\n",
    "ax.scatter(_df.x1, _df.x2, _df.y, alpha=0.2)\n",
    "ax.set_xlabel('x1', fontsize=15)\n",
    "ax.set_ylabel('x2', fontsize=15)\n",
    "ax.set_zlabel('y', fontsize=15)\n",
    "plt.close()\n",
    "data_fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tHFbv-URm10P"
   },
   "source": [
    "# Exercise 6.2: Create a small neural network\n",
    "\n",
    "Please complete the class `Net` by completing the `__init__` and `forward` methods.\n",
    "Your parameters should contain a float tensor with two elements, one for $x_1$ and one for $x_2$.\n",
    "Please initialize each of those parameters to the value 3.0.\n",
    "The `forward` method should multiply your inputs by these parameters.\n",
    "In this case, let's ignore the bias parameter, since we didn't use it in creating our dataset.\n",
    "\n",
    "<!-- startquestion -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1658778338859,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "nNXD3MLroaYn"
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.params = ...\n",
    "\n",
    "    def forward(self, x):\n",
    "        return ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "executionInfo": {
     "elapsed": 193,
     "status": "ok",
     "timestamp": 1658778339048,
     "user": {
      "displayName": "Mike Frantz",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "VE3VOPrCn27s"
   },
   "outputs": [],
   "source": [
    "assert Net()(torch.tensor([[1., 1.]])) == 6 # Check the outputs against a known input\n",
    "assert (next(Net().parameters()).data == torch.tensor([3., 3.])).all() # Check that the parameters are all 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o1SJCtEhPJeC"
   },
   "outputs": [],
   "source": [
    "def get_parameters_over_training(optimizer, opt_kwargs={'lr':1e-3}):\n",
    "    \"\"\"\n",
    "    This function is used to learn about different optimizers.\n",
    "    The optimizer argument is a class from torch.optim,\n",
    "    and the opt_kwargs is a dictionary of keyword arguments to pass\n",
    "    to the optimizer as it's instantiated.\n",
    "\n",
    "    This function returns a tuple of (params_history, losses)\n",
    "    where params_history is a numpy array where each column represents\n",
    "    the history of a parameter as the model's trained.\n",
    "    The losses are just the loss at every step.\n",
    "    \"\"\"\n",
    "    # Instantiate our network\n",
    "    model = Net()\n",
    "    # Instantiate the optimizer\n",
    "    opt = optimizer(model.parameters(), **opt_kwargs)\n",
    "    # Create placeholders for the params history and losses\n",
    "    params_history = []\n",
    "    losses = []\n",
    "    # Instantiate the loss\n",
    "    loss = nn.MSELoss()\n",
    "    # Train the model for 50 epochs\n",
    "    for epoch in range(50):\n",
    "        for x_b, y_b in ds:\n",
    "            # Make predictions\n",
    "            preds = model(x_b)\n",
    "            # Log the parameters\n",
    "            c = model.params.clone().detach().numpy()\n",
    "            params_history.append(c)\n",
    "            # Calculate and record the loss\n",
    "            l = loss(preds, y_b)\n",
    "            losses.append(l.clone().detach().numpy())\n",
    "            # Calculate the gradients\n",
    "            l.backward()\n",
    "            # Apply the optimizer update and zero the gradients\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "        if epoch % 10 == 0:\n",
    "            print(l.detach().numpy(), c)\n",
    "    return np.vstack(params_history), np.hstack(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NJB6G989qt59"
   },
   "outputs": [],
   "source": [
    "def loss_fig(smooth=True):\n",
    "    \"\"\"\n",
    "    Plots a contour map of the loss.\n",
    "    If smooth = True, the landscape will look like nice circles.\n",
    "    While this is a nicer visualization, the actual loss landscape\n",
    "    is hinted at when smooth = False.\n",
    "\n",
    "    Returns a pyplot.Figure\n",
    "    \"\"\"\n",
    "    p1r = np.linspace(2, 10)\n",
    "    p2r = np.linspace(4, -10)\n",
    "    X, Y = np.meshgrid(p1r, p2r)\n",
    "    Z = np.zeros_like(X)\n",
    "    rows, cols = X.shape\n",
    "    for i in range(rows):\n",
    "        for j in range(cols):\n",
    "            preds = f(xs, X[i,j], Y[i,j], add_noise=not smooth)\n",
    "            Z[i,j] = float(nn.MSELoss()(preds, ys).numpy())\n",
    "    fig, ax = plt.subplots(figsize=(12, 12))\n",
    "    CS = ax.contour(X, Y, Z, levels=[.1, 1, 2, 4, 8, 16, 32, 64], colors='k')\n",
    "    ax.clabel(CS, inline=True, fontsize=10)\n",
    "    ax.set_xlabel('Param 1')\n",
    "    ax.set_ylabel('Param 2')\n",
    "    plt.close()\n",
    "    return fig\n",
    "\n",
    "def plot_losses(losses, n_steps=None):\n",
    "    \"\"\"\n",
    "    Plots a list of losses, assuming it's MSE loss.\n",
    "    Returns a plt.Figure.\n",
    "    \"\"\"\n",
    "    if n_steps:\n",
    "        ls = pd.Series(losses[n_steps])\n",
    "    else:\n",
    "        ls = pd.Series(losses)\n",
    "    fig, ax = plt.subplots(figsize=(12,8))\n",
    "    ax.plot(ls, label='MSE')\n",
    "    ax.plot(ls.rolling(100).mean(), label='MSE rolling mean')\n",
    "    ax.set_xlabel('Step')\n",
    "    ax.set_ylabel('Mean squared error')\n",
    "    ax.legend()\n",
    "    plt.close()\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nUdqDJilxAtK"
   },
   "outputs": [],
   "source": [
    "# Write our own sgd optimizer\n",
    "class SGD(optim.Optimizer):\n",
    "    def __init__(self, params, lr):\n",
    "        # Store the params and defaults in the Optimizer parent class\n",
    "        super().__init__(params, defaults=dict(lr=lr))\n",
    "        self.lr = lr\n",
    "\n",
    "    def step(self):\n",
    "        # For each param group\n",
    "        for group in self.param_groups:\n",
    "            # Fetch the parameters\n",
    "            for p in group['params']:\n",
    "                # Apply SGD update rule\n",
    "                # Note: sub_ is an in-place subtraction\n",
    "                p.data.sub_(p.grad * self.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9295,
     "status": "ok",
     "timestamp": 1624738634405,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "S8AAg8YSP2iK",
    "outputId": "801a920b-fc4f-4f5d-9ad5-c0fab70b2d41"
   },
   "outputs": [],
   "source": [
    "sgd_params, sgd_losses = get_parameters_over_training(SGD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 709,
     "status": "ok",
     "timestamp": 1624738635111,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "y94-qRhfrW9G",
    "outputId": "7b5af8e7-1f9d-45f7-ec1f-e41d7accbc28"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(sgd_params[:,0], sgd_params[:,1], label='Our SGD')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9457,
     "status": "ok",
     "timestamp": 1624738644566,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "5Emqf4NNxAtL",
    "outputId": "27868c5a-c40b-4c2c-9d5c-7261ca050d32"
   },
   "outputs": [],
   "source": [
    "sgd_params, sgd_losses = get_parameters_over_training(optim.SGD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 692,
     "status": "ok",
     "timestamp": 1624738645254,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "_DW2XYB7xAtL",
    "outputId": "6620be8c-6896-4f68-cf2a-ba4f01b733a7"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(sgd_params[:,0], sgd_params[:,1], label='SGD')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rO5qT-QkxAtM"
   },
   "outputs": [],
   "source": [
    "class SGDWithMomentum(optim.Optimizer):\n",
    "    def __init__(self, params, lr, momentum=0.9):\n",
    "        defaults=dict(\n",
    "            lr=lr,\n",
    "            momentum=momentum\n",
    "        )\n",
    "        super().__init__(params, defaults=defaults)\n",
    "        self.lr = lr\n",
    "        self.momentum = momentum\n",
    "        \n",
    "    \n",
    "    def step(self):\n",
    "        # Iterate through all the parameters in the model\n",
    "        for group in self.param_groups:\n",
    "            for p in group['params']:\n",
    "                # if the parameter doesn't ahve an \"avg\" attribute,\n",
    "                # let's create one as a tensor of zeros.\n",
    "                if not hasattr(p, 'avg'):\n",
    "                    p.avg = torch.zeros_like(p.grad.data)\n",
    "                # Let's update the moving average with the new update\n",
    "                p.avg = self.momentum * p.avg - self.lr * p.grad\n",
    "                # Finally, let's perform the parameter update step.\n",
    "                p.data.add_(p.avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9873,
     "status": "ok",
     "timestamp": 1624738655125,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "Cs1BZbkcxAtN",
    "outputId": "5dd57b51-e18a-4da4-aa2c-1e6c106b26c7"
   },
   "outputs": [],
   "source": [
    "sgd_mom_params, sgd_mom_losses = get_parameters_over_training(SGDWithMomentum, opt_kwargs={'lr':1e-3, 'momentum':0.9})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 642,
     "status": "ok",
     "timestamp": 1624738655754,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "1CvFhhvUxAtN",
    "outputId": "c1a733e2-8dce-4a15-a2f4-081695d3a3ec"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(sgd_mom_params[:,0], sgd_mom_params[:,1], label='Our SGD with momentum')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 30798,
     "status": "ok",
     "timestamp": 1624738686549,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "-9qtVjBAxAtO",
    "outputId": "a9b8d6f6-1e85-4d93-b9bf-a20b5055af05"
   },
   "outputs": [],
   "source": [
    "sgd_small_mom_params, sgd_small_mom_losses = get_parameters_over_training(optim.SGD, opt_kwargs={'lr':1e-3, 'momentum':0.1})\n",
    "sgd_med_mom_params, sgd_med_mom_losses = get_parameters_over_training(optim.SGD, opt_kwargs={'lr':1e-3, 'momentum':0.5})\n",
    "sgd_large_mom_params, sgd_large_mom_losses = get_parameters_over_training(optim.SGD, opt_kwargs={'lr':1e-3, 'momentum':0.9})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 855,
     "status": "ok",
     "timestamp": 1624738687391,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "tbssPokmxAtO",
    "outputId": "fd1e3225-6a86-47e4-c8f0-40f4778f085e"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(sgd_small_mom_params[:,0], sgd_small_mom_params[:,1], label='SGD (momentum=0.1)')\n",
    "ax.plot(sgd_med_mom_params[:,0], sgd_med_mom_params[:,1], label='SGD (momentum=0.5)')\n",
    "ax.plot(sgd_large_mom_params[:,0], sgd_large_mom_params[:,1], label='SGD (momentum=0.9)')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q0jlJ2UjxAtP"
   },
   "outputs": [],
   "source": [
    "class RMSProp(optim.Optimizer):\n",
    "    def __init__(self, params, lr, alpha=0.9, eps=1e-8):\n",
    "        defaults=dict(\n",
    "            lr=lr,\n",
    "            alpha=alpha\n",
    "        )\n",
    "        super().__init__(params, defaults=defaults)\n",
    "        self.lr = lr\n",
    "        self.alpha = alpha\n",
    "        self.eps = eps\n",
    "        \n",
    "    \n",
    "    def step(self):\n",
    "        # Iterate through the parameters\n",
    "        for group in self.param_groups:\n",
    "            for p in group['params']:\n",
    "                # If the parameter doesn't have a squared gradient attribute,\n",
    "                # let's add one\n",
    "                if not hasattr(p, 'sq_grad'):\n",
    "                    p.sq_grad = torch.zeros_like(p.grad.data)\n",
    "                # Update the squared gradient attribute\n",
    "                p.sq_grad = self.alpha * p.sq_grad + (1 - self.alpha) * p.grad.data**2\n",
    "                # Using the squared gradient attribute, calculate the update\n",
    "                update = p.grad * self.lr / p.sq_grad.add(self.eps).sqrt()\n",
    "                # Apply the update to the parameters\n",
    "                p.data.sub_(update)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 11338,
     "status": "ok",
     "timestamp": 1624738698726,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "xo2Fj7uzxAtP",
    "outputId": "bae13540-a596-4acf-c4d2-6db225f253b3"
   },
   "outputs": [],
   "source": [
    "rms_params, rms_losses = get_parameters_over_training(RMSProp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 641,
     "status": "ok",
     "timestamp": 1624738699352,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "Z1DBXd-AxAtP",
    "outputId": "4244bbb3-f58e-4551-b12d-bc338fdb88bd"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(rms_params[:,0], rms_params[:,1], label='Our RMSProp')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 11333,
     "status": "ok",
     "timestamp": 1624738710682,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "zikG1-wNxAtQ",
    "outputId": "3b071ec6-f22e-46f8-a945-85cfdd9a5a96"
   },
   "outputs": [],
   "source": [
    "rms_params, rms_losses = get_parameters_over_training(optim.RMSprop, opt_kwargs={'lr':1e-3})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 656,
     "status": "ok",
     "timestamp": 1624738711334,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "qVxWe2RXxAtQ",
    "outputId": "ad8231dd-9be2-463a-da7b-3b123310aee7"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "# ax.plot(sgd_params[:,0], sgd_params[:,1], label='SGD')\n",
    "ax.plot(rms_params[:,0], rms_params[:,1], label='RMSProp')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ql6CVRR5xAtQ"
   },
   "outputs": [],
   "source": [
    "class Adam(optim.Optimizer):\n",
    "    def __init__(self, params, lr, grad_mom=0.9, rms_mom=0.99, eps=1e-8):\n",
    "        defaults=dict(\n",
    "            lr=lr,\n",
    "            grad_mom=grad_mom,\n",
    "            rms_mom=rms_mom,\n",
    "            eps=eps\n",
    "        )\n",
    "        super().__init__(params, defaults=defaults)\n",
    "        self.lr = lr\n",
    "        self.grad_mom = grad_mom\n",
    "        self.rms_mom = rms_mom\n",
    "        self.eps = eps\n",
    "        self.step_no = 0\n",
    "        \n",
    "    def step(self):\n",
    "        # Iterate through the parameters\n",
    "        for group in self.param_groups:\n",
    "            for p in group['params']:\n",
    "                # If the parameters dont' have records of the average or squared\n",
    "                # average, let's add them\n",
    "                if not hasattr(p, 'avg'):\n",
    "                    p.avg = torch.zeros_like(p.grad.data)\n",
    "                    p.sq_avg = torch.zeros_like(p.grad.data)\n",
    "                # First, let's calculate the moving average update\n",
    "                p.avg = self.grad_mom * p.avg + (1 - self.grad_mom) * p.grad.data\n",
    "                ma_update = p.avg / (1 - self.grad_mom**self.step_no)\n",
    "                \n",
    "                # Next, let's calculate the RMS update\n",
    "                p.sq_avg = self.rms_mom * p.sq_avg + (1 - self.rms_mom) * p.grad.data**2\n",
    "                rms_update = p.sq_avg / (1 - self.rms_mom**self.step_no)\n",
    "                \n",
    "                # Finally, we'll combine the grad and sq grad updates\n",
    "                # into the final update, then apply the update to our params.\n",
    "                update = self.lr * p.avg / p.sq_avg.add(self.eps).sqrt()\n",
    "                p.data.sub_(update)\n",
    "                self.step_no += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 13175,
     "status": "ok",
     "timestamp": 1624738724506,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "4ut6ycjpvS_7",
    "outputId": "48a31119-f519-4cae-8b02-c4627cc53535"
   },
   "outputs": [],
   "source": [
    "adam_params, adam_losses = get_parameters_over_training(Adam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 718,
     "status": "ok",
     "timestamp": 1624738725213,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "0KBXSdod0gcO",
    "outputId": "dd62571d-897b-4158-e64c-7b9c2c6eeb17"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True) # Try this with smooth=False so we can see that the actual loss landscape is bumpy\n",
    "ax = fig.gca()\n",
    "ax.plot(adam_params[:,0], adam_params[:,1], label='Our Adam')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 12337,
     "status": "ok",
     "timestamp": 1624738737547,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "6FMBtgLLxAtR",
    "outputId": "172e335b-682a-44fd-92fc-6066c8ddfb62"
   },
   "outputs": [],
   "source": [
    "adam_params, adam_losses = get_parameters_over_training(optim.Adam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 718
    },
    "executionInfo": {
     "elapsed": 580,
     "status": "ok",
     "timestamp": 1624738738122,
     "user": {
      "displayName": "Mike Frantz",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GiGYCwZXGfGg-CE4_33PQcH9SW8fN6MyRkenN8-IA=s64",
      "userId": "08482506601269125384"
     },
     "user_tz": 420
    },
    "id": "uV3ilkD_xAtR",
    "outputId": "189f5e5c-b7ca-44d5-cbdc-36f539f084fa"
   },
   "outputs": [],
   "source": [
    "fig = loss_fig(True)\n",
    "ax = fig.gca()\n",
    "ax.plot(adam_params[:,0], adam_params[:,1], label='Adam')\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Kcf3fwmDU95T"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "06_Optimizers.ipynb",
   "provenance": []
  },
  "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}