assignment2/TensorFlow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What's this TensorFlow business?\n",
    "\n",
    "You've written a lot of code in this assignment to provide a whole host of neural network functionality. Dropout, Batch Norm, and 2D convolutions are some of the workhorses of deep learning in computer vision. You've also worked hard to make your code efficient and vectorized.\n",
    "\n",
    "For the last part of this assignment, though, we're going to leave behind your beautiful codebase and instead migrate to one of two popular deep learning frameworks: in this instance, TensorFlow (or PyTorch, if you switch over to that notebook)\n",
    "\n",
    "#### What is it?\n",
    "TensorFlow is a system for executing computational graphs over Tensor objects, with native support for performing backpropogation for its Variables. In it, we work with Tensors which are n-dimensional arrays analogous to the numpy ndarray.\n",
    "\n",
    "#### Why?\n",
    "\n",
    "* Our code will now run on GPUs! Much faster training. Writing your own modules to run on GPUs is beyond the scope of this class, unfortunately.\n",
    "* We want you to be ready to use one of these frameworks for your project so you can experiment more efficiently than if you were writing every feature you want to use by hand. \n",
    "* We want you to stand on the shoulders of giants! TensorFlow and PyTorch are both excellent frameworks that will make your lives a lot easier, and now that you understand their guts, you are free to use them :) \n",
    "* We want you to be exposed to the sort of deep learning code you might run into in academia or industry. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How will I learn TensorFlow?\n",
    "\n",
    "TensorFlow has many excellent tutorials available, including those from [Google themselves](https://www.tensorflow.org/get_started/get_started).\n",
    "\n",
    "Otherwise, this notebook will walk you through much of what you need to do to train models in TensorFlow. See the end of the notebook for some links to helpful tutorials if you want to learn more or need further clarification on topics that aren't fully explained here.\n",
    "\n",
    "\n",
    "# Table of Contents\n",
    "\n",
    "This notebook has 5 parts. We will walk through TensorFlow at three different levels of abstraction, which should help you better understand it and prepare you for working on your project.\n",
    "\n",
    "1. Preparation: load the CIFAR-10 dataset.\n",
    "2. Barebone TensorFlow: we will work directly with low-level TensorFlow graphs. \n",
    "3. Keras Model API: we will use `tf.keras.Model` to define arbitrary neural network architecture. \n",
    "4. Keras Sequential API: we will use `tf.keras.Sequential` to define a linear feed-forward network very conveniently. \n",
    "5. CIFAR-10 open-ended challenge: please implement your own network to get as high accuracy as possible on CIFAR-10. You can experiment with any layer, optimizer, hyperparameters or other advanced features. \n",
    "\n",
    "Here is a table of comparison:\n",
    "\n",
    "| API           | Flexibility | Convenience |\n",
    "|---------------|-------------|-------------|\n",
    "| Barebone      | High        | Low         |\n",
    "| `tf.keras.Model`     | High        | Medium      |\n",
    "| `tf.keras.Sequential` | Low         | High        |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part I: Preparation\n",
    "\n",
    "First, we load the CIFAR-10 dataset. This might take a few minutes to download the first time you run it, but after that the files should be cached on disk and loading should be faster.\n",
    "\n",
    "In previous parts of the assignment we used CS231N-specific code to download and read the CIFAR-10 dataset; however the `tf.keras.datasets` package in TensorFlow provides prebuilt utility functions for loading many common datasets.\n",
    "\n",
    "For the purposes of this assignment we will still write our own code to preprocess the data and iterate through it in minibatches. The `tf.data` package in TensorFlow provides tools for automating this process, but working with this package adds extra complication and is beyond the scope of this notebook. However using `tf.data` can be much more efficient than the simple approach used in this notebook, so you should consider using it for your project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/weidong/anaconda/envs/py36/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import math\n",
    "import timeit\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data shape:  (49000, 32, 32, 3)\n",
      "Train labels shape:  (49000,) int32\n",
      "Validation data shape:  (1000, 32, 32, 3)\n",
      "Validation labels shape:  (1000,)\n",
      "Test data shape:  (10000, 32, 32, 3)\n",
      "Test labels shape:  (10000,)\n"
     ]
    }
   ],
   "source": [
    "def load_cifar10(num_training=49000, num_validation=1000, num_test=10000):\n",
    "    \"\"\"\n",
    "    Fetch the CIFAR-10 dataset from the web and perform preprocessing to prepare\n",
    "    it for the two-layer neural net classifier. These are the same steps as\n",
    "    we used for the SVM, but condensed to a single function.\n",
    "    \"\"\"\n",
    "    # Load the raw CIFAR-10 dataset and use appropriate data types and shapes\n",
    "    cifar10 = tf.keras.datasets.cifar10.load_data()\n",
    "    (X_train, y_train), (X_test, y_test) = cifar10\n",
    "    X_train = np.asarray(X_train, dtype=np.float32)\n",
    "    y_train = np.asarray(y_train, dtype=np.int32).flatten()\n",
    "    X_test = np.asarray(X_test, dtype=np.float32)\n",
    "    y_test = np.asarray(y_test, dtype=np.int32).flatten()\n",
    "\n",
    "    # Subsample the data\n",
    "    mask = range(num_training, num_training + num_validation)\n",
    "    X_val = X_train[mask]\n",
    "    y_val = y_train[mask]\n",
    "    mask = range(num_training)\n",
    "    X_train = X_train[mask]\n",
    "    y_train = y_train[mask]\n",
    "    mask = range(num_test)\n",
    "    X_test = X_test[mask]\n",
    "    y_test = y_test[mask]\n",
    "\n",
    "    # Normalize the data: subtract the mean pixel and divide by std\n",
    "    mean_pixel = X_train.mean(axis=(0, 1, 2), keepdims=True)\n",
    "    std_pixel = X_train.std(axis=(0, 1, 2), keepdims=True)\n",
    "    X_train = (X_train - mean_pixel) / std_pixel\n",
    "    X_val = (X_val - mean_pixel) / std_pixel\n",
    "    X_test = (X_test - mean_pixel) / std_pixel\n",
    "\n",
    "    return X_train, y_train, X_val, y_val, X_test, y_test\n",
    "\n",
    "\n",
    "# Invoke the above function to get our data.\n",
    "NHW = (0, 1, 2)\n",
    "X_train, y_train, X_val, y_val, X_test, y_test = load_cifar10()\n",
    "print('Train data shape: ', X_train.shape)\n",
    "print('Train labels shape: ', y_train.shape, y_train.dtype)\n",
    "print('Validation data shape: ', X_val.shape)\n",
    "print('Validation labels shape: ', y_val.shape)\n",
    "print('Test data shape: ', X_test.shape)\n",
    "print('Test labels shape: ', y_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparation: Dataset object\n",
    "\n",
    "For our own convenience we'll define a lightweight `Dataset` class which lets us iterate over data and labels. This is not the most flexible or most efficient way to iterate through data, but it will serve our purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(object):\n",
    "    def __init__(self, X, y, batch_size, shuffle=False):\n",
    "        \"\"\"\n",
    "        Construct a Dataset object to iterate over data X and labels y\n",
    "        \n",
    "        Inputs:\n",
    "        - X: Numpy array of data, of any shape\n",
    "        - y: Numpy array of labels, of any shape but with y.shape[0] == X.shape[0]\n",
    "        - batch_size: Integer giving number of elements per minibatch\n",
    "        - shuffle: (optional) Boolean, whether to shuffle the data on each epoch\n",
    "        \"\"\"\n",
    "        assert X.shape[0] == y.shape[0], 'Got different numbers of data and labels'\n",
    "        self.X, self.y = X, y\n",
    "        self.batch_size, self.shuffle = batch_size, shuffle\n",
    "\n",
    "    def __iter__(self):\n",
    "        N, B = self.X.shape[0], self.batch_size\n",
    "        idxs = np.arange(N)\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(idxs)\n",
    "        return iter((self.X[i:i+B], self.y[i:i+B]) for i in range(0, N, B))\n",
    "\n",
    "\n",
    "train_dset = Dataset(X_train, y_train, batch_size=64, shuffle=True)\n",
    "val_dset = Dataset(X_val, y_val, batch_size=64, shuffle=False)\n",
    "test_dset = Dataset(X_test, y_test, batch_size=64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 (64, 32, 32, 3) (64,)\n",
      "1 (64, 32, 32, 3) (64,)\n",
      "2 (64, 32, 32, 3) (64,)\n",
      "3 (64, 32, 32, 3) (64,)\n",
      "4 (64, 32, 32, 3) (64,)\n",
      "5 (64, 32, 32, 3) (64,)\n",
      "6 (64, 32, 32, 3) (64,)\n"
     ]
    }
   ],
   "source": [
    "# We can iterate through a dataset like this:\n",
    "for t, (x, y) in enumerate(train_dset):\n",
    "    print(t, x.shape, y.shape)\n",
    "    if t > 5: break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can optionally **use GPU by setting the flag to True below**. It's not neccessary to use a GPU for this assignment; if you are working on Google Cloud then we recommend that you do not use a GPU, as it will be significantly more expensive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device:  /device:GPU:0\n"
     ]
    }
   ],
   "source": [
    "# Set up some global variables\n",
    "USE_GPU = True\n",
    "\n",
    "if USE_GPU:\n",
    "    device = '/device:GPU:0'\n",
    "else:\n",
    "    device = '/cpu:0'\n",
    "\n",
    "# Constant to control how often we print when training models\n",
    "print_every = 100\n",
    "\n",
    "print('Using device: ', device)\n",
    "\n",
    "# Limit GPU memory allocation\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part II: Barebone TensorFlow\n",
    "TensorFlow ships with various high-level APIs which make it very convenient to define and train neural networks; we will cover some of these constructs in Part III and Part IV of this notebook. In this section we will start by building a model with basic TensorFlow constructs to help you better understand what's going on under the hood of the higher-level APIs.\n",
    "\n",
    "TensorFlow is primarily a framework for working with **static computational graphs**. Nodes in the computational graph are Tensors which will hold n-dimensional arrays when the graph is run; edges in the graph represent functions that will operate on Tensors when the graph is run to actually perform useful computation.\n",
    "\n",
    "This means that a typical TensorFlow program is written in two distinct phases:\n",
    "\n",
    "1. Build a computational graph that describes the computation that you want to perform. This stage doesn't actually perform any computation; it just builds up a symbolic representation of your computation. This stage will typically define one or more `placeholder` objects that represent inputs to the computational graph.\n",
    "2. Run the computational graph many times. Each time the graph is run you will specify which parts of the graph you want to compute, and pass a `feed_dict` dictionary that will give concrete values to any `placeholder`s in the graph.\n",
    "\n",
    "### TensorFlow warmup: Flatten Function\n",
    "\n",
    "We can see this in action by defining a simple `flatten` function that will reshape image data for use in a fully-connected network.\n",
    "\n",
    "In TensorFlow, data for convolutional feature maps is typically stored in a Tensor of shape N x H x W x C where:\n",
    "\n",
    "- N is the number of datapoints (minibatch size)\n",
    "- H is the height of the feature map\n",
    "- W is the width of the feature map\n",
    "- C is the number of channels in the feature map\n",
    "\n",
    "This is the right way to represent the data when we are doing something like a 2D convolution, that needs spatial understanding of where the intermediate features are relative to each other. When we use fully connected affine layers to process the image, however, we want each datapoint to be represented by a single vector -- it's no longer useful to segregate the different channels, rows, and columns of the data. So, we use a \"flatten\" operation to collapse the `H x W x C` values per representation into a single long vector. The flatten function below first reads in the value of N from a given batch of data, and then returns a \"view\" of that data. \"View\" is analogous to numpy's \"reshape\" method: it reshapes x's dimensions to be N x ??, where ?? is allowed to be anything (in this case, it will be H x W x C, but we don't need to specify that explicitly). \n",
    "\n",
    "**NOTE**: TensorFlow and PyTorch differ on the default Tensor layout; TensorFlow uses N x H x W x C but PyTorch uses N x C x H x W."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten(x):\n",
    "    \"\"\"    \n",
    "    Input:\n",
    "    - TensorFlow Tensor of shape (N, D1, ..., DM)\n",
    "    \n",
    "    Output:\n",
    "    - TensorFlow Tensor of shape (N, D1 * ... * DM)\n",
    "    \"\"\"\n",
    "    N = tf.shape(x)[0]\n",
    "    return tf.reshape(x, (N, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  <class 'tensorflow.python.framework.ops.Tensor'> Tensor(\"Placeholder:0\", dtype=float32, device=/device:GPU:0)\n",
      "x_flat:  <class 'tensorflow.python.framework.ops.Tensor'> Tensor(\"Reshape:0\", shape=(?, ?), dtype=float32, device=/device:GPU:0)\n",
      "\n",
      "x_np:\n",
      " [[[ 0  1  2  3]\n",
      "  [ 4  5  6  7]\n",
      "  [ 8  9 10 11]]\n",
      "\n",
      " [[12 13 14 15]\n",
      "  [16 17 18 19]\n",
      "  [20 21 22 23]]] \n",
      "\n",
      "x_flat_np:\n",
      " [[ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11.]\n",
      " [12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23.]] \n",
      "\n",
      "x_np:\n",
      " [[[ 0  1]\n",
      "  [ 2  3]\n",
      "  [ 4  5]]\n",
      "\n",
      " [[ 6  7]\n",
      "  [ 8  9]\n",
      "  [10 11]]] \n",
      "\n",
      "x_flat_np:\n",
      " [[ 0.  1.  2.  3.  4.  5.]\n",
      " [ 6.  7.  8.  9. 10. 11.]]\n"
     ]
    }
   ],
   "source": [
    "def test_flatten():\n",
    "    # Clear the current TensorFlow graph.\n",
    "    tf.reset_default_graph()\n",
    "    \n",
    "    # Stage I: Define the TensorFlow graph describing our computation.\n",
    "    # In this case the computation is trivial: we just want to flatten\n",
    "    # a Tensor using the flatten function defined above.\n",
    "    \n",
    "    # Our computation will have a single input, x. We don't know its\n",
    "    # value yet, so we define a placeholder which will hold the value\n",
    "    # when the graph is run. We then pass this placeholder Tensor to\n",
    "    # the flatten function; this gives us a new Tensor which will hold\n",
    "    # a flattened view of x when the graph is run. The tf.device\n",
    "    # context manager tells TensorFlow whether to place these Tensors\n",
    "    # on CPU or GPU.\n",
    "    with tf.device(device):\n",
    "        x = tf.placeholder(tf.float32)\n",
    "        x_flat = flatten(x)\n",
    "    \n",
    "    # At this point we have just built the graph describing our computation,\n",
    "    # but we haven't actually computed anything yet. If we print x and x_flat\n",
    "    # we see that they don't hold any data; they are just TensorFlow Tensors\n",
    "    # representing values that will be computed when the graph is run.\n",
    "    print('x: ', type(x), x)\n",
    "    print('x_flat: ', type(x_flat), x_flat)\n",
    "    print()\n",
    "    \n",
    "    # We need to use a TensorFlow Session object to actually run the graph.\n",
    "    with tf.Session(config=config) as sess:\n",
    "        # Construct concrete values of the input data x using numpy\n",
    "        x_np = np.arange(24).reshape((2, 3, 4))\n",
    "        print('x_np:\\n', x_np, '\\n')\n",
    "    \n",
    "        # Run our computational graph to compute a concrete output value.\n",
    "        # The first argument to sess.run tells TensorFlow which Tensor\n",
    "        # we want it to compute the value of; the feed_dict specifies\n",
    "        # values to plug into all placeholder nodes in the graph. The\n",
    "        # resulting value of x_flat is returned from sess.run as a\n",
    "        # numpy array.\n",
    "        x_flat_np = sess.run(x_flat, feed_dict={x: x_np})\n",
    "        print('x_flat_np:\\n', x_flat_np, '\\n')\n",
    "\n",
    "        # We can reuse the same graph to perform the same computation\n",
    "        # with different input data\n",
    "        x_np = np.arange(12).reshape((2, 3, 2))\n",
    "        print('x_np:\\n', x_np, '\\n')\n",
    "        x_flat_np = sess.run(x_flat, feed_dict={x: x_np})\n",
    "        print('x_flat_np:\\n', x_flat_np)\n",
    "test_flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Two-Layer Network\n",
    "We will now implement our first neural network with TensorFlow: a fully-connected ReLU network with two hidden layers and no biases on the CIFAR10 dataset. For now we will use only low-level TensorFlow operators to define the network; later we will see how to use the higher-level abstractions provided by `tf.keras` to simplify the process.\n",
    "\n",
    "We will define the forward pass of the network in the function `two_layer_fc`; this will accept TensorFlow Tensors for the inputs and weights of the network, and return a TensorFlow Tensor for the scores. It's important to keep in mind that calling the `two_layer_fc` function **does not** perform any computation; instead it just sets up the computational graph for the forward computation. To actually run the network we need to enter a TensorFlow Session and feed data to the computational graph.\n",
    "\n",
    "After defining the network architecture in the `two_layer_fc` function, we will test the implementation by setting up and running a computational graph, feeding zeros to the network and checking the shape of the output.\n",
    "\n",
    "It's important that you read and understand this implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def two_layer_fc(x, params):\n",
    "    \"\"\"\n",
    "    A fully-connected neural network; the architecture is:\n",
    "    fully-connected layer -> ReLU -> fully connected layer.\n",
    "    Note that we only need to define the forward pass here; TensorFlow will take\n",
    "    care of computing the gradients for us.\n",
    "    \n",
    "    The input to the network will be a minibatch of data, of shape\n",
    "    (N, d1, ..., dM) where d1 * ... * dM = D. The hidden layer will have H units,\n",
    "    and the output layer will produce scores for C classes.\n",
    "\n",
    "    Inputs:\n",
    "    - x: A TensorFlow Tensor of shape (N, d1, ..., dM) giving a minibatch of\n",
    "      input data.\n",
    "    - params: A list [w1, w2] of TensorFlow Tensors giving weights for the\n",
    "      network, where w1 has shape (D, H) and w2 has shape (H, C).\n",
    "    \n",
    "    Returns:\n",
    "    - scores: A TensorFlow Tensor of shape (N, C) giving classification scores\n",
    "      for the input data x.\n",
    "    \"\"\"\n",
    "    w1, w2 = params  # Unpack the parameters\n",
    "    x = flatten(x)   # Flatten the input; now x has shape (N, D)\n",
    "    h = tf.nn.relu(tf.matmul(x, w1)) # Hidden layer: h has shape (N, H)\n",
    "    scores = tf.matmul(h, w2)        # Compute scores of shape (N, C)\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 10)\n"
     ]
    }
   ],
   "source": [
    "def two_layer_fc_test():\n",
    "    # TensorFlow's default computational graph is essentially a hidden global\n",
    "    # variable. To avoid adding to this default graph when you rerun this cell,\n",
    "    # we clear the default graph before constructing the graph we care about.\n",
    "    tf.reset_default_graph()\n",
    "    hidden_layer_size = 42\n",
    "\n",
    "    # Scoping our computational graph setup code under a tf.device context\n",
    "    # manager lets us tell TensorFlow where we want these Tensors to be\n",
    "    # placed.\n",
    "    with tf.device(device):\n",
    "        # Set up a placehoder for the input of the network, and constant\n",
    "        # zero Tensors for the network weights. Here we declare w1 and w2\n",
    "        # using tf.zeros instead of tf.placeholder as we've seen before - this\n",
    "        # means that the values of w1 and w2 will be stored in the computational\n",
    "        # graph itself and will persist across multiple runs of the graph; in\n",
    "        # particular this means that we don't have to pass values for w1 and w2\n",
    "        # using a feed_dict when we eventually run the graph.\n",
    "        x = tf.placeholder(tf.float32)\n",
    "        w1 = tf.zeros((32 * 32 * 3, hidden_layer_size))\n",
    "        w2 = tf.zeros((hidden_layer_size, 10))\n",
    "        \n",
    "        # Call our two_layer_fc function to set up the computational\n",
    "        # graph for the forward pass of the network.\n",
    "        scores = two_layer_fc(x, [w1, w2])\n",
    "    \n",
    "    # Use numpy to create some concrete data that we will pass to the\n",
    "    # computational graph for the x placeholder.\n",
    "    x_np = np.zeros((64, 32, 32, 3))\n",
    "    with tf.Session(config=config) as sess:\n",
    "        # The calls to tf.zeros above do not actually instantiate the values\n",
    "        # for w1 and w2; the following line tells TensorFlow to instantiate\n",
    "        # the values of all Tensors (like w1 and w2) that live in the graph.\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        \n",
    "        # Here we actually run the graph, using the feed_dict to pass the\n",
    "        # value to bind to the placeholder for x; we ask TensorFlow to compute\n",
    "        # the value of the scores Tensor, which it returns as a numpy array.\n",
    "        scores_np = sess.run(scores, feed_dict={x: x_np})\n",
    "        print(scores_np.shape)\n",
    "\n",
    "two_layer_fc_test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Three-Layer ConvNet\n",
    "Here you will complete the implementation of the function `three_layer_convnet` which will perform the forward pass of a three-layer convolutional network. The network should have the following architecture:\n",
    "\n",
    "1. A convolutional layer (with bias) with `channel_1` filters, each with shape `KW1 x KH1`, and zero-padding of two\n",
    "2. ReLU nonlinearity\n",
    "3. A convolutional layer (with bias) with `channel_2` filters, each with shape `KW2 x KH2`, and zero-padding of one\n",
    "4. ReLU nonlinearity\n",
    "5. Fully-connected layer with bias, producing scores for `C` classes.\n",
    "\n",
    "**HINT**: For convolutions: https://www.tensorflow.org/api_docs/python/tf/nn/conv2d; be careful with padding!\n",
    "\n",
    "**HINT**: For biases: https://www.tensorflow.org/performance/xla/broadcasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def three_layer_convnet(x, params):\n",
    "    \"\"\"\n",
    "    A three-layer convolutional network with the architecture described above.\n",
    "    \n",
    "    Inputs:\n",
    "    - x: A TensorFlow Tensor of shape (N, H, W, 3) giving a minibatch of images\n",
    "    - params: A list of TensorFlow Tensors giving the weights and biases for the\n",
    "      network; should contain the following:\n",
    "      - conv_w1: TensorFlow Tensor of shape (KH1, KW1, 3, channel_1) giving\n",
    "        weights for the first convolutional layer.\n",
    "      - conv_b1: TensorFlow Tensor of shape (channel_1,) giving biases for the\n",
    "        first convolutional layer.\n",
    "      - conv_w2: TensorFlow Tensor of shape (KH2, KW2, channel_1, channel_2)\n",
    "        giving weights for the second convolutional layer\n",
    "      - conv_b2: TensorFlow Tensor of shape (channel_2,) giving biases for the\n",
    "        second convolutional layer.\n",
    "      - fc_w: TensorFlow Tensor giving weights for the fully-connected layer.\n",
    "        Can you figure out what the shape should be?\n",
    "      - fc_b: TensorFlow Tensor giving biases for the fully-connected layer.\n",
    "        Can you figure out what the shape should be?\n",
    "    \"\"\"\n",
    "    conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b = params\n",
    "    scores = None\n",
    "    ############################################################################\n",
    "    # TODO: Implement the forward pass for the three-layer ConvNet.            #\n",
    "    ############################################################################\n",
    "    x = tf.nn.conv2d(x, conv_w1, strides=[1, 1, 1, 1], padding='SAME') + conv_b1\n",
    "    x = tf.nn.relu(x)\n",
    "    x = tf.nn.conv2d(x, conv_w2, strides=[1, 1, 1, 1], padding='SAME') + conv_b2\n",
    "    x = tf.nn.relu(x)\n",
    "    scores = tf.matmul(flatten(x), fc_w) + fc_b\n",
    "    ############################################################################\n",
    "    #                              END OF YOUR CODE                            #\n",
    "    ############################################################################\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After defing the forward pass of the three-layer ConvNet above, run the following cell to test your implementation. Like the two-layer network, we use the `three_layer_convnet` function to set up the computational graph, then run the graph on a batch of zeros just to make sure the function doesn't crash, and produces outputs of the correct shape.\n",
    "\n",
    "When you run this function, `scores_np` should have shape `(64, 10)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scores_np has shape:  (64, 10)\n"
     ]
    }
   ],
   "source": [
    "def three_layer_convnet_test():\n",
    "    tf.reset_default_graph()\n",
    "\n",
    "    with tf.device(device):\n",
    "        x = tf.placeholder(tf.float32)\n",
    "        conv_w1 = tf.zeros((5, 5, 3, 6))\n",
    "        conv_b1 = tf.zeros((6,))\n",
    "        conv_w2 = tf.zeros((3, 3, 6, 9))\n",
    "        conv_b2 = tf.zeros((9,))\n",
    "        fc_w = tf.zeros((32 * 32 * 9, 10))\n",
    "        fc_b = tf.zeros((10,))\n",
    "        params = [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b]\n",
    "        scores = three_layer_convnet(x, params)\n",
    "\n",
    "    # Inputs to convolutional layers are 4-dimensional arrays with shape\n",
    "    # [batch_size, height, width, channels]\n",
    "    x_np = np.zeros((64, 32, 32, 3))\n",
    "    \n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        scores_np = sess.run(scores, feed_dict={x: x_np})\n",
    "        print('scores_np has shape: ', scores_np.shape)\n",
    "\n",
    "with tf.device('/cpu:0'):\n",
    "    three_layer_convnet_test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Training Step\n",
    "We now define the `training_step` function which sets up the part of the computational graph that performs a single training step. This will take three basic steps:\n",
    "\n",
    "1. Compute the loss\n",
    "2. Compute the gradient of the loss with respect to all network weights\n",
    "3. Make a weight update step using (stochastic) gradient descent.\n",
    "\n",
    "Note that the step of updating the weights is itself an operation in the computational graph - the calls to `tf.assign_sub` in `training_step` return TensorFlow operations that mutate the weights when they are executed. There is an important bit of subtlety here - when we call `sess.run`, TensorFlow does not execute all operations in the computational graph; it only executes the minimal subset of the graph necessary to compute the outputs that we ask TensorFlow to produce. As a result, naively computing the loss would not cause the weight update operations to execute, since the operations needed to compute the loss do not depend on the output of the weight update. To fix this problem, we insert a **control dependency** into the graph, adding a duplicate `loss` node to the graph that does depend on the outputs of the weight update operations; this is the object that we actually return from the `training_step` function. As a result, asking TensorFlow to evaluate the value of the `loss` returned from `training_step` will also implicitly update the weights of the network using that minibatch of data.\n",
    "\n",
    "We need to use a few new TensorFlow functions to do all of this:\n",
    "- For computing the cross-entropy loss we'll use `tf.nn.sparse_softmax_cross_entropy_with_logits`: https://www.tensorflow.org/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits\n",
    "- For averaging the loss across a minibatch of data we'll use `tf.reduce_mean`:\n",
    "https://www.tensorflow.org/api_docs/python/tf/reduce_mean\n",
    "- For computing gradients of the loss with respect to the weights we'll use `tf.gradients`:  https://www.tensorflow.org/api_docs/python/tf/gradients\n",
    "- We'll mutate the weight values stored in a TensorFlow Tensor using `tf.assign_sub`: https://www.tensorflow.org/api_docs/python/tf/assign_sub\n",
    "- We'll add a control dependency to the graph using `tf.control_dependencies`: https://www.tensorflow.org/api_docs/python/tf/control_dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_step(scores, y, params, learning_rate):\n",
    "    \"\"\"\n",
    "    Set up the part of the computational graph which makes a training step.\n",
    "\n",
    "    Inputs:\n",
    "    - scores: TensorFlow Tensor of shape (N, C) giving classification scores for\n",
    "      the model.\n",
    "    - y: TensorFlow Tensor of shape (N,) giving ground-truth labels for scores;\n",
    "      y[i] == c means that c is the correct class for scores[i].\n",
    "    - params: List of TensorFlow Tensors giving the weights of the model\n",
    "    - learning_rate: Python scalar giving the learning rate to use for gradient\n",
    "      descent step.\n",
    "      \n",
    "    Returns:\n",
    "    - loss: A TensorFlow Tensor of shape () (scalar) giving the loss for this\n",
    "      batch of data; evaluating the loss also performs a gradient descent step\n",
    "      on params (see above).\n",
    "    \"\"\"\n",
    "    # First compute the loss; the first line gives losses for each example in\n",
    "    # the minibatch, and the second averages the losses acros the batch\n",
    "    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=scores)\n",
    "    loss = tf.reduce_mean(losses)\n",
    "\n",
    "    # Compute the gradient of the loss with respect to each parameter of the the\n",
    "    # network. This is a very magical function call: TensorFlow internally\n",
    "    # traverses the computational graph starting at loss backward to each element\n",
    "    # of params, and uses backpropagation to figure out how to compute gradients;\n",
    "    # it then adds new operations to the computational graph which compute the\n",
    "    # requested gradients, and returns a list of TensorFlow Tensors that will\n",
    "    # contain the requested gradients when evaluated.\n",
    "    grad_params = tf.gradients(loss, params)\n",
    "    \n",
    "    # Make a gradient descent step on all of the model parameters.\n",
    "    new_weights = []   \n",
    "    for w, grad_w in zip(params, grad_params):\n",
    "        new_w = tf.assign_sub(w, learning_rate * grad_w)\n",
    "        new_weights.append(new_w)\n",
    "\n",
    "    # Insert a control dependency so that evaluting the loss causes a weight\n",
    "    # update to happen; see the discussion above.\n",
    "    with tf.control_dependencies(new_weights):\n",
    "        return tf.identity(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Training Loop\n",
    "Now we set up a basic training loop using low-level TensorFlow operations. We will train the model using stochastic gradient descent without momentum. The `training_step` function sets up the part of the computational graph that performs the training step, and the function `train_part2` iterates through the training data, making training steps on each minibatch, and periodically evaluates accuracy on the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_part2(model_fn, init_fn, learning_rate):\n",
    "    \"\"\"\n",
    "    Train a model on CIFAR-10.\n",
    "    \n",
    "    Inputs:\n",
    "    - model_fn: A Python function that performs the forward pass of the model\n",
    "      using TensorFlow; it should have the following signature:\n",
    "      scores = model_fn(x, params) where x is a TensorFlow Tensor giving a\n",
    "      minibatch of image data, params is a list of TensorFlow Tensors holding\n",
    "      the model weights, and scores is a TensorFlow Tensor of shape (N, C)\n",
    "      giving scores for all elements of x.\n",
    "    - init_fn: A Python function that initializes the parameters of the model.\n",
    "      It should have the signature params = init_fn() where params is a list\n",
    "      of TensorFlow Tensors holding the (randomly initialized) weights of the\n",
    "      model.\n",
    "    - learning_rate: Python float giving the learning rate to use for SGD.\n",
    "    \"\"\"\n",
    "    # First clear the default graph\n",
    "    tf.reset_default_graph()\n",
    "    is_training = tf.placeholder(tf.bool, name='is_training')\n",
    "    # Set up the computational graph for performing forward and backward passes,\n",
    "    # and weight updates.\n",
    "    with tf.device(device):\n",
    "        # Set up placeholders for the data and labels\n",
    "        x = tf.placeholder(tf.float32, [None, 32, 32, 3])\n",
    "        y = tf.placeholder(tf.int32, [None])\n",
    "        params = init_fn()           # Initialize the model parameters\n",
    "        scores = model_fn(x, params) # Forward pass of the model\n",
    "        loss = training_step(scores, y, params, learning_rate)\n",
    "\n",
    "    # Now we actually run the graph many times using the training data\n",
    "    with tf.Session(config=config) as sess:\n",
    "        # Initialize variables that will live in the graph\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        for t, (x_np, y_np) in enumerate(train_dset):\n",
    "            # Run the graph on a batch of training data; recall that asking\n",
    "            # TensorFlow to evaluate loss will cause an SGD step to happen.\n",
    "            feed_dict = {x: x_np, y: y_np}\n",
    "            loss_np = sess.run(loss, feed_dict=feed_dict)\n",
    "            \n",
    "            # Periodically print the loss and check accuracy on the val set\n",
    "            if t % print_every == 0:\n",
    "                print('Iteration %d, loss = %.4f' % (t, loss_np))\n",
    "                check_accuracy(sess, val_dset, x, scores, is_training)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Check Accuracy\n",
    "When training the model we will use the following function to check the accuracy of our model on the training or validation sets. Note that this function accepts a TensorFlow Session object as one of its arguments; this is needed since the function must actually run the computational graph many times on the data that it loads from the dataset `dset`.\n",
    "\n",
    "Also note that we reuse the same computational graph both for taking training steps and for evaluating the model; however since the `check_accuracy` function never evalutes the `loss` value in the computational graph, the part of the graph that updates the weights of the graph do not execute on the validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_accuracy(sess, dset, x, scores, is_training=None):\n",
    "    \"\"\"\n",
    "    Check accuracy on a classification model.\n",
    "    \n",
    "    Inputs:\n",
    "    - sess: A TensorFlow Session that will be used to run the graph\n",
    "    - dset: A Dataset object on which to check accuracy\n",
    "    - x: A TensorFlow placeholder Tensor where input images should be fed\n",
    "    - scores: A TensorFlow Tensor representing the scores output from the\n",
    "      model; this is the Tensor we will ask TensorFlow to evaluate.\n",
    "      \n",
    "    Returns: Nothing, but prints the accuracy of the model\n",
    "    \"\"\"\n",
    "    num_correct, num_samples = 0, 0\n",
    "    for x_batch, y_batch in dset:\n",
    "        feed_dict = {x: x_batch, is_training: 0}\n",
    "        scores_np = sess.run(scores, feed_dict=feed_dict)\n",
    "        y_pred = scores_np.argmax(axis=1)\n",
    "        num_samples += x_batch.shape[0]\n",
    "        num_correct += (y_pred == y_batch).sum()\n",
    "    acc = float(num_correct) / num_samples\n",
    "    print('Got %d / %d correct (%.2f%%)' % (num_correct, num_samples, 100 * acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Initialization\n",
    "We'll use the following utility method to initialize the weight matrices for our models using Kaiming's normalization method.\n",
    "\n",
    "[1] He et al, *Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification\n",
    "*, ICCV 2015, https://arxiv.org/abs/1502.01852"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kaiming_normal(shape):\n",
    "    if len(shape) == 2:\n",
    "        fan_in, fan_out = shape[0], shape[1]\n",
    "    elif len(shape) == 4:\n",
    "        fan_in, fan_out = np.prod(shape[:3]), shape[3]\n",
    "    return tf.random_normal(shape) * np.sqrt(2.0 / fan_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Train a Two-Layer Network\n",
    "We are finally ready to use all of the pieces defined above to train a two-layer fully-connected network on CIFAR-10.\n",
    "\n",
    "We just need to define a function to initialize the weights of the model, and call `train_part2`.\n",
    "\n",
    "Defining the weights of the network introduces another important piece of TensorFlow API: `tf.Variable`. A TensorFlow Variable is a Tensor whose value is stored in the graph and persists across runs of the computational graph; however unlike constants defined with `tf.zeros` or `tf.random_normal`, the values of a Variable can be mutated as the graph runs; these mutations will persist across graph runs. Learnable parameters of the network are usually stored in Variables.\n",
    "\n",
    "You don't need to tune any hyperparameters, but you should achieve accuracies above 40% after one epoch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 2.9162\n",
      "Got 132 / 1000 correct (13.20%)\n",
      "Iteration 100, loss = 1.8897\n",
      "Got 390 / 1000 correct (39.00%)\n",
      "Iteration 200, loss = 1.4826\n",
      "Got 394 / 1000 correct (39.40%)\n",
      "Iteration 300, loss = 1.7932\n",
      "Got 380 / 1000 correct (38.00%)\n",
      "Iteration 400, loss = 1.7440\n",
      "Got 423 / 1000 correct (42.30%)\n",
      "Iteration 500, loss = 1.8032\n",
      "Got 422 / 1000 correct (42.20%)\n",
      "Iteration 600, loss = 1.8155\n",
      "Got 422 / 1000 correct (42.20%)\n",
      "Iteration 700, loss = 1.9840\n",
      "Got 450 / 1000 correct (45.00%)\n"
     ]
    }
   ],
   "source": [
    "def two_layer_fc_init():\n",
    "    \"\"\"\n",
    "    Initialize the weights of a two-layer network, for use with the\n",
    "    two_layer_network function defined above.\n",
    "    \n",
    "    Inputs: None\n",
    "    \n",
    "    Returns: A list of:\n",
    "    - w1: TensorFlow Variable giving the weights for the first layer\n",
    "    - w2: TensorFlow Variable giving the weights for the second layer\n",
    "    \"\"\"\n",
    "    hidden_layer_size = 4000\n",
    "    w1 = tf.Variable(kaiming_normal((3 * 32 * 32, 4000)))\n",
    "    w2 = tf.Variable(kaiming_normal((4000, 10)))\n",
    "    return [w1, w2]\n",
    "\n",
    "learning_rate = 1e-2\n",
    "train_part2(two_layer_fc, two_layer_fc_init, learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones TensorFlow: Train a three-layer ConvNet\n",
    "We will now use TensorFlow to train a three-layer ConvNet on CIFAR-10.\n",
    "\n",
    "You need to implement the `three_layer_convnet_init` function. Recall that the architecture of the network is:\n",
    "\n",
    "1. Convolutional layer (with bias) with 32 5x5 filters, with zero-padding 2\n",
    "2. ReLU\n",
    "3. Convolutional layer (with bias) with 16 3x3 filters, with zero-padding 1\n",
    "4. ReLU\n",
    "5. Fully-connected layer (with bias) to compute scores for 10 classes\n",
    "\n",
    "You don't need to do any hyperparameter tuning, but you should see accuracies above 43% after one epoch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 2.6499\n",
      "Got 90 / 1000 correct (9.00%)\n",
      "Iteration 100, loss = 1.8992\n",
      "Got 366 / 1000 correct (36.60%)\n",
      "Iteration 200, loss = 1.5991\n",
      "Got 389 / 1000 correct (38.90%)\n",
      "Iteration 300, loss = 1.7262\n",
      "Got 383 / 1000 correct (38.30%)\n",
      "Iteration 400, loss = 1.7509\n",
      "Got 418 / 1000 correct (41.80%)\n",
      "Iteration 500, loss = 1.7163\n",
      "Got 449 / 1000 correct (44.90%)\n",
      "Iteration 600, loss = 1.6371\n",
      "Got 463 / 1000 correct (46.30%)\n",
      "Iteration 700, loss = 1.5887\n",
      "Got 456 / 1000 correct (45.60%)\n"
     ]
    }
   ],
   "source": [
    "def three_layer_convnet_init():\n",
    "    \"\"\"\n",
    "    Initialize the weights of a Three-Layer ConvNet, for use with the\n",
    "    three_layer_convnet function defined above.\n",
    "    \n",
    "    Inputs: None\n",
    "    \n",
    "    Returns a list containing:\n",
    "    - conv_w1: TensorFlow Variable giving weights for the first conv layer\n",
    "    - conv_b1: TensorFlow Variable giving biases for the first conv layer\n",
    "    - conv_w2: TensorFlow Variable giving weights for the second conv layer\n",
    "    - conv_b2: TensorFlow Variable giving biases for the second conv layer\n",
    "    - fc_w: TensorFlow Variable giving weights for the fully-connected layer\n",
    "    - fc_b: TensorFlow Variable giving biases for the fully-connected layer\n",
    "    \"\"\"\n",
    "    # params = None\n",
    "    ############################################################################\n",
    "    # TODO: Initialize the parameters of the three-layer network.              #\n",
    "    ############################################################################\n",
    "    conv_w1 = tf.Variable(kaiming_normal((5, 5, 3, 32)))\n",
    "    conv_b1 = tf.Variable(tf.zeros((32, )))\n",
    "    conv_w2 = tf.Variable(kaiming_normal((3, 3, 32, 16)))\n",
    "    conv_b2 = tf.Variable(tf.zeros((16, )))\n",
    "    fc_w = tf.Variable(kaiming_normal((32 * 32 * 16, 10)))\n",
    "    fc_b = tf.Variable(tf.zeros((10, )))\n",
    "    params = [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b]\n",
    "    ############################################################################\n",
    "    #                             END OF YOUR CODE                             #\n",
    "    ############################################################################\n",
    "    return params\n",
    "\n",
    "learning_rate = 3e-3\n",
    "train_part2(three_layer_convnet, three_layer_convnet_init, learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part III: Keras Model API\n",
    "Implementing a neural network using the low-level TensorFlow API is a good way to understand how TensorFlow works, but it's a little inconvenient - we had to manually keep track of all Tensors holding learnable parameters, and we had to use a control dependency to implement the gradient descent update step. This was fine for a small network, but could quickly become unweildy for a large complex model.\n",
    "\n",
    "Fortunately TensorFlow provides higher-level packages such as `tf.keras` and `tf.layers` which make it easy to build models out of modular, object-oriented layers; `tf.train` allows you to easily train these models using a variety of different optimization algorithms.\n",
    "\n",
    "In this part of the notebook we will define neural network models using the `tf.keras.Model` API. To implement your own model, you need to do the following:\n",
    "\n",
    "1. Define a new class which subclasses `tf.keras.model`. Give your class an intuitive name that describes it, like `TwoLayerFC` or `ThreeLayerConvNet`.\n",
    "2. In the initializer `__init__()` for your new class, define all the layers you need as class attributes. The `tf.layers` package provides many common neural-network layers, like `tf.layers.Dense` for fully-connected layers and `tf.layers.Conv2D` for convolutional layers. Under the hood, these layers will construct `Variable` Tensors for any learnable parameters. **Warning**: Don't forget to call `super().__init__()` as the first line in your initializer!\n",
    "3. Implement the `call()` method for your class; this implements the forward pass of your model, and defines the *connectivity* of your network. Layers defined in `__init__()` implement `__call__()` so they can be used as function objects that transform input Tensors into output Tensors. Don't define any new layers in `call()`; any layers you want to use in the forward pass should be defined in `__init__()`.\n",
    "\n",
    "After you define your `tf.keras.Model` subclass, you can instantiate it and use it like the model functions from Part II.\n",
    "\n",
    "### Module API: Two-Layer Network\n",
    "\n",
    "Here is a concrete example of using the `tf.keras.Model` API to define a two-layer network. There are a few new bits of API to be aware of here:\n",
    "\n",
    "We use an `Initializer` object to set up the initial values of the learnable parameters of the layers; in particular `tf.variance_scaling_initializer` gives behavior similar to the Kaiming initialization method we used in Part II. You can read more about it here: https://www.tensorflow.org/api_docs/python/tf/variance_scaling_initializer\n",
    "\n",
    "We construct `tf.layers.Dense` objects to represent the two fully-connected layers of the model. In addition to multiplying their input by a weight matrix and adding a bias vector, these layer can also apply a nonlinearity for you. For the first layer we specify a ReLU activation function by passing `activation=tf.nn.relu` to the constructor; the second layer does not apply any activation function.\n",
    "\n",
    "Unfortunately the `flatten` function we defined in Part II is not compatible with the `tf.keras.Model` API; fortunately we can use `tf.layers.flatten` to perform the same operation. The issue with our `flatten` function from Part II has to do with static vs dynamic shapes for Tensors, which is beyond the scope of this notebook; you can read more about the distinction [in the documentation](https://www.tensorflow.org/programmers_guide/faq#tensor_shapes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 10)\n"
     ]
    }
   ],
   "source": [
    "class TwoLayerFC(tf.keras.Model):\n",
    "    def __init__(self, hidden_size, num_classes):\n",
    "        super().__init__()        \n",
    "        initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "        self.fc1 = tf.layers.Dense(hidden_size, activation=tf.nn.relu,\n",
    "                                   kernel_initializer=initializer)\n",
    "        self.fc2 = tf.layers.Dense(num_classes,\n",
    "                                   kernel_initializer=initializer)\n",
    "    def call(self, x, training=None):\n",
    "        x = tf.layers.flatten(x)\n",
    "        x = self.fc1(x)\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "def test_TwoLayerFC():\n",
    "    \"\"\" A small unit test to exercise the TwoLayerFC model above. \"\"\"\n",
    "    tf.reset_default_graph()\n",
    "    input_size, hidden_size, num_classes = 50, 42, 10\n",
    "\n",
    "    # As usual in TensorFlow, we first need to define our computational graph.\n",
    "    # To this end we first construct a TwoLayerFC object, then use it to construct\n",
    "    # the scores Tensor.\n",
    "    model = TwoLayerFC(hidden_size, num_classes)\n",
    "    with tf.device(device):\n",
    "        x = tf.zeros((64, input_size))\n",
    "        scores = model(x)\n",
    "\n",
    "    # Now that our computational graph has been defined we can run the graph\n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        scores_np = sess.run(scores)\n",
    "        print(scores_np.shape)\n",
    "        \n",
    "test_TwoLayerFC()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Funtional API: Two-Layer Network\n",
    "The `tf.layers` package provides two different higher-level APIs for defining neural network models. In the example above we used the **object-oriented API**, where each layer of the neural network is represented as a Python object (like `tf.layers.Dense`). Here we showcase the **functional API**, where each layer is a Python function (like `tf.layers.dense`) which inputs and outputs TensorFlow Tensors, and which internally sets up Tensors in the computational graph to hold any learnable weights.\n",
    "\n",
    "To construct a network, one needs to pass the input tensor to the first layer, and construct the subsequent layers sequentially. Here's an example of how to construct the same two-layer nework with the functional API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 10)\n"
     ]
    }
   ],
   "source": [
    "def two_layer_fc_functional(inputs, hidden_size, num_classes):     \n",
    "    initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "    flattened_inputs = tf.layers.flatten(inputs)\n",
    "    fc1_output = tf.layers.dense(flattened_inputs, hidden_size, activation=tf.nn.relu,\n",
    "                                 kernel_initializer=initializer)\n",
    "    scores = tf.layers.dense(fc1_output, num_classes,\n",
    "                             kernel_initializer=initializer)\n",
    "    return scores\n",
    "\n",
    "def test_two_layer_fc_functional():\n",
    "    \"\"\" A small unit test to exercise the TwoLayerFC model above. \"\"\"\n",
    "    tf.reset_default_graph()\n",
    "    input_size, hidden_size, num_classes = 50, 42, 10\n",
    "\n",
    "    # As usual in TensorFlow, we first need to define our computational graph.\n",
    "    # To this end we first construct a two layer network graph by calling the\n",
    "    # two_layer_network() function. This function constructs the computation\n",
    "    # graph and outputs the score tensor.\n",
    "    with tf.device(device):\n",
    "        x = tf.zeros((64, input_size))\n",
    "        scores = two_layer_fc_functional(x, hidden_size, num_classes)\n",
    "\n",
    "    # Now that our computational graph has been defined we can run the graph\n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        scores_np = sess.run(scores)\n",
    "        print(scores_np.shape)\n",
    "        \n",
    "test_two_layer_fc_functional()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Model API: Three-Layer ConvNet\n",
    "Now it's your turn to implement a three-layer ConvNet using the `tf.keras.Model` API. Your model should have the same architecture used in Part II:\n",
    "\n",
    "1. Convolutional layer with 5 x 5 kernels, with zero-padding of 2\n",
    "2. ReLU nonlinearity\n",
    "3. Convolutional layer with 3 x 3 kernels, with zero-padding of 1\n",
    "4. ReLU nonlinearity\n",
    "5. Fully-connected layer to give class scores\n",
    "\n",
    "You should initialize the weights of your network using the same initialization method as was used in the two-layer network above.\n",
    "\n",
    "**Hint**: Refer to the documentation for `tf.layers.Conv2D` and `tf.layers.Dense`:\n",
    "\n",
    "https://www.tensorflow.org/api_docs/python/tf/layers/Conv2D\n",
    "\n",
    "https://www.tensorflow.org/api_docs/python/tf/layers/Dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ThreeLayerConvNet(tf.keras.Model):\n",
    "    def __init__(self, channel_1, channel_2, num_classes):\n",
    "        super().__init__()\n",
    "        ########################################################################\n",
    "        # TODO: Implement the __init__ method for a three-layer ConvNet. You   #\n",
    "        # should instantiate layer objects to be used in the forward pass.     #\n",
    "        ########################################################################\n",
    "        initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "        self.conv1 = tf.layers.Conv2D(filters=channel_1, kernel_size=(5, 5), strides=(1, 1), padding='same',\n",
    "                                      activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer)\n",
    "        self.conv2 = tf.layers.Conv2D(filters=channel_1, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                                      activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer)\n",
    "        self.fc = tf.layers.Dense(num_classes, use_bias=True, kernel_initializer=initializer)\n",
    "        ########################################################################\n",
    "        #                           END OF YOUR CODE                           #\n",
    "        ########################################################################\n",
    "        \n",
    "    def call(self, x, training=None):\n",
    "        scores = None\n",
    "        ########################################################################\n",
    "        # TODO: Implement the forward pass for a three-layer ConvNet. You      #\n",
    "        # should use the layer objects defined in the __init__ method.         #\n",
    "        ########################################################################\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        scores = self.fc(tf.layers.flatten(x))\n",
    "        ########################################################################\n",
    "        #                           END OF YOUR CODE                           #\n",
    "        ########################################################################        \n",
    "        return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you complete the implementation of the `ThreeLayerConvNet` above you can run the following to ensure that your implementation does not crash and produces outputs of the expected shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 10)\n"
     ]
    }
   ],
   "source": [
    "def test_ThreeLayerConvNet():\n",
    "    tf.reset_default_graph()\n",
    "    \n",
    "    channel_1, channel_2, num_classes = 12, 8, 10\n",
    "    model = ThreeLayerConvNet(channel_1, channel_2, num_classes)\n",
    "    with tf.device(device):\n",
    "        x = tf.zeros((64, 3, 32, 32))\n",
    "        scores = model(x)\n",
    "    \n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        scores_np = sess.run(scores)\n",
    "        print(scores_np.shape)\n",
    "\n",
    "test_ThreeLayerConvNet()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Model API: Training Loop\n",
    "We need to implement a slightly different training loop when using the `tf.keras.Model` API. Instead of computing gradients and updating the weights of the model manually, we use an `Optimizer` object from the `tf.train` package which takes care of these details for us. You can read more about `Optimizer`s here: https://www.tensorflow.org/api_docs/python/tf/train/Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_part34(model_init_fn, optimizer_init_fn, num_epochs=1):\n",
    "    \"\"\"\n",
    "    Simple training loop for use with models defined using tf.keras. It trains\n",
    "    a model for one epoch on the CIFAR-10 training set and periodically checks\n",
    "    accuracy on the CIFAR-10 validation set.\n",
    "    \n",
    "    Inputs:\n",
    "    - model_init_fn: A function that takes no parameters; when called it\n",
    "      constructs the model we want to train: model = model_init_fn()\n",
    "    - optimizer_init_fn: A function which takes no parameters; when called it\n",
    "      constructs the Optimizer object we will use to optimize the model:\n",
    "      optimizer = optimizer_init_fn()\n",
    "    - num_epochs: The number of epochs to train for\n",
    "    \n",
    "    Returns: Nothing, but prints progress during trainingn\n",
    "    \"\"\"\n",
    "    tf.reset_default_graph()    \n",
    "    with tf.device(device):\n",
    "        # Construct the computational graph we will use to train the model. We\n",
    "        # use the model_init_fn to construct the model, declare placeholders for\n",
    "        # the data and labels\n",
    "        x = tf.placeholder(tf.float32, [None, 32, 32, 3])\n",
    "        y = tf.placeholder(tf.int32, [None])\n",
    "        \n",
    "        # We need a place holder to explicitly specify if the model is in the training\n",
    "        # phase or not. This is because a number of layers behaves differently in\n",
    "        # training and in testing, e.g., dropout and batch normalization.\n",
    "        # We pass this variable to the computation graph through feed_dict as shown below.\n",
    "        is_training = tf.placeholder(tf.bool, name='is_training')\n",
    "        \n",
    "        # Use the model function to build the forward pass.\n",
    "        scores = model_init_fn(x, is_training)\n",
    "\n",
    "        # Compute the loss like we did in Part II\n",
    "        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=scores)\n",
    "        loss = tf.reduce_mean(loss)\n",
    "\n",
    "        # Use the optimizer_fn to construct an Optimizer, then use the optimizer\n",
    "        # to set up the training step. Asking TensorFlow to evaluate the\n",
    "        # train_op returned by optimizer.minimize(loss) will cause us to make a\n",
    "        # single update step using the current minibatch of data.\n",
    "        \n",
    "        # Note that we use tf.control_dependencies to force the model to run\n",
    "        # the tf.GraphKeys.UPDATE_OPS at each training step. tf.GraphKeys.UPDATE_OPS\n",
    "        # holds the operators that update the states of the network.\n",
    "        # For example, the tf.layers.batch_normalization function adds the running mean\n",
    "        # and variance update operators to tf.GraphKeys.UPDATE_OPS.\n",
    "        optimizer = optimizer_init_fn()\n",
    "        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
    "        with tf.control_dependencies(update_ops):\n",
    "            train_op = optimizer.minimize(loss)\n",
    "\n",
    "    # Now we can run the computational graph many times to train the model.\n",
    "    # When we call sess.run we ask it to evaluate train_op, which causes the\n",
    "    # model to update.\n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        t = 0\n",
    "        for epoch in range(num_epochs):\n",
    "            print('Starting epoch %d' % epoch)\n",
    "            for x_np, y_np in train_dset:\n",
    "                feed_dict = {x: x_np, y: y_np, is_training:1}\n",
    "                loss_np, _ = sess.run([loss, train_op], feed_dict=feed_dict)\n",
    "                if t % print_every == 0:\n",
    "                    print('Iteration %d, loss = %.4f' % (t, loss_np))\n",
    "                    check_accuracy(sess, val_dset, x, scores, is_training=is_training)\n",
    "                    print()\n",
    "                t += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Model API: Train a Two-Layer Network\n",
    "We can now use the tools defined above to train a two-layer network on CIFAR-10. We define the `model_init_fn` and `optimizer_init_fn` that construct the model and optimizer respectively when called. Here we want to train the model using stochastic gradient descent with no momentum, so we construct a `tf.train.GradientDescentOptimizer` function; you can [read about it here](https://www.tensorflow.org/api_docs/python/tf/train/GradientDescentOptimizer).\n",
    "\n",
    "You don't need to tune any hyperparameters here, but you should achieve accuracies above 40% after one epoch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.8653\n",
      "Got 111 / 1000 correct (11.10%)\n",
      "\n",
      "Iteration 100, loss = 1.8670\n",
      "Got 382 / 1000 correct (38.20%)\n",
      "\n",
      "Iteration 200, loss = 1.4273\n",
      "Got 387 / 1000 correct (38.70%)\n",
      "\n",
      "Iteration 300, loss = 1.7106\n",
      "Got 395 / 1000 correct (39.50%)\n",
      "\n",
      "Iteration 400, loss = 1.7743\n",
      "Got 432 / 1000 correct (43.20%)\n",
      "\n",
      "Iteration 500, loss = 1.6999\n",
      "Got 452 / 1000 correct (45.20%)\n",
      "\n",
      "Iteration 600, loss = 1.8557\n",
      "Got 426 / 1000 correct (42.60%)\n",
      "\n",
      "Iteration 700, loss = 1.8305\n",
      "Got 440 / 1000 correct (44.00%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hidden_size, num_classes = 4000, 10\n",
    "learning_rate = 1e-2\n",
    "\n",
    "def model_init_fn(inputs, is_training):\n",
    "    return TwoLayerFC(hidden_size, num_classes)(inputs)\n",
    "\n",
    "def optimizer_init_fn():\n",
    "    return tf.train.GradientDescentOptimizer(learning_rate)\n",
    "\n",
    "train_part34(model_init_fn, optimizer_init_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Model API: Train a Two-Layer Network (functional API)\n",
    "Similarly, we train the two-layer network constructed using the functional API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.6408\n",
      "Got 123 / 1000 correct (12.30%)\n",
      "\n",
      "Iteration 100, loss = 1.9099\n",
      "Got 370 / 1000 correct (37.00%)\n",
      "\n",
      "Iteration 200, loss = 1.4287\n",
      "Got 402 / 1000 correct (40.20%)\n",
      "\n",
      "Iteration 300, loss = 1.7444\n",
      "Got 385 / 1000 correct (38.50%)\n",
      "\n",
      "Iteration 400, loss = 1.7191\n",
      "Got 421 / 1000 correct (42.10%)\n",
      "\n",
      "Iteration 500, loss = 1.6728\n",
      "Got 437 / 1000 correct (43.70%)\n",
      "\n",
      "Iteration 600, loss = 1.7512\n",
      "Got 420 / 1000 correct (42.00%)\n",
      "\n",
      "Iteration 700, loss = 1.8484\n",
      "Got 448 / 1000 correct (44.80%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hidden_size, num_classes = 4000, 10\n",
    "learning_rate = 1e-2\n",
    "\n",
    "def model_init_fn(inputs, is_training):\n",
    "    return two_layer_fc_functional(inputs, hidden_size, num_classes)\n",
    "\n",
    "def optimizer_init_fn():\n",
    "    return tf.train.GradientDescentOptimizer(learning_rate)\n",
    "\n",
    "train_part34(model_init_fn, optimizer_init_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Model API: Train a Three-Layer ConvNet\n",
    "Here you should use the tools we've defined above to train a three-layer ConvNet on CIFAR-10. Your ConvNet should use 32 filters in the first convolutional layer and 16 filters in the second layer.\n",
    "\n",
    "To train the model you should use gradient descent with Nesterov momentum 0.9. \n",
    "\n",
    "**HINT**: https://www.tensorflow.org/api_docs/python/tf/train/MomentumOptimizer\n",
    "\n",
    "You don't need to perform any hyperparameter tuning, but you should achieve accuracies above 45% after training for one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.5871\n",
      "Got 100 / 1000 correct (10.00%)\n",
      "\n",
      "Iteration 100, loss = 1.5666\n",
      "Got 442 / 1000 correct (44.20%)\n",
      "\n",
      "Iteration 200, loss = 1.2672\n",
      "Got 509 / 1000 correct (50.90%)\n",
      "\n",
      "Iteration 300, loss = 1.3184\n",
      "Got 522 / 1000 correct (52.20%)\n",
      "\n",
      "Iteration 400, loss = 1.2229\n",
      "Got 528 / 1000 correct (52.80%)\n",
      "\n",
      "Iteration 500, loss = 1.3870\n",
      "Got 555 / 1000 correct (55.50%)\n",
      "\n",
      "Iteration 600, loss = 1.4371\n",
      "Got 563 / 1000 correct (56.30%)\n",
      "\n",
      "Iteration 700, loss = 1.2691\n",
      "Got 583 / 1000 correct (58.30%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 3e-3\n",
    "channel_1, channel_2, num_classes = 32, 16, 10\n",
    "\n",
    "def model_init_fn(inputs, is_training):\n",
    "    model = None\n",
    "    ############################################################################\n",
    "    # TODO: Complete the implementation of model_fn.                           #\n",
    "    ############################################################################\n",
    "    model = ThreeLayerConvNet(channel_1, channel_2, num_classes)\n",
    "    ############################################################################\n",
    "    #                           END OF YOUR CODE                               #\n",
    "    ############################################################################\n",
    "    return model(inputs)\n",
    "\n",
    "def optimizer_init_fn():\n",
    "    optimizer = None\n",
    "    ############################################################################\n",
    "    # TODO: Complete the implementation of model_fn.                           #\n",
    "    ############################################################################\n",
    "    optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=0.9, use_nesterov=True)\n",
    "    ############################################################################\n",
    "    #                           END OF YOUR CODE                               #\n",
    "    ############################################################################\n",
    "    return optimizer\n",
    "\n",
    "train_part34(model_init_fn, optimizer_init_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part IV: Keras Sequential API\n",
    "In Part III we introduced the `tf.keras.Model` API, which allows you to define models with any number of learnable layers and with arbitrary connectivity between layers.\n",
    "\n",
    "However for many models you don't need such flexibility - a lot of models can be expressed as a sequential stack of layers, with the output of each layer fed to the next layer as input. If your model fits this pattern, then there is an even easier way to define your model: using `tf.keras.Sequential`. You don't need to write any custom classes; you simply call the `tf.keras.Sequential` constructor with a list containing a sequence of layer objects.\n",
    "\n",
    "One complication with `tf.keras.Sequential` is that you must define the shape of the input to the model by passing a value to the `input_shape` of the first layer in your model.\n",
    "\n",
    "### Keras Sequential API: Two-Layer Network\n",
    "Here we rewrite the two-layer fully-connected network using `tf.keras.Sequential`, and train it using the training loop defined above.\n",
    "\n",
    "You don't need to perform any hyperparameter tuning here, but you should see accuracies above 40% after training for one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.3987\n",
      "Got 151 / 1000 correct (15.10%)\n",
      "\n",
      "Iteration 100, loss = 1.7430\n",
      "Got 391 / 1000 correct (39.10%)\n",
      "\n",
      "Iteration 200, loss = 1.4660\n",
      "Got 414 / 1000 correct (41.40%)\n",
      "\n",
      "Iteration 300, loss = 1.7606\n",
      "Got 384 / 1000 correct (38.40%)\n",
      "\n",
      "Iteration 400, loss = 1.7892\n",
      "Got 432 / 1000 correct (43.20%)\n",
      "\n",
      "Iteration 500, loss = 1.7433\n",
      "Got 440 / 1000 correct (44.00%)\n",
      "\n",
      "Iteration 600, loss = 1.7601\n",
      "Got 427 / 1000 correct (42.70%)\n",
      "\n",
      "Iteration 700, loss = 1.7620\n",
      "Got 460 / 1000 correct (46.00%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 1e-2\n",
    "\n",
    "def model_init_fn(inputs, is_training):\n",
    "    input_shape = (32, 32, 3)\n",
    "    hidden_layer_size, num_classes = 4000, 10\n",
    "    initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "    layers = [\n",
    "        tf.layers.Flatten(input_shape=input_shape),\n",
    "        tf.layers.Dense(hidden_layer_size, activation=tf.nn.relu,\n",
    "                        kernel_initializer=initializer),\n",
    "        tf.layers.Dense(num_classes, kernel_initializer=initializer),\n",
    "    ]\n",
    "    model = tf.keras.Sequential(layers)\n",
    "    return model(inputs)\n",
    "\n",
    "def optimizer_init_fn():\n",
    "    return tf.train.GradientDescentOptimizer(learning_rate)\n",
    "\n",
    "train_part34(model_init_fn, optimizer_init_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras Sequential API: Three-Layer ConvNet\n",
    "Here you should use `tf.keras.Sequential` to reimplement the same three-layer ConvNet architecture used in Part II and Part III. As a reminder, your model should have the following architecture:\n",
    "\n",
    "1. Convolutional layer with 16 5x5 kernels, using zero padding of 2\n",
    "2. ReLU nonlinearity\n",
    "3. Convolutional layer with 32 3x3 kernels, using zero padding of 1\n",
    "4. ReLU nonlinearity\n",
    "5. Fully-connected layer giving class scores\n",
    "\n",
    "You should initialize the weights of the model using a `tf.variance_scaling_initializer` as above.\n",
    "\n",
    "You should train the model using Nesterov momentum 0.9.\n",
    "\n",
    "You don't need to perform any hyperparameter search, but you should achieve accuracy above 45% after training for one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.4997\n",
      "Got 113 / 1000 correct (11.30%)\n",
      "\n",
      "Iteration 100, loss = 1.6991\n",
      "Got 407 / 1000 correct (40.70%)\n",
      "\n",
      "Iteration 200, loss = 1.4965\n",
      "Got 466 / 1000 correct (46.60%)\n",
      "\n",
      "Iteration 300, loss = 1.5634\n",
      "Got 496 / 1000 correct (49.60%)\n",
      "\n",
      "Iteration 400, loss = 1.5800\n",
      "Got 497 / 1000 correct (49.70%)\n",
      "\n",
      "Iteration 500, loss = 1.5451\n",
      "Got 508 / 1000 correct (50.80%)\n",
      "\n",
      "Iteration 600, loss = 1.5406\n",
      "Got 518 / 1000 correct (51.80%)\n",
      "\n",
      "Iteration 700, loss = 1.4586\n",
      "Got 521 / 1000 correct (52.10%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def model_init_fn(inputs, is_training):\n",
    "    model = None\n",
    "    ############################################################################\n",
    "    # TODO: Construct a three-layer ConvNet using tf.keras.Sequential.         #\n",
    "    ############################################################################\n",
    "    input_shape = (32, 32, 3)\n",
    "    initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "    layers = [\n",
    "        tf.layers.Conv2D(input_shape=input_shape, filters=16, kernel_size=(5, 5), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.Flatten(),\n",
    "        tf.layers.Dense(10, use_bias=True, kernel_initializer=initializer)\n",
    "    ]\n",
    "    model = tf.keras.Sequential(layers)\n",
    "    ############################################################################\n",
    "    #                            END OF YOUR CODE                              #\n",
    "    ############################################################################\n",
    "    return model(inputs)\n",
    "\n",
    "learning_rate = 5e-4\n",
    "def optimizer_init_fn():\n",
    "    optimizer = None\n",
    "    ############################################################################\n",
    "    # TODO: Complete the implementation of model_fn.                           #\n",
    "    ############################################################################\n",
    "    optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=0.9, use_nesterov=True)\n",
    "    ############################################################################\n",
    "    #                           END OF YOUR CODE                               #\n",
    "    ############################################################################\n",
    "    return optimizer\n",
    "\n",
    "train_part34(model_init_fn, optimizer_init_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part V: CIFAR-10 open-ended challenge\n",
    "\n",
    "In this section you can experiment with whatever ConvNet architecture you'd like on CIFAR-10.\n",
    "\n",
    "You should experiment with architectures, hyperparameters, loss functions, regularization, or anything else you can think of to train a model that achieves **at least 70%** accuracy on the **validation** set within 10 epochs. You can use the `check_accuracy` and `train` functions from above, or you can implement your own training loop.\n",
    "\n",
    "Describe what you did at the end of the notebook.\n",
    "\n",
    "### Some things you can try:\n",
    "- **Filter size**: Above we used 5x5 and 3x3; is this optimal?\n",
    "- **Number of filters**: Above we used 16 and 32 filters. Would more or fewer do better?\n",
    "- **Pooling**: We didn't use any pooling above. Would this improve the model?\n",
    "- **Normalization**: Would your model be improved with batch normalization, layer normalization, group normalization, or some other normalization strategy?\n",
    "- **Network architecture**: The ConvNet above has only three layers of trainable parameters. Would a deeper model do better?\n",
    "- **Global average pooling**: Instead of flattening after the final convolutional layer, would global average pooling do better? This strategy is used for example in Google's Inception network and in Residual Networks.\n",
    "- **Regularization**: Would some kind of regularization improve performance? Maybe weight decay or dropout?\n",
    "\n",
    "### WARNING: Batch Normalization / Dropout\n",
    "Batch Normalization and Dropout **WILL NOT WORK CORRECTLY** if you use the `train_part34()` function with the object-oriented `tf.keras.Model` or `tf.keras.Sequential` APIs; if you want to use these layers with this training loop then you **must use the tf.layers functional API**.\n",
    "\n",
    "We wrote `train_part34()` to explicitly demonstrate how TensorFlow works; however there are some subtleties that make it tough to handle the object-oriented batch normalization layer in a simple training loop. In practice both `tf.keras` and `tf` provide higher-level APIs which handle the training loop for you, such as [keras.fit](https://keras.io/models/sequential/) and [tf.Estimator](https://www.tensorflow.org/programmers_guide/estimators), both of which will properly handle batch normalization when using the object-oriented API.\n",
    "\n",
    "### Tips for training\n",
    "For each network architecture that you try, you should tune the learning rate and other hyperparameters. When doing this there are a couple important things to keep in mind:\n",
    "\n",
    "- If the parameters are working well, you should see improvement within a few hundred iterations\n",
    "- Remember the coarse-to-fine approach for hyperparameter tuning: start by testing a large range of hyperparameters for just a few training iterations to find the combinations of parameters that are working at all.\n",
    "- Once you have found some sets of parameters that seem to work, search more finely around these parameters. You may need to train for more epochs.\n",
    "- You should use the validation set for hyperparameter search, and save your test set for evaluating your architecture on the best parameters as selected by the validation set.\n",
    "\n",
    "### Going above and beyond\n",
    "If you are feeling adventurous there are many other features you can implement to try and improve your performance. You are **not required** to implement any of these, but don't miss the fun if you have time!\n",
    "\n",
    "- Alternative optimizers: you can try Adam, Adagrad, RMSprop, etc.\n",
    "- Alternative activation functions such as leaky ReLU, parametric ReLU, ELU, or MaxOut.\n",
    "- Model ensembles\n",
    "- Data augmentation\n",
    "- New Architectures\n",
    "  - [ResNets](https://arxiv.org/abs/1512.03385) where the input from the previous layer is added to the output.\n",
    "  - [DenseNets](https://arxiv.org/abs/1608.06993) where inputs into previous layers are concatenated together.\n",
    "  - [This blog has an in-depth overview](https://chatbotslife.com/resnets-highwaynets-and-densenets-oh-my-9bb15918ee32)\n",
    "  \n",
    "### Have fun and happy training! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0\n",
      "Iteration 0, loss = 2.6499\n",
      "Got 115 / 1000 correct (11.50%)\n",
      "\n",
      "Iteration 500, loss = 1.4037\n",
      "Got 530 / 1000 correct (53.00%)\n",
      "\n",
      "Starting epoch 1\n",
      "Iteration 1000, loss = 1.0195\n",
      "Got 589 / 1000 correct (58.90%)\n",
      "\n",
      "Iteration 1500, loss = 0.9765\n",
      "Got 643 / 1000 correct (64.30%)\n",
      "\n",
      "Starting epoch 2\n",
      "Iteration 2000, loss = 0.7247\n",
      "Got 646 / 1000 correct (64.60%)\n",
      "\n",
      "Starting epoch 3\n",
      "Iteration 2500, loss = 0.7839\n",
      "Got 669 / 1000 correct (66.90%)\n",
      "\n",
      "Iteration 3000, loss = 0.8779\n",
      "Got 696 / 1000 correct (69.60%)\n",
      "\n",
      "Starting epoch 4\n",
      "Iteration 3500, loss = 0.7538\n",
      "Got 711 / 1000 correct (71.10%)\n",
      "\n",
      "Starting epoch 5\n",
      "Iteration 4000, loss = 0.7768\n",
      "Got 711 / 1000 correct (71.10%)\n",
      "\n",
      "Iteration 4500, loss = 0.8019\n",
      "Got 713 / 1000 correct (71.30%)\n",
      "\n",
      "Starting epoch 6\n",
      "Iteration 5000, loss = 0.5387\n",
      "Got 729 / 1000 correct (72.90%)\n",
      "\n",
      "Starting epoch 7\n",
      "Iteration 5500, loss = 0.4169\n",
      "Got 722 / 1000 correct (72.20%)\n",
      "\n",
      "Iteration 6000, loss = 0.5624\n",
      "Got 715 / 1000 correct (71.50%)\n",
      "\n",
      "Starting epoch 8\n",
      "Iteration 6500, loss = 0.3362\n",
      "Got 740 / 1000 correct (74.00%)\n",
      "\n",
      "Starting epoch 9\n",
      "Iteration 7000, loss = 0.5272\n",
      "Got 728 / 1000 correct (72.80%)\n",
      "\n",
      "Iteration 7500, loss = 0.4886\n",
      "Got 723 / 1000 correct (72.30%)\n",
      "\n",
      "CPU times: user 29.9 s, sys: 3.63 s, total: 33.5 s\n",
      "Wall time: 27.9 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "def model_init_fn(inputs, is_training):\n",
    "    model = None\n",
    "    ############################################################################\n",
    "    # TODO: Construct a model that performs well on CIFAR-10                   #\n",
    "    ############################################################################\n",
    "    input_shape = (32, 32, 3)\n",
    "    initializer = tf.variance_scaling_initializer(scale=2.0)\n",
    "    layers = [\n",
    "        tf.layers.Conv2D(input_shape=input_shape, filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'),\n",
    "        tf.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "                         activation=tf.nn.relu, use_bias=True, kernel_initializer=initializer),\n",
    "        tf.layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'),\n",
    "        tf.layers.Flatten(),\n",
    "        tf.layers.Dense(10, use_bias=True, kernel_initializer=initializer)\n",
    "    ]\n",
    "    model = tf.keras.Sequential(layers)\n",
    "    ############################################################################\n",
    "    #                            END OF YOUR CODE                              #\n",
    "    ############################################################################\n",
    "    return model(inputs)\n",
    "\n",
    "def optimizer_init_fn():\n",
    "    optimizer = None\n",
    "    ############################################################################\n",
    "    # TODO: Construct an optimizer that performs well on CIFAR-10              #\n",
    "    ############################################################################\n",
    "    learning_rate = 5e-4\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate)\n",
    "    ############################################################################\n",
    "    #                            END OF YOUR CODE                              #\n",
    "    ############################################################################\n",
    "    return optimizer\n",
    "\n",
    "device = '/gpu:0'\n",
    "print_every = 500\n",
    "num_epochs = 10\n",
    "train_part34(model_init_fn, optimizer_init_fn, num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Describe what you did \n",
    "\n",
    "In the cell below you should write an explanation of what you did, any additional features that you implemented, and/or any graphs that you made in the process of training and evaluating your network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I used 4 convolutional layers with 32 3x3 kernels and relu activation, and added a max pooling layer with 2x2 kernel after every 2 convolutional layers."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (py36)",
   "language": "python",
   "name": "py36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}