2024_11_04_Tyriek_Homework3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Tcraigs1/ai-science-training-series/blob/main/2024_11_04_Tyriek_Homework3.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T4Jh7pqpqYob"
      },
      "source": [
        "# Convolutional Neural Networks\n",
        "\n",
        "Notebook by Corey Adams, some modifications by Bethany Lusch\n",
        "\n",
        "Up until transformers, convolutions were *the* state of the art in computer vision.  In many ways and applications they still are!\n",
        "\n",
        "Large Language Models, which are what we'll focus on the rest of the series after this lecture, are really good at ordered, *tokenized data.  But there is lots of data that isn't _implicitly_ ordered like `images`, and their more general cousins `graphs`.\n",
        "\n",
        "Today's lecture focuses on computer vision models, and particularly on convolutional neural networks.  There are a ton of applications you can do with these, and not nearly enough time to get into them.  Check out the extra references file to see some publications to get you started if you want to learn more.\n",
        "\n",
        "Tip: this notebook is much faster on the GPU!\n",
        "\n",
        "\n",
        "## Convolutional Networks: A brief historical context\n",
        "\n",
        "![ImageNet Accuracy by Yearh](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/ImageNet.png?raw=1)\n",
        "\n",
        "[reference](https://www.researchgate.net/publication/332452649_A_Roadmap_for_Foundational_Research_on_Artificial_Intelligence_in_Medical_Imaging_From_the_2018_NIHRSNAACRThe_Academy_Workshop)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RTcCPYxqYoe"
      },
      "outputs": [],
      "source": [
        "import torch, torchvision"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-C_fK1gqYof"
      },
      "source": [
        "# Convolutional Building Blocks\n",
        "\n",
        "We're going to go through some examples of building blocks for convolutional networks.  To help illustate some of these, let's use an image for examples:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B7IZkjPmqYog"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "# wget line useful in Google Colab\n",
        "! wget https://raw.githubusercontent.com/argonne-lcf/ai-science-training-series/main/03_advanced_neural_networks/ALCF-Staff.jpg\n",
        "alcf_image = Image.open(\"ALCF-Staff.jpg\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ar-8Q0SLqYog"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "figure = plt.figure(figsize=(20,20))\n",
        "plt.imshow(alcf_image)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s5_qX1iFqYog"
      },
      "source": [
        "## Convolutions\n",
        "\n",
        "Convolutions are a restriction of - and a specialization of - dense linear layers.  A convolution of an image produces another image, and each output pixel is a function of only it's local neighborhood of points.  This is called an _inductive bias_ and is a big reason why convolutions work for image data: neighboring pixels are correlated and you can operate on just those pixels at a time.\n",
        "\n",
        "See examples of convolutions [here](https://github.com/vdumoulin/conv_arithmetic)\n",
        "\n",
        "![image-2.png](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/conv_eqn.png?raw=1)\n",
        "\n",
        "![image.png](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/conv.png?raw=1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "96_lB-GtqYoh"
      },
      "outputs": [],
      "source": [
        "# Let's apply a convolution to the ALCF Staff photo:\n",
        "alcf_tensor = torchvision.transforms.ToTensor()(alcf_image)\n",
        "\n",
        "# Reshape the tensor to have a batch size of 1:\n",
        "alcf_tensor = alcf_tensor.reshape((1,) + alcf_tensor.shape)\n",
        "\n",
        "# Create a random convolution:\n",
        "# shape is: (channels_in, channels_out, kernel_x, kernel_y)\n",
        "conv_random = torch.rand((3,3,15,15))\n",
        "\n",
        "alcf_rand = torch.nn.functional.conv2d(alcf_tensor, conv_random)\n",
        "alcf_rand = (1./alcf_rand.max()) * alcf_rand\n",
        "print(alcf_rand.shape)\n",
        "alcf_rand = alcf_rand.reshape(alcf_rand.shape[1:])\n",
        "\n",
        "print(alcf_tensor.shape)\n",
        "\n",
        "rand_image = alcf_rand.permute((1,2,0)).cpu()\n",
        "\n",
        "figure = plt.figure(figsize=(20,20))\n",
        "\n",
        "plt.imshow(rand_image)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dh7iDQUeqYoh"
      },
      "source": [
        "## Normalization\n",
        "\n",
        "![Batch Norm](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/batch_norm.png?raw=1)\n",
        "Reference: [Normalizations](https://arxiv.org/pdf/1903.10520.pdf)\n",
        "\n",
        "Normalization is the act of transforming the mean and moment of your data to standard values (usually 0.0 and 1.0).  It's particularly useful in machine learning since it stabilizes training, and allows higher learning rates.\n",
        "\n",
        "![Batch Normalization accelerates training](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/batch_norm_effect.png?raw=1)\n",
        "\n",
        "Reference: [Batch Norm](https://arxiv.org/pdf/1502.03167.pdf)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IzIL-5uhqYoi"
      },
      "outputs": [],
      "source": [
        "# Let's apply a normalization to the ALCF Staff photo:\n",
        "alcf_tensor = torchvision.transforms.ToTensor()(alcf_image)\n",
        "\n",
        "# Reshape the tensor to have a batch size of 1:\n",
        "alcf_tensor = alcf_tensor.reshape((1,) + alcf_tensor.shape)\n",
        "\n",
        "\n",
        "alcf_rand = torch.nn.functional.normalize(alcf_tensor)\n",
        "alcf_rand = alcf_rand.reshape(alcf_rand.shape[1:])\n",
        "\n",
        "print(alcf_tensor.shape)\n",
        "\n",
        "rand_image = alcf_rand.permute((1,2,0)).cpu()\n",
        "\n",
        "figure = plt.figure(figsize=(20,20))\n",
        "\n",
        "plt.imshow(rand_image)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Otq0ASS0qYoi"
      },
      "source": [
        "\n",
        "## Downsampling (And upsampling)\n",
        "\n",
        "Downsampling is a critical component of convolutional and many vision models.  Because of the local-only nature of convolutional filters, learning large-range features can be too slow for convergence.  Downsampling of layers can bring information from far away closer, effectively changing what it means to be \"local\" as the input to a convolution.\n",
        "\n",
        "![Convolutional Pooling](conv_pooling.png \"Pooling\")\n",
        "\n",
        "[Reference](https://www.researchgate.net/publication/333593451_Application_of_Transfer_Learning_Using_Convolutional_Neural_Network_Method_for_Early_Detection_of_Terry's_Nail)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZrgdRzYDqYoi"
      },
      "outputs": [],
      "source": [
        "# Let's apply a normalization to the ALCF Staff photo:\n",
        "alcf_tensor = torchvision.transforms.ToTensor()(alcf_image)\n",
        "\n",
        "# Reshape the tensor to have a batch size of 1:\n",
        "alcf_tensor = alcf_tensor.reshape((1,) + alcf_tensor.shape)\n",
        "\n",
        "\n",
        "alcf_rand = torch.nn.functional.max_pool2d(alcf_tensor, 2)\n",
        "alcf_rand = alcf_rand.reshape(alcf_rand.shape[1:])\n",
        "\n",
        "print(alcf_tensor.shape)\n",
        "\n",
        "rand_image = alcf_rand.permute((1,2,0)).cpu()\n",
        "\n",
        "figure = plt.figure(figsize=(20,20))\n",
        "\n",
        "plt.imshow(rand_image)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vyJy6VWBqYoi"
      },
      "source": [
        "## Residual Connections\n",
        "\n",
        "One issue, quickly encountered when making convolutional networks deeper and deeper, is the \"Vanishing Gradients\" problem.  As layers were stacked on top of each other, the size of updates dimished at the earlier layers of a convolutional network.  The paper \"Deep Residual Learning for Image Recognition\" solved this by introduction \"residual connections\" as skip layers.\n",
        "\n",
        "\n",
        "Reference: [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)\n",
        "\n",
        "![Residual Layer](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/residual_layer.png?raw=1)\n",
        "\n",
        "\n",
        "Compare the performance of the models before and after the introduction of these layers:\n",
        "\n",
        "![Resnet Performance vs. Plain network performance](https://github.com/Tcraigs1/ai-science-training-series/blob/main/03_advanced_neural_networks/resnet_comparison.png?raw=1)\n",
        "\n",
        "If you have time to read only one paper on computer vision, make it this one!  Resnet was the first model to beat human accuracy on ImageNet and is one of the most impactful papers in AI ever published."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BufELvkQqYoj"
      },
      "source": [
        "# Building a ConvNet\n",
        "\n",
        "In this section we'll build and apply a conv net to the mnist dataset.  The layers here are loosely based off of the ConvNext architecture.  Why?  Because we're getting into LLM's soon, and this ConvNet uses LLM features.  ConvNext is an update to the ResNet architecture that outperforms it.\n",
        "\n",
        "[ConvNext](https://arxiv.org/abs/2201.03545)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v4wEu8u9qYoj"
      },
      "source": [
        "The dataset here is CIFAR-10 - slightly harder than MNIST but still relatively easy and computationally tractable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NSWJKnCRqYoj"
      },
      "outputs": [],
      "source": [
        "from torchvision.transforms import v2\n",
        "training_data = torchvision.datasets.CIFAR10(\n",
        "    # Polaris: root=\"/lus/eagle/projects/datasets/CIFAR-10/\",\n",
        "    # Polaris: download=False,\n",
        "    root=\"data\",\n",
        "    download=True,\n",
        "    train=True,\n",
        "    transform=v2.Compose([\n",
        "        v2.ToTensor(),\n",
        "        v2.RandomHorizontalFlip(),\n",
        "        v2.RandomResizedCrop(size=32, scale=[0.85,1.0], antialias=False),\n",
        "        v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n",
        "    ])\n",
        ")\n",
        "\n",
        "test_data = torchvision.datasets.CIFAR10(\n",
        "    # Polaris: root=\"/lus/eagle/projects/datasets/CIFAR-10/\",\n",
        "    # Polaris: download=False,\n",
        "    root=\"data\",\n",
        "    download=True,\n",
        "    train=False,\n",
        "    transform=torchvision.transforms.ToTensor()\n",
        ")\n",
        "\n",
        "training_data, validation_data = torch.utils.data.random_split(training_data, [0.8, 0.2], generator=torch.Generator().manual_seed(55))\n",
        "\n",
        "batch_size = 128\n",
        "\n",
        "# The dataloader makes our dataset iterable\n",
        "train_dataloader = torch.utils.data.DataLoader(training_data,\n",
        "    batch_size=batch_size,\n",
        "    pin_memory=True,\n",
        "    shuffle=True,\n",
        "    num_workers=2)\n",
        "val_dataloader = torch.utils.data.DataLoader(validation_data,\n",
        "    batch_size=batch_size,\n",
        "    pin_memory=True,\n",
        "    shuffle=False,\n",
        "    num_workers=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_itU5xJtqYoj"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_gdB8jq3qYoj"
      },
      "outputs": [],
      "source": [
        "batch, (X, Y) = next(enumerate(train_dataloader))\n",
        "plt.imshow(X[0].cpu().permute((1,2,0))); plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kgyGNQHNqYoj"
      },
      "source": [
        "This code below is important as our models get bigger: this is wrapping the pytorch data loaders to put the data onto the GPU!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "crIEnwa1qYoj"
      },
      "outputs": [],
      "source": [
        "dev = torch.device(\n",
        "    \"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "\n",
        "\n",
        "def preprocess(x, y):\n",
        "    # CIFAR-10 is *color* images so 3 layers!\n",
        "    return x.view(-1, 3, 32, 32).to(dev), y.to(dev)\n",
        "\n",
        "\n",
        "class WrappedDataLoader:\n",
        "    def __init__(self, dl, func):\n",
        "        self.dl = dl\n",
        "        self.func = func\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.dl)\n",
        "\n",
        "    def __iter__(self):\n",
        "        for b in self.dl:\n",
        "            yield (self.func(*b))\n",
        "\n",
        "\n",
        "train_dataloader = WrappedDataLoader(train_dataloader, preprocess)\n",
        "val_dataloader = WrappedDataLoader(val_dataloader, preprocess)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N0i9p0RHqYoj"
      },
      "outputs": [],
      "source": [
        "from torch import nn\n",
        "\n",
        "\n",
        "class Downsampler(nn.Module):\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, shape, stride=2):\n",
        "        super(Downsampler, self).__init__()\n",
        "\n",
        "        self.norm = nn.LayerNorm([in_channels, *shape])\n",
        "\n",
        "        self.downsample = nn.Conv2d(\n",
        "            in_channels=in_channels,\n",
        "            out_channels=out_channels,\n",
        "            kernel_size = stride,\n",
        "            stride = stride,\n",
        "        )\n",
        "\n",
        "    def forward(self, inputs):\n",
        "\n",
        "\n",
        "        return self.downsample(self.norm(inputs))\n",
        "\n",
        "\n",
        "\n",
        "class ConvNextBlock(nn.Module):\n",
        "    \"\"\"This block of operations is loosely based on this paper:\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "    def __init__(self, in_channels, shape):\n",
        "        super(ConvNextBlock, self).__init__()\n",
        "\n",
        "        # Depthwise, seperable convolution with a large number of output filters:\n",
        "        self.conv1 = nn.Conv2d(in_channels=in_channels,\n",
        "                                     out_channels=in_channels,\n",
        "                                     groups=in_channels,\n",
        "                                     kernel_size=[7,7],\n",
        "                                     padding='same' )\n",
        "\n",
        "        self.norm = nn.LayerNorm([in_channels, *shape])\n",
        "\n",
        "        # Two more convolutions:\n",
        "        self.conv2 = nn.Conv2d(in_channels=in_channels,\n",
        "                                     out_channels=4*in_channels,\n",
        "                                     kernel_size=1)\n",
        "\n",
        "        self.conv3 = nn.Conv2d(in_channels=4*in_channels,\n",
        "                                     out_channels=in_channels,\n",
        "                                     kernel_size=1\n",
        "                                     )\n",
        "\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        x = self.conv1(inputs)\n",
        "\n",
        "        # The normalization layer:\n",
        "        x = self.norm(x)\n",
        "\n",
        "        x = self.conv2(x)\n",
        "\n",
        "        # The non-linear activation layer:\n",
        "        x = torch.nn.functional.gelu(x)\n",
        "\n",
        "        x = self.conv3(x)\n",
        "\n",
        "        # This makes it a residual network:\n",
        "        return x + inputs\n",
        "\n",
        "\n",
        "class Classifier(nn.Module):\n",
        "\n",
        "\n",
        "    def __init__(self, n_initial_filters, n_stages, blocks_per_stage):\n",
        "        super(Classifier, self).__init__()\n",
        "\n",
        "        # This is a downsampling convolution that will produce patches of output.\n",
        "\n",
        "        # This is similar to what vision transformers do to tokenize the images.\n",
        "        self.stem = nn.Conv2d(in_channels=3,\n",
        "                                    out_channels=n_initial_filters,\n",
        "                                    kernel_size=1,\n",
        "                                    stride=1)\n",
        "\n",
        "        current_shape = [32, 32]\n",
        "\n",
        "        self.norm1 = nn.LayerNorm([n_initial_filters,*current_shape])\n",
        "        # self.norm1 = WrappedLayerNorm()\n",
        "\n",
        "        current_n_filters = n_initial_filters\n",
        "\n",
        "        self.layers = nn.Sequential()\n",
        "        for i, n_blocks in enumerate(range(n_stages)):\n",
        "            # Add a convnext block series:\n",
        "            for _ in range(blocks_per_stage):\n",
        "                self.layers.append(ConvNextBlock(in_channels=current_n_filters, shape=current_shape))\n",
        "            # Add a downsampling layer:\n",
        "            if i != n_stages - 1:\n",
        "                # Skip downsampling if it's the last layer!\n",
        "                self.layers.append(Downsampler(\n",
        "                    in_channels=current_n_filters,\n",
        "                    out_channels=2*current_n_filters,\n",
        "                    shape = current_shape,\n",
        "                    )\n",
        "                )\n",
        "                # Double the number of filters:\n",
        "                current_n_filters = 2*current_n_filters\n",
        "                # Cut the shape in half:\n",
        "                current_shape = [ cs // 2 for cs in current_shape]\n",
        "\n",
        "\n",
        "\n",
        "        self.head = nn.Sequential(\n",
        "            nn.Flatten(),\n",
        "            nn.LayerNorm(current_n_filters),\n",
        "            nn.Linear(current_n_filters, 10)\n",
        "        )\n",
        "        # self.norm2 = nn.InstanceNorm2d(current_n_filters)\n",
        "        # # This brings it down to one channel / class\n",
        "        # self.bottleneck = nn.Conv2d(in_channels=current_n_filters, out_channels=10,\n",
        "        #                                   kernel_size=1, stride=1)\n",
        "\n",
        "    def forward(self, inputs):\n",
        "\n",
        "        x = self.stem(inputs)\n",
        "        # Apply a normalization after the initial patching:\n",
        "        x = self.norm1(x)\n",
        "\n",
        "        # Apply the main chunk of the network:\n",
        "        x = self.layers(x)\n",
        "\n",
        "        # Normalize and readout:\n",
        "        x = nn.functional.avg_pool2d(x, x.shape[2:])\n",
        "        x = self.head(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "\n",
        "        # x = self.norm2(x)\n",
        "        # x = self.bottleneck(x)\n",
        "\n",
        "        # # Average pooling of the remaining spatial dimensions (and reshape) makes this label-like:\n",
        "        # return nn.functional.avg_pool2d(x, kernel_size=x.shape[-2:]).reshape((-1,10))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "smUimDPdqYok"
      },
      "outputs": [],
      "source": [
        "!pip install torchinfo # if not on Polaris"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c6xG3pRzqYok"
      },
      "outputs": [],
      "source": [
        "model = Classifier(32, 4, 2).to(device=dev)\n",
        "\n",
        "from torchinfo import summary\n",
        "\n",
        "print(summary(model, input_size=(batch_size, 3, 32, 32)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9WExbWCpqYok"
      },
      "outputs": [],
      "source": [
        "def evaluate(dataloader, model, loss_fn, val_bar):\n",
        "    # Set the model to evaluation mode - some NN pieces behave differently during training\n",
        "    # Unnecessary in this situation but added for best practices\n",
        "    model.eval()\n",
        "    size = len(dataloader)\n",
        "    num_batches = len(dataloader)\n",
        "    loss, correct = 0, 0\n",
        "\n",
        "    # We can save computation and memory by not calculating gradients here - we aren't optimizing\n",
        "    with torch.no_grad():\n",
        "        # loop over all of the batches\n",
        "        for X, y in dataloader:\n",
        "\n",
        "            pred = model(X)\n",
        "            loss += loss_fn(pred, y).item()\n",
        "            # how many are correct in this batch? Tracking for accuracy\n",
        "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "            val_bar.update()\n",
        "\n",
        "    loss /= num_batches\n",
        "    correct /= (size*batch_size)\n",
        "\n",
        "    accuracy = 100*correct\n",
        "    return accuracy, loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SX27RvNoqYok"
      },
      "outputs": [],
      "source": [
        "def train_one_epoch(dataloader, model, loss_fn, optimizer, progress_bar):\n",
        "    model.train()\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        # forward pass\n",
        "        pred = model(X)\n",
        "        loss = loss_fn(pred, y)\n",
        "\n",
        "        # backward pass calculates gradients\n",
        "        loss.backward()\n",
        "\n",
        "        # take one step with these gradients\n",
        "        optimizer.step()\n",
        "\n",
        "        # resets the gradients\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        progress_bar.update()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oiAlXpYgqYok"
      },
      "outputs": [],
      "source": [
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, weight_decay=0.01)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hD_Sm_kPqYok"
      },
      "outputs": [],
      "source": [
        "from tqdm.notebook import tqdm\n",
        "\n",
        "\n",
        "epochs = 30\n",
        "for j in range(epochs):\n",
        "    with tqdm(total=len(train_dataloader), position=0, leave=True, desc=f\"Train Epoch {j}\") as train_bar:\n",
        "        train_one_epoch(train_dataloader, model, loss_fn, optimizer, train_bar)\n",
        "\n",
        "    # checking on the training & validation loss & accuracy\n",
        "    # for training data - only once every 5 epochs (takes a while)\n",
        "    if j % 5 == 0:\n",
        "        with tqdm(total=len(train_dataloader), position=0, leave=True, desc=f\"Validate (train) Epoch {j}\") as train_eval:\n",
        "            acc, loss = evaluate(train_dataloader, model, loss_fn, train_eval)\n",
        "            print(f\"Epoch {j}: training loss: {loss:.3f}, accuracy: {acc:.3f}\")\n",
        "\n",
        "    with tqdm(total=len(val_dataloader), position=0, leave=True, desc=f\"Validate Epoch {j}\") as val_bar:\n",
        "\n",
        "        acc_val, loss_val = evaluate(val_dataloader, model, loss_fn, val_bar)\n",
        "        print(f\"Epoch {j}: validation loss: {loss_val:.3f}, accuracy: {acc_val:.3f}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M8D1T4C1qYok"
      },
      "source": [
        "# Homework 1:\n",
        "\n",
        "In this notebook, we've learned about some basic convolutional networks and trained one on CIFAR-10 images.  It did ... OK.  There is significant overfitting of this model.  There are some ways to address that, but we didn't have time to get into that in this session.\n",
        "\n",
        "Meanwhile, your homework (part 1) for this week is to try to train the model again but with a different architecture.  Change one or more of the following:\n",
        "- The number of convolutions between downsampling\n",
        "- The number of filters in each layer\n",
        "- The initial \"patchify\" layer\n",
        "- Another hyper-parameter of your choosing\n",
        "\n",
        "\n",
        "And compare your final validation accuracy to the accuracy shown here.  Can you beat the validation accuracy shown?\n",
        "\n",
        "For full credit on the homework, you need to show (via text, or make a plot) the training and validation data sets' performance (loss and accuracy) for all the epochs you train.  You also need to explain, in several sentences, what you changed in the network and why you think it makes a difference."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torch import nn\n",
        "\n",
        "\n",
        "class Downsampler(nn.Module):\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, shape, stride=2):\n",
        "        super(Downsampler, self).__init__()\n",
        "\n",
        "        self.norm = nn.LayerNorm([in_channels, *shape])\n",
        "\n",
        "        self.downsample = nn.Conv2d(\n",
        "            in_channels=in_channels,\n",
        "            out_channels=out_channels,\n",
        "            kernel_size = stride,\n",
        "            stride = stride,\n",
        "        )\n",
        "\n",
        "    def forward(self, inputs):\n",
        "\n",
        "\n",
        "        return self.downsample(self.norm(inputs))\n",
        "\n",
        "\n",
        "\n",
        "class ConvNextBlock(nn.Module):\n",
        "    \"\"\"This block of operations is loosely based on this paper:\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "    def __init__(self, in_channels, shape):\n",
        "        super(ConvNextBlock, self).__init__()\n",
        "\n",
        "        # Depthwise, seperable convolution with a large number of output filters:\n",
        "        self.conv1 = nn.Conv2d(in_channels=in_channels,\n",
        "                                     out_channels=in_channels,\n",
        "                                     groups=in_channels,\n",
        "                                     kernel_size=[9,9], # Changed to 9x9 for hw\n",
        "                                     padding='same' )\n",
        "\n",
        "        self.norm = nn.LayerNorm([in_channels, *shape])\n",
        "\n",
        "        # Two more convolutions:\n",
        "        self.conv2 = nn.Conv2d(in_channels=in_channels,\n",
        "                                     out_channels=4*in_channels,\n",
        "                                     kernel_size=1)\n",
        "\n",
        "        self.conv3 = nn.Conv2d(in_channels=4*in_channels,\n",
        "                                     out_channels=in_channels,\n",
        "                                     kernel_size=1\n",
        "                                     )\n",
        "        # Adding two more for hw\n",
        "        self.conv4 = nn.Conv2d(in_channels=in_channels,\n",
        "                                     out_channels=4*in_channels,\n",
        "                                     kernel_size=1)\n",
        "\n",
        "        self.conv5 = nn.Conv2d(in_channels=4*in_channels,\n",
        "                                     out_channels=in_channels,\n",
        "                                     kernel_size=1\n",
        "                                     )\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        x = self.conv1(inputs)\n",
        "\n",
        "        # The normalization layer:\n",
        "        x = self.norm(x)\n",
        "\n",
        "        x = self.conv2(x)\n",
        "\n",
        "        # The non-linear activation layer:\n",
        "        x = torch.nn.functional.gelu(x)\n",
        "\n",
        "        x = self.conv3(x)\n",
        "\n",
        "        # This makes it a residual network:\n",
        "        return x + inputs\n",
        "\n",
        "\n",
        "class Classifier(nn.Module):\n",
        "\n",
        "\n",
        "    def __init__(self, n_initial_filters, n_stages, blocks_per_stage):\n",
        "        super(Classifier, self).__init__()\n",
        "\n",
        "        # This is a downsampling convolution that will produce patches of output.\n",
        "\n",
        "        # This is similar to what vision transformers do to tokenize the images.\n",
        "        self.stem = nn.Conv2d(in_channels=3,\n",
        "                                    out_channels=n_initial_filters,\n",
        "                                    kernel_size=1,\n",
        "                                    stride=1)\n",
        "\n",
        "        current_shape = [32, 32]\n",
        "\n",
        "        self.norm1 = nn.LayerNorm([n_initial_filters,*current_shape])\n",
        "        # self.norm1 = WrappedLayerNorm()\n",
        "\n",
        "        current_n_filters = n_initial_filters\n",
        "\n",
        "        self.layers = nn.Sequential()\n",
        "        for i, n_blocks in enumerate(range(n_stages)):\n",
        "            # Add a convnext block series:\n",
        "            for _ in range(blocks_per_stage):\n",
        "                self.layers.append(ConvNextBlock(in_channels=current_n_filters, shape=current_shape))\n",
        "            # Add a downsampling layer:\n",
        "            if i != n_stages - 1:\n",
        "                # Skip downsampling if it's the last layer!\n",
        "                self.layers.append(Downsampler(\n",
        "                    in_channels=current_n_filters,\n",
        "                    out_channels=2*current_n_filters,\n",
        "                    shape = current_shape,\n",
        "                    )\n",
        "                )\n",
        "                # Double the number of filters:\n",
        "                current_n_filters = 2*current_n_filters\n",
        "                # Cut the shape in half:\n",
        "                current_shape = [ cs // 2 for cs in current_shape]\n",
        "\n",
        "\n",
        "\n",
        "        self.head = nn.Sequential(\n",
        "            nn.Flatten(),\n",
        "            nn.LayerNorm(current_n_filters),\n",
        "            nn.Linear(current_n_filters, 10)\n",
        "        )\n",
        "        # self.norm2 = nn.InstanceNorm2d(current_n_filters)\n",
        "        # # This brings it down to one channel / class\n",
        "        # self.bottleneck = nn.Conv2d(in_channels=current_n_filters, out_channels=10,\n",
        "        #                                   kernel_size=1, stride=1)\n",
        "\n",
        "    def forward(self, inputs):\n",
        "\n",
        "        x = self.stem(inputs)\n",
        "        # Apply a normalization after the initial patching:\n",
        "        x = self.norm1(x)\n",
        "\n",
        "        # Apply the main chunk of the network:\n",
        "        x = self.layers(x)\n",
        "\n",
        "        # Normalize and readout:\n",
        "        x = nn.functional.avg_pool2d(x, x.shape[2:])\n",
        "        x = self.head(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "\n",
        "        # x = self.norm2(x)\n",
        "        # x = self.bottleneck(x)\n",
        "\n",
        "        # # Average pooling of the remaining spatial dimensions (and reshape) makes this label-like:\n",
        "        # return nn.functional.avg_pool2d(x, kernel_size=x.shape[-2:]).reshape((-1,10))"
      ],
      "metadata": {
        "id": "iUOy5tKs-Lj3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Classifier(64, 6, 2).to(device=dev)\n",
        "\n",
        "from torchinfo import summary\n",
        "\n",
        "print(summary(model, input_size=(batch_size, 3, 32, 32)))"
      ],
      "metadata": {
        "id": "_DvONwkc-Oka"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(dataloader, model, loss_fn, val_bar):\n",
        "    # Set the model to evaluation mode - some NN pieces behave differently during training\n",
        "    # Unnecessary in this situation but added for best practices\n",
        "    model.eval()\n",
        "    size = len(dataloader)\n",
        "    num_batches = len(dataloader)\n",
        "    loss, correct = 0, 0\n",
        "\n",
        "    # We can save computation and memory by not calculating gradients here - we aren't optimizing\n",
        "    with torch.no_grad():\n",
        "        # loop over all of the batches\n",
        "        for X, y in dataloader:\n",
        "\n",
        "            pred = model(X)\n",
        "            loss += loss_fn(pred, y).item()\n",
        "            # how many are correct in this batch? Tracking for accuracy\n",
        "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "            val_bar.update()\n",
        "\n",
        "    loss /= num_batches\n",
        "    correct /= (size*batch_size)\n",
        "\n",
        "    accuracy = 100*correct\n",
        "    return accuracy, loss"
      ],
      "metadata": {
        "id": "2B32XDYM-Rk8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_one_epoch(dataloader, model, loss_fn, optimizer, progress_bar):\n",
        "    model.train()\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        # forward pass\n",
        "        pred = model(X)\n",
        "        loss = loss_fn(pred, y)\n",
        "\n",
        "        # backward pass calculates gradients\n",
        "        loss.backward()\n",
        "\n",
        "        # take one step with these gradients\n",
        "        optimizer.step()\n",
        "\n",
        "        # resets the gradients\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        progress_bar.update()"
      ],
      "metadata": {
        "id": "CDNfpXJY-UyY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, weight_decay=0.01)"
      ],
      "metadata": {
        "id": "nTeB2eAr-Xsb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm.notebook import tqdm\n",
        "\n",
        "\n",
        "epochs = 30\n",
        "for j in range(epochs):\n",
        "    with tqdm(total=len(train_dataloader), position=0, leave=True, desc=f\"Train Epoch {j}\") as train_bar:\n",
        "        train_one_epoch(train_dataloader, model, loss_fn, optimizer, train_bar)\n",
        "\n",
        "    # checking on the training & validation loss & accuracy\n",
        "    # for training data - only once every 5 epochs (takes a while)\n",
        "    if j % 5 == 0:\n",
        "        with tqdm(total=len(train_dataloader), position=0, leave=True, desc=f\"Validate (train) Epoch {j}\") as train_eval:\n",
        "            acc, loss = evaluate(train_dataloader, model, loss_fn, train_eval)\n",
        "            print(f\"Epoch {j}: training loss: {loss:.3f}, accuracy: {acc:.3f}\")\n",
        "\n",
        "    with tqdm(total=len(val_dataloader), position=0, leave=True, desc=f\"Validate Epoch {j}\") as val_bar:\n",
        "\n",
        "        acc_val, loss_val = evaluate(val_dataloader, model, loss_fn, val_bar)\n",
        "        print(f\"Epoch {j}: validation loss: {loss_val:.3f}, accuracy: {acc_val:.3f}\")"
      ],
      "metadata": {
        "id": "ijYgwPH_-aON"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.4"
    },
    "colab": {
      "provenance": [],
      "name": "2024-11-04_Tyriek_Craigs_Hw3.ipynb",
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}