code/08-feature-extraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** <amirhosseinheydari78@gmail.com> - 📍 **Origin:** [mr-pylin/pytorch-workshop](https://github.com/mr-pylin/pytorch-workshop)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- [Dependencies](#toc1_)    \n",
    "- [Load Dataset](#toc2_)    \n",
    "- [Model](#toc3_)    \n",
    "- [Set up remaining Hyperparameters](#toc4_)    \n",
    "- [Train Loop](#toc5_)    \n",
    "- [Feature Extraction](#toc6_)    \n",
    "  - [model.feature_extractor.0](#toc6_1_)    \n",
    "  - [model.feature_extractor.5](#toc6_2_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=1\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc1_'></a>[Dependencies](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader\n",
    "from torchinfo import summary\n",
    "from torchmetrics.classification import MulticlassAccuracy\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision.models.feature_extraction import create_feature_extractor\n",
    "from torchvision.transforms import v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a seed for deterministic results\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if cuda is available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# log\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# update paths as needed based on your project structure\n",
    "DATASET_DIR = r\"../datasets\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc2_'></a>[Load Dataset](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = v2.Compose(\n",
    "    [\n",
    "        v2.ToImage(),\n",
    "        v2.ToDtype(dtype=torch.float32, scale=True),\n",
    "        v2.Normalize(mean=(0.5,), std=(0.5,)),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = CIFAR10(DATASET_DIR, train=True, transform=transform, download=False)\n",
    "testset = CIFAR10(DATASET_DIR, train=False, transform=transform, download=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = np.array(trainset.classes)\n",
    "num_classes = len(classes)\n",
    "num_trainset, height, width, depth = trainset.data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "\n",
    "trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "testloader = DataLoader(trainset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc3_'></a>[Model](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(nn.Module):\n",
    "    def __init__(self, layers: list[int]):\n",
    "        super().__init__()\n",
    "\n",
    "        feature_extractor_layers = []\n",
    "        for i, o in zip(layers, layers[1:]):\n",
    "            feature_extractor_layers.append(nn.Conv2d(i, out_channels=o, kernel_size=3))\n",
    "            feature_extractor_layers.append(nn.BatchNorm2d(o))\n",
    "            feature_extractor_layers.append(nn.ReLU())\n",
    "            feature_extractor_layers.append(nn.MaxPool2d(kernel_size=2))\n",
    "\n",
    "        self.feature_extractor = nn.Sequential(*feature_extractor_layers)\n",
    "        self.flatten = nn.Flatten(start_dim=1)\n",
    "        self.classifier = nn.Linear(1152, len(classes))\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.feature_extractor(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialize the model\n",
    "model = CustomModel(layers=[depth, 16, 32]).to(device)\n",
    "\n",
    "# log\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(model, input_size=(batch_size, *testset.data.transpose(0, 3, 1, 2).shape[1:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc4_'></a>[Set up remaining Hyperparameters](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.01\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = Adam(params=model.parameters(), lr=lr)\n",
    "num_epochs = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc5_'></a>[Train Loop](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_acc_per_epoch = []\n",
    "train_loss_per_epoch = []\n",
    "train_acc = MulticlassAccuracy(num_classes=len(testset.classes), top_k=1).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "\n",
    "    # train loop\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "\n",
    "    for x, y in trainloader:\n",
    "\n",
    "        # send data to GPU\n",
    "        x, y_true = x.to(device), y.to(device)\n",
    "\n",
    "        # forward\n",
    "        y_pred = model(x)\n",
    "        loss = criterion(y_pred, y_true)\n",
    "\n",
    "        # backward\n",
    "        loss.backward()\n",
    "\n",
    "        # update parameters\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # store loss and accuracy per iteration\n",
    "        train_loss += loss.item() * len(x)\n",
    "        train_acc.update(y_pred, y_true)\n",
    "\n",
    "    # store loss and accuracy per epoch\n",
    "    train_loss_per_epoch.append(train_loss / len(trainset))\n",
    "    train_acc_per_epoch.append(train_acc.compute().item())\n",
    "    train_acc.reset()\n",
    "\n",
    "    # log\n",
    "    print(\n",
    "        f\"epoch {epoch+1:0{len(str(num_epochs))}}/{num_epochs} -> train[loss: {train_loss_per_epoch[epoch]:7.5f} - acc: {train_acc_per_epoch[epoch]*100:5.2f}%]\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc6_'></a>[Feature Extraction](#toc0_)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print(f\"{name:<26} - requires_grad: {param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = {\"feature_extractor.0\": \"conv1\", \"feature_extractor.5\": \"conv2\"}\n",
    "\n",
    "feature_extractor = create_feature_extractor(model, return_nodes=nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frog = transform(testset.data[0]).to(device)[None, :, :, :]\n",
    "\n",
    "feature_maps = feature_extractor(frog)\n",
    "feature_maps.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(8, 4), layout=\"compressed\")\n",
    "axs[0].imshow(testset.data[0])\n",
    "axs[0].axis(\"off\")\n",
    "axs[0].set_title(\"Orignal\")\n",
    "axs[1].imshow(frog.detach().cpu()[0].permute(1, 2, 0).clip(0, 1))\n",
    "axs[1].axis(\"off\")\n",
    "axs[1].set_title(\"Transformed\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc6_1_'></a>[model.feature_extractor.0](#toc0_)\n",
    "\n",
    "- Feature maps : 16x30x30\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "total_rows = 4\n",
    "total_cols = 4\n",
    "fig, axs = plt.subplots(nrows=total_rows, ncols=total_cols, figsize=(total_cols * 2, total_rows * 2), layout=\"compressed\")\n",
    "fig.suptitle(\"model.feature_extractor.0 feature maps\")\n",
    "for row in range(total_rows):\n",
    "    for col in range(total_cols):\n",
    "        axs[row, col].imshow(feature_maps[\"conv1\"][0, row * total_cols + col].detach().cpu(), cmap=\"gray\")\n",
    "        axs[row, col].axis(\"off\")\n",
    "        axs[row, col].set(title=row * total_cols + col)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc6_2_'></a>[model.feature_extractor.5](#toc0_)\n",
    "\n",
    "- Feature maps : 32x13x13\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "total_rows = 4\n",
    "total_cols = 8\n",
    "fig, axs = plt.subplots(nrows=total_rows, ncols=total_cols, figsize=(total_cols * 2, total_rows * 2), layout=\"compressed\")\n",
    "fig.suptitle(\"model.feature_extractor.0 feature maps\")\n",
    "for row in range(total_rows):\n",
    "    for col in range(total_cols):\n",
    "        axs[row, col].imshow(feature_maps[\"conv2\"][0, row * total_cols + col].detach().cpu(), cmap=\"gray\")\n",
    "        axs[row, col].axis(\"off\")\n",
    "        axs[row, col].set(title=row * total_cols + col)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "author_email": "AmirhosseinHeydari78@gmail.com",
  "author_github": "https://github.com/mr-pylin",
  "author_name": "Amirhossein Heydari",
  "kernelspec": {
   "display_name": "pytorch-workshop-U_zYfVTd-py3.12",
   "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.12.8"
  },
  "origin_repo": "https://github.com/mr-pylin/pytorch-workshop"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}