code/09-transfer-learning.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",
    "- [Pre-Processing](#toc2_)    \n",
    "  - [Load Dataset](#toc2_1_)    \n",
    "  - [Normalization](#toc2_2_)    \n",
    "  - [DataLoader](#toc2_3_)    \n",
    "- [Transfer Learning](#toc3_)    \n",
    "  - [Extract a subset of Pretrained Model](#toc3_1_)    \n",
    "  - [Freeze all transferred layers](#toc3_2_)    \n",
    "  - [Transfer the model](#toc3_3_)    \n",
    "- [Set up remaining Hyperparameters](#toc4_)    \n",
    "- [Train Loop](#toc5_)    \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": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\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 import ResNet50_Weights, resnet50\n",
    "from torchvision.transforms import v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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": 25,
   "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>[Pre-Processing](#toc0_)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_1_'></a>[Load Dataset](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial transforms\n",
    "transforms = v2.Compose(\n",
    "    [\n",
    "        v2.ToImage(),\n",
    "        v2.ToDtype(torch.float32, scale=True),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = CIFAR10(DATASET_DIR, train=True, transform=transforms, download=False)\n",
    "testset = CIFAR10(DATASET_DIR, train=False, transform=transforms, download=False)\n",
    "\n",
    "# log\n",
    "print(\"trainset:\")\n",
    "print(f\"    -> trainset.data.shape    : {trainset.data.shape}\")\n",
    "print(f\"    -> trainset.data.dtype    : {trainset.data.dtype}\")\n",
    "print(f\"    -> type(trainset.data)    : {type(trainset.data)}\")\n",
    "print(f\"    -> type(trainset.targets) : {type(trainset.targets)}\")\n",
    "print(\"-\" * 50)\n",
    "print(\"testset:\")\n",
    "print(f\"    -> testset.data.shape     : {testset.data.shape}\")\n",
    "print(f\"    -> testset.data.dtype     : {testset.data.dtype}\")\n",
    "print(f\"    -> type(testset.data)     : {type(testset.data)}\")\n",
    "print(f\"    -> type(testset.targets)  : {type(testset.targets)}\")\n",
    "print(\"-\" * 50)\n",
    "print(f\"classes : {trainset.classes}\")\n",
    "print(f\"trainset distribution : {torch.unique(torch.tensor(trainset.targets), return_counts=True)[1]}\")\n",
    "print(f\"testset  distribution : {torch.unique(torch.tensor(testset.targets), return_counts=True)[1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig, axs = plt.subplots(nrows=4, ncols=8, figsize=(12, 6), layout=\"compressed\")\n",
    "for i in range(4):\n",
    "    for j in range(8):\n",
    "        axs[i, j].imshow(trainset.data[i * 8 + j], cmap=\"gray\")\n",
    "        axs[i, j].set_title(trainset.classes[trainset.targets[i * 8 + j]])\n",
    "        axs[i, j].axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_2_'></a>[Normalization](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = next(iter(DataLoader(trainset, batch_size=len(trainset))))[0]\n",
    "\n",
    "train_mean = data.mean(axis=(0, 2, 3))\n",
    "train_std = data.std(axis=(0, 2, 3))\n",
    "\n",
    "del data\n",
    "\n",
    "# log\n",
    "print(f\"train mean per channel : {train_mean}\")\n",
    "print(f\"train std  per channel : {train_std}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms.transforms.append(v2.Normalize(mean=train_mean, std=train_std))\n",
    "\n",
    "# log\n",
    "print(f\"trainset.transforms:\\n{trainset.transforms}\\n\")\n",
    "print(f\"testset.transforms:\\n{testset.transforms}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(\"before applying transform:\")\n",
    "print(f\"    -> type(testset.data[0]) : {type(testset.data[0])}\")\n",
    "print(f\"    -> testset.data[0].dtype : {testset.data[0].dtype}\")\n",
    "print(f\"    -> testset.data[0].shape : {testset.data[0].shape}\")\n",
    "print(\"-\" * 50)\n",
    "print(\"after applying transform:\")\n",
    "print(f\"    -> type(testset[0][0])   : {type(testset[0][0])}\")\n",
    "print(f\"    -> testset[0][0].dtype   : {testset[0][0].dtype}\")\n",
    "print(f\"    -> testset[0][0].shape   : {testset[0][0].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_3_'></a>[DataLoader](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc3_'></a>[Transfer Learning](#toc0_)\n",
    "\n",
    "- resnet50 pretrained on IMAGENET1K_V1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)\n",
    "\n",
    "# log\n",
    "resnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(resnet, input_size=(1, *trainset[0][0].shape), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_1_'></a>[Extract a subset of Pretrained Model](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_extractor = nn.Sequential(*list(resnet.children())[:-1])\n",
    "\n",
    "# log\n",
    "for name, param in feature_extractor.named_parameters():\n",
    "    print(f\"{name:<23} - requires_grad: {param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_2_'></a>[Freeze all transferred layers](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in resnet.named_parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# log\n",
    "for name, param in resnet.named_parameters():\n",
    "    print(f\"{name:<28} - requires_grad : {param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_3_'></a>[Transfer the model](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.feature_extractor = feature_extractor\n",
    "        self.flatten = nn.Flatten(start_dim=1)\n",
    "        self.classifier = nn.Sequential(nn.Linear(2048, 256), nn.ReLU(), nn.Linear(256, 10))\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 = MyModel().to(device)\n",
    "\n",
    "# log\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print(f\"{name:<41} - requires_grad : {param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(model, input_size=(1, *trainset[0][0].shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc4_'></a>[Set up remaining Hyperparameters](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.001\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = Adam(params=model.parameters(), lr=lr)\n",
    "num_epochs = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc5_'></a>[Train Loop](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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",
    "    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",
    "    )"
   ]
  }
 ],
 "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
}