docs/tutorials/layers_weightnormalization.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 Layers: WeightNormalization\n",
    "\n",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.tensorflow.org/addons/tutorials/layers_weightnormalization\"><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/layers_weightnormalization.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/layers_weightnormalization.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/layers_weightnormalization.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 Weight Normalization layer and how it can improve convergence.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KR01t9v_fxbT"
   },
   "source": [
    "# WeightNormalization\n",
    "\n",
    "A Simple Reparameterization to Accelerate Training of Deep Neural Networks:\n",
    "\n",
    "Tim Salimans, Diederik P. Kingma (2016)\n",
    "\n",
    "> By reparameterizing the weights in this way you improve the conditioning of the optimization problem and speed up convergence of stochastic gradient descent. Our reparameterization is inspired by batch normalization but does not introduce any dependencies between the examples in a minibatch. This means that our method can also be applied successfully to recurrent models such as LSTMs and to noise-sensitive applications such as deep reinforcement learning or generative models, for which batch normalization is less well suited. Although our method is much simpler, it still provides much of the speed-up of full batch normalization. In addition, the computational overhead of our method is lower, permitting more optimization steps to be taken in the same amount of time.\n",
    "\n",
    "> https://arxiv.org/abs/1602.07868 \n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/seanpmorgan/tf-weightnorm/master/static/wrapped-graph.png\" width=\"80%\"><br><br>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MUXex9ctTuDB"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CyWHXw9mQ6mp"
   },
   "outputs": [],
   "source": [
    "!pip install -U tensorflow-addons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OywLbs7EXiE_"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_addons as tfa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KQMhhq1qXiFF"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ULWHqMAnTVZD"
   },
   "outputs": [],
   "source": [
    "# Hyper Parameters\n",
    "batch_size = 32\n",
    "epochs = 10\n",
    "num_classes = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fhM0ieDpSnKh"
   },
   "source": [
    "## Build Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7XZXnBYgRPSk"
   },
   "outputs": [],
   "source": [
    "# Standard ConvNet\n",
    "reg_model = tf.keras.Sequential(\n",
    "    [\n",
    "        tf.keras.layers.Conv2D(6, 5, activation=\"relu\"),\n",
    "        tf.keras.layers.MaxPooling2D(2, 2),\n",
    "        tf.keras.layers.Conv2D(16, 5, activation=\"relu\"),\n",
    "        tf.keras.layers.MaxPooling2D(2, 2),\n",
    "        tf.keras.layers.Flatten(),\n",
    "        tf.keras.layers.Dense(120, activation=\"relu\"),\n",
    "        tf.keras.layers.Dense(84, activation=\"relu\"),\n",
    "        tf.keras.layers.Dense(num_classes, activation=\"softmax\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UZd6V90eR4Gm"
   },
   "outputs": [],
   "source": [
    "# WeightNorm ConvNet\n",
    "wn_model = tf.keras.Sequential(\n",
    "    [\n",
    "        tfa.layers.WeightNormalization(tf.keras.layers.Conv2D(6, 5, activation=\"relu\")),\n",
    "        tf.keras.layers.MaxPooling2D(2, 2),\n",
    "        tfa.layers.WeightNormalization(\n",
    "            tf.keras.layers.Conv2D(16, 5, activation=\"relu\")\n",
    "        ),\n",
    "        tf.keras.layers.MaxPooling2D(2, 2),\n",
    "        tf.keras.layers.Flatten(),\n",
    "        tfa.layers.WeightNormalization(tf.keras.layers.Dense(120, activation=\"relu\")),\n",
    "        tfa.layers.WeightNormalization(tf.keras.layers.Dense(84, activation=\"relu\")),\n",
    "        tfa.layers.WeightNormalization(\n",
    "            tf.keras.layers.Dense(num_classes, activation=\"softmax\")\n",
    "        ),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AA5dti8AS2Y7"
   },
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L8Isjc7W8MEn"
   },
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
    "\n",
    "# Convert class vectors to binary class matrices.\n",
    "y_train = tf.keras.utils.to_categorical(y_train, num_classes)\n",
    "y_test = tf.keras.utils.to_categorical(y_test, num_classes)\n",
    "\n",
    "x_train = x_train.astype(\"float32\")\n",
    "x_test = x_test.astype(\"float32\")\n",
    "x_train /= 255\n",
    "x_test /= 255"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cH1CG9E7S34C"
   },
   "source": [
    "## Train Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EvNKxfaI7vSm"
   },
   "outputs": [],
   "source": [
    "reg_model.compile(\n",
    "    optimizer=\"adam\", loss=\"categorical_crossentropy\", metrics=[\"accuracy\"]\n",
    ")\n",
    "\n",
    "reg_history = reg_model.fit(\n",
    "    x_train,\n",
    "    y_train,\n",
    "    batch_size=batch_size,\n",
    "    epochs=epochs,\n",
    "    validation_data=(x_test, y_test),\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "esmMh-5g7wmp"
   },
   "outputs": [],
   "source": [
    "wn_model.compile(\n",
    "    optimizer=\"adam\", loss=\"categorical_crossentropy\", metrics=[\"accuracy\"]\n",
    ")\n",
    "\n",
    "wn_history = wn_model.fit(\n",
    "    x_train,\n",
    "    y_train,\n",
    "    batch_size=batch_size,\n",
    "    epochs=epochs,\n",
    "    validation_data=(x_test, y_test),\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yujf2YRbwX55"
   },
   "outputs": [],
   "source": [
    "reg_accuracy = reg_history.history[\"accuracy\"]\n",
    "wn_accuracy = wn_history.history[\"accuracy\"]\n",
    "\n",
    "plt.plot(\n",
    "    np.linspace(0, epochs, epochs), reg_accuracy, color=\"red\", label=\"Regular ConvNet\"\n",
    ")\n",
    "\n",
    "plt.plot(\n",
    "    np.linspace(0, epochs, epochs),\n",
    "    wn_accuracy,\n",
    "    color=\"blue\",\n",
    "    label=\"WeightNorm ConvNet\",\n",
    ")\n",
    "\n",
    "plt.title(\"WeightNorm Accuracy Comparison\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "layers_weightnormalization.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}