CPEN455_2023W02_Assignment_1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/SimonX13/DeepLearning/blob/main/CPEN455_2023W02_Assignment_1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9AS6zKmYSwVD"
      },
      "source": [
        "# Programming Assignment 1: Convolution and Back-Propagation\n",
        "\n",
        "**UBC CPEN 455: Deep Learning, 2023 Winter Term 2**\n",
        "\n",
        "**Created By Renjie Liao**\n",
        "\n",
        "**Date: Feb. 19, 2024**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HhakGhuGTBzQ"
      },
      "source": [
        "---\n",
        "# Setup\n",
        "\n",
        "We will use PyTorch to implement this assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PoqjO9gIKMMV"
      },
      "outputs": [],
      "source": [
        "# Imports\n",
        "import pdb\n",
        "import torch\n",
        "import numpy as np\n",
        "from torch import nn\n",
        "from math import pi\n",
        "import matplotlib.pyplot as plt\n",
        "import torch.nn.functional as F\n",
        "from torchvision import datasets, transforms, utils\n",
        "\n",
        "torch.set_default_dtype(torch.float64)\n",
        "\n",
        "## load MNIST images\n",
        "B = 5 # batch size\n",
        "train_set = datasets.MNIST('./data',\n",
        "                            train=True,\n",
        "                            download=True,\n",
        "                            transform=transforms.ToTensor())\n",
        "\n",
        "loader = torch.utils.data.DataLoader(train_set, batch_size=B)\n",
        "\n",
        "## load a batch of MNIST images as a PyTorch tensor (shape: B x C x H x W)\n",
        "# B: batch size\n",
        "# C: number of channels\n",
        "# H: height of images\n",
        "# W: width of images\n",
        "img, label = next(iter(loader)) # img shape: B x C x H x W, label shape: B X 1\n",
        "\n",
        "## create a random filter (shape: D x C x K x K)\n",
        "K = 3 # kernel size\n",
        "P = 1 # padding size\n",
        "C = 1 # channel size\n",
        "D = 2 # number of filters\n",
        "filter = torch.randn(D, C, K, K) # filter shape: D x C x K x K"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6t6YO5QlYhbX"
      },
      "source": [
        "---\n",
        "# Q1 [60Pts]: 2D convolution and its gradient\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "koMGWjFvWWOP"
      },
      "source": [
        "## 1.1 [30Pts]  Implement 2D convolution:\n",
        "\n",
        "Discrete convolution can be implemented in multiple ways, e.g., matrix multiplication in spatial/Fourier domains.\n",
        "\n",
        "First, let us take a look at 1D convolution in spatial domain. Suppose we have a 1D signal with $n$ elements $x_1, x_2, \\dots, x_n$ and a 1D filter with $m$ weights $h_1, h_2, \\dots, h_m$. Note that we typically use filters with odd sizes for the ease of indexing.\n",
        "\n",
        "The (discrete) convolution with zero-padding and stride 1 is defined as:\n",
        "\n",
        "\\begin{align}\n",
        "    y = h \\ast x = \\sum_{i=1}^{n} \\sum_{j=1}^{m} h_j x_{i - \\lfloor m/2 \\rfloor - 1 + j},\n",
        "\\end{align}\n",
        "where padded values $x_{-\\lfloor m/2 \\rfloor + 1}, \\dots, x_{0}, x_{n+1}, \\dots, x_{n - \\lfloor m/2 \\rfloor - 1 + m}$ are all zeros.\n",
        "\n",
        "If you forget about the concepts of padding and stride, take a look at [this guide](https://arxiv.org/pdf/1603.07285.pdf) or [these pictures](https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md).\n",
        "\n",
        "In the 1D case, we can illustrate the two matrix multiplication views of spatial convolution as below.\n",
        "\n",
        "1.   **im2col**:\n",
        "The key idea is to first extract the spatial windows from the signal $x$ for individual convolutions and then perform convolutions (i.e., dot product with the filter).\n",
        "If we put each window as a column in a matrix (the right one in RHS below), then we can perform convolution via the following matrix multiplication (N.B.: the products between the filter and individual columns can be done in parallel).\n",
        "\n",
        "\\begin{align}\n",
        "    y^\\top = (h \\ast x)^{\\top} = \\begin{bmatrix}\n",
        "                h_m & h_{m-1} & \\cdots & h_3 & h_2 & h_1\n",
        "            \\end{bmatrix}\n",
        "            \\begin{bmatrix}\n",
        "                x_{m - \\lfloor m/2 \\rfloor} & x_{m - \\lfloor m/2 \\rfloor + 1} & \\cdots & x_m & x_{m+1} & \\cdots & 0 & 0 \\\\\n",
        "                \\vdots & \\vdots & \\cdots & x_{m-1} & x_m & \\cdots & \\vdots & \\vdots \\\\\n",
        "                x_1 & x_2 & \\cdots & \\vdots & x_{m-1} & \\cdots  & x_n & 0 \\\\\n",
        "                0 & x_1 & \\cdots & \\vdots & \\vdots & \\cdots  & x_{n-1} & x_n \\\\\n",
        "                \\vdots & 0 & \\cdots & \\vdots & \\vdots & \\cdots & \\vdots & \\vdots \\\\                        \n",
        "                0 & 0 & \\cdots & x_1 & x_2 & \\cdots & x_{n - \\lfloor m/2 \\rfloor+1} & x_{n - \\lfloor m/2 \\rfloor}\n",
        "            \\end{bmatrix}.\n",
        "\\end{align}\n",
        "\n",
        "2.   **filter2row**: The key idea is to convert the filter and the signal to a sparse cyclic matrix and a vector respectively.\n",
        "Then the convolution is simply the matrix multiplication between the filter and the signal.\n",
        "\n",
        "\\begin{align}\n",
        "        y = h \\ast x =\n",
        "            \\begin{bmatrix}\n",
        "                h_{\\lfloor m/2 \\rfloor + 1} & h_{\\lfloor m/2 \\rfloor + 2} & \\cdots & h_m & 0 & \\cdots & \\cdots & \\cdots & \\cdots & 0 \\\\\n",
        "                h_{\\lfloor m/2 \\rfloor} & h_{\\lfloor m/2 \\rfloor + 1} & \\cdots & h_{m-1} & h_m & 0 & \\cdots & \\cdots & \\cdots & 0 \\\\\n",
        "                \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots \\\\\n",
        "                h_1 & h_2 & \\cdots & \\cdots & \\cdots & \\cdots & h_m & 0 & \\cdots & 0 \\\\\n",
        "                0 & h_1 & h_2 & \\cdots & \\cdots & \\cdots & \\cdots & h_m & \\cdots & 0 \\\\\n",
        "                \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots \\\\\n",
        "                0 & 0 & \\cdots & \\cdots & \\cdots & 0 & h_1 & h_2 & \\cdots & h_{\\lfloor m/2 \\rfloor + 2} \\\\                \n",
        "                0 & 0 & \\cdots & \\cdots & & \\cdots \\cdots & 0 & h_1 & \\cdots & h_{\\lfloor m/2 \\rfloor + 1}\n",
        "            \\end{bmatrix}\n",
        "            \\begin{bmatrix}\n",
        "                x_1 \\\\\n",
        "                x_2 \\\\\n",
        "                x_3 \\\\\n",
        "                \\vdots \\\\\n",
        "                x_n\n",
        "            \\end{bmatrix}\n",
        "\\end{align}\n",
        "\n",
        "**Task**:\n",
        "Implement the 2D convolution in the spatial domain via matrix multiplication following the above two views: **im2col** and **filter2row**.\n",
        "The starter code is provided below.\n",
        "You just need to fill in the missing parts of function ***conv2d_im2col*** and ***conv2d_filter2row***.\n",
        "If your implementation is correct, the ***unit_test*** will output:\n",
        "\n",
        "*Your implementation of xxx is correct!*\n",
        "\n",
        "Otherwise, it will output:\n",
        "\n",
        "*Your implementation of xxx is wrong!*\n",
        "\n",
        "**N.B.**: we assume the strides along height and width are the same and the kernel is square\n",
        "\n",
        "**Hint**: you can reduce the 2D case to 1D and follow the above construction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-IRaDddEdgT0",
        "outputId": "f26c3e3a-6807-4a43-a2c2-d698cca1a855"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of conv2d_im2col is correct!\n",
            "Your implementation of conv2d_filter2row is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following two functions\n",
        "\n",
        "def conv2d_im2col(img, filter, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter: filters, shape D x C x K x K\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   out: convoluted images, shape B x D x H x W\n",
        "    # cols = im2col(img, kernel_size, stride, padding)\n",
        "    # # Reshape the filter into a row per filter\n",
        "    # filter_row = filter.reshape(num_filters, -1)\n",
        "    # # Perform matrix multiplication and then reshape the output\n",
        "    # out = torch.matmul(filter_row, cols)\n",
        "    # out = out.reshape(img.shape[0], num_filters, (img.shape[2] - kernel_size + 2 * padding) // stride + 1, (img.shape[3] - kernel_size + 2 * padding) // stride + 1)\n",
        "    # return out\n",
        "      B, C, H, W = img.shape\n",
        "\n",
        "      out_H = (H + 2 * padding - kernel_size) // stride + 1\n",
        "      out_W = (W + 2 * padding - kernel_size) // stride + 1\n",
        "      H_padded = H + 2 * padding\n",
        "      H_padded = W + 2 * padding\n",
        "      img_padded = torch.zeros((B, C, H_padded, H_padded))\n",
        "      img_padded[:, :, padding:padding + H, padding:padding + W] = img\n",
        "      out = torch.zeros((B, num_filters, out_H, out_W))\n",
        "      filter_row = filter.view(num_filters, -1).t()\n",
        "\n",
        "      for x in range(out_H):\n",
        "          for y in range(out_W):\n",
        "              img_trans = img_padded[:, :, x * stride: x * stride + kernel_size, y * stride: y * stride + kernel_size]\n",
        "              img_trans_col = img_trans.reshape(B, -1)\n",
        "              for z in range(num_filters):\n",
        "                  out[:, z, x, y] = torch.matmul(img_trans_col, filter_row[:, z])\n",
        "      return out\n",
        "\n",
        "\n",
        "def conv2d_filter2row(img, filter, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter: filters, shape D x C x K x K\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   out: convoluted images, shape B x D x H x W\n",
        "    B, C, H, W = img.shape\n",
        "\n",
        "    out_H = (H + 2 * padding - kernel_size) // stride + 1\n",
        "    out_W = (W + 2 * padding - kernel_size) // stride + 1\n",
        "\n",
        "    H_padded = H + 2 * padding\n",
        "    W_padded = W + 2 * padding\n",
        "    img_padded = torch.zeros((B, C, H_padded, W_padded))\n",
        "\n",
        "    img_padded[ :, :, padding:padding + H, padding:padding + W ] = img\n",
        "\n",
        "    out = torch.zeros((B, num_filters, out_H, out_W))\n",
        "\n",
        "    for b in range(B):\n",
        "        for f in range(num_filters):\n",
        "            filter_matrix = torch.zeros((out_H * out_W, C*kernel_size*kernel_size))\n",
        "            for c in range(C):\n",
        "                filter_row = filter[f, c].flatten()\n",
        "                for i in range(out_H *out_W):\n",
        "                    start_row = (i // out_W) * stride\n",
        "                    start_col = (i % out_W) * stride\n",
        "                    for k in range(kernel_size * kernel_size):\n",
        "                        filter_matrix[i, c *  kernel_size * kernel_size + k] = filter_row[k]\n",
        "                        img_col = img_padded[b, c, start_row + k//kernel_size, start_col + k%kernel_size]\n",
        "                        out[b, f, i // out_W, i%out_W] = out[b, f, i // out_W, i%out_W] + img_col*filter_row[k]\n",
        "\n",
        "    return out\n",
        "\n",
        "\n",
        "\n",
        "def unit_test_conv2d(img, filter, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    # call your implemented \"im2col\" conv2D\n",
        "    y_im2col = conv2d_im2col(img, filter, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "\n",
        "    # ground truth conv2D\n",
        "    y_gt = F.conv2d(img, filter, stride=stride, padding=padding)\n",
        "\n",
        "    diff = (y_im2col - y_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of conv2d_im2col is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of conv2d_im2col is wrong!\")\n",
        "\n",
        "    # call your implemented \"im2col\" conv2D\n",
        "    y_filter2row = conv2d_filter2row(img, filter, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "\n",
        "    diff = (y_filter2row - y_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of conv2d_filter2row is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of conv2d_filter2row is wrong!\")\n",
        "\n",
        "\n",
        "unit_test_conv2d(img, filter, channel_size=C, num_filters=D, kernel_size=K, stride=1, padding=P)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0l4zm6mEXCxO"
      },
      "source": [
        "## 1.2 [20Pts] Implement the gradient of 2D convolution\n",
        "\n",
        "We now turn to the gradient of 2D convolution.\n",
        "In particular, given a batch of images $x$ with the shape $B \\times C \\times H \\times W$ and filters $h$ with shape $D \\times C \\times K \\times K$, we can view the convolution (zero-padding and stride 1) as a function\n",
        "\\begin{align}\n",
        "    y = f(h, x),\n",
        "\\end{align}\n",
        "that would produce an output tensor $y$ with shape $B \\times D \\times H \\times W$.\n",
        "\n",
        "If we vectorize $x$, $h$, and $y$, then the Jacobian matrix $∇f = [\\frac{\\partial y}{\\partial h}, \\frac{\\partial y}{\\partial x}]$ would be of shape $BDHW \\times (DCKK + BCHW)$.\n",
        "In practice, we almost never need to compute the Jacobian matrix directly as it is unnecessary for back-propagation.\n",
        "Instead, we often need to compute the product between the transposed Jacobian and a vector (a.k.a. vector-Jacobian product), i.e., ${\\frac{\\partial y}{\\partial h}}^{\\top} v$ and ${\\frac{\\partial y}{\\partial x}}^{\\top} v$.\n",
        "For example, the vector $v$ could the gradient of some loss $\\ell$ (scalar) w.r.t. the output above (i.e., $\\frac{\\partial \\ell}{\\partial y}$).\n",
        "\n",
        "**Task**:\n",
        "Given input images $x$, filters $h$, output $y$, and a vector $v$, implement the gradients ${\\frac{\\partial y}{\\partial h}}^{\\top} v$ and ${\\frac{\\partial y}{\\partial x}}^{\\top} v$ in the function below.\n",
        "\n",
        "**N.B.**: The function needs to return the gradients in the original shapes, i.e., ${\\frac{\\partial y}{\\partial h}}^{\\top} v$ should have the same shape as $h$ ($D \\times C \\times K \\times K$) and ${\\frac{\\partial y}{\\partial x}}^{\\top} v$ should have the same shape as $x$ ($B \\times C \\times H \\times W$)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vtz-vk95d9WR",
        "outputId": "8fb5f83b-a3fb-4205-bbde-b949136aba7d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of grad_img is correct!\n",
            "Your implementation of grad_filter is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following functions\n",
        "def grad_conv2d(img, filter, out, grad_out, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter: filters, shape D x C x K x K\n",
        "    #   out: convoluted images, shape B x D x H x W\n",
        "    #   grad_out: gradient w.r.t. output, shape B x D x H x W\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   grad_img: gradient w.r.t. img, shape B x C x H x W\n",
        "    #   grad_filer: gradient w.r.t. filter, shape D x C x K x K\n",
        "\n",
        "\n",
        "    # Initialize gradients with zeros\n",
        "    grad_img = torch.zeros_like(img)\n",
        "    grad_filter = torch.zeros_like(filter)\n",
        "\n",
        "    # Manually pad the image\n",
        "    #if padding > 0:\n",
        "    img_padded = torch.zeros(img.size(0), img.size(1), img.size(2) + 2 * padding, img.size(3) + 2 * padding)\n",
        "    img_padded[:, :, padding:-padding, padding:-padding] = img\n",
        "    #else:\n",
        "        #img_padded = img.clone()\n",
        "\n",
        "    B_O, C_O, H_O, W_O = out.shape\n",
        "\n",
        "    # Calculate the gradient w.r.t the filter\n",
        "    for b in range(B_O):  # iterate over the batch\n",
        "        for d in range(C_O):  # iterate over the number of filters\n",
        "            for i in range(H_O):  # iterate over the output height\n",
        "                for j in range(W_O):  # iterate over the output width\n",
        "\n",
        "                    # Get the window of the image corresponding to the current output\n",
        "                    window = img_padded[b, :, i*stride:i*stride+kernel_size, j*stride:j*stride+kernel_size]\n",
        "                    # Accumulate the gradients for the filter\n",
        "                    grad_filter[d] += grad_out[b, d, i, j] * window\n",
        "\n",
        "    B_I, C_I, H_I, W_I = grad_out.shape\n",
        "\n",
        "    # Calculate the gradient w.r.t the input image\n",
        "    for b in range(B_I):  # iterate over the batch\n",
        "        for d in range(C_I):  # iterate over the number of filters\n",
        "            for i in range(H_I):  # iterate over the output height\n",
        "                for j in range(W_I):  # iterate over the output width\n",
        "\n",
        "                    # Compute the start and end indices for the current slice of the input gradient\n",
        "                    i_start = max(i * stride - padding, 0)\n",
        "                    i_end = min(i * stride + kernel_size - padding, img.size(2))\n",
        "                    j_start = max(j * stride - padding, 0)\n",
        "                    j_end = min(j * stride + kernel_size - padding, img.size(3))\n",
        "\n",
        "                    # Compute the start and end indices for the filter\n",
        "                    f_i_start = max(padding - i * stride, 0)\n",
        "                    f_i_end = kernel_size - max((i * stride + kernel_size) - (img.size(2) + padding), 0)\n",
        "                    f_j_start = max(padding - j * stride, 0)\n",
        "                    f_j_end = kernel_size - max((j * stride + kernel_size) - (img.size(3) + padding), 0)\n",
        "\n",
        "                    # Update the input gradient\n",
        "                    grad_img[b, :, i_start:i_end, j_start:j_end] += filter[d, :, f_i_start:f_i_end, f_j_start:f_j_end] * grad_out[b, d, i, j]\n",
        "\n",
        "    return grad_img, grad_filter\n",
        "\n",
        "\n",
        "\n",
        "    # Compute the gradient w.r.t. the filter\n",
        "    for i in range(out.size(2)):  # Loop over the output height\n",
        "        for j in range(out.size(3)):  # Loop over the output width\n",
        "            # Select the current slice of the input image\n",
        "            h_start, w_start = i * stride, j * stride\n",
        "            h_end, w_end = h_start + kernel_size, w_start + kernel_size\n",
        "            img_slice = img_padded[:, :, h_start:h_end, w_start:w_end]\n",
        "\n",
        "            # Select the corresponding slice of the gradient of the output\n",
        "            grad_out_slice = grad_out[:, :, i, j].unsqueeze(1)\n",
        "\n",
        "            # Add contribution to the gradient of the filter\n",
        "            for b in range(img.size(0)):  # Loop over the batch size\n",
        "                grad_filter += (img_slice[b].unsqueeze(0) * grad_out_slice[b]).sum(0).unsqueeze(0)\n",
        "\n",
        "    return grad_img, grad_filter\n",
        "\n",
        "\n",
        "def unit_test_grad_conv2d(img, filter, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    filter.requires_grad = True\n",
        "    img.requires_grad = True\n",
        "\n",
        "    ### ground truth conv2D\n",
        "    img_out = F.conv2d(img, filter, stride=stride, padding=padding)\n",
        "\n",
        "    # create a random vector v\n",
        "    v = torch.randn_like(img_out)\n",
        "\n",
        "    # call your implemented \"grad_conv2d\" function\n",
        "    grad_img, grad_filter = grad_conv2d(img, filter, img_out, v, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "\n",
        "    # compute ground-truth gradients\n",
        "    grad_img_gt = torch.autograd.grad(img_out, img, grad_outputs=v, retain_graph=True)[0]\n",
        "    grad_filter_gt = torch.autograd.grad(img_out, filter, grad_outputs=v, retain_graph=True)[0]\n",
        "    #pdb.set_trace()\n",
        "\n",
        "    diff = (grad_img - grad_img_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_img is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_img is wrong!\")\n",
        "\n",
        "    diff = (grad_filter - grad_filter_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_filter is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_filter is wrong!\")\n",
        "\n",
        "unit_test_grad_conv2d(img, filter, channel_size=C, num_filters=D, kernel_size=K, stride=1, padding=P)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QI9J5CDCUSCc"
      },
      "source": [
        "---\n",
        "## 1.3 [10Pts]: Implement gradient checking via the finite difference approximation\n",
        "\n",
        "We verify the correctness of the implementation of gradient operators by calling PyTorch's autograd function.\n",
        "However, PyTorch's autograd function just calls the gradient operators implemented by the PyTorch team.\n",
        "How do they verify the correctness of their implementation?\n",
        "\n",
        "The answer is **finite difference approximation**.\n",
        "Following the setup in 1.2, given a batch of images $x$ with the shape $B \\times C \\times H \\times W$ and filters $h$ with shape $D \\times C \\times K \\times K$, we have the convolution (zero-padding and stride 1)\n",
        "\\begin{align}\n",
        "    y = f(h, x).\n",
        "\\end{align}\n",
        "\n",
        "Again, mentally vectorize $x$ and $y$ would help us understand the math.\n",
        "Given any vector $v$ with the same shape as $y$, we are interested in computing ${\\frac{\\partial y}{\\partial h}}^{\\top} v$ and ${\\frac{\\partial y}{\\partial x}}^{\\top} v$.\n",
        "These two gradients are equivalent to ${\\frac{\\partial \\ell}{\\partial h}}$ and ${\\frac{\\partial \\ell}{\\partial x}}$ where\n",
        "\\begin{align}\n",
        "    \\ell(h, x) = y^{\\top}v = f(h, x)^{\\top} v.\n",
        "\\end{align}\n",
        "Note that here $\\ell$ becomes a scalar.\n",
        "Based on Talyor's theorem, we have\n",
        "\\begin{align}\n",
        "    d^{\\top} \\frac{\\partial \\ell}{\\partial h} = \\lim_{\\epsilon → 0} \\frac{\\ell(h + \\epsilon \\cdot d, x) - \\ell(h - \\epsilon \\cdot d, x)}{2 ϵ},\n",
        "\\end{align}\n",
        "where $d$ could be any direction vector and $ϵ$ is a scalar.\n",
        "For our purpose, we just need to set $d$ to be the unit vector to compute the per-dimension value of $\\frac{\\partial \\ell}{\\partial h}$.\n",
        "Specifically, if we set $d$ as the $i$-th unit vector $e_i$, i.e., $d[i] = 1$ and $d[j] = 0, \\forall j \\neq i$, we can then compute\n",
        "\\begin{align}\n",
        "    \\frac{\\partial \\ell}{\\partial h}[i] &= \\lim_{\\epsilon → 0} \\frac{\\ell(h + \\epsilon \\cdot e_i, x) - \\ell(h - \\epsilon \\cdot e_i, x)}{2 ϵ} \\\\\n",
        "    & ≈ \\frac{\\ell(h + \\epsilon \\cdot e_i, x) - \\ell(h - \\epsilon \\cdot e_i, x)}{2 ϵ}.\n",
        "\\end{align}\n",
        "\n",
        "**Task**: Implement the finite-difference based gradient checker for ${\\frac{\\partial \\ell}{\\partial h}}$ and ${\\frac{\\partial \\ell}{\\partial x}}$.\n",
        "\n",
        "\n",
        "**N.B.**: For efficiency consideration in the unit test, you can use F.conv2d to compute the convolution in your implementation of *grad_checker*. This assignment is to let you understand how to implement finte-difference. But in pratice, if we want to verify our implementation of conv2d, then we should use our conv2d instead of F.conv2d from PyTorch.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qgeHgGn1JbBn",
        "outputId": "35c7e881-73b3-4a59-eeb1-2e10ae785b93"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of grad_img is correct!\n",
            "Your implementation of grad_filter is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following functions\n",
        "def grad_checker(img, filter, conv, grad_out, epsilon=1.0e-5, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter: filters, shape D x C x K x K\n",
        "    #   conv: convolution function\n",
        "    #   out: convoluted images, shape B x D x H x W\n",
        "    #   grad_out: gradient w.r.t. output, shape B x D x H x W\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   grad_img: gradient w.r.t. img, shape B x C x H x W\n",
        "    #   grad_filer: gradient w.r.t. filter, shape D x C x K x K\n",
        "    B, C, H, W = img.size()\n",
        "    D = filter.size(0)\n",
        "\n",
        "    # Initialize gradients with zeros\n",
        "    grad_img = torch.zeros_like(img)\n",
        "    grad_filter = torch.zeros_like(filter)\n",
        "\n",
        "    # Loop over each element in the filter\n",
        "    for i in range(D):\n",
        "        for j in range(channel_size):\n",
        "            for k in range(kernel_size):\n",
        "                for l in range(kernel_size):\n",
        "                    # Perturb the filter element up and down\n",
        "                    filter_d = filter.clone()\n",
        "                    filter_d[i, j, k, l] += epsilon\n",
        "                    out_d = conv(img, filter_d, stride=stride, padding=padding)\n",
        "                    loss_d = (grad_out * out_d).sum()\n",
        "\n",
        "                    filter_d[i, j, k, l] -= 2 * epsilon\n",
        "                    out_d = conv(img, filter_d, stride=stride, padding=padding)\n",
        "                    loss_d -= (grad_out * out_d).sum()\n",
        "\n",
        "                    # Compute the gradient\n",
        "                    grad_filter[i, j, k, l] = loss_d / (2 * epsilon)\n",
        "\n",
        "    # Loop over each element in the image\n",
        "    for b in range(B):\n",
        "        for c in range(C):\n",
        "            for h in range(H):\n",
        "                for w in range(W):\n",
        "                    # Perturb the image pixel up and down\n",
        "                    img_d = img.clone()\n",
        "                    img_d[b, c, h, w] += epsilon\n",
        "                    out_d = conv(img_d, filter, stride=stride, padding=padding)\n",
        "                    loss_d = (grad_out * out_d).sum()\n",
        "\n",
        "                    img_d[b, c, h, w] -= 2 * epsilon\n",
        "                    out_d = conv(img_d, filter, stride=stride, padding=padding)\n",
        "                    loss_d -= (grad_out * out_d).sum()\n",
        "\n",
        "                    # Compute the gradient\n",
        "                    grad_img[b, c, h, w] = loss_d / (2 * epsilon)\n",
        "\n",
        "    return grad_img, grad_filter\n",
        "\n",
        "\n",
        "def unit_test_grad_checker(img, filter, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    epsilon = 1.0e-5\n",
        "    filter.requires_grad = True\n",
        "    img.requires_grad = True\n",
        "\n",
        "    ### ground truth conv2D\n",
        "    img_out = F.conv2d(img, filter, stride=stride, padding=padding)\n",
        "\n",
        "    # create a random vector v\n",
        "    v = torch.randn_like(img_out)\n",
        "\n",
        "    # call your implemented \"grad_checker\" function\n",
        "    grad_img, grad_filter = grad_checker(img, filter, F.conv2d, v, epsilon=epsilon, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "\n",
        "    # compute ground-truth gradients\n",
        "    grad_img_gt = torch.autograd.grad(img_out, img, grad_outputs=v, retain_graph=True)[0]\n",
        "    grad_filter_gt = torch.autograd.grad(img_out, filter, grad_outputs=v, retain_graph=True)[0]\n",
        "    #pdb.set_trace()\n",
        "\n",
        "    diff = (grad_img - grad_img_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_img is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_img is wrong!\")\n",
        "\n",
        "    diff = (grad_filter - grad_filter_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_filter is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_filter is wrong!\")\n",
        "\n",
        "unit_test_grad_checker(img, filter, channel_size=C, num_filters=D, kernel_size=K, stride=1, padding=P)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rJ3DtRlwUlzY"
      },
      "source": [
        "---\n",
        "#Q2 [5Pts]: Implement ReLU and its gradient\n",
        "\n",
        "**Task**: Implement ReLU operator, i.e., $f(x) = max(x, 0)$, and its gradient operator ${\\frac{\\partial f}{\\partial x}}^{\\top} v$ for any given tensor $v$ that is of the same shape as $x$.\n",
        "\n",
        "**N.B.**: For simplicity, we can assume the input $x$ is of shape $B \\times C \\times H \\times W$ as before."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_ruiFoTsJb9w",
        "outputId": "122b4077-5ce6-4c4b-d402-577a40f34ff8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of func_relu is correct!\n",
            "Your implementation of grad_relu is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following functions\n",
        "def func_relu(x):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   x: input, shape B x C x H x W\n",
        "    #\n",
        "    # Returns:\n",
        "    #   y: output, shape B x C x H x W\n",
        "    y = torch.max(x, torch.zeros_like(x))\n",
        "    return y\n",
        "\n",
        "\n",
        "def grad_relu(x, y, grad_out):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   x: input, shape B x C x H x W\n",
        "    #   y: output, shape B x C x H x W\n",
        "    #   grad_out: gradient w.r.t. output y, shape B x D x H x W\n",
        "    #\n",
        "    # Returns:\n",
        "    #   grad_x: gradient w.r.t. x, shape B x C x H x W\n",
        "    grad_x = torch.zeros_like(x)\n",
        "    grad_x[x > 0] = 1\n",
        "    return grad_x * grad_out\n",
        "\n",
        "\n",
        "def unit_test_relu(x):\n",
        "    x.requires_grad = True\n",
        "\n",
        "    # call your implemented \"func_relu\" function\n",
        "    y = func_relu(x)\n",
        "\n",
        "    # ground truth ReLU\n",
        "    y_gt = F.relu(x)\n",
        "\n",
        "    diff = (y - y_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of func_relu is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of func_relu is wrong!\")\n",
        "\n",
        "    # create a random vector v\n",
        "    v = torch.randn_like(y)\n",
        "\n",
        "    # call your implemented \"grad_relu\" function\n",
        "    grad_x = grad_relu(x, y, v)\n",
        "\n",
        "    # compute ground-truth gradients\n",
        "    grad_x_gt = torch.autograd.grad(y_gt, x, grad_outputs=v, retain_graph=True)[0]\n",
        "\n",
        "    diff = (grad_x - grad_x_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_relu is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_relu is wrong!\")\n",
        "\n",
        "unit_test_relu(torch.randn_like(img))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jiVr7-xJU066"
      },
      "source": [
        "---\n",
        "#Q3 [20Pts]: Implement Batch-normalization (BN) for convolution and its gradient\n",
        "\n",
        "Given a batch of input images $x$ with shape $B \\times C \\times H \\times W$, we compute a single mean and a single standard deviation per channel as below,\n",
        "\\begin{align}\n",
        "    \\mu[c] &= \\frac{1}{BHW} \\sum_{i=1}^{B} \\sum_{m=1}^{H} \\sum_{n=1}^{W} x[i, c, m, n] \\\\\n",
        "    \\sigma^2[c] &= \\frac{1}{BHW} \\sum_{i=1}^{B} \\sum_{m=1}^{H} \\sum_{n=1}^{W} (x[i, c, m, n] - \\mu[c])^2.\n",
        "\\end{align}\n",
        "Then we perform BN, $y = f(x, \\beta, \\gamma)$, as,\n",
        "\\begin{align}\n",
        "    y[i,c,m,n] &= \\gamma[c] \\frac{x[i,c,m,n] - \\mu[c]}{\\sqrt{\\sigma^2[c] + \\epsilon}} + \\beta[c],\n",
        "\\end{align}\n",
        "where $\\gamma$ and $\\beta$ are learnable parameters are of shape $C$.\n",
        "$ϵ$ is a constant.\n",
        "\n",
        "**Task**: For simplicity, we fix the learnable parameters as $\\gamma = 1$ and $\\beta = 0$.\n",
        "Implement BN for convolution and its gradient operators ${\\frac{\\partial f}{\\partial x}}^{\\top} v$ for any $v$ that is compatible with the matrix multiplication.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3Sou_6LzJcxS",
        "outputId": "1e58d38d-7fa8-4797-fdbf-386fcd60fb80"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of func_batch_norm is correct!\n",
            "Your implementation of grad_batch_norm is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following functions\n",
        "def func_batch_norm(x, epsilon=1.0e-5):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   x: input, shape B x C x H x W\n",
        "    #   epsilon: constant, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   y: output, shape B x C x H x W\n",
        "    mu = x.mean(dim=(0, 2, 3), keepdim=True)\n",
        "    var = x.var(dim=(0, 2, 3), keepdim=True, unbiased=False)\n",
        "\n",
        "    # Normalize the batch\n",
        "    x_normalized = (x - mu) / torch.sqrt(var + epsilon)\n",
        "\n",
        "    # Apply scale and shift (with gamma=1 and beta=0)\n",
        "    gamma = torch.ones(x.shape[1], dtype=x.dtype, device=x.device).view(1, -1, 1, 1)\n",
        "    beta = torch.zeros(x.shape[1], dtype=x.dtype, device=x.device).view(1, -1, 1, 1)\n",
        "    y = gamma * x_normalized + beta\n",
        "\n",
        "    return y\n",
        "\n",
        "\n",
        "def grad_batch_norm(x, y, grad_out, epsilon=1.0e-5):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   x: input, shape B x C x H x W\n",
        "    #   y: output, shape B x C x H x W\n",
        "    #   grad_out: gradient w.r.t. output y, shape B x D x H x W\n",
        "    #   epsilon: constant, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   grad_x: gradient w.r.t. x, shape B x C x H x W\n",
        "    B, C, H, W = x.shape\n",
        "    mu = x.mean(dim=(0, 2, 3), keepdim=True)\n",
        "    var = x.var(dim=(0, 2, 3), keepdim=True, unbiased=False)\n",
        "    x_normalized = (x - mu) / torch.sqrt(var + epsilon)\n",
        "\n",
        "    # Partial derivatives needed for gradient computation\n",
        "    grad_x_normalized = grad_out\n",
        "    grad_var = (-0.5 * ((var + epsilon) ** (-1.5)) * (grad_x_normalized * (x - mu)).sum(dim=(0, 2, 3), keepdim=True)).sum(dim=1, keepdim=True)\n",
        "    grad_mu = (-1.0 / torch.sqrt(var + epsilon) * grad_x_normalized.sum(dim=(0, 2, 3), keepdim=True)).sum(dim=1, keepdim=True)\n",
        "\n",
        "    # Compute the gradient with respect to x\n",
        "    grad_x = (1.0 / torch.sqrt(var + epsilon) * grad_x_normalized) + \\\n",
        "             (2.0 / (B * H * W) * grad_var * (x - mu)) + \\\n",
        "             (1.0 / (B * H * W) * grad_mu)\n",
        "\n",
        "    return grad_x\n",
        "\n",
        "\n",
        "def unit_test_batch_norm(x):\n",
        "    x.requires_grad = True\n",
        "    epsilon = 1e-5\n",
        "\n",
        "    # call your implemented \"func_batch_norm\" function\n",
        "    y = func_batch_norm(x, epsilon=epsilon)\n",
        "\n",
        "    # ground truth ReLU\n",
        "    BN_gt = nn.BatchNorm2d(x.shape[1], eps=epsilon, momentum=1.0, affine=False, track_running_stats=False)\n",
        "    y_gt = BN_gt(x)\n",
        "\n",
        "    diff = (y - y_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of func_batch_norm is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of func_batch_norm is wrong!\")\n",
        "\n",
        "    # create a random vector v\n",
        "    v = torch.randn_like(y)\n",
        "\n",
        "    # call your implemented \"grad_batch_norm\" function\n",
        "    grad_x = grad_batch_norm(x, y, v, epsilon=epsilon)\n",
        "\n",
        "    # compute ground-truth gradients\n",
        "    grad_x_gt = torch.autograd.grad(y_gt, x, grad_outputs=v, retain_graph=True)[0]\n",
        "\n",
        "    diff = (grad_x - grad_x_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_batch_norm is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_batch_norm is wrong!\")\n",
        "\n",
        "unit_test_batch_norm(torch.randn_like(img))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2m-KuDI7U2zb"
      },
      "source": [
        "---\n",
        "#Q4 [15Pts]: Implement a simple CNN and back-propagation (BP)\n",
        "\n",
        "Now we are ready to build a deep CNN and learn it with back-propagation.\n",
        "In particular, let us build a simple CNN with following architecture:\n",
        "\n",
        "Conv $→$ BN $→$ ReLU $→$ Conv $→$ BN $→$ ReLU $→$ Linear.\n",
        "\n",
        "Here, for all layers, the convolutions are the same as before (i.e., kernel size $3 \\times 3$, zero-padding, number of filters $D = 2$, and stride 1), the BNs are without learnable $\\gamma$ and $\\beta$, and the last linear layer would map whatever input dimension to $10$ classes in MNIST.\n",
        "\n",
        "**Task**: Implement the above CNN, compute the cross-entropy loss, and compute gradient of the loss w.r.t. filter weights.\n",
        "\n",
        "**N.B.**: You can use F.cross_entropy provided by PyTorch. But for other operators like Conv, BN, and ReLU and their gradietns, you should use your previous implementations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OnpM4Yj_JdVO",
        "outputId": "86b2a9de-2b6f-40d0-cd15-ff4493d18f5b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your implementation of grad_filter_1 is wrong!\n",
            "Your implementation of grad_filter_2 is wrong!\n",
            "Your implementation of grad_weight is correct!\n"
          ]
        }
      ],
      "source": [
        "## implement the following two functions\n",
        "def CNN(img, filter_1, filter_2, weight, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter_1: filters at 1st layer, shape D x C x K x K\n",
        "    #   filter_2: filters at 2nd layer, shape D x C x K x K\n",
        "    #   weight: weights of linear readout layer, shape ? x 10\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   out: logits, shape B x 10\n",
        "    B, C, H, W = img.shape\n",
        "\n",
        "    ### 1st layer\n",
        "    conv1_out = conv2d_filter2row(img, filter_1, C, num_filters, kernel_size, stride, padding)\n",
        "    bn1_out = func_batch_norm(conv1_out)\n",
        "    relu1_out = func_relu(bn1_out)\n",
        "\n",
        "    ### 2nd layer\n",
        "    conv2_out = conv2d_filter2row(relu1_out, filter_2, num_filters, num_filters, kernel_size, stride, padding)\n",
        "    bn2_out = func_batch_norm(conv2_out)\n",
        "    relu2_out = func_relu(bn2_out)\n",
        "\n",
        "    ### linear readout\n",
        "    relu2_flat = relu2_out.view(B, -1)\n",
        "    out = torch.matmul(relu2_flat, weight)\n",
        "    return out\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def grad_CNN(img, filter_1, filter_2, weight, grad_loss, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    ### Fill in this function ###\n",
        "    # Args:\n",
        "    #   img: images, shape B x C x H x W\n",
        "    #   filter_1: filters at 1st layer, shape D x C x K x K\n",
        "    #   filter_2: filters at 2nd layer, shape D x C x K x K\n",
        "    #   weight: weights of linear readout layer, shape ? x 10\n",
        "    #   grad_loss: gradient of loss w.r.t. logits, shape B x 10\n",
        "    #   channel_size: number of channels, scalar (C)\n",
        "    #   num_filters: number of filters, scalar (D)\n",
        "    #   kernel_size: kernel size, scalar (K)\n",
        "    #   stride: stride size, scalar\n",
        "    #   padding: padding size, scalar\n",
        "    #\n",
        "    # Returns:\n",
        "    #   grad_filter_1: filters, shape D x C x K x K\n",
        "    #   grad_filter_2: filters, shape D x C x K x K\n",
        "    #   grad_weight: weight, shape ? x 10\n",
        "\n",
        "\n",
        "    B, C, H, W = img.shape\n",
        "    conv1_out = conv2d_filter2row(img, filter_1, C, num_filters, kernel_size, stride=stride, padding=padding)\n",
        "    bn1_out = func_batch_norm(conv1_out)\n",
        "    relu1_out = func_relu(bn1_out)\n",
        "\n",
        "    # 2nd layer: Convolution, Batch Norm, ReLU\n",
        "    conv2_out = conv2d_filter2row(relu1_out, filter_2, num_filters, num_filters, kernel_size, stride=stride, padding=padding)\n",
        "    bn2_out = func_batch_norm(conv2_out)\n",
        "    relu2_out = func_relu(bn2_out)\n",
        "\n",
        "    # Flatten the output for the linear layer\n",
        "    flat_out = relu2_out.view(relu2_out.size(0), -1)\n",
        "\n",
        "    # Linear readout\n",
        "    logits = torch.matmul(flat_out, weight)\n",
        "\n",
        "     # Backward pass\n",
        "    # Gradient through linear readout layer\n",
        "    grad_flat_out = torch.matmul(grad_loss, weight.t())\n",
        "    grad_weight = torch.matmul(flat_out.t(), grad_loss)\n",
        "\n",
        "    # Unflatten the gradient to the shape of relu2_out for further backpropagation\n",
        "    grad_relu2_out = grad_flat_out.view_as(relu2_out)\n",
        "\n",
        "    # Gradients through the second ReLU\n",
        "    grad_bn2_out = grad_relu(relu2_out, bn2_out, grad_relu2_out)\n",
        "\n",
        "    # Gradients through the second batch normalization\n",
        "    grad_conv2_out = grad_batch_norm(conv2_out, bn2_out, grad_bn2_out)\n",
        "\n",
        "    # Gradients through the second convolutional layer\n",
        "    grad_relu1_out, grad_filter_2 = grad_conv2d(relu1_out, filter_2, conv2_out, grad_conv2_out, num_filters, num_filters, kernel_size, stride, padding)\n",
        "\n",
        "    # Gradients through the first ReLU\n",
        "    grad_bn1_out = grad_relu(relu1_out, bn1_out, grad_relu1_out)\n",
        "\n",
        "    # Gradients through the first batch normalization\n",
        "    grad_conv1_out = grad_batch_norm(conv1_out, bn1_out, grad_bn1_out)\n",
        "\n",
        "    # Gradients through the first convolutional layer\n",
        "    grad_conv1_out, grad_filter_1 = grad_conv2d(img, filter_1, conv1_out, grad_conv1_out, C, num_filters, kernel_size, stride, padding)\n",
        "\n",
        "    return grad_filter_1, grad_filter_2, grad_weight\n",
        "\n",
        "\n",
        "\n",
        "def unit_test_CNN(img, label, filter_1, filter_2, weight, channel_size=1, num_filters=1, kernel_size=3, stride=1, padding=1):\n",
        "    # call your implemented \"CNN\"\n",
        "    img.requires_grad_()\n",
        "    filter_1.requires_grad_()\n",
        "    filter_2.requires_grad_()\n",
        "    weight.requires_grad_()\n",
        "    y = CNN(img, filter_1, filter_2, weight, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "    y.requires_grad_()\n",
        "\n",
        "    # compute loss function\n",
        "    loss = F.cross_entropy(y, label).mean()\n",
        "    loss.requires_grad_()\n",
        "\n",
        "    # compute gradient of loss w.r.t. logits\n",
        "    grad_loss = torch.autograd.grad(loss, y, retain_graph=True)[0]\n",
        "\n",
        "    # call your implemented \"grad_batch_norm\" function\n",
        "    grad_filter_1, grad_filter_2, grad_weight = grad_CNN(img, filter_1, filter_2, weight, grad_loss, channel_size=channel_size, num_filters=num_filters, kernel_size=kernel_size, stride=stride, padding=padding)\n",
        "\n",
        "    # compute ground-truth gradients\n",
        "    grad_filter_1_gt = torch.autograd.grad(loss, filter_1, retain_graph=True)[0]\n",
        "    grad_filter_2_gt = torch.autograd.grad(loss, filter_2, retain_graph=True)[0]\n",
        "    grad_weight_gt = torch.autograd.grad(loss, weight, retain_graph=True)[0]\n",
        "\n",
        "    diff = (grad_filter_1 - grad_filter_1_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_filter_1 is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_filter_1 is wrong!\")\n",
        "\n",
        "    diff = (grad_filter_2 - grad_filter_2_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_filter_2 is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_filter_2 is wrong!\")\n",
        "\n",
        "    diff = (grad_weight - grad_weight_gt).norm()\n",
        "    if diff < 1.0e-5:\n",
        "        print(\"Your implementation of grad_weight is correct!\")\n",
        "    else:\n",
        "        print(\"Your implementation of grad_weight is wrong!\")\n",
        "\n",
        "\n",
        "filter_1 = torch.randn(D, C, K, K) # filter shape: D x C x K x K\n",
        "filter_2 = torch.randn(D, D, K, K) # filter shape: D x C x K x K\n",
        "\n",
        "### compute the correct shape and then replace None with it ###\n",
        "weight = torch.randn(D*img.shape[2]*img.shape[3], 10) # weight of the last linear layer\n",
        "\n",
        "unit_test_CNN(img, label, filter_1, filter_2, weight, channel_size=C, num_filters=D, kernel_size=K, stride=1, padding=P)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "9AS6zKmYSwVD"
      ],
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3.11.5 64-bit",
      "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.11.5"
    },
    "vscode": {
      "interpreter": {
        "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}