CS231n/assignment2/PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What's this PyTorch 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, PyTorch (or TensorFlow, if you switch over to that notebook). \n",
    "\n",
    "### What is PyTorch?\n",
    "\n",
    "PyTorch is a system for executing dynamic computational graphs over Tensor objects that behave similarly as numpy ndarray. It comes with a powerful automatic differentiation engine that removes the need for manual back-propagation. \n",
    "\n",
    "### Why?\n",
    "\n",
    "* Our code will now run on GPUs! Much faster training. When using a framework like PyTorch or TensorFlow you can harness the power of the GPU for your own custom neural network architectures without having to write CUDA code directly (which is beyond the scope of this class).\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.\n",
    "\n",
    "### PyTorch versions\n",
    "This notebook assumes that you are using **PyTorch version 0.4**. Prior to this version, Tensors had to be wrapped in Variable objects to be used in autograd; however Variables have now been deprecated. In addition 0.4 also separates a Tensor's datatype from its device, and uses numpy-style factories for constructing Tensors rather than directly invoking Tensor constructors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How will I learn PyTorch?\n",
    "\n",
    "Justin Johnson has made an excellent [tutorial](https://github.com/jcjohnson/pytorch-examples) for PyTorch. \n",
    "\n",
    "You can also find the detailed [API doc](http://pytorch.org/docs/stable/index.html) here. If you have other questions that are not addressed by the API docs, the [PyTorch forum](https://discuss.pytorch.org/) is a much better place to ask than StackOverflow.\n",
    "\n",
    "\n",
    "# Table of Contents\n",
    "\n",
    "This assignment has 5 parts. You will learn PyTorch on different levels of abstractions, which will help you understand it better and prepare you for the final project. \n",
    "\n",
    "1. Preparation: we will use CIFAR-10 dataset.\n",
    "2. Barebones PyTorch: we will work directly with the lowest-level PyTorch Tensors. \n",
    "3. PyTorch Module API: we will use `nn.Module` to define arbitrary neural network architecture. \n",
    "4. PyTorch Sequential API: we will use `nn.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",
    "| `nn.Module`     | High        | Medium      |\n",
    "| `nn.Sequential` | Low         | High        |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part I. Preparation\n",
    "\n",
    "First, we load the CIFAR-10 dataset. This might take a couple minutes the first time you do it, but the files should stay cached after that.\n",
    "\n",
    "In previous parts of the assignment we had to write our own code to download the CIFAR-10 dataset, preprocess it, and iterate through it in minibatches; PyTorch provides convenient tools to automate this process for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import sampler\n",
    "\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as T\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "NUM_TRAIN = 49000\n",
    "\n",
    "# The torchvision.transforms package provides tools for preprocessing data\n",
    "# and for performing data augmentation; here we set up a transform to\n",
    "# preprocess the data by subtracting the mean RGB value and dividing by the\n",
    "# standard deviation of each RGB value; we've hardcoded the mean and std.\n",
    "transform = T.Compose([\n",
    "                T.ToTensor(),\n",
    "                T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n",
    "            ])\n",
    "\n",
    "# We set up a Dataset object for each split (train / val / test); Datasets load\n",
    "# training examples one at a time, so we wrap each Dataset in a DataLoader which\n",
    "# iterates through the Dataset and forms minibatches. We divide the CIFAR-10\n",
    "# training set into train and val sets by passing a Sampler object to the\n",
    "# DataLoader telling how it should sample from the underlying Dataset.\n",
    "cifar10_train = dset.CIFAR10('./cs231n/datasets', train=True, download=True,\n",
    "                             transform=transform)\n",
    "loader_train = DataLoader(cifar10_train, batch_size=64, \n",
    "                          sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN)))\n",
    "\n",
    "cifar10_val = dset.CIFAR10('./cs231n/datasets', train=True, download=True,\n",
    "                           transform=transform)\n",
    "loader_val = DataLoader(cifar10_val, batch_size=64, \n",
    "                        sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN, 50000)))\n",
    "\n",
    "cifar10_test = dset.CIFAR10('./cs231n/datasets', train=False, download=True, \n",
    "                            transform=transform)\n",
    "loader_test = DataLoader(cifar10_test, batch_size=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have an option to **use GPU by setting the flag to True below**. It is not necessary to use GPU for this assignment. Note that if your computer does not have CUDA enabled, `torch.cuda.is_available()` will return False and this notebook will fallback to CPU mode.\n",
    "\n",
    "The global variables `dtype` and `device` will control the data types throughout this assignment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using device: cuda\n"
     ]
    }
   ],
   "source": [
    "USE_GPU = True\n",
    "\n",
    "dtype = torch.float32 # we will be using float throughout this tutorial\n",
    "\n",
    "if USE_GPU and torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "\n",
    "# Constant to control how frequently we print train loss\n",
    "print_every = 100\n",
    "\n",
    "print('using device:', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part II. Barebones PyTorch\n",
    "\n",
    "PyTorch ships with high-level APIs to help us define model architectures conveniently, which we will cover in Part II of this tutorial. In this section, we will start with the barebone PyTorch elements to understand the autograd engine better. After this exercise, you will come to appreciate the high-level model API more.\n",
    "\n",
    "We will start with a simple fully-connected ReLU network with two hidden layers and no biases for CIFAR classification. \n",
    "This implementation computes the forward pass using operations on PyTorch Tensors, and uses PyTorch autograd to compute gradients. It is important that you understand every line, because you will write a harder version after the example.\n",
    "\n",
    "When we create a PyTorch Tensor with `requires_grad=True`, then operations involving that Tensor will not just compute values; they will also build up a computational graph in the background, allowing us to easily backpropagate through the graph to compute gradients of some Tensors with respect to a downstream loss. Concretely if x is a Tensor with `x.requires_grad == True` then after backpropagation `x.grad` will be another Tensor holding the gradient of x with respect to the scalar loss at the end."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Tensors: Flatten Function\n",
    "A PyTorch Tensor is conceptionally similar to a numpy array: it is an n-dimensional grid of numbers, and like numpy PyTorch provides many functions to efficiently operate on Tensors. As a simple example, we provide a `flatten` function below which reshapes image data for use in a fully-connected neural network.\n",
    "\n",
    "Recall that image data is typically stored in a Tensor of shape N x C x H x W, where:\n",
    "\n",
    "* N is the number of datapoints\n",
    "* C is the number of channels\n",
    "* H is the height of the intermediate feature map in pixels\n",
    "* W is the height of the intermediate feature map in pixels\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 `C x H x W` values per representation into a single long vector. The flatten function below first reads in the N, C, H, and W values 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 C x H x W, but we don't need to specify that explicitly). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before flattening:  tensor([[[[  0.,   1.],\n",
      "          [  2.,   3.],\n",
      "          [  4.,   5.]]],\n",
      "\n",
      "\n",
      "        [[[  6.,   7.],\n",
      "          [  8.,   9.],\n",
      "          [ 10.,  11.]]]])\n",
      "After flattening:  tensor([[  0.,   1.,   2.,   3.,   4.,   5.],\n",
      "        [  6.,   7.,   8.,   9.,  10.,  11.]])\n"
     ]
    }
   ],
   "source": [
    "def flatten(x):\n",
    "    N = x.shape[0] # read in N, C, H, W\n",
    "    return x.view(N, -1)  # \"flatten\" the C * H * W values into a single vector per image\n",
    "\n",
    "def test_flatten():\n",
    "    x = torch.arange(12).view(2, 1, 3, 2)\n",
    "    print('Before flattening: ', x)\n",
    "    print('After flattening: ', flatten(x))\n",
    "\n",
    "test_flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones PyTorch: Two-Layer Network\n",
    "\n",
    "Here we define a function `two_layer_fc` which performs the forward pass of a two-layer fully-connected ReLU network on a batch of image data. After defining the forward pass we check that it doesn't crash and that it produces outputs of the right shape by running zeros through the network.\n",
    "\n",
    "You don't have to write any code here, but it's important that you read and understand the implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 10])\n"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F  # useful stateless functions\n",
    "\n",
    "def two_layer_fc(x, params):\n",
    "    \"\"\"\n",
    "    A fully-connected neural networks; the architecture is:\n",
    "    NN is fully connected -> ReLU -> fully connected layer.\n",
    "    Note that this function only defines the forward pass; \n",
    "    PyTorch will take care of the backward pass 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 PyTorch Tensor of shape (N, d1, ..., dM) giving a minibatch of\n",
    "      input data.\n",
    "    - params: A list [w1, w2] of PyTorch Tensors giving weights for the network;\n",
    "      w1 has shape (D, H) and w2 has shape (H, C).\n",
    "    \n",
    "    Returns:\n",
    "    - scores: A PyTorch Tensor of shape (N, C) giving classification scores for\n",
    "      the input data x.\n",
    "    \"\"\"\n",
    "    # first we flatten the image\n",
    "    x = flatten(x)  # shape: [batch_size, C x H x W]\n",
    "    \n",
    "    w1, w2 = params\n",
    "    \n",
    "    # Forward pass: compute predicted y using operations on Tensors. Since w1 and\n",
    "    # w2 have requires_grad=True, operations involving these Tensors will cause\n",
    "    # PyTorch to build a computational graph, allowing automatic computation of\n",
    "    # gradients. Since we are no longer implementing the backward pass by hand we\n",
    "    # don't need to keep references to intermediate values.\n",
    "    # you can also use `.clamp(min=0)`, equivalent to F.relu()\n",
    "    x = F.relu(x.mm(w1))\n",
    "    x = x.mm(w2)\n",
    "    return x\n",
    "    \n",
    "\n",
    "def two_layer_fc_test():\n",
    "    hidden_layer_size = 42\n",
    "    x = torch.zeros((64, 50), dtype=dtype)  # minibatch size 64, feature dimension 50\n",
    "    w1 = torch.zeros((50, hidden_layer_size), dtype=dtype)\n",
    "    w2 = torch.zeros((hidden_layer_size, 10), dtype=dtype)\n",
    "    scores = two_layer_fc(x, [w1, w2])\n",
    "    print(scores.size())  # you should see [64, 10]\n",
    "\n",
    "two_layer_fc_test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones PyTorch: Three-Layer ConvNet\n",
    "\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. Like above, we can immediately test our implementation by passing zeros through the 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: http://pytorch.org/docs/stable/nn.html#torch.nn.functional.conv2d; pay attention to the shapes of convolutional filters!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def three_layer_convnet(x, params):\n",
    "    \"\"\"\n",
    "    Performs the forward pass of a three-layer convolutional network with the\n",
    "    architecture defined above.\n",
    "\n",
    "    Inputs:\n",
    "    - x: A PyTorch Tensor of shape (N, 3, H, W) giving a minibatch of images\n",
    "    - params: A list of PyTorch Tensors giving the weights and biases for the\n",
    "      network; should contain the following:\n",
    "      - conv_w1: PyTorch Tensor of shape (channel_1, 3, KH1, KW1) giving weights\n",
    "        for the first convolutional layer\n",
    "      - conv_b1: PyTorch Tensor of shape (channel_1,) giving biases for the first\n",
    "        convolutional layer\n",
    "      - conv_w2: PyTorch Tensor of shape (channel_2, channel_1, KH2, KW2) giving\n",
    "        weights for the second convolutional layer\n",
    "      - conv_b2: PyTorch Tensor of shape (channel_2,) giving biases for the second\n",
    "        convolutional layer\n",
    "      - fc_w: PyTorch Tensor giving weights for the fully-connected layer. Can you\n",
    "        figure out what the shape should be?\n",
    "      - fc_b: PyTorch Tensor giving biases for the fully-connected layer. Can you\n",
    "        figure out what the shape should be?\n",
    "    \n",
    "    Returns:\n",
    "    - scores: PyTorch Tensor of shape (N, C) giving classification scores for x\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",
    "    rc1 = F.conv2d(x, conv_w1, conv_b1, padding=2)\n",
    "    rr1 = F.relu(rc1)\n",
    "    rc2 = F.conv2d(rr1, conv_w2, conv_b2, padding=1)\n",
    "    rr2 = F.relu(rc2)\n",
    "    rr2 = flatten(rr2)\n",
    "    scores = F.linear(rr2, fc_w.t(), fc_b) # 这里w要转置，看api文档\n",
    "    ################################################################################\n",
    "    #                                 END OF YOUR CODE                             #\n",
    "    ################################################################################\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After defining the forward pass of the ConvNet above, run the following cell to test your implementation.\n",
    "\n",
    "When you run this function, scores should have shape (64, 10)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 10])\n"
     ]
    }
   ],
   "source": [
    "def three_layer_convnet_test():\n",
    "    x = torch.zeros((64, 3, 32, 32), dtype=dtype)  # minibatch size 64, image size [3, 32, 32]\n",
    "\n",
    "    conv_w1 = torch.zeros((6, 3, 5, 5), dtype=dtype)  # [out_channel, in_channel, kernel_H, kernel_W]\n",
    "    conv_b1 = torch.zeros((6,))  # out_channel\n",
    "    conv_w2 = torch.zeros((9, 6, 3, 3), dtype=dtype)  # [out_channel, in_channel, kernel_H, kernel_W]\n",
    "    conv_b2 = torch.zeros((9,))  # out_channel\n",
    "\n",
    "    # you must calculate the shape of the tensor after two conv layers, before the fully-connected layer\n",
    "    fc_w = torch.zeros((9 * 32 * 32, 10))\n",
    "    fc_b = torch.zeros(10)\n",
    "\n",
    "    scores = three_layer_convnet(x, [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b])\n",
    "    print(scores.size())  # you should see [64, 10]\n",
    "three_layer_convnet_test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones PyTorch: Initialization\n",
    "Let's write a couple utility methods to initialize the weight matrices for our models.\n",
    "\n",
    "- `random_weight(shape)` initializes a weight tensor with the Kaiming normalization method.\n",
    "- `zero_weight(shape)` initializes a weight tensor with all zeros. Useful for instantiating bias parameters.\n",
    "\n",
    "The `random_weight` function uses the Kaiming normal initialization method, described in:\n",
    "\n",
    "He et al, *Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification*, ICCV 2015, https://arxiv.org/abs/1502.01852"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.7929,  0.8114, -1.0011, -0.8142, -0.7674],\n",
       "        [ 0.5590,  1.5965, -1.1593,  1.9211, -0.5401],\n",
       "        [-1.0938, -0.0085,  1.1711, -0.5517,  0.0749]], device='cuda:0')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def random_weight(shape):\n",
    "    \"\"\"\n",
    "    Create random Tensors for weights; setting requires_grad=True means that we\n",
    "    want to compute gradients for these Tensors during the backward pass.\n",
    "    We use Kaiming normalization: sqrt(2 / fan_in)\n",
    "    \"\"\"\n",
    "    if len(shape) == 2:  # FC weight\n",
    "        fan_in = shape[0]\n",
    "    else:\n",
    "        fan_in = np.prod(shape[1:]) # conv weight [out_channel, in_channel, kH, kW]\n",
    "    # randn is standard normal distribution generator. \n",
    "    w = torch.randn(shape, device=device, dtype=dtype) * np.sqrt(2. / fan_in)\n",
    "    w.requires_grad = True # 这里非常重要，标识w需要算梯度\n",
    "    return w\n",
    "\n",
    "def zero_weight(shape):\n",
    "    return torch.zeros(shape, device=device, dtype=dtype, requires_grad=True)\n",
    "\n",
    "# create a weight of shape [3 x 5]\n",
    "# you should see the type `torch.cuda.FloatTensor` if you use GPU. \n",
    "# Otherwise it should be `torch.FloatTensor`\n",
    "random_weight((3, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barebones PyTorch: 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.\n",
    "\n",
    "When checking accuracy we don't need to compute any gradients; as a result we don't need PyTorch to build a computational graph for us when we compute scores. To prevent a graph from being built we scope our computation under a `torch.no_grad()` context manager."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_accuracy_part2(loader, model_fn, params):\n",
    "    \"\"\"\n",
    "    Check the accuracy of a classification model.\n",
    "    \n",
    "    Inputs:\n",
    "    - loader: A DataLoader for the data split we want to check\n",
    "    - model_fn: A function that performs the forward pass of the model,\n",
    "      with the signature scores = model_fn(x, params)\n",
    "    - params: List of PyTorch Tensors giving parameters of the model\n",
    "    \n",
    "    Returns: Nothing, but prints the accuracy of the model\n",
    "    \"\"\"\n",
    "    split = 'val' if loader.dataset.train else 'test'\n",
    "    print('Checking accuracy on the %s set' % split)\n",
    "    num_correct, num_samples = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for x, y in loader:\n",
    "            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU\n",
    "            y = y.to(device=device, dtype=torch.int64)\n",
    "            scores = model_fn(x, params)\n",
    "            _, preds = scores.max(1)\n",
    "            num_correct += (preds == y).sum()\n",
    "            num_samples += preds.size(0)\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 PyTorch: Training Loop\n",
    "We can now set up a basic training loop to train our network. We will train the model using stochastic gradient descent without momentum. We will use `torch.functional.cross_entropy` to compute the loss; you can [read about it here](http://pytorch.org/docs/stable/nn.html#cross-entropy).\n",
    "\n",
    "The training loop takes as input the neural network function, a list of initialized parameters (`[w1, w2]` in our example), and learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_part2(model_fn, params, 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",
    "      It should have the signature scores = model_fn(x, params) where x is a\n",
    "      PyTorch Tensor of image data, params is a list of PyTorch Tensors giving\n",
    "      model weights, and scores is a PyTorch Tensor of shape (N, C) giving\n",
    "      scores for the elements in x.\n",
    "    - params: List of PyTorch Tensors giving weights for the model\n",
    "    - learning_rate: Python scalar giving the learning rate to use for SGD\n",
    "    \n",
    "    Returns: Nothing\n",
    "    \"\"\"\n",
    "    for t, (x, y) in enumerate(loader_train):\n",
    "        # Move the data to the proper device (GPU or CPU)\n",
    "        x = x.to(device=device, dtype=dtype)\n",
    "        y = y.to(device=device, dtype=torch.long)\n",
    "\n",
    "        # Forward pass: compute scores and loss\n",
    "        scores = model_fn(x, params)\n",
    "        loss = F.cross_entropy(scores, y)\n",
    "\n",
    "        # Backward pass: PyTorch figures out which Tensors in the computational\n",
    "        # graph has requires_grad=True and uses backpropagation to compute the\n",
    "        # gradient of the loss with respect to these Tensors, and stores the\n",
    "        # gradients in the .grad attribute of each Tensor.\n",
    "        loss.backward()\n",
    "\n",
    "        # Update parameters. We don't want to backpropagate through the\n",
    "        # parameter updates, so we scope the updates under a torch.no_grad()\n",
    "        # context manager to prevent a computational graph from being built.\n",
    "        with torch.no_grad():\n",
    "            for w in params:\n",
    "                w -= learning_rate * w.grad\n",
    "\n",
    "                # Manually zero the gradients after running the backward pass\n",
    "                w.grad.zero_()\n",
    "\n",
    "        if t % print_every == 0:\n",
    "            print('Iteration %d, loss = %.4f' % (t, loss.item()))\n",
    "            check_accuracy_part2(loader_val, model_fn, params)\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BareBones PyTorch: Train a Two-Layer Network\n",
    "Now we are ready to run the training loop. We need to explicitly allocate tensors for the fully connected weights, `w1` and `w2`. \n",
    "\n",
    "Each minibatch of CIFAR has 64 examples, so the tensor shape is `[64, 3, 32, 32]`. \n",
    "\n",
    "After flattening, `x` shape should be `[64, 3 * 32 * 32]`. This will be the size of the first dimension of `w1`. \n",
    "The second dimension of `w1` is the hidden layer size, which will also be the first dimension of `w2`. \n",
    "\n",
    "Finally, the output of the network is a 10-dimensional vector that represents the probability distribution over 10 classes. \n",
    "\n",
    "You don't need to tune any hyperparameters but you should see accuracies above 40% after training for one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 3.9007\n",
      "Checking accuracy on the val set\n",
      "Got 123 / 1000 correct (12.30%)\n",
      "\n",
      "Iteration 100, loss = 2.9128\n",
      "Checking accuracy on the val set\n",
      "Got 326 / 1000 correct (32.60%)\n",
      "\n",
      "Iteration 200, loss = 1.8723\n",
      "Checking accuracy on the val set\n",
      "Got 394 / 1000 correct (39.40%)\n",
      "\n",
      "Iteration 300, loss = 2.0363\n",
      "Checking accuracy on the val set\n",
      "Got 390 / 1000 correct (39.00%)\n",
      "\n",
      "Iteration 400, loss = 1.6701\n",
      "Checking accuracy on the val set\n",
      "Got 411 / 1000 correct (41.10%)\n",
      "\n",
      "Iteration 500, loss = 2.2864\n",
      "Checking accuracy on the val set\n",
      "Got 446 / 1000 correct (44.60%)\n",
      "\n",
      "Iteration 600, loss = 1.5476\n",
      "Checking accuracy on the val set\n",
      "Got 408 / 1000 correct (40.80%)\n",
      "\n",
      "Iteration 700, loss = 1.6752\n",
      "Checking accuracy on the val set\n",
      "Got 434 / 1000 correct (43.40%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hidden_layer_size = 4000\n",
    "learning_rate = 1e-2\n",
    "\n",
    "w1 = random_weight((3 * 32 * 32, hidden_layer_size))\n",
    "w2 = random_weight((hidden_layer_size, 10))\n",
    "\n",
    "train_part2(two_layer_fc, [w1, w2], learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BareBones PyTorch: Training a ConvNet\n",
    "\n",
    "In the below you should use the functions defined above to train a three-layer convolutional network on CIFAR. The network should have the following architecture:\n",
    "\n",
    "1. Convolutional layer (with bias) with 32 5x5 filters, with zero-padding of 2\n",
    "2. ReLU\n",
    "3. Convolutional layer (with bias) with 16 3x3 filters, with zero-padding of 1\n",
    "4. ReLU\n",
    "5. Fully-connected layer (with bias) to compute scores for 10 classes\n",
    "\n",
    "You should initialize your weight matrices using the `random_weight` function defined above, and you should initialize your bias vectors using the `zero_weight` function above.\n",
    "\n",
    "You don't need to tune any hyperparameters, but if everything works correctly you should achieve an accuracy above 42% after one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 4.0809\n",
      "Checking accuracy on the val set\n",
      "Got 86 / 1000 correct (8.60%)\n",
      "\n",
      "Iteration 100, loss = 1.8272\n",
      "Checking accuracy on the val set\n",
      "Got 362 / 1000 correct (36.20%)\n",
      "\n",
      "Iteration 200, loss = 1.6330\n",
      "Checking accuracy on the val set\n",
      "Got 414 / 1000 correct (41.40%)\n",
      "\n",
      "Iteration 300, loss = 1.7572\n",
      "Checking accuracy on the val set\n",
      "Got 431 / 1000 correct (43.10%)\n",
      "\n",
      "Iteration 400, loss = 1.5868\n",
      "Checking accuracy on the val set\n",
      "Got 466 / 1000 correct (46.60%)\n",
      "\n",
      "Iteration 500, loss = 1.5742\n",
      "Checking accuracy on the val set\n",
      "Got 466 / 1000 correct (46.60%)\n",
      "\n",
      "Iteration 600, loss = 1.5070\n",
      "Checking accuracy on the val set\n",
      "Got 457 / 1000 correct (45.70%)\n",
      "\n",
      "Iteration 700, loss = 1.7960\n",
      "Checking accuracy on the val set\n",
      "Got 493 / 1000 correct (49.30%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 3e-3\n",
    "\n",
    "channel_1 = 32\n",
    "channel_2 = 16\n",
    "\n",
    "conv_w1 = None\n",
    "conv_b1 = None\n",
    "conv_w2 = None\n",
    "conv_b2 = None\n",
    "fc_w = None\n",
    "fc_b = None\n",
    "\n",
    "################################################################################\n",
    "# TODO: Initialize the parameters of a three-layer ConvNet.                    #\n",
    "################################################################################\n",
    "conv_w1 = random_weight((32,3,5,5))\n",
    "conv_b1 = zero_weight((32,))\n",
    "conv_w2 = random_weight((16,32,3,3))\n",
    "conv_b2 = zero_weight((16,))\n",
    "fc_w = random_weight((16*32*32,10))\n",
    "fc_b = zero_weight((10,))\n",
    "################################################################################\n",
    "#                                 END OF YOUR CODE                             #\n",
    "################################################################################\n",
    "\n",
    "params = [conv_w1, conv_b1, conv_w2, conv_b2, fc_w, fc_b]\n",
    "train_part2(three_layer_convnet, params, learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part III. PyTorch Module API\n",
    "\n",
    "Barebone PyTorch requires that we track all the parameter tensors by hand. This is fine for small networks with a few tensors, but it would be extremely inconvenient and error-prone to track tens or hundreds of tensors in larger networks.\n",
    "\n",
    "PyTorch provides the `nn.Module` API for you to define arbitrary network architectures, while tracking every learnable parameters for you. In Part II, we implemented SGD ourselves. PyTorch also provides the `torch.optim` package that implements all the common optimizers, such as RMSProp, Adagrad, and Adam. It even supports approximate second-order methods like L-BFGS! You can refer to the [doc](http://pytorch.org/docs/master/optim.html) for the exact specifications of each optimizer.\n",
    "\n",
    "To use the Module API, follow the steps below:\n",
    "\n",
    "1. Subclass `nn.Module`. Give your network class an intuitive name like `TwoLayerFC`. \n",
    "\n",
    "2. In the constructor `__init__()`, define all the layers you need as class attributes. Layer objects like `nn.Linear` and `nn.Conv2d` are themselves `nn.Module` subclasses and contain learnable parameters, so that you don't have to instantiate the raw tensors yourself. `nn.Module` will track these internal parameters for you. Refer to the [doc](http://pytorch.org/docs/master/nn.html) to learn more about the dozens of builtin layers. **Warning**: don't forget to call the `super().__init__()` first!\n",
    "\n",
    "3. In the `forward()` method, define the *connectivity* of your network. You should use the attributes defined in `__init__` as function calls that take tensor as input and output the \"transformed\" tensor. Do *not* create any new layers with learnable parameters in `forward()`! All of them must be declared upfront in `__init__`. \n",
    "\n",
    "After you define your Module subclass, you can instantiate it as an object and call it just like the NN forward function in part II.\n",
    "\n",
    "### Module API: Two-Layer Network\n",
    "Here is a concrete example of a 2-layer fully connected network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 10])\n"
     ]
    }
   ],
   "source": [
    "class TwoLayerFC(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_classes):\n",
    "        super().__init__()\n",
    "        # assign layer objects to class attributes\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        # nn.init package contains convenient initialization methods\n",
    "        # http://pytorch.org/docs/master/nn.html#torch-nn-init \n",
    "        nn.init.kaiming_normal_(self.fc1.weight)\n",
    "        self.fc2 = nn.Linear(hidden_size, num_classes)\n",
    "        nn.init.kaiming_normal_(self.fc2.weight)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # forward always defines connectivity\n",
    "        x = flatten(x)\n",
    "        scores = self.fc2(F.relu(self.fc1(x)))\n",
    "        return scores\n",
    "\n",
    "def test_TwoLayerFC():\n",
    "    input_size = 50\n",
    "    x = torch.zeros((64, input_size), dtype=dtype)  # minibatch size 64, feature dimension 50\n",
    "    model = TwoLayerFC(input_size, 42, 10)\n",
    "    scores = model(x)\n",
    "    print(scores.size())  # you should see [64, 10]\n",
    "test_TwoLayerFC()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Module API: Three-Layer ConvNet\n",
    "It's your turn to implement a 3-layer ConvNet followed by a fully connected layer. The network architecture should be the same as in Part II:\n",
    "\n",
    "1. Convolutional layer with `channel_1` 5x5 filters with zero-padding of 2\n",
    "2. ReLU\n",
    "3. Convolutional layer with `channel_2` 3x3 filters with zero-padding of 1\n",
    "4. ReLU\n",
    "5. Fully-connected layer to `num_classes` classes\n",
    "\n",
    "You should initialize the weight matrices of the model using the Kaiming normal initialization method.\n",
    "\n",
    "**HINT**: http://pytorch.org/docs/stable/nn.html#conv2d\n",
    "\n",
    "After you implement the three-layer ConvNet, the `test_ThreeLayerConvNet` function will run your implementation; it should print `(64, 10)` for the shape of the output scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 10])\n"
     ]
    }
   ],
   "source": [
    "class ThreeLayerConvNet(nn.Module):\n",
    "    def __init__(self, in_channel, channel_1, channel_2, num_classes):\n",
    "        super().__init__()\n",
    "        ########################################################################\n",
    "        # TODO: Set up the layers you need for a three-layer ConvNet with the  #\n",
    "        # architecture defined above.                                          #\n",
    "        ########################################################################\n",
    "        self.conv1 = nn.Conv2d(in_channel, channel_1, (5,5), padding=2)\n",
    "        nn.init.kaiming_normal_(self.conv1.weight)\n",
    "        self.conv2 = nn.Conv2d(channel_1, channel_2, (3,3), padding=1)\n",
    "        nn.init.kaiming_normal_(self.conv2.weight)\n",
    "        self.fc = nn.Linear(channel_2*32*32,num_classes)\n",
    "        nn.init.kaiming_normal_(self.fc.weight)\n",
    "        ########################################################################\n",
    "        #                          END OF YOUR CODE                            #       \n",
    "        ########################################################################\n",
    "\n",
    "    def forward(self, x):\n",
    "        scores = None\n",
    "        ########################################################################\n",
    "        # TODO: Implement the forward function for a 3-layer ConvNet. you      #\n",
    "        # should use the layers you defined in __init__ and specify the        #\n",
    "        # connectivity of those layers in forward()                            #\n",
    "        ########################################################################\n",
    "        x = F.relu(self.conv2(F.relu(self.conv1(x))))\n",
    "        x = flatten(x)\n",
    "        scores = self.fc(x)\n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################\n",
    "        return scores\n",
    "\n",
    "\n",
    "def test_ThreeLayerConvNet():\n",
    "    x = torch.zeros((64, 3, 32, 32), dtype=dtype)  # minibatch size 64, image size [3, 32, 32]\n",
    "    model = ThreeLayerConvNet(in_channel=3, channel_1=12, channel_2=8, num_classes=10)\n",
    "    scores = model(x)\n",
    "    print(scores.size())  # you should see [64, 10]\n",
    "test_ThreeLayerConvNet()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Module API: Check Accuracy\n",
    "Given the validation or test set, we can check the classification accuracy of a neural network. \n",
    "\n",
    "This version is slightly different from the one in part II. You don't manually pass in the parameters anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_accuracy_part34(loader, model):\n",
    "    if loader.dataset.train:\n",
    "        print('Checking accuracy on validation set')\n",
    "    else:\n",
    "        print('Checking accuracy on test set')   \n",
    "    num_correct = 0\n",
    "    num_samples = 0\n",
    "    model.eval()  # set model to evaluation mode\n",
    "    with torch.no_grad():\n",
    "        for x, y in loader:\n",
    "            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU\n",
    "            y = y.to(device=device, dtype=torch.long)\n",
    "            scores = model(x)\n",
    "            _, preds = scores.max(1)\n",
    "            num_correct += (preds == y).sum()\n",
    "            num_samples += preds.size(0)\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": [
    "### Module API: Training Loop\n",
    "We also use a slightly different training loop. Rather than updating the values of the weights ourselves, we use an Optimizer object from the `torch.optim` package, which abstract the notion of an optimization algorithm and provides implementations of most of the algorithms commonly used to optimize neural networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_part34(model, optimizer, epochs=1):\n",
    "    \"\"\"\n",
    "    Train a model on CIFAR-10 using the PyTorch Module API.\n",
    "    \n",
    "    Inputs:\n",
    "    - model: A PyTorch Module giving the model to train.\n",
    "    - optimizer: An Optimizer object we will use to train the model\n",
    "    - epochs: (Optional) A Python integer giving the number of epochs to train for\n",
    "    \n",
    "    Returns: Nothing, but prints model accuracies during training.\n",
    "    \"\"\"\n",
    "    model = model.to(device=device)  # move the model parameters to CPU/GPU\n",
    "    for e in range(epochs):\n",
    "        for t, (x, y) in enumerate(loader_train):\n",
    "            model.train()  # put model to training mode\n",
    "            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU\n",
    "            y = y.to(device=device, dtype=torch.long)\n",
    "\n",
    "            scores = model(x)\n",
    "            loss = F.cross_entropy(scores, y)\n",
    "\n",
    "            # Zero out all of the gradients for the variables which the optimizer\n",
    "            # will update.\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # This is the backwards pass: compute the gradient of the loss with\n",
    "            # respect to each  parameter of the model.\n",
    "            loss.backward()\n",
    "\n",
    "            # Actually update the parameters of the model using the gradients\n",
    "            # computed by the backwards pass.\n",
    "            optimizer.step()\n",
    "\n",
    "            if t % print_every == 0:\n",
    "                print('Iteration %d, loss = %.4f' % (t, loss.item()))\n",
    "                check_accuracy_part34(loader_val, model)\n",
    "                print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Module API: Train a Two-Layer Network\n",
    "Now we are ready to run the training loop. In contrast to part II, we don't explicitly allocate parameter tensors anymore.\n",
    "\n",
    "Simply pass the input size, hidden layer size, and number of classes (i.e. output size) to the constructor of `TwoLayerFC`. \n",
    "\n",
    "You also need to define an optimizer that tracks all the learnable parameters inside `TwoLayerFC`.\n",
    "\n",
    "You don't need to tune any hyperparameters, but you should see model accuracies above 40% after training for one epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 3.5496\n",
      "Checking accuracy on validation set\n",
      "Got 151 / 1000 correct (15.10)\n",
      "\n",
      "Iteration 100, loss = 1.9458\n",
      "Checking accuracy on validation set\n",
      "Got 285 / 1000 correct (28.50)\n",
      "\n",
      "Iteration 200, loss = 2.0304\n",
      "Checking accuracy on validation set\n",
      "Got 408 / 1000 correct (40.80)\n",
      "\n",
      "Iteration 300, loss = 1.7830\n",
      "Checking accuracy on validation set\n",
      "Got 343 / 1000 correct (34.30)\n",
      "\n",
      "Iteration 400, loss = 1.4227\n",
      "Checking accuracy on validation set\n",
      "Got 414 / 1000 correct (41.40)\n",
      "\n",
      "Iteration 500, loss = 1.8738\n",
      "Checking accuracy on validation set\n",
      "Got 415 / 1000 correct (41.50)\n",
      "\n",
      "Iteration 600, loss = 1.5694\n",
      "Checking accuracy on validation set\n",
      "Got 447 / 1000 correct (44.70)\n",
      "\n",
      "Iteration 700, loss = 1.6504\n",
      "Checking accuracy on validation set\n",
      "Got 442 / 1000 correct (44.20)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hidden_layer_size = 4000\n",
    "learning_rate = 1e-2\n",
    "model = TwoLayerFC(3 * 32 * 32, hidden_layer_size, 10)\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "train_part34(model, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Module API: Train a Three-Layer ConvNet\n",
    "You should now use the Module API to train a three-layer ConvNet on CIFAR. This should look very similar to training the two-layer network! You don't need to tune any hyperparameters, but you should achieve above above 45% after training for one epoch.\n",
    "\n",
    "You should train the model using stochastic gradient descent without momentum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 3.0433\n",
      "Checking accuracy on validation set\n",
      "Got 109 / 1000 correct (10.90)\n",
      "\n",
      "Iteration 100, loss = 2.0086\n",
      "Checking accuracy on validation set\n",
      "Got 356 / 1000 correct (35.60)\n",
      "\n",
      "Iteration 200, loss = 1.8688\n",
      "Checking accuracy on validation set\n",
      "Got 399 / 1000 correct (39.90)\n",
      "\n",
      "Iteration 300, loss = 1.6527\n",
      "Checking accuracy on validation set\n",
      "Got 421 / 1000 correct (42.10)\n",
      "\n",
      "Iteration 400, loss = 1.5756\n",
      "Checking accuracy on validation set\n",
      "Got 442 / 1000 correct (44.20)\n",
      "\n",
      "Iteration 500, loss = 1.6130\n",
      "Checking accuracy on validation set\n",
      "Got 466 / 1000 correct (46.60)\n",
      "\n",
      "Iteration 600, loss = 1.6782\n",
      "Checking accuracy on validation set\n",
      "Got 452 / 1000 correct (45.20)\n",
      "\n",
      "Iteration 700, loss = 1.4245\n",
      "Checking accuracy on validation set\n",
      "Got 475 / 1000 correct (47.50)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 3e-3\n",
    "channel_1 = 32\n",
    "channel_2 = 16\n",
    "\n",
    "model = None\n",
    "optimizer = None\n",
    "################################################################################\n",
    "# TODO: Instantiate your ThreeLayerConvNet model and a corresponding optimizer #\n",
    "################################################################################\n",
    "model = ThreeLayerConvNet(3, channel_1, channel_2, 10)\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
    "################################################################################\n",
    "#                                 END OF YOUR CODE                             \n",
    "################################################################################\n",
    "\n",
    "train_part34(model, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part IV. PyTorch Sequential API\n",
    "\n",
    "Part III introduced the PyTorch Module API, which allows you to define arbitrary learnable layers and their connectivity. \n",
    "\n",
    "For simple models like a stack of feed forward layers, you still need to go through 3 steps: subclass `nn.Module`, assign layers to class attributes in `__init__`, and call each layer one by one in `forward()`. Is there a more convenient way? \n",
    "\n",
    "Fortunately, PyTorch provides a container Module called `nn.Sequential`, which merges the above steps into one. It is not as flexible as `nn.Module`, because you cannot specify more complex topology than a feed-forward stack, but it's good enough for many use cases.\n",
    "\n",
    "### Sequential API: Two-Layer Network\n",
    "Let's see how to rewrite our two-layer fully connected network example with `nn.Sequential`, and train it using the training loop defined above.\n",
    "\n",
    "Again, you don't need to tune any hyperparameters here, but you shoud achieve above 40% accuracy after one epoch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 2.3414\n",
      "Checking accuracy on validation set\n",
      "Got 119 / 1000 correct (11.90)\n",
      "\n",
      "Iteration 100, loss = 1.9493\n",
      "Checking accuracy on validation set\n",
      "Got 388 / 1000 correct (38.80)\n",
      "\n",
      "Iteration 200, loss = 1.6697\n",
      "Checking accuracy on validation set\n",
      "Got 398 / 1000 correct (39.80)\n",
      "\n",
      "Iteration 300, loss = 1.8498\n",
      "Checking accuracy on validation set\n",
      "Got 409 / 1000 correct (40.90)\n",
      "\n",
      "Iteration 400, loss = 2.0913\n",
      "Checking accuracy on validation set\n",
      "Got 429 / 1000 correct (42.90)\n",
      "\n",
      "Iteration 500, loss = 1.5562\n",
      "Checking accuracy on validation set\n",
      "Got 437 / 1000 correct (43.70)\n",
      "\n",
      "Iteration 600, loss = 1.6216\n",
      "Checking accuracy on validation set\n",
      "Got 424 / 1000 correct (42.40)\n",
      "\n",
      "Iteration 700, loss = 1.5994\n",
      "Checking accuracy on validation set\n",
      "Got 447 / 1000 correct (44.70)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# We need to wrap `flatten` function in a module in order to stack it\n",
    "# in nn.Sequential\n",
    "class Flatten(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return flatten(x)\n",
    "\n",
    "hidden_layer_size = 4000\n",
    "learning_rate = 1e-2\n",
    "\n",
    "model = nn.Sequential(\n",
    "    Flatten(),\n",
    "    nn.Linear(3 * 32 * 32, hidden_layer_size),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(hidden_layer_size, 10),\n",
    ")\n",
    "\n",
    "# you can use Nesterov momentum in optim.SGD\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate,\n",
    "                     momentum=0.9, nesterov=True)\n",
    "\n",
    "train_part34(model, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sequential API: Three-Layer ConvNet\n",
    "Here you should use `nn.Sequential` to define and train a three-layer ConvNet with the same architecture we used in Part III:\n",
    "\n",
    "1. Convolutional layer (with bias) with 32 5x5 filters, with zero-padding of 2\n",
    "2. ReLU\n",
    "3. Convolutional layer (with bias) with 16 3x3 filters, with zero-padding of 1\n",
    "4. ReLU\n",
    "5. Fully-connected layer (with bias) to compute scores for 10 classes\n",
    "\n",
    "You should initialize your weight matrices using the `random_weight` function defined above, and you should initialize your bias vectors using the `zero_weight` function above.\n",
    "\n",
    "You should optimize your model using stochastic gradient descent with Nesterov momentum 0.9.\n",
    "\n",
    "Again, you don't need to tune any hyperparameters but you should see accuracy above 55% after one epoch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 2.2852\n",
      "Checking accuracy on validation set\n",
      "Got 124 / 1000 correct (12.40)\n",
      "\n",
      "Iteration 100, loss = 1.5513\n",
      "Checking accuracy on validation set\n",
      "Got 444 / 1000 correct (44.40)\n",
      "\n",
      "Iteration 200, loss = 1.4262\n",
      "Checking accuracy on validation set\n",
      "Got 469 / 1000 correct (46.90)\n",
      "\n",
      "Iteration 300, loss = 1.1712\n",
      "Checking accuracy on validation set\n",
      "Got 490 / 1000 correct (49.00)\n",
      "\n",
      "Iteration 400, loss = 1.3891\n",
      "Checking accuracy on validation set\n",
      "Got 523 / 1000 correct (52.30)\n",
      "\n",
      "Iteration 500, loss = 1.3876\n",
      "Checking accuracy on validation set\n",
      "Got 509 / 1000 correct (50.90)\n",
      "\n",
      "Iteration 600, loss = 1.2339\n",
      "Checking accuracy on validation set\n",
      "Got 572 / 1000 correct (57.20)\n",
      "\n",
      "Iteration 700, loss = 1.4657\n",
      "Checking accuracy on validation set\n",
      "Got 587 / 1000 correct (58.70)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "channel_1 = 32\n",
    "channel_2 = 16\n",
    "learning_rate = 1e-2\n",
    "\n",
    "model = None\n",
    "optimizer = None\n",
    "\n",
    "################################################################################\n",
    "# TODO: Rewrite the 2-layer ConvNet with bias from Part III with the           #\n",
    "# Sequential API.                                                              #\n",
    "################################################################################\n",
    "model = nn.Sequential(\n",
    "    nn.Conv2d(3, channel_1, (5,5), padding=2),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(channel_1, channel_2, (3,3), padding=1),\n",
    "    nn.ReLU(),\n",
    "    Flatten(),\n",
    "    nn.Linear(channel_2*32*32,10),\n",
    ")\n",
    "\n",
    "# you can use Nesterov momentum in optim.SGD\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate,\n",
    "                     momentum=0.9, nesterov=True)\n",
    "################################################################################\n",
    "#                                 END OF YOUR CODE                             \n",
    "################################################################################\n",
    "\n",
    "train_part34(model, optimizer)"
   ]
  },
  {
   "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",
    "Now it's your job to experiment with architectures, hyperparameters, loss functions, and optimizers to train a model that achieves **at least 70%** accuracy on the CIFAR-10 **validation** set within 10 epochs. You can use the check_accuracy and train functions from above. You can use either `nn.Module` or `nn.Sequential` API. \n",
    "\n",
    "Describe what you did at the end of this notebook.\n",
    "\n",
    "Here are the official API documentation for each component. One note: what we call in the class \"spatial batch norm\" is called \"BatchNorm2D\" in PyTorch.\n",
    "\n",
    "* Layers in torch.nn package: http://pytorch.org/docs/stable/nn.html\n",
    "* Activations: http://pytorch.org/docs/stable/nn.html#non-linear-activations\n",
    "* Loss functions: http://pytorch.org/docs/stable/nn.html#loss-functions\n",
    "* Optimizers: http://pytorch.org/docs/stable/optim.html\n",
    "\n",
    "\n",
    "### Things you might try:\n",
    "- **Filter size**: Above we used 5x5; would smaller filters be more efficient?\n",
    "- **Number of filters**: Above we used 32 filters. Do more or fewer do better?\n",
    "- **Pooling vs Strided Convolution**: Do you use max pooling or just stride convolutions?\n",
    "- **Batch normalization**: Try adding spatial batch normalization after convolution layers and vanilla batch normalization after affine layers. Do your networks train faster?\n",
    "- **Network architecture**: The network above has two layers of trainable parameters. Can you do better with a deep network? Good architectures to try include:\n",
    "    - [conv-relu-pool]xN -> [affine]xM -> [softmax or SVM]\n",
    "    - [conv-relu-conv-relu-pool]xN -> [affine]xM -> [softmax or SVM]\n",
    "    - [batchnorm-relu-conv]xN -> [affine]xM -> [softmax or SVM]\n",
    "- **Global Average Pooling**: Instead of flattening and then having multiple affine layers, perform convolutions until your image gets small (7x7 or so) and then perform an average pooling operation to get to a 1x1 image picture (1, 1 , Filter#), which is then reshaped into a (Filter#) vector. This is used in [Google's Inception Network](https://arxiv.org/abs/1512.00567) (See Table 1 for their architecture).\n",
    "- **Regularization**: Add l2 weight regularization, or perhaps use Dropout.\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": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 2.5105\n",
      "Checking accuracy on validation set\n",
      "Got 103 / 1000 correct (10.30)\n",
      "\n",
      "Iteration 300, loss = 0.8683\n",
      "Checking accuracy on validation set\n",
      "Got 636 / 1000 correct (63.60)\n",
      "\n",
      "Iteration 600, loss = 1.1072\n",
      "Checking accuracy on validation set\n",
      "Got 679 / 1000 correct (67.90)\n",
      "\n",
      "Iteration 0, loss = 1.0412\n",
      "Checking accuracy on validation set\n",
      "Got 707 / 1000 correct (70.70)\n",
      "\n",
      "Iteration 300, loss = 0.9487\n",
      "Checking accuracy on validation set\n",
      "Got 737 / 1000 correct (73.70)\n",
      "\n",
      "Iteration 600, loss = 0.5790\n",
      "Checking accuracy on validation set\n",
      "Got 756 / 1000 correct (75.60)\n",
      "\n",
      "Iteration 0, loss = 0.7634\n",
      "Checking accuracy on validation set\n",
      "Got 747 / 1000 correct (74.70)\n",
      "\n",
      "Iteration 300, loss = 0.6576\n",
      "Checking accuracy on validation set\n",
      "Got 773 / 1000 correct (77.30)\n",
      "\n",
      "Iteration 600, loss = 0.7913\n",
      "Checking accuracy on validation set\n",
      "Got 794 / 1000 correct (79.40)\n",
      "\n",
      "Iteration 0, loss = 0.4503\n",
      "Checking accuracy on validation set\n",
      "Got 794 / 1000 correct (79.40)\n",
      "\n",
      "Iteration 300, loss = 0.4867\n",
      "Checking accuracy on validation set\n",
      "Got 793 / 1000 correct (79.30)\n",
      "\n",
      "Iteration 600, loss = 0.6074\n",
      "Checking accuracy on validation set\n",
      "Got 767 / 1000 correct (76.70)\n",
      "\n",
      "Iteration 0, loss = 0.3161\n",
      "Checking accuracy on validation set\n",
      "Got 808 / 1000 correct (80.80)\n",
      "\n",
      "Iteration 300, loss = 0.4668\n",
      "Checking accuracy on validation set\n",
      "Got 787 / 1000 correct (78.70)\n",
      "\n",
      "Iteration 600, loss = 0.3349\n",
      "Checking accuracy on validation set\n",
      "Got 819 / 1000 correct (81.90)\n",
      "\n",
      "Iteration 0, loss = 0.5020\n",
      "Checking accuracy on validation set\n",
      "Got 803 / 1000 correct (80.30)\n",
      "\n",
      "Iteration 300, loss = 0.2995\n",
      "Checking accuracy on validation set\n",
      "Got 821 / 1000 correct (82.10)\n",
      "\n",
      "Iteration 600, loss = 0.2754\n",
      "Checking accuracy on validation set\n",
      "Got 836 / 1000 correct (83.60)\n",
      "\n",
      "Iteration 0, loss = 0.2857\n",
      "Checking accuracy on validation set\n",
      "Got 819 / 1000 correct (81.90)\n",
      "\n",
      "Iteration 300, loss = 0.2925\n",
      "Checking accuracy on validation set\n",
      "Got 814 / 1000 correct (81.40)\n",
      "\n",
      "Iteration 600, loss = 0.3658\n",
      "Checking accuracy on validation set\n",
      "Got 812 / 1000 correct (81.20)\n",
      "\n",
      "Iteration 0, loss = 0.1379\n",
      "Checking accuracy on validation set\n",
      "Got 828 / 1000 correct (82.80)\n",
      "\n",
      "Iteration 300, loss = 0.2591\n",
      "Checking accuracy on validation set\n",
      "Got 817 / 1000 correct (81.70)\n",
      "\n",
      "Iteration 600, loss = 0.1676\n",
      "Checking accuracy on validation set\n",
      "Got 816 / 1000 correct (81.60)\n",
      "\n",
      "Iteration 0, loss = 0.1325\n",
      "Checking accuracy on validation set\n",
      "Got 825 / 1000 correct (82.50)\n",
      "\n",
      "Iteration 300, loss = 0.2111\n",
      "Checking accuracy on validation set\n",
      "Got 823 / 1000 correct (82.30)\n",
      "\n",
      "Iteration 600, loss = 0.2078\n",
      "Checking accuracy on validation set\n",
      "Got 830 / 1000 correct (83.00)\n",
      "\n",
      "Iteration 0, loss = 0.1824\n",
      "Checking accuracy on validation set\n",
      "Got 811 / 1000 correct (81.10)\n",
      "\n",
      "Iteration 300, loss = 0.0935\n",
      "Checking accuracy on validation set\n",
      "Got 808 / 1000 correct (80.80)\n",
      "\n",
      "Iteration 600, loss = 0.1654\n",
      "Checking accuracy on validation set\n",
      "Got 822 / 1000 correct (82.20)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "################################################################################\n",
    "# TODO:                                                                        #         \n",
    "# Experiment with any architectures, optimizers, and hyperparameters.          #\n",
    "# Achieve AT LEAST 70% accuracy on the *validation set* within 10 epochs.      #\n",
    "#                                                                              #\n",
    "# Note that you can use the check_accuracy function to evaluate on either      #\n",
    "# the test set or the validation set, by passing either loader_test or         #\n",
    "# loader_val as the second argument to check_accuracy. You should not touch    #\n",
    "# the test set until you have finished your architecture and  hyperparameter   #\n",
    "# tuning, and only run the test set once at the end to report a final value.   #\n",
    "################################################################################\n",
    "model = None\n",
    "optimizer = None\n",
    "\n",
    "c,c1,c2,num_class = 3,128,256,10\n",
    "print_every = 300\n",
    "learning_rate = 1e-2\n",
    "\n",
    "class Flatten(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return flatten(x)\n",
    "    \n",
    "class PrintShape(nn.Module):\n",
    "    def forward(self, x):\n",
    "        print(x.shape)\n",
    "        return x\n",
    "\n",
    "model = nn.Sequential(\n",
    "    nn.Conv2d(c, c1, (5,5), padding=2),\n",
    "    nn.BatchNorm2d(c1),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(c1, c2, (3,3), padding=1),\n",
    "    nn.BatchNorm2d(c2),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d((2,2),2),\n",
    "    nn.Conv2d(c2, c2, (3,3), padding=1),\n",
    "    nn.BatchNorm2d(c2),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(c2, c1, (5,5), padding=2),\n",
    "    nn.BatchNorm2d(c1),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d((2,2),2),\n",
    "    Flatten(),\n",
    "    nn.Linear(c1*8*8,10),\n",
    "    nn.BatchNorm1d(10),\n",
    ")\n",
    "\n",
    "# you can use Nesterov momentum in optim.SGD\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate,\n",
    "                     momentum=0.9, nesterov=True)\n",
    "################################################################################\n",
    "#                                 END OF YOUR CODE                             \n",
    "################################################################################\n",
    "\n",
    "# You should get at least 70% accuracy\n",
    "train_part34(model, optimizer, epochs=10)"
   ]
  },
  {
   "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": [
    "TODO: Describe what you did"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test set -- run this only once\n",
    "\n",
    "Now that we've gotten a result we're happy with, we test our final model on the test set (which you should store in best_model). Think about how this compares to your validation set accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking accuracy on test set\n",
      "Got 8025 / 10000 correct (80.25)\n"
     ]
    }
   ],
   "source": [
    "best_model = model\n",
    "check_accuracy_part34(loader_test, best_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}