docs/tutorials/optimizers_lazyadam.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0\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": "MfBg1C5NB3X0"
      },
      "source": [
        "# TensorFlow Addons Optimizers: LazyAdam\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/addons/tutorials/optimizers_lazyadam\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/addons/blob/master/docs/tutorials/optimizers_lazyadam.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/addons/blob/master/docs/tutorials/optimizers_lazyadam.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/addons/docs/tutorials/optimizers_lazyadam.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHxb-dlhMIzW"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook will demonstrate how to use the lazy adam optimizer from the Addons package.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQwBbFVAyHJ_"
      },
      "source": [
        "## LazyAdam\n",
        "\n",
        "> LazyAdam is a variant of the Adam optimizer that handles sparse updates more efficiently.\n",
        "    The original Adam algorithm maintains two moving-average accumulators for\n",
        "    each trainable variable; the accumulators are updated at every step.\n",
        "    This class provides lazier handling of gradient updates for sparse\n",
        "    variables.  It only updates moving-average accumulators for sparse variable\n",
        "    indices that appear in the current batch, rather than updating the\n",
        "    accumulators for all indices. Compared with the original Adam optimizer,\n",
        "    it can provide large improvements in model training throughput for some\n",
        "    applications. However, it provides slightly different semantics than the\n",
        "    original Adam algorithm, and may lead to different empirical results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cHAOyeOVx-k3"
      },
      "outputs": [],
      "source": [
        "!pip install -U tensorflow-addons"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "42ztALK4ZdyZ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ys65MwOLKnXq"
      },
      "outputs": [],
      "source": [
        "# Hyperparameters\n",
        "batch_size=64\n",
        "epochs=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KR01t9v_fxbT"
      },
      "source": [
        "## Build the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "djpoAvfWNyL5"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Dense(64, input_shape=(784,), activation='relu', name='dense_1'),\n",
        "    tf.keras.layers.Dense(64, activation='relu', name='dense_2'),\n",
        "    tf.keras.layers.Dense(10, activation='softmax', name='predictions'),\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0_D7CZqkv_Hj"
      },
      "source": [
        "## Prepare the Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U0bS3SyowBoB"
      },
      "outputs": [],
      "source": [
        "# Load MNIST dataset as NumPy arrays\n",
        "dataset = {}\n",
        "num_validation = 10000\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "# Preprocess the data\n",
        "x_train = x_train.reshape(-1, 784).astype('float32') / 255\n",
        "x_test = x_test.reshape(-1, 784).astype('float32') / 255"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HYE-BxhOzFQp"
      },
      "source": [
        "## Train and Evaluate\n",
        "\n",
        "Simply replace typical keras optimizers with the new tfa optimizer "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NxfYhtiSzHf-"
      },
      "outputs": [],
      "source": [
        "# Compile the model\n",
        "model.compile(\n",
        "    optimizer=tfa.optimizers.LazyAdam(0.001),  # Utilize TFA optimizer\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "# Train the network\n",
        "history = model.fit(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    batch_size=batch_size,\n",
        "    epochs=epochs)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Y--0tK69SXf"
      },
      "outputs": [],
      "source": [
        "# Evaluate the network\n",
        "print('Evaluate on test data:')\n",
        "results = model.evaluate(x_test, y_test, batch_size=128, verbose = 2)\n",
        "print('Test loss = {0}, Test acc: {1}'.format(results[0], results[1]))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "optimizers_lazyadam.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}