week5_network_blocks.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install lightning\n",
        "!rm logs -rf"
      ],
      "metadata": {
        "id": "HGzzvPk4NmHI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GugF_VS8NEgB"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets import MNIST\n",
        "import lightning as pl\n",
        "from torchmetrics import Accuracy\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "\n",
        "# Define the LightningModule\n",
        "class MNISTClassifier(pl.LightningModule):\n",
        "    def __init__(self, learning_rate=0.001):\n",
        "        super(MNISTClassifier, self).__init__()\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Linear(784, 128),\n",
        "            nn.Sigmoid(),\n",
        "            nn.Linear(128, 128),\n",
        "            nn.Sigmoid(),\n",
        "            nn.Linear(128, 10),\n",
        "        )\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "        self.learning_rate = learning_rate\n",
        "        self.val_logits = []\n",
        "        self.val_labels = []\n",
        "        self.val_accuracy = Accuracy('multiclass', num_classes=10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.layers(x)\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        x, y = batch\n",
        "        logits = self(x)\n",
        "        loss = self.loss_fn(logits, y)\n",
        "        self.log('train_loss', loss)\n",
        "        return loss\n",
        "\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        x, y = batch\n",
        "        logits = self(x)\n",
        "        loss = self.loss_fn(logits, y)\n",
        "\n",
        "        # Store logits and labels for later calculation\n",
        "        self.val_logits.append(logits)\n",
        "        self.val_labels.append(y)\n",
        "\n",
        "    def on_validation_epoch_end(self):\n",
        "        # Concatenate all stored logits and labels\n",
        "        val_logits = torch.cat(self.val_logits, dim=0)\n",
        "\n",
        "        # Calculate the average loss over the entire validation set\n",
        "        val_labels = torch.cat(self.val_labels, dim=0)\n",
        "        avg_val_loss = self.loss_fn(val_logits, val_labels)\n",
        "\n",
        "        # Calculate accuracy using torchmetrics\n",
        "        avg_val_acc = self.val_accuracy(val_logits, val_labels)\n",
        "\n",
        "        # Log the average validation loss and accuracy\n",
        "        self.log('val_loss', avg_val_loss, prog_bar=True)\n",
        "        self.log('val_acc', avg_val_acc, prog_bar=True)\n",
        "\n",
        "        # Clear the lists for the next epoch\n",
        "        self.val_logits = []\n",
        "        self.val_labels = []\n",
        "        self.val_accuracy.reset()\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        return torch.optim.SGD(self.parameters(), lr=self.learning_rate, momentum=0.9)\n",
        "\n",
        "def train_and_log(model, model_name):\n",
        "    tb_logger = TensorBoardLogger(save_dir=\"logs/\", name=model_name)\n",
        "\n",
        "    trainer = pl.Trainer(max_epochs=5,\n",
        "                        accelerator=\"gpu\",\n",
        "                        devices=1,\n",
        "                        log_every_n_steps=1000,\n",
        "                        logger=tb_logger)\n",
        "\n",
        "    # Train the model\n",
        "    trainer.fit(model, train_loader, val_loader)\n",
        "\n",
        "# Data preparation\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5,), (0.5,)),\n",
        "    transforms.Lambda(lambda x: torch.flatten(x))\n",
        "])\n",
        "\n",
        "train_dataset = MNIST('./mnist', train=True, transform=transform, download=True)\n",
        "val_dataset = MNIST('./mnist', train=False, transform=transform, download=True)\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True)\n",
        "val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=4, pin_memory=True)\n",
        "\n",
        "# Instantiate the model\n",
        "model = MNISTClassifier()\n",
        "\n",
        "train_and_log(model, 'fcc')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Data preparation\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5,), (0.5,)),\n",
        "])\n",
        "\n",
        "train_dataset = MNIST('./mnist', train=True, transform=transform, download=True)\n",
        "val_dataset = MNIST('./mnist', train=False, transform=transform, download=True)\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True)\n",
        "val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=4, pin_memory=True)"
      ],
      "metadata": {
        "id": "fXJRuLUjOjpi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### exercise: Convolutional network\n"
      ],
      "metadata": {
        "id": "7s6rsQ3oPHv5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "\n",
        "class MNISTCNN(MNISTClassifier):\n",
        "    def __init__(self, learning_rate=0.001):\n",
        "        super(MNISTClassifier, self).__init__()\n",
        "\n",
        "        # edit the follwing to creat layers\n",
        "        # conv 3x3, 1 -> 32 + ReLU\n",
        "        self.stem = nn.Sequential(nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),\n",
        "                                 nn.ReLU()\n",
        "                                 )\n",
        "        # conv 3x3, 32 -> 32 + ReLU\n",
        "        self.layer0 = nn.Sequential(nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), nn.ReLU())\n",
        "        # conv 3x3, 32 -> 32 + ReLU\n",
        "        self.layer1 = nn.Sequential(nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), nn.ReLU())\n",
        "        # make a conv 1x1, 32 -> 10, no activation function\n",
        "        self.fc = nn.Sequential(nn.Conv2d(32, 10, kernel_size=1, stride=1, padding=0))\n",
        "        # END OF edit the follwing to creat layers\n",
        "\n",
        "        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "        self.learning_rate = learning_rate\n",
        "        self.val_logits = []\n",
        "        self.val_labels = []\n",
        "        self.val_accuracy = Accuracy('multiclass', num_classes=10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.stem(x)\n",
        "        x = self.layer0(x)\n",
        "        x = self.max_pool(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.max_pool(x)\n",
        "        x = self.fc(x)\n",
        "        x = x.mean([2, 3])\n",
        "        return x\n",
        "\n",
        "model = MNISTCNN(0.001)\n",
        "train_and_log(model, 'cnn')"
      ],
      "metadata": {
        "id": "lICmQOeGNUw5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MNISTBNCNN(MNISTCNN):\n",
        "    def __init__(self, learning_rate=0.001):\n",
        "        super(MNISTClassifier, self).__init__()\n",
        "\n",
        "        # edit the follwing to creat layers\n",
        "        # add nn.BatchNorm2d right after conv2d\n",
        "        self.stem = nn.Sequential(nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),\n",
        "                                 ...\n",
        "                                 )\n",
        "        # conv 3x3, 32 -> 32 + ReLU\n",
        "        # add nn.BatchNorm2d right after conv2d\n",
        "        self.layer0 = nn.Sequential(...)\n",
        "\n",
        "        # conv 3x3, 32 -> 32 + ReLU\n",
        "        # add nn.BatchNorm2d right after conv2d\n",
        "        self.layer1 = nn.Sequential(...)\n",
        "\n",
        "        # make a conv 1x1, 32 -> 10, no activation function\n",
        "        self.fc = nn.Sequential(...)\n",
        "        # END OF edit the follwing to creat layers\n",
        "\n",
        "        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "        self.learning_rate = learning_rate\n",
        "        self.val_logits = []\n",
        "        self.val_labels = []\n",
        "        self.val_accuracy = Accuracy('multiclass', num_classes=10)\n",
        "\n",
        "model = MNISTBNCNN(0.001)\n",
        "train_and_log(model, 'cnn_batchnorm')"
      ],
      "metadata": {
        "id": "o9PJlxKPFNZT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**exercise**: skip connection\n",
        "\n",
        "hint:\n",
        "x = x + f(x)"
      ],
      "metadata": {
        "id": "huuqssPWPOhn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "\n",
        "class MNISTSKIP(MNISTCNN):\n",
        "    def forward(self, x):\n",
        "        # edit the follwing to make skip connection\n",
        "        x = self.stem(x)\n",
        "        x = self.layer0(x) + x\n",
        "        x = self.max_pool(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.max_pool(x)\n",
        "        x = self.fc(x)\n",
        "        # ENDOF edit the follwing to make skip connection\n",
        "        x = x.mean([2, 3])\n",
        "        return x\n",
        "\n",
        "model = MNISTSKIP(0.001)\n",
        "train_and_log(model, 'skip')"
      ],
      "metadata": {
        "id": "I8Jp89HvPN3o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#!pip install tensorboard\n",
        "%load_ext tensorboard\n",
        "\n",
        "from pytorch_lightning import Trainer\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "%tensorboard --logdir logs/"
      ],
      "metadata": {
        "id": "-dVl4WJmTzlP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "GwtNo_vpdBGy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "zEWhvM1Hc9bw"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}