labs/08_frameworks/A_PyTorch_introduction__autograd_in_action.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to PyTorch\n",
    "\n",
    "PyTorch is a popular deep learning framework that was introduced in 2016. It has a slightly lower-level API than Keras, but it is very easy to use when it comes to defining models with custom loss functions or exotic architectures with a structure that depends on the input data (e.g. for NLP). It is very popular among researchers and kagglers - a little less in the industry (yet) compared to TensorFlow / Keras.\n",
    "\n",
    "http://pytorch.org/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Principles\n",
    "\n",
    "PyTorch essentially provides the user with two main utilities:\n",
    "\n",
    "- It allows to perform Tensor/vector computation on the GPU with an API similar (but still not 100% compatible) to Numpy.\n",
    "\n",
    "- It records all computation to be able to backpropagate through them. That is, provided a sequence of operations that starts from a tensor $\\theta$ to define a scalar $g(\\theta)$, it is able to compute \n",
    "$\\nabla_\\theta g(\\theta)$ exactly, with only one function call.\n",
    "\n",
    "Typically, $\\theta$ will be a parameter of a neural network and $g$ a loss function such as $\\ell(f_{\\theta}(x), y)$ for supervised learning. \n",
    "\n",
    "In PyTorch, every node in the computation graph is defined on the fly while executing the forward pass, from within the Python interpreter: any NumPy code can be ported to PyTorch easily, and all flow control operation (e.g. loops, if/else, etc.) can be kept untouched.\n",
    "\n",
    "PyTorch takes care of recording everything it needs to do the backpropagation, *on the fly*.\n",
    "\n",
    "Note that recent versions of TensorFlow (2.0 and later) now come with [eager execution](https://www.tensorflow.org/guide/eager) enabled by default to make TensorFlow use the define-by-run semantics of PyTorch (and Chainer)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "You can check the version number of the currently installed pytorch package with: `import torch; print(torch.__version__)`.\n",
    "\n",
    "If you do not have a CUDA-based GPU on your machine please feel free to use the CPU-only version as those introductory notebooks do not require a GPU.\n",
    "\n",
    "Up-to-date installation instructions (including pip-based installation instructions) can be found on the official home page of the project: http://pytorch.org.\n",
    "\n",
    "If pytorch version is too old and was previously installed from another conda channel, conda upgrade might not work. You need to `conda remove pytorch` first and then reinstall it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the device\n",
    "PyTorch can run on either CPU or GPU. The new API (v0.4.0) lets us define it in a nice way "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Two options: 'cuda' (for GPU) or 'cpu' (if no GPU available)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "print(\"Device type set:\", \"GPU\" if device.type == \"cuda\" else \"CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A first example\n",
    "\n",
    "We will define a vector $x \\in \\mathbb{R}^n$ and computes its norms using PyTorch. For this, we will define our first `Tensor` -- the Tensor object is the central element of PyTorch. It is very similar to a numpy array.\n",
    "\n",
    "Let us fill this vector with random values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "\n",
    "n = 3\n",
    "x = torch.rand(n, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`x` is a tensor, which can be manipulated roughly as a Numpy array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's do a small API tour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x * 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can go back and forth from numpy to PyTorch tensors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_np = x.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "A = torch.from_numpy(np.ones((n, n)))\n",
    "\n",
    "A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that all tensors are typed and that you can only do operations with tensors of the same type. Mixing tensor types will lead to a runtime error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    torch.matmul(A, x)\n",
    "except RuntimeError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explicit type conversion are required. Note while numpy uses the `dtype` attribute to store the type of the data items, this information is part of the type of the Tensor itself in PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.matmul(A.float(), x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Python 3.5 and later it's possible to use the `@` operator syntax for matrix-matrix or matrix-vector multiplication operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.float() @ x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now define a norm function that returns the euclidean norm of a vector $x \\in \\mathbb{R}^n$:\n",
    "\n",
    "$$f(x) = || x ||_2$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return torch.sqrt(torch.sum(x ** 2, axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now interested in computing $\\nabla_x f(x) = \\frac{x}{|| x ||_2}$. \n",
    "\n",
    "\n",
    "**Exercise**:\n",
    "\n",
    "- Show that if $f(x) = || x ||_2$ then $\\nabla_x f(x) = \\frac{x}{|| x ||_2}$.\n",
    "\n",
    "\n",
    "Assume we are too lazy to derive the analytical form of the gradient manually. Instead we will use the `autograd` facilities of PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Did we require autograd for our tensor x?\n",
    "x.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.requires_grad_()\n",
    "x.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now compute the norm `f(x)` for our specific vector `x`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm = f(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This result is a PyTorch scalar tensor (that also stores the necessary provenance information to be able to compute a gradient). To convert into a regular Python scalar value we can use the `.item()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the gradient of this scalar variable with respect to all the tensor variables that were used to compute its value.\n",
    "\n",
    "The following `.backward()` call will assign `.grad` attributes to all `Tensor` variables used in the computation of $f$ and that are marked with the `requires_grad` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The gradient $\\nabla_{x} f(x)$ can be found in `x.grad`, which is also a `Tensor`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us compare it to the result of the evaluation of the analytical expression of the derivative of f(x) given x:  $\\nabla_x f(x) = \\frac{x}{|| x ||_2}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_grad = x / f(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works ! You now know everything you need to know to use PyTorch. Note that, similar to Keras, PyTorch comes with a number of predefined functions that are useful in network definition. Check out http://pytorch.org/docs/ and tutorials for an overview of the tools you can use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercises**\n",
    "\n",
    "- Write a function the pytorch code for function `g` that computes the cosine similarity of two variable vectors with float entries $\\mathbf{x}$ and $\\mathbf{y}$:\n",
    "\n",
    "$$g(\\mathbf{x}, \\mathbf{y}) = \\frac{\\mathbf{x}^T \\mathbf{y}}{|| \\mathbf{x} ||_2 || \\mathbf{y} ||_2 }$$\n",
    "\n",
    "- Use `torch.autograd` to compute the derivatives with respect to $\\mathbf{x} \\in \\mathbb{R}^3$ and $\\mathbf{y} \\in \\mathbb{R}^3$ for some random values of your choice;\n",
    "\n",
    "- What is the expected value of the cosine similarity of two colinear vectors?\n",
    "\n",
    "- What is the expected value of the gradients of the cosine function with respect to each of those input tensors?\n",
    "\n",
    "- Compute $\\nabla_x g(x, y)$ and $\\nabla_y g(x, y)$ for some choice of $\\mathbf{x} = \\alpha \\cdot \\mathbf{y}$ with any $\\alpha \\in \\mathbb{R}$. Check that you can get the expected result with Pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def g(x, y):\n",
    "    # TODO: fix the following code to implement the cosine\n",
    "    # similarity function instead.\n",
    "    return torch.sum(x) + torch.sum(y)\n",
    "\n",
    "\n",
    "x = torch.tensor([0, 1, 2], dtype=torch.float32, requires_grad=True)\n",
    "y = torch.tensor([3, 0.9, 2.2], dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "cosine = g(x, y)\n",
    "cosine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/cosine_autograd.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/cosine_autograd_colinear.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's reinitialize our two variables to non colinear values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([0, 1, 2], dtype=torch.float32, requires_grad=True)\n",
    "y = torch.tensor([3, 0.9, 2.2], dtype=torch.float32, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute the following cells many times (use `ctrl-enter`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosine = g(x, y)\n",
    "print('cosine:', cosine)\n",
    "cosine.backward()\n",
    "with torch.no_grad():\n",
    "    # Do not record the following operations for future\n",
    "    # calls to cosine.backward().\n",
    "    x.add_(0.5 * x.grad)\n",
    "    y.add_(0.5 * y.grad)\n",
    "print(\"x:\", x)\n",
    "print(\"y:\", y)\n",
    "print(\"x/y:\", x / y)\n",
    "print(\"x.grad:\", x.grad)\n",
    "print(\"y.grad:\", y.grad)\n",
    "x.grad.zero_()\n",
    "y.grad.zero_();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do you observe?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing gradient descent methods\n",
    "\n",
    "In this notebook, we will use PyTorch to compare the different gradient methods and a toy 2D examples: we will try to find the minimum of the difference of two Gaussians. PyTorch provides a convenient wrapper named `nn.Module` to define parametrized functions, that we will use along this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "\n",
    "class Norm(nn.Module):\n",
    "\n",
    "    def __init__(self, p=2.):\n",
    "        super(Norm, self).__init__()\n",
    "        self.p = Parameter(torch.tensor([p], dtype=torch.float32))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        p_sum = torch.sum(torch.pow(x, self.p), dim=0)\n",
    "        return torch.pow(p_sum, 1 / self.p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above definition, we treat the attribute `p` of the \"p-norm\" (a generalization of the euclidean norm) as a parameter that can be differentiated upon:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "\n",
    "x = torch.rand(n)\n",
    "norm = Norm(p=3.)\n",
    "v = norm(x)\n",
    "v.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can access $\\nabla_p(x \\to || x ||_p)$ in `norm.p.grad`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm.p.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define a Gaussian operator, along with a generic Gaussian combination. We will not consider the gradient w.r.t the parameters of these modules, hence we specify `requires_grad=False`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Gaussian(nn.Module):\n",
    "\n",
    "    def __init__(self, precision, mean):\n",
    "        super(Gaussian, self).__init__()\n",
    "\n",
    "        assert precision.shape == (2, 2)\n",
    "        assert mean.shape == (2,)\n",
    "\n",
    "        self.precision = Parameter(precision, requires_grad=False)\n",
    "        self.mean = Parameter(mean, requires_grad=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Compute the (unormalized) likelihood of x given a Gaussian.\n",
    "        \n",
    "        x can be a 2D-vector or a batch of 2D-vectors.\n",
    "        \"\"\"\n",
    "        xc = x - self.mean\n",
    "        if len(xc.shape) == 1:\n",
    "            # Reshape xc to work as a mini-batch of one vector.\n",
    "            xc = xc.view(1, -1)\n",
    "        return torch.exp(-.5 * (torch.sum((xc @ self.precision) * xc, dim=1)))\n",
    "\n",
    "\n",
    "class GaussianCombination(nn.Module):\n",
    "\n",
    "    def __init__(self, precisions, means, weights):\n",
    "        super(GaussianCombination, self).__init__()\n",
    "        assert len(precisions) == len(means) == len(weights)\n",
    "        self.gaussians = nn.ModuleList()\n",
    "        for precision, mean in zip(precisions, means):\n",
    "            self.gaussians.append(Gaussian(precision, mean))\n",
    "        self.weights = weights\n",
    "\n",
    "    def forward(self, x):\n",
    "        return sum(w * g(x) for g, w in zip(self.gaussians, self.weights))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define $f(x) = \\exp(-(x- m_1)^T P_1 (x - m_1)) - \\exp(-(x- m_2)^T P_2 (x - m_2))$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1 = torch.tensor([[1, 0], [0, 4]], dtype=torch.float32)\n",
    "m1 = torch.tensor([0, 1], dtype=torch.float)\n",
    "w1 = 1\n",
    "p2 = torch.tensor([[1, -2], [-2, 10]], dtype=torch.float32)\n",
    "m2 = torch.tensor([0, -2], dtype=torch.float32)\n",
    "w2 = -1\n",
    "\n",
    "f = GaussianCombination([p1, p2], [m1, m2], [w1, w2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a plotting function to visualize $f$. Note the small boilerplate to interface PyTorch with Numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def plot_function(f, ax):\n",
    "    x_max, y_max, x_min, y_min = 3, 3, -3, -3\n",
    "    x = np.linspace(x_min, x_max, 100, dtype=np.float32)\n",
    "    y = np.linspace(y_min, y_max, 100, dtype=np.float32)\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    samples = np.concatenate((X[:, :, None], Y[:, :, None]), axis=2)\n",
    "    samples = samples.reshape(-1, 2)\n",
    "    samples = torch.from_numpy(samples).requires_grad_()\n",
    "    Z = f(samples).data.numpy()\n",
    "    Z = Z.reshape(100, 100)\n",
    "    CS = ax.contour(X, Y, Z)\n",
    "    ax.clabel(CS, inline=1, fontsize=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "plot_function(f, ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now try to minimize $f$ using gradient descent, with optional flavors. For this, we define a minimize function that performs gradient descent, along with a helper class `GradientDescent` that will perform the updates given the gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GradientDescent:\n",
    "    \n",
    "    def __init__(self, params, lr=0.1):\n",
    "        self.params = params\n",
    "        self.lr = lr\n",
    "\n",
    "    def step(self):\n",
    "        with torch.no_grad():\n",
    "            for param in self.params:\n",
    "                param -= self.lr * param.grad\n",
    "\n",
    "    def zero_grad(self):\n",
    "        for param in self.params:\n",
    "            if param.grad is not None:\n",
    "                param.grad.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def minimize(f, optimizer, max_iter=500, verbose=False):\n",
    "    if hasattr(optimizer, 'params'):\n",
    "        [iterate] = optimizer.params\n",
    "    else:\n",
    "        # For pytorch optimizers.\n",
    "        assert len(optimizer.param_groups) == 1\n",
    "        [iterate] = optimizer.param_groups[0]['params']\n",
    "    iterate_record = []\n",
    "\n",
    "    for i in range(max_iter):\n",
    "        # iterate.grad may be non zero. We zero it first:\n",
    "        optimizer.zero_grad()\n",
    "        value = f(iterate)\n",
    "        \n",
    "        # Compute the gradient of f with respect to the parameters:\n",
    "        value.backward()\n",
    "        # iterate.grad now holds $\\nabla_x f(x)$\n",
    "        if float(torch.sum(iterate.grad ** 2)) < 1e-6:\n",
    "            if verbose:\n",
    "                print('Converged at iteration %i: '\n",
    "                      'f(x) = %e, x = [%e, %e]'\n",
    "                      % (i, value, iterate[0], iterate[1]))\n",
    "            break\n",
    "        \n",
    "        # We store the trajectory of the iterates\n",
    "        iterate_record.append(iterate.data.clone()[None, :])\n",
    "        if verbose and i % 10 == 0:\n",
    "            print('Iteration %i: f(x) = %e, x = [%e, %e]'\n",
    "                  % (i, value, iterate[0], iterate[1]))\n",
    "\n",
    "        # Perform the parameter update step using the gradient\n",
    "        # values:\n",
    "        optimizer.step()\n",
    "    return torch.cat(iterate_record, dim=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the minimization algorithm and plot it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The extra dimension marked with `None` is used to make it \n",
    "# possible\n",
    "init = torch.tensor([0.8, 0.8], dtype=torch.float32)\n",
    "optimizer = GradientDescent([init.clone().requires_grad_()], lr=0.1)\n",
    "iterate_rec = minimize(f, optimizer, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_trace(iterate_rec, ax, label='', tags=True):\n",
    "    iterate_rec = iterate_rec.numpy()\n",
    "    n_steps = len(iterate_rec)\n",
    "    line = ax.plot(iterate_rec[:, 0], iterate_rec[:, 1], linestyle=':',\n",
    "                   marker='o', markersize=2,\n",
    "                   label=label + \" (%d steps)\" % n_steps)\n",
    "    color = plt.getp(line[0], 'color')\n",
    "    bbox_props = dict(boxstyle=\"square,pad=0.3\", ec=color, fc='white',\n",
    "                      lw=1)\n",
    "    if tags:\n",
    "        # Put tags every 10 iterates. The small randomization makes it\n",
    "        # easier to spot tag overlap when the steps are very small.\n",
    "        for i in range(0, len(iterate_rec), 10):\n",
    "            ax.annotate(i, xy=(iterate_rec[i, 0], iterate_rec[i, 1]),\n",
    "                        xycoords='data',\n",
    "                        xytext=(5 + np.random.uniform(-2, 2),\n",
    "                                5 + np.random.uniform(-2, 2)),\n",
    "                        textcoords='offset points',\n",
    "                        bbox=bbox_props)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))\n",
    "plot_function(f, ax1)\n",
    "plot_function(f, ax2)\n",
    "plot_trace(iterate_rec, ax1, label='gradient_descent')\n",
    "plot_trace(iterate_rec, ax2, label='gradient_descent', tags=False)\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot([f(x).numpy() for x in iterate_rec])\n",
    "plt.xlabel('optimization step')\n",
    "plt.ylabel('objective function value');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercices\n",
    "\n",
    "- Try to move the initialization point to the other side of the yellow mountain, for instance at position `[0.8, 1.2]`. What do you observe? How to do you explain this?\n",
    "\n",
    "- Put back the init to `[0.8, 0.8]` and implement the step method of `MomemtumGradientDescent` in the following cell,\n",
    "- Check that it behaves as `GradientDescent` for `momemtum=0`\n",
    "- Can you find a value of `momentum` that makes it converge faster than gradient descent on for this objective function? Try with different values of `lr` and `momentum`.\n",
    "- Try to use [torch.optim.Adam](http://pytorch.org/docs/master/optim.html#torch.optim.Adam) in the minimization loop.\n",
    "- Compare the three trajectories.\n",
    "\n",
    "\n",
    "For the momentum let's use the following 2-stage mathematical description of the update step:\n",
    "\n",
    "$$ \\mathbf{v} \\leftarrow \\mu \\cdot \\mathbf{v} + \\nabla_\\mathbf{x}\\ell(\\mathbf{x}) $$\n",
    "$$ \\mathbf{x} \\leftarrow \\eta \\cdot \\mathbf{v} $$\n",
    "\n",
    "where:\n",
    "\n",
    "- $\\eta$ is the learning rate coefficient (`lr` in Python code);\n",
    "- $\\mu$ is the momentum coefficient (set to 0.9 by default);\n",
    "- $\\nabla_\\mathbf{x}\\ell(\\mathbf{x})$ is the gradient of the loss function for the current value of the parameters $\\mathbf{x}$;\n",
    "- $\\mathbf{v}$ is the tensor of velocities and as the same shape as the parameters tensor $\\theta$. $v$ is initialized to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MomentumGradientDescent(GradientDescent):\n",
    "\n",
    "    def __init__(self, params, lr=0.1, momentum=.9):\n",
    "        super(MomentumGradientDescent, self).__init__(params, lr)\n",
    "        self.momentum = momentum\n",
    "        self.velocities = [torch.zeros_like(param, requires_grad=False)\n",
    "                           for param in params]\n",
    "\n",
    "    def step(self):\n",
    "        # TODO: implement me!\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/momentum_optimizer.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.1\n",
    "init = torch.FloatTensor([0.8, 0.8])\n",
    "optimizer = GradientDescent([init.clone().requires_grad_()], lr=lr)\n",
    "iterate_rec_gd = minimize(f, optimizer)\n",
    "\n",
    "optimizer = MomentumGradientDescent([init.clone().requires_grad_()],\n",
    "                                    lr=lr, momentum=0.9)\n",
    "iterate_rec_mom = minimize(f, optimizer)\n",
    "\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(ncols=2, figsize=(16, 8))\n",
    "plot_function(f, ax0)\n",
    "plot_trace(iterate_rec_gd, ax0, label='gradient descent', tags=False)\n",
    "plot_trace(iterate_rec_mom, ax0, label='momentum', tags=False)\n",
    "ax0.legend();\n",
    "\n",
    "ax1.plot([f(x).numpy() for x in iterate_rec_gd], label='gradient descent')\n",
    "ax1.plot([f(x).numpy() for x in iterate_rec_mom], label='momentum')\n",
    "ax1.set_xlabel('optimization step')\n",
    "ax1.set_ylabel('objective function value')\n",
    "ax1.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do you observe ? Try changing the momentum and the initialization to compare optimization traces."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}