site/en/tutorials/eager/custom_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Custom training: basics",
      "version": "0.3.2",
      "views": {},
      "default_view": {},
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "metadata": {
        "id": "5rmpybwysXGV",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2018 The TensorFlow Authors."
      ]
    },
    {
      "metadata": {
        "id": "m8y3rGtQsYP2",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "cellView": "form"
      },
      "cell_type": "code",
      "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."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "hrXv0rU9sIma",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Custom training: basics"
      ]
    },
    {
      "metadata": {
        "id": "7S0BwJ_8sLu7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\"><td>\n",
        "<a target=\"_blank\"  href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/eager/custom_training.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "</td><td>\n",
        "<a target=\"_blank\"  href=\"https://github.com/tensorflow/docs/blob/master/site/en/tutorials/eager/custom_training.ipynb\"><img width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a></td></table>"
      ]
    },
    {
      "metadata": {
        "id": "k2o3TTG4TFpt",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "In the previous tutorial we covered the TensorFlow APIs for automatic differentiation, a basic building block for machine learning.\n",
        "In this tutorial we will use the TensorFlow primitives introduced in the prior tutorials to do some simple machine learning.\n",
        "\n",
        "TensorFlow also includes a higher-level neural networks API (`tf.keras`) which provides useful abstractions to reduce boilerplate. We strongly recommend those higher level APIs for people working with neural networks. However, in this short tutorial we cover neural network training from first principles to establish a strong foundation."
      ]
    },
    {
      "metadata": {
        "id": "3LXMVuV0VhDr",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ]
    },
    {
      "metadata": {
        "id": "PJ64L90aVir3",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "tf.enable_eager_execution()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "eMAWbDJFVmMk",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Variables\n",
        "\n",
        "Tensors in TensorFlow are immutable stateless objects. Machine learning models, however, need to have changing state: as your model trains, the same code to compute predictions should behave differently over time (hopefully with a lower loss!). To represent this state which needs to change over the course of your computation, you can choose to rely on the fact that Python is a stateful programming language:\n"
      ]
    },
    {
      "metadata": {
        "id": "VkJwtLS_Jbn8",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "# Using python state\n",
        "x = tf.zeros([10, 10])\n",
        "x += 2  # This is equivalent to x = x + 2, which does not mutate the original\n",
        "        # value of x\n",
        "print(x)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "wfneTXy7JcUz",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "TensorFlow, however, has stateful operations built in, and these are often more pleasant to use than low-level Python representations of your state. To represent weights in a model, for example, it's often convenient and efficient to use TensorFlow variables.\n",
        "\n",
        "A Variable is an object which stores a value and, when used in a TensorFlow computation, will implicitly read from this stored value. There are operations (`tf.assign_sub`, `tf.scatter_update`, etc) which manipulate the value stored in a TensorFlow variable."
      ]
    },
    {
      "metadata": {
        "id": "itxmrMil6DQi",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "v = tf.Variable(1.0)\n",
        "assert v.numpy() == 1.0\n",
        "\n",
        "# Re-assign the value\n",
        "v.assign(3.0)\n",
        "assert v.numpy() == 3.0\n",
        "\n",
        "# Use `v` in a TensorFlow operation like tf.square() and reassign\n",
        "v.assign(tf.square(v))\n",
        "assert v.numpy() == 9.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "-paSaeq1JzwC",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Computations using Variables are automatically traced when computing gradients. For Variables representing embeddings TensorFlow will do sparse updates by default, which are more computation and memory efficient.\n",
        "\n",
        "Using Variables is also a way to quickly let a reader of your code know that this piece of state is mutable."
      ]
    },
    {
      "metadata": {
        "id": "BMiFcDzE7Qu3",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Example: Fitting a linear model\n",
        "\n",
        "Let's now put the few concepts we have so far ---`Tensor`, `GradientTape`, `Variable` --- to build and train a simple model. This typically involves a few steps:\n",
        "\n",
        "1. Define the model.\n",
        "2. Define a loss function.\n",
        "3. Obtain training data.\n",
        "4. Run through the training data and use an \"optimizer\" to adjust the variables to fit the data.\n",
        "\n",
        "In this tutorial, we'll walk through a trivial example of a simple linear model: `f(x) = x * W + b`, which has two variables - `W` and `b`. Furthermore, we'll synthesize data such that a well trained model would have `W = 3.0` and `b = 2.0`."
      ]
    },
    {
      "metadata": {
        "id": "gFzH64Jn9PIm",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define the model\n",
        "\n",
        "Let's define a simple class to encapsulate the variables and the computation."
      ]
    },
    {
      "metadata": {
        "id": "_WRu7Pze7wk8",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "class Model(object):\n",
        "  def __init__(self):\n",
        "    # Initialize variable to (5.0, 0.0)\n",
        "    # In practice, these should be initialized to random values.\n",
        "    self.W = tf.Variable(5.0)\n",
        "    self.b = tf.Variable(0.0)\n",
        "    \n",
        "  def __call__(self, x):\n",
        "    return self.W * x + self.b\n",
        "  \n",
        "model = Model()\n",
        "\n",
        "assert model(3.0).numpy() == 15.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xa6j_yXa-j79",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define a loss function\n",
        "\n",
        "A loss function measures how well the output of a model for a given input matches the desired output. Let's use the standard L2 loss."
      ]
    },
    {
      "metadata": {
        "id": "Y0ysUFGY924U",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "def loss(predicted_y, desired_y):\n",
        "  return tf.reduce_mean(tf.square(predicted_y - desired_y))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "qutT_fkl_CBc",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Obtain training data\n",
        "\n",
        "Let's synthesize the training data with some noise."
      ]
    },
    {
      "metadata": {
        "id": "gxPTb-kt_N5m",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "TRUE_W = 3.0\n",
        "TRUE_b = 2.0\n",
        "NUM_EXAMPLES = 1000\n",
        "\n",
        "inputs  = tf.random_normal(shape=[NUM_EXAMPLES])\n",
        "noise   = tf.random_normal(shape=[NUM_EXAMPLES])\n",
        "outputs = inputs * TRUE_W + TRUE_b + noise"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "-50nq-wPBsAW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Before we train the model let's visualize where the model stands right now. We'll plot the model's predictions in red and the training data in blue."
      ]
    },
    {
      "metadata": {
        "id": "_eb83LtrB4nt",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.scatter(inputs, outputs, c='b')\n",
        "plt.scatter(inputs, model(inputs), c='r')\n",
        "plt.show()\n",
        "\n",
        "print('Current loss: '),\n",
        "print(loss(model(inputs), outputs).numpy())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sSDP-yeq_4jE",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define a training loop\n",
        "\n",
        "We now have our network and our training data. Let's train it, i.e., use the training data to update the model's variables (`W` and `b`) so that the loss goes down using [gradient descent](https://en.wikipedia.org/wiki/Gradient_descent). There are many variants of the gradient descent scheme that are captured in `tf.train.Optimizer` implementations. We'd highly recommend using those implementations, but in the spirit of building from first principles, in this particular example we will implement the basic math ourselves."
      ]
    },
    {
      "metadata": {
        "id": "MBIACgdnA55X",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "def train(model, inputs, outputs, learning_rate):\n",
        "  with tf.GradientTape() as t:\n",
        "    current_loss = loss(model(inputs), outputs)\n",
        "  dW, db = t.gradient(current_loss, [model.W, model.b])\n",
        "  model.W.assign_sub(learning_rate * dW)\n",
        "  model.b.assign_sub(learning_rate * db)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "RwWPaJryD2aN",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Finally, let's repeatedly run through the training data and see how `W` and `b` evolve."
      ]
    },
    {
      "metadata": {
        "id": "XdfkR223D9dW",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "cell_type": "code",
      "source": [
        "model = Model()\n",
        "\n",
        "# Collect the history of W-values and b-values to plot later\n",
        "Ws, bs = [], []\n",
        "epochs = range(10)\n",
        "for epoch in epochs:\n",
        "  Ws.append(model.W.numpy())\n",
        "  bs.append(model.b.numpy())\n",
        "  current_loss = loss(model(inputs), outputs)\n",
        "\n",
        "  train(model, inputs, outputs, learning_rate=0.1)\n",
        "  print('Epoch %2d: W=%1.2f b=%1.2f, loss=%2.5f' %\n",
        "        (epoch, Ws[-1], bs[-1], current_loss))\n",
        "\n",
        "# Let's plot it all\n",
        "plt.plot(epochs, Ws, 'r',\n",
        "         epochs, bs, 'b')\n",
        "plt.plot([TRUE_W] * len(epochs), 'r--',\n",
        "         [TRUE_b] * len(epochs), 'b--')\n",
        "plt.legend(['W', 'b', 'true W', 'true_b'])\n",
        "plt.show()\n",
        "  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "vPnIVuaSJwWz",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Next Steps\n",
        "\n",
        "In this tutorial we covered `Variable`s and built and trained a simple linear model using the TensorFlow primitives discussed so far.\n",
        "\n",
        "In theory, this is pretty much all you need to use TensorFlow for your machine learning research.\n",
        "In practice, particularly for neural networks, the higher level APIs like `tf.keras` will be much more convenient since it provides higher level building blocks (called \"layers\"), utilities to save and restore state, a suite of loss functions, a suite of optimization strategies etc. \n",
        "\n",
        "The [next tutorial](TODO) will cover these higher level APIs."
      ]
    }
  ]
}