{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "FhGuhbZ6M5tl" }, "source": [ "##### Copyright 2022 The TensorFlow Authors." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "AwOEIRJC6Une" }, "outputs": [], "source": [ "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "EIdT9iu_Z4Rb" }, "source": [ "# Optimizers with Core APIs" ] }, { "cell_type": "markdown", "metadata": { "id": "bBIlTPscrIT9" }, "source": [ "\n", " \n", " \n", " \n", " \n", "
\n", " View on TensorFlow.org\n", " \n", " Run in Google Colab\n", " \n", " View source on GitHub\n", " \n", " Download notebook\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "SjAxxRpBzVYg" }, "source": [ "## Introduction\n", "\n", "This notebook introduces the process of creating custom optimizers with the [TensorFlow Core low-level APIs](https://www.tensorflow.org/guide/core). Visit the [Core APIs overview](https://www.tensorflow.org/guide/core) to learn more about TensorFlow Core and its intended use cases. \n", "\n", "The [Keras optimizers](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers) module is the recommended optimization toolkit for many general training purposes. It includes a variety of prebuilt optimiziers as well as subclassing functionality for customization. The Keras optimizers are also compatible with custom layers, models, and training loops built with the Core APIs. These prebuilt and customizable optimizers are suitable for most cases, but the Core APIs allow for complete control over the optimization process. For example, techniques such as Sharpness-Aware Minimization (SAM) require the model and optimizer to be coupled, which does not fit the traditional definition of ML optimizers. This guide walks through the process of building custom optimizers from scratch with the Core APIs, giving you the power to have full control over the structure, implementation, and behavior of your optimizers." ] }, { "cell_type": "markdown", "metadata": { "id": "nBmqYyodNRd_" }, "source": [ "## Optimizers overview\n", "\n", "An optimizer is an algorithm used to minimize a loss function with respect to a model's trainable parameters. The most straightforward optimization technique is gradient descent, which iteratively updates a model's parameters by taking a step in the direction of its loss function's steepest descent. Its step size is directly proportional to the size of the gradient, which can be problematic when the gradient is either too large or too small. There are many other gradient-based optimizers such as Adam, Adagrad, and RMSprop that leverage various mathematical properties of gradients for memory efficiency and fast convergence." ] }, { "cell_type": "markdown", "metadata": { "id": "nchsZfwEVtVs" }, "source": [ "## Setup" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "d9idwpXCltUl" }, "outputs": [], "source": [ "import matplotlib\n", "from matplotlib import pyplot as plt\n", "# Preset Matplotlib figure sizes.\n", "matplotlib.rcParams['figure.figsize'] = [9, 6]" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "9xQKvCJ85kCQ" }, "outputs": [], "source": [ "import tensorflow as tf\n", "print(tf.__version__)\n", "# set random seed for reproducible results \n", "tf.random.set_seed(22)" ] }, { "cell_type": "markdown", "metadata": { "id": "0UmF5aU3MnwX" }, "source": [ "## Gradient descent\n", "\n", "The basic optimizer class should have an initialization method and a function to update a list of variables given a list of gradients. Start by implementing the basic gradient descent optimizer which updates each variable by subtracting its gradient scaled by a learning rate." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "MWjmUmeOQFFN" }, "outputs": [], "source": [ "class GradientDescent(tf.Module):\n", "\n", " def __init__(self, learning_rate=1e-3):\n", " # Initialize parameters\n", " self.learning_rate = learning_rate\n", " self.title = f\"Gradient descent optimizer: learning rate={self.learning_rate}\"\n", "\n", " def apply_gradients(self, grads, vars):\n", " # Update variables\n", " for grad, var in zip(grads, vars):\n", " var.assign_sub(self.learning_rate*grad)" ] }, { "cell_type": "markdown", "metadata": { "id": "ZSekgBHDRzmp" }, "source": [ "To test this optimizer, create a sample loss function to minimize with respect to a single variable, $x$. Compute its gradient function and solve for its minimizing parameter value:\n", "\n", "$$L = 2x^4 + 3x^3 + 2$$\n", "\n", "$$\\frac{dL}{dx} = 8x^3 + 9x^2$$\n", "\n", "$\\frac{dL}{dx}$ is 0 at $x = 0$, which is a saddle point and at $x = - \\frac{9}{8}$, which is the global minimum. Therefore, the loss function is optimized at $x^\\star = - \\frac{9}{8}$." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VCtJaUo6Ry8V" }, "outputs": [], "source": [ "x_vals = tf.linspace(-2, 2, 201)\n", "x_vals = tf.cast(x_vals, tf.float32)\n", "\n", "def loss(x):\n", " return 2*(x**4) + 3*(x**3) + 2\n", "\n", "def grad(f, x):\n", " with tf.GradientTape() as tape:\n", " tape.watch(x)\n", " result = f(x)\n", " return tape.gradient(result, x)\n", "\n", "plt.plot(x_vals, loss(x_vals), c='k', label = \"Loss function\")\n", "plt.plot(x_vals, grad(loss, x_vals), c='tab:blue', label = \"Gradient function\")\n", "plt.plot(0, loss(0), marker=\"o\", c='g', label = \"Inflection point\")\n", "plt.plot(-9/8, loss(-9/8), marker=\"o\", c='r', label = \"Global minimum\")\n", "plt.legend()\n", "plt.ylim(0,5)\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"loss\")\n", "plt.title(\"Sample loss function and gradient\");" ] }, { "cell_type": "markdown", "metadata": { "id": "fLlIBJ9yuwhE" }, "source": [ "Write a function to test the convergence of an optimizer with a single variable loss function. Assume that convergence has been achieved when the updated parameter's value at timestep $t$ is the same as its value held at timestep $t-1$. Terminate the test after a set number of iterations and also keep track of any exploding gradients during the process. In order to truly challenge the optimization algorithm, initialize the parameter poorly. In the above example, $x = 2$ is a good choice since it involves an steep gradient and also leads into an inflection point." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "SLQTc41ouv0F" }, "outputs": [], "source": [ "def convergence_test(optimizer, loss_fn, grad_fn=grad, init_val=2., max_iters=2000):\n", " # Function for optimizer convergence test\n", " print(optimizer.title)\n", " print(\"-------------------------------\")\n", " # Initializing variables and structures\n", " x_star = tf.Variable(init_val)\n", " param_path = []\n", " converged = False\n", "\n", " for iter in range(1, max_iters + 1):\n", " x_grad = grad_fn(loss_fn, x_star)\n", "\n", " # Case for exploding gradient\n", " if tf.math.is_nan(x_grad):\n", " print(f\"Gradient exploded at iteration {iter}\\n\")\n", " return []\n", "\n", " # Updating the variable and storing its old-version\n", " x_old = x_star.numpy()\n", " optimizer.apply_gradients([x_grad], [x_star])\n", " param_path.append(x_star.numpy())\n", "\n", " # Checking for convergence\n", " if x_star == x_old:\n", " print(f\"Converged in {iter} iterations\\n\")\n", " converged = True\n", " break\n", " \n", " # Print early termination message\n", " if not converged:\n", " print(f\"Exceeded maximum of {max_iters} iterations. Test terminated.\\n\")\n", " return param_path" ] }, { "cell_type": "markdown", "metadata": { "id": "vK-7_TsmyAgI" }, "source": [ "Test the convergence of the gradient descent optimizer for the following learning rates: 1e-3, 1e-2, 1e-1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "lWRn8c91mqB0" }, "outputs": [], "source": [ "param_map_gd = {}\n", "learning_rates = [1e-3, 1e-2, 1e-1]\n", "for learning_rate in learning_rates:\n", " param_map_gd[learning_rate] = (convergence_test(\n", " GradientDescent(learning_rate=learning_rate), loss_fn=loss))" ] }, { "cell_type": "markdown", "metadata": { "id": "TydrGHF5y6iI" }, "source": [ "Visualize the path of the parameters over a contour plot of the loss function." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "piffzGHI_u5G" }, "outputs": [], "source": [ "def viz_paths(param_map, x_vals, loss_fn, title, max_iters=2000):\n", " # Creating a controur plot of the loss function\n", " t_vals = tf.range(1., max_iters + 100.)\n", " t_grid, x_grid = tf.meshgrid(t_vals, x_vals)\n", " loss_grid = tf.math.log(loss_fn(x_grid))\n", " plt.pcolormesh(t_vals, x_vals, loss_grid, vmin=0, shading='nearest')\n", " colors = ['r', 'w', 'c']\n", " # Plotting the parameter paths over the contour plot\n", " for i, learning_rate in enumerate(param_map):\n", " param_path = param_map[learning_rate]\n", " if len(param_path) > 0:\n", " x_star = param_path[-1]\n", " plt.plot(t_vals[:len(param_path)], param_path, c=colors[i])\n", " plt.plot(len(param_path), x_star, marker='o', c=colors[i], \n", " label = f\"x*: learning rate={learning_rate}\")\n", " plt.xlabel(\"Iterations\")\n", " plt.ylabel(\"Parameter value\")\n", " plt.legend()\n", " plt.title(f\"{title} parameter paths\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Ssyj2sO4BcNY" }, "outputs": [], "source": [ "viz_paths(param_map_gd, x_vals, loss, \"Gradient descent\")" ] }, { "cell_type": "markdown", "metadata": { "id": "MmM-5eDLFnmC" }, "source": [ "Gradient descent seems to get stuck at the inflection point when using smaller learning rates. Increasing the learning rate can encourage faster movement around the plateau region due to a larger step size; however, this comes at the risk of having exploding gradients in early iterations when the loss function is extremely steep." ] }, { "cell_type": "markdown", "metadata": { "id": "m5CDeXN8S1SF" }, "source": [ "## Gradient descent with momentum\n", "\n", "Gradient descent with momentum not only uses the gradient to update a variable but also involves the change in position of a variable based on its previous update. The momentum parameter determines the level of influence the update at timestep $t-1$ has on the update at timestep $t$. Accumulating momentum helps to move variables past plataeu regions faster than basic gradient descent. The momentum update rule is as follows:\n", "\n", "$$\\Delta_x^{[t]} = lr \\cdot L^\\prime(x^{[t-1]}) + p \\cdot \\Delta_x^{[t-1]}$$\n", "\n", "$$x^{[t]} = x^{[t-1]} - \\Delta_x^{[t]}$$\n", "\n", "where\n", "\n", "* $x$: the variable being optimized\n", "* $\\Delta_x$: change in $x$ \n", "* $lr$: learning rate\n", "* $L^\\prime(x)$: gradient of the loss function with respect to x\n", "* $p$: momentum parameter" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "rOBY8Tz4S0dX" }, "outputs": [], "source": [ "class Momentum(tf.Module):\n", "\n", " def __init__(self, learning_rate=1e-3, momentum=0.7):\n", " # Initialize parameters\n", " self.learning_rate = learning_rate\n", " self.momentum = momentum\n", " self.change = 0.\n", " self.title = f\"Gradient descent optimizer: learning rate={self.learning_rate}\"\n", "\n", " def apply_gradients(self, grads, vars):\n", " # Update variables \n", " for grad, var in zip(grads, vars):\n", " curr_change = self.learning_rate*grad + self.momentum*self.change\n", " var.assign_sub(curr_change)\n", " self.change = curr_change" ] }, { "cell_type": "markdown", "metadata": { "id": "t_nDu38gW6Fu" }, "source": [ "Test the convergence of the momentum optimizer for the following learning rates: 1e-3, 1e-2, 1e-1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tA6oQL-sW2xg" }, "outputs": [], "source": [ "param_map_mtm = {}\n", "learning_rates = [1e-3, 1e-2, 1e-1]\n", "for learning_rate in learning_rates:\n", " param_map_mtm[learning_rate] = (convergence_test(\n", " Momentum(learning_rate=learning_rate),\n", " loss_fn=loss, grad_fn=grad))" ] }, { "cell_type": "markdown", "metadata": { "id": "wz_LV0EPYE6k" }, "source": [ "Visualize the path of the parameters over a contour plot of the loss function." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "qbW1eEKaX3T9" }, "outputs": [], "source": [ "viz_paths(param_map_mtm, x_vals, loss, \"Momentum\")" ] }, { "cell_type": "markdown", "metadata": { "id": "4bEFnhPRTBXh" }, "source": [ "## Adaptive moment estimation (Adam)\n", "\n", "The Adaptive Moment Estimation (Adam) algorithm is an efficient and highly generalizable optimization technique that leverages two key gradient descent methedologies: momentum, and root mean square propogation (RMSP). Momentum helps accelerate gradient descent by using the first moment (sum of gradients) along with a decay parameter. RMSP is similar; however, it leverages the second moment (sum of gradients squared). \n", "\n", "The Adam algorithm combines both the first and second moment to provide a more generalizable update rule. The sign of a variable, $x$, can be determined by computing $\\frac{x}{\\sqrt{x^2}}$. The Adam optimizer uses this fact to calculate an update step which is effectively a smoothed sign. Instead of calculating $\\frac{x}{\\sqrt{x^2}}$, the optimizer calculates a smoothed version of $x$ (first moment) and $x^2$ (second moment) for each variable update. \n" ] }, { "cell_type": "markdown", "metadata": { "id": "WjgyqRiZ7XhA" }, "source": [ "**Adam algorithm**\n", "\n", "$\\beta_1 \\gets 0.9 \\; \\triangleright \\text{literature value}$\n", "\n", "$\\beta_2 \\gets 0.999 \\; \\triangleright \\text{literature value}$\n", "\n", "$lr \\gets \\text{1e-3} \\; \\triangleright \\text{configurable learning rate}$\n", "\n", "$\\epsilon \\gets \\text{1e-7} \\; \\triangleright \\text{prevents divide by 0 error}$\n", "\n", "$V_{dv} \\gets \\vec {\\underset{n\\times1}{0}} \\;\\triangleright \\text{stores momentum updates for each variable}$\n", "\n", "$S_{dv} \\gets \\vec {\\underset{n\\times1}{0}} \\; \\triangleright \\text{stores RMSP updates for each variable}$\n", "\n", "$t \\gets 1$\n", "\n", "$\\text{On iteration } t:$\n", "\n", "$\\;\\;\\;\\; \\text{For} (\\frac{dL}{dv}, v) \\text{ in gradient variable pairs}:$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; V_{dv\\_i} = \\beta_1V_{dv\\_i} + (1 - \\beta_1)\\frac{dL}{dv} \\; \\triangleright \\text{momentum update}$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; S_{dv\\_i} = \\beta_2V_{dv\\_i} + (1 - \\beta_2)(\\frac{dL}{dv})^2 \\; \\triangleright \\text{RMSP update}$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; v_{dv}^{bc} = \\frac{V_{dv\\_i}}{(1-\\beta_1)^t} \\; \\triangleright \\text{momentum bias correction}$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; s_{dv}^{bc} = \\frac{S_{dv\\_i}}{(1-\\beta_2)^t} \\; \\triangleright \\text{RMSP bias correction}$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; v = v - lr\\frac{v_{dv}^{bc}}{\\sqrt{s_{dv}^{bc}} + \\epsilon} \\; \\triangleright \\text{parameter update}$\n", "\n", "$\\;\\;\\;\\;\\;\\;\\;\\; t = t + 1$\n", "\n", "**End of algorithm**\n", "\n", "Given that $V_{dv}$ and $S_{dv}$ are initialized to 0 and that $\\beta_1$ and $\\beta_2$ are close to 1, the momentum and RMSP updates are naturally biased towards 0; therefore, the variables can benefit from bias correction. Bias correction also helps to control the osccilation of weights as they approach the global minimum." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "hm5vffRJRsEc" }, "outputs": [], "source": [ "class Adam(tf.Module):\n", " \n", " def __init__(self, learning_rate=1e-3, beta_1=0.9, beta_2=0.999, ep=1e-7):\n", " # Initialize the Adam parameters\n", " self.beta_1 = beta_1\n", " self.beta_2 = beta_2\n", " self.learning_rate = learning_rate\n", " self.ep = ep\n", " self.t = 1.\n", " self.v_dvar, self.s_dvar = [], []\n", " self.title = f\"Adam: learning rate={self.learning_rate}\"\n", " self.built = False\n", "\n", " def apply_gradients(self, grads, vars):\n", " # Set up moment and RMSprop slots for each variable on the first call\n", " if not self.built:\n", " for var in vars:\n", " v = tf.Variable(tf.zeros(shape=var.shape))\n", " s = tf.Variable(tf.zeros(shape=var.shape))\n", " self.v_dvar.append(v)\n", " self.s_dvar.append(s)\n", " self.built = True\n", " # Perform Adam updates\n", " for i, (d_var, var) in enumerate(zip(grads, vars)):\n", " # Moment calculation\n", " self.v_dvar[i] = self.beta_1*self.v_dvar[i] + (1-self.beta_1)*d_var\n", " # RMSprop calculation\n", " self.s_dvar[i] = self.beta_2*self.s_dvar[i] + (1-self.beta_2)*tf.square(d_var)\n", " # Bias correction\n", " v_dvar_bc = self.v_dvar[i]/(1-(self.beta_1**self.t))\n", " s_dvar_bc = self.s_dvar[i]/(1-(self.beta_2**self.t))\n", " # Update model variables\n", " var.assign_sub(self.learning_rate*(v_dvar_bc/(tf.sqrt(s_dvar_bc) + self.ep)))\n", " # Increment the iteration counter\n", " self.t += 1." ] }, { "cell_type": "markdown", "metadata": { "id": "UWN4Qus7flUO" }, "source": [ "Test the performance of the Adam optimizer with the same learning rates used with the gradient descent examples. " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "GXHCxtemFBpR" }, "outputs": [], "source": [ "param_map_adam = {}\n", "learning_rates = [1e-3, 1e-2, 1e-1]\n", "for learning_rate in learning_rates:\n", " param_map_adam[learning_rate] = (convergence_test(\n", " Adam(learning_rate=learning_rate), loss_fn=loss))" ] }, { "cell_type": "markdown", "metadata": { "id": "jgpUcs_xXEjX" }, "source": [ "Visualize the path of the parameters over a contour plot of the loss function." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ctvOUmlzFK8s" }, "outputs": [], "source": [ "viz_paths(param_map_adam, x_vals, loss, \"Adam\")" ] }, { "cell_type": "markdown", "metadata": { "id": "_oGScF8zJcY4" }, "source": [ "In this particular example, the Adam optimizer has slower convergence compared to traditional gradient descent when using small learning rates. However, the algorithm successfully moves past the plataeu region and converges to the global minimum when a larger learning rate. Exploding gradients are no longer an issue due to Adam's dynamic scaling of learning rates when encountering large gradients." ] }, { "cell_type": "markdown", "metadata": { "id": "VFLfEH4ManbW" }, "source": [ "## Conclusion\n", "\n", "This notebook introduced the basics of writing and comparing optimizers with the [TensorFlow Core APIs](https://www.tensorflow.org/guide/core). Although prebuilt optimizers like Adam are generalizable, they may not always be the best choice for every model or dataset. Having fine-grained control over the optimization process can help streamline ML training workflows and improve overall performance. Refer to the following documentation for more examples of custom optimizers:\n", "\n", "* This Adam optimizer is used in the [Multilayer perceptrons](https://www.tensorflow.org/guide/core/mlp_core) tutorial and the [Distributed training]()\n", "* [Model Garden](https://blog.tensorflow.org/2020/03/introducing-model-garden-for-tensorflow-2.html) has a variety of [custom optimizers](https://github.com/tensorflow/models/tree/master/official/modeling/optimization) written with the Core APIs.\n" ] } ], "metadata": { "colab": { "collapsed_sections": [], "name": "optimizers_core.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }