Improving_Neural_Network_Training_in_TensorFlow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Improving_Neural_Network_Training_in_TensorFlow.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyM/LgnQd5gIwKPGzdmVEyWI",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lukmanr/codenext/blob/master/Improving_Neural_Network_Training_in_TensorFlow.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ta2WuhvgIjRa",
        "colab_type": "text"
      },
      "source": [
        "# Improving Neural Network Training in TensorFlow\n",
        "\n",
        "This notebook shows you how to improve training for the basic Neural Network we saw in the previous notebook. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tLQWJal3_Hv0",
        "colab_type": "text"
      },
      "source": [
        "![two layer MNIST network](https://i.imgur.com/5LOj7mtl.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0jzreYS1JH4E",
        "colab_type": "text"
      },
      "source": [
        "Let's start with the code from the basic Neural Network in TensorFlow we did last time.  We are going to \"refactor\", or reorganize, the code a little.  We'll put all the code that loads the MNIST training data into one method, ```load_training_data```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uxLvKF-kYASA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "\n",
        "# Load the MNIST data set. The training set and test set are split\n",
        "# automatically. We pre-process the data to \"normalize\" the pixels (make all \n",
        "# the pixel values between 0 and 1), reshape the data into 2-D tensors of\n",
        "# shape 1 x 784, and one-hot encode the labels.\n",
        "def load_training_data():\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "  # convert the integer pixel values to floats\n",
        "  x_train = x_train / 255.0\n",
        "  x_test = x_test / 255.0\n",
        "\n",
        "  # reshape the images to be 2-D tensors of 1 x 784 pixels\n",
        "  x_train = x_train.reshape([-1, 1, 784])\n",
        "  x_test = x_test.reshape([-1, 1, 784])\n",
        "\n",
        "  # reshape the labels to be one hot encoded tensors\n",
        "  y_train = tf.one_hot(y_train, depth=10)\n",
        "  y_test = tf.one_hot(y_test, depth=10)\n",
        "\n",
        "  return (x_train, y_train), (x_test, y_test)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1PAIkeTAJKoi",
        "colab_type": "text"
      },
      "source": [
        "The code that defines the neural network will be the same, for now.  Remember from last time:  the network is a simple function that takes the inputs X as argument.  It computes the input to each neuron by multiplying the inputs times the weights and adding the biases.  Then it applies the \"softmax\" activation function to the inputs, to compute the output for each of the 10 neurons."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "63CkAiRATFN4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# weights W[784, 10]   784=28*28\n",
        "W = tf.Variable(tf.zeros([784, 10]))\n",
        "\n",
        "# biases b[10]\n",
        "b = tf.Variable(tf.zeros([10]))\n",
        "\n",
        "# The neural network\n",
        "def neural_network(X):\n",
        "  Inputs = tf.matmul(X, W) + b\n",
        "  Y = tf.nn.softmax(Inputs)\n",
        "  return Y"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oGxF-8xZMkGj",
        "colab_type": "text"
      },
      "source": [
        "The loss function is the same sum of squares function we have seen before."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QsmqazsridHB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# The loss function\n",
        "def loss(Y, Y_l):\n",
        "  return tf.reduce_sum(tf.square(Y - Y_l))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4k5lJp1-MuEA",
        "colab_type": "text"
      },
      "source": [
        "Here is the training loop.  We'll put that into a method as well, to make it easier to call repeatedly.  The method take as required parameters the training images and labels.  The method also takes ```num_epochs```, ```learning_rate```, and ```eval_steps``` as hyperparameters, specifying the default values of each of those parameters.  We've also modified it to plot the average loss values every ```eval_steps``` if the plot flag is set to True. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oT7AiAqHlHWl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def nn_training_loop(x_train,\n",
        "                     y_train,\n",
        "                     num_epochs = 1, \n",
        "                     learning_rate = 0.001, \n",
        "                     eval_steps = 500,\n",
        "                     plot = False):\n",
        "\n",
        "  counter = 0\n",
        "  total_loss = 0.0\n",
        "  avg_losses = []\n",
        "\n",
        "  # the outer training loop:  repeat for num_epochs\n",
        "  for e in range(num_epochs):\n",
        "\n",
        "      # the inner training loop: train on one image and label from the data set\n",
        "      for image, Y_l in zip(x_train, y_train):\n",
        "\n",
        "          # use this to count how many training loops we've executed\n",
        "          counter += 1\n",
        "\n",
        "          # convert the image to float tensor\n",
        "          X = tf.Variable(image, dtype=tf.float32)\n",
        "\n",
        "          # we wrap this 'with' statement around the next two lines, to tell \n",
        "          # TensorFlow to auto-compute the gradients\n",
        "          with tf.GradientTape() as tape:\n",
        "              # get the output of the neural net\n",
        "              Y = neural_network(X)\n",
        "\n",
        "              # compute the loss function \n",
        "              current_loss = loss(Y, Y_l)\n",
        "\n",
        "          # compute the gradients of the weights and biases with respect to the\n",
        "          # loss function\n",
        "          dW, db = tape.gradient(current_loss, [W, b])\n",
        "\n",
        "          # update the weights and biases. Remember we need to reverse the sign; \n",
        "          # ie. we want the add the negative of the gradient (times the \n",
        "          # learning rate) to the weights and biases.\n",
        "          W.assign(W - learning_rate * dW)\n",
        "          b.assign(b - learning_rate * db)\n",
        "\n",
        "          # every eval_steps, print the average loss since the last eval\n",
        "          total_loss += current_loss.numpy()\n",
        "          if (counter % eval_steps == 0):\n",
        "              avg_loss = total_loss / eval_steps\n",
        "              avg_losses.append(avg_loss)\n",
        "              if plot:\n",
        "                plt.plot(avg_losses)\n",
        "                plt.show()\n",
        "              else:\n",
        "                print(\"train step =\", counter, \" avg loss =\", avg_loss)\n",
        "              total_loss = 0.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "62U-FpKEbZwH",
        "colab_type": "text"
      },
      "source": [
        "Now call the function to execute the loop."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FX1Qn8FqWqp5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "(x_train, y_train), (x_test, y_test) = load_training_data()\n",
        "\n",
        "nn_training_loop(x_train, y_train)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HNYj3vNPJ7-d",
        "colab_type": "text"
      },
      "source": [
        "Now let's define a method to compute the average test loss.  We just enclose the code from the previous notebook into a method:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UfAmU7EFHmkK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def test_loss(x_test, y_test):\n",
        "    # First convert the test set images and labels into Tensors\n",
        "    x_test = tf.Variable(x_test, dtype=tf.float32)\n",
        "    y_test = tf.Variable(y_test, dtype=tf.float32)\n",
        "\n",
        "    # The labels need to be the same shape as the neural network output:\n",
        "    # a vector of [1, 10].  The \"-1\" means, whatever dimension is needed\n",
        "    # to get to the right shape.  In this case there are 10000 test images\n",
        "    # so that \"-1\" is going to become 10000.\n",
        "    y_test = tf.reshape(y_test, [-1, 1, 10])\n",
        "\n",
        "    # Now compute the loss on the entire test set.  First compute the\n",
        "    # neural network output on the entire test set:\n",
        "    Y_test = neural_network(x_test)\n",
        "\n",
        "    # Then use that output, along with the test labels, to compute the\n",
        "    # loss for the entire test set:\n",
        "    test_loss = loss(Y_test, y_test)\n",
        "\n",
        "    # divide by y_test.shape[0] which is the number of test images\n",
        "    average_test_loss = test_loss / y_test.shape[0]\n",
        "\n",
        "    # output the numpy value of the tensor (i.e. just the normal python value)\n",
        "    return average_test_loss.numpy()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Vja_61FIL-G",
        "colab_type": "text"
      },
      "source": [
        "Now let's call that method to compute the test loss.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x6_kk96tIKVp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "t_l = test_loss(x_test, y_test)\n",
        "print(t_l)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sz9hQ--YbgUX",
        "colab_type": "text"
      },
      "source": [
        "After one epoch the network gets to about 80% accuracy on this problem. Let's make some improvements to this Neural Network! "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3OmXWdLlVuta",
        "colab_type": "text"
      },
      "source": [
        "## Improvements to training\n",
        "We are going to try these things to improve the training of this neural network:\n",
        "\n",
        "1.   Use **Stochastic Gradient Descent**:  train in batches.\n",
        "2.   Use **a better loss** function (the \"cross entropy loss\").\n",
        "3.   Add some **\"hidden layers\"**.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqrsKP_yN1Q4",
        "colab_type": "text"
      },
      "source": [
        "First let's create a new method to create data for training.  We use the TensorFlow data module to create a TensorFlow Dataset directly from the training and test data.  TensorFlow's Dataset makes it easy to define randomly sampled batches of data.  Datasets also make it easy to define batches, which we will apply in the training loop method.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M_La_NOrWpTT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def training_dataset():\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "  # reshape the images to be 2-D tensors of 1 x 784 pixels\n",
        "  x_train = x_train.reshape([-1, 784])\n",
        "  x_train = tf.convert_to_tensor(x_train, dtype=tf.float32) / 255.\n",
        "  x_test = x_test.reshape([-1, 784])\n",
        "  x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) / 255.\n",
        "\n",
        "  # convert the labels to one hot encoded\n",
        "  y_train = tf.one_hot(y_train, depth=10)\n",
        "  y_test = tf.one_hot(y_test, depth=10)\n",
        "\n",
        "  # create a TensorFlow dataset for train and test sets\n",
        "  train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "  test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "\n",
        "  # shuffle (randomly select) the dataset\n",
        "  train_dataset = train_dataset.shuffle(10000).repeat()\n",
        "\n",
        "  return train_dataset, test_dataset"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "db-b1rDX-xpx",
        "colab_type": "text"
      },
      "source": [
        "We make a couple minor modifications to the training loop, to make use of the Dataset, and apply the batch size."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pBRxy-6yb-Dh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def nn_dataset_training_loop(train_data,\n",
        "                             batch_size = 100,\n",
        "                             num_epochs = 1, \n",
        "                             learning_rate = 0.001, \n",
        "                             eval_steps = 500,\n",
        "                             plot = False):\n",
        "\n",
        "  counter = 0\n",
        "  total_loss = 0.0\n",
        "  avg_losses = []\n",
        "\n",
        "  # set batch size in training data\n",
        "  train_data = train_data.batch(batch_size)\n",
        "\n",
        "  # the outer training loop:  repeat for num_epochs\n",
        "  for e in range(num_epochs):\n",
        "\n",
        "      # the inner training loop: train on one image and label from the data set\n",
        "      for X, Y_l in train_data:\n",
        "\n",
        "          # use this to count how many training loops we've executed\n",
        "          counter += 1          \n",
        "\n",
        "          # we wrap this 'with' statement around the next two lines, to tell \n",
        "          # TensorFlow to auto-compute the gradients                  \n",
        "          with tf.GradientTape() as tape:\n",
        "              # get the output of the neural net\n",
        "              Y = neural_network(X)\n",
        "\n",
        "              # compute the loss function \n",
        "              current_loss = loss(Y, Y_l)\n",
        "\n",
        "          # compute the gradients of the weights and biases with respect to the\n",
        "          # loss function\n",
        "          dW, db = tape.gradient(current_loss, [W, b])\n",
        "\n",
        "          # update the weights and biases. Remember we need to reverse the sign; \n",
        "          # ie. we want the add the negative of the gradient (times the \n",
        "          # learning rate) to the weights and biases.\n",
        "          W.assign(W - learning_rate * dW)\n",
        "          b.assign(b - learning_rate * db)\n",
        "\n",
        "          # every eval_steps, print the average loss since the last eval\n",
        "          total_loss += current_loss.numpy()\n",
        "          if (counter % eval_steps == 0):\n",
        "              avg_loss = total_loss / eval_steps\n",
        "              avg_losses.append(avg_loss)\n",
        "              if plot:\n",
        "                plt.plot(avg_losses)\n",
        "                plt.show()\n",
        "              else:\n",
        "                print(\"train step =\", counter, \" avg loss =\", avg_loss)\n",
        "              total_loss = 0.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rJSvqVQq_O3b",
        "colab_type": "text"
      },
      "source": [
        "Let's execute the new batch training scheme and see what we get."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sfuITEwBHus6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_data, test_data = training_dataset()\n",
        "\n",
        "nn_dataset_training_loop(train_data)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4RUNpjVBMSjI",
        "colab_type": "text"
      },
      "source": [
        "### A better loss function\n",
        "\n",
        "For a problem in which we are trying to classify a set of objects, a better loss function than \"squared error\" is \"cross entropy\".  It computes the difference between two probability distributions.  In this case, the output of the 10 neurons is a probablity distribution, and the label represents another probability distribution."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JkED7j_YI1nu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# The loss function\n",
        "def loss(Y, Y_l):\n",
        "  loss = tf.keras.losses.CategoricalCrossentropy()\n",
        "  return loss(Y_l, Y)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PmNL-9OLNejA",
        "colab_type": "text"
      },
      "source": [
        "Let's try training with the new loss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dWUL9_3BNhdj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "nn_dataset_training_loop(train_data)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeO94bCjNiGy",
        "colab_type": "text"
      },
      "source": [
        "Finally we will add some layers to the network.  We'll pick an arbitrary number of unit for the hidden layers. We define separate weights and biases for each layer, then define a new neural network function that computes the activations of each layer, feeding them into the next one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-7hRwXGyH2rk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# number of units in hidden layers\n",
        "n_hidden = 20\n",
        "\n",
        "# weights/biases for first layer\n",
        "W1 = tf.Variable(tf.zeros([784, n_hidden]))\n",
        "b1 = tf.Variable(tf.zeros([n_hidden]))\n",
        "\n",
        "# weights/biases for second layer\n",
        "W2 = tf.Variable(tf.zeros([n_hidden, n_hidden]))\n",
        "b2 = tf.Variable(tf.zeros([n_hidden]))\n",
        "\n",
        "# weights/biases for final layer\n",
        "W3 = tf.Variable(tf.zeros([n_hidden, 10]))\n",
        "b3 = tf.Variable(tf.zeros([10]))\n",
        "\n",
        "# The neural network\n",
        "def multilayer_neural_network(X):\n",
        "  input1 = tf.matmul(X, W1) + b1\n",
        "  output1 = tf.nn.relu(input1)\n",
        "\n",
        "  input2 = tf.matmul(output1, W2) + b2\n",
        "  output2 = tf.nn.relu(input2)\n",
        "\n",
        "  input3 = tf.matmul(output2, W3) + b3\n",
        "  Y = tf.nn.softmax(input3)\n",
        "  return Y"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fDT_ziLreUi-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def multi_layer_nn_training_loop(train_data,\n",
        "                                 batch_size = 100,\n",
        "                                 num_epochs = 1, \n",
        "                                 learning_rate = 0.001, \n",
        "                                 eval_steps = 500,\n",
        "                                 plot = False):\n",
        "\n",
        "  counter = 0\n",
        "  total_loss = 0.0\n",
        "  avg_losses = []\n",
        "\n",
        "  # set batch size in training data\n",
        "  train_data = train_data.batch(batch_size)\n",
        "\n",
        "  # the outer training loop:  repeat for num_epochs\n",
        "  for e in range(num_epochs):\n",
        "\n",
        "      # the inner training loop: train on one image and label from the data set\n",
        "      for X, Y_l in train_data:\n",
        "\n",
        "          # use this to count how many training loops we've executed\n",
        "          counter += 1\n",
        "\n",
        "          # we wrap this 'with' statement around the next two lines, to tell \n",
        "          # TensorFlow to auto-compute the gradients\n",
        "          with tf.GradientTape() as tape:\n",
        "              # get the output of the neural net\n",
        "              Y = multilayer_neural_network(X)\n",
        "\n",
        "              # compute the loss function \n",
        "              current_loss = loss(Y, Y_l)\n",
        "\n",
        "          # compute the gradients of the weights and biases with respect to the\n",
        "          # loss function\n",
        "          dW1, db1, dW2, db2, dW3, db3 = tape.gradient(current_loss, [W1, b1, W2, b2, W3, b3])\n",
        "\n",
        "          # update the weights and biases. Remember we need to reverse the sign; \n",
        "          # ie. we want the add the negative of the gradient (times the \n",
        "          # learning rate) to the weights and biases.\n",
        "          W1.assign(W1 - learning_rate * dW1)\n",
        "          b1.assign(b1 - learning_rate * db1)\n",
        "\n",
        "          W2.assign(W2 - learning_rate * dW2)\n",
        "          b2.assign(b2 - learning_rate * db2)\n",
        "\n",
        "          W3.assign(W3 - learning_rate * dW3)\n",
        "          b3.assign(b3 - learning_rate * db3)\n",
        "\n",
        "          # every eval_steps, print the average loss since the last eval\n",
        "          total_loss += current_loss.numpy()\n",
        "          if (counter % eval_steps == 0):\n",
        "              avg_loss = total_loss / eval_steps\n",
        "              avg_losses.append(avg_loss)\n",
        "              if (plot):\n",
        "                plt.plot(avg_losses)\n",
        "                plt.show()\n",
        "              else:\n",
        "                print(\"train step =\", counter, \" avg loss =\", avg_loss)\n",
        "              total_loss = 0.0\n",
        "\n",
        "  test_results(test_data)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "leZOaQ4ee2lk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_data, test_data = training_dataset()\n",
        "\n",
        "multi_layer_nn_training_loop(train_data)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}