transfer-learning-1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9dd78e0a-ba36-4a81-91ca-47528b8aed5d",
   "metadata": {},
   "source": [
    "# Steps to Implement Transfer Learning for Image Classification in PyTorch\n",
    "\n",
    "Transfer learning for image classification is essentially reusing a pre-trained neural network to improve the result on a different dataset. Follow the steps to implement Transfer Learning for Image Classification.\n",
    "\n",
    "1. Choose a pre-trained model (ResNet, VGG, etc.) based on your task.\n",
    "2. Modify the model by potentially replacing the final classification layer to match the number of classes in your new dataset.\n",
    "3. Freeze the pre-trained layers (make their weights non-trainable) to prevent them from being updated during training on the new dataset. This is especially useful when you have a small dataset.\n",
    "4. Preprocess your data, including resizing images and normalization.\n",
    "5. Optionally, perform data augmentation to increase the size and diversity of your dataset.\n",
    "6. Define the new model architecture by adding the new classifier on top of the pre-trained model.\n",
    "7. Compile the model by specifying the loss function, optimizer, and metrics.\n",
    "8. Train the model on your new dataset. Freezing the pre-trained layers might require fewer training epochs compared to training from scratch.\n",
    "9. Fine-tuning: You can further train the model by unfreezing some or all of the pre-trained layers.\n",
    "10. Evaluate the model’s performance on a validation or test dataset to assess its accuracy and generalization capabilities.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df101f7f-b9fe-4e05-8c09-445868a09071",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import neccessary libraries\n",
    "\n",
    "import torch\n",
    "from torchsummary import summary\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.transforms import ToTensor, Normalize\n",
    "from torchvision.datasets import MNIST, CIFAR10\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1775dec-f9ae-497a-9a4d-033db218fb4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.models as models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d21aba-6236-48da-bbc8-4c1f6fefc0b3",
   "metadata": {},
   "source": [
    "#### Step 1: Choose a pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f3f99c-cdcb-46e0-8ca0-3e80e24175b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the pre-trained ResNet-50 model\n",
    "model = models.resnet50(pretrained=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9ebc29-a1c5-4bc6-89e0-48e73f8b4c84",
   "metadata": {},
   "source": [
    "#### Step 2: Modify the pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab895788-a5a7-48e2-8ea4-ce88ccde11c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModifiedResNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ModifiedResNet, self).__init__()\n",
    "        self.resnet = torch.hub.load('pytorch/vision', 'resnet50', pretrained=True)\n",
    "        num_classes = 10  # MNIST has 10 classes\n",
    "        self.resnet.fc = nn.Linear(pretrained_model.fc.in_features, num_classes)  # Change the final fully connected layer for 10 classes\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.resnet(x)\n",
    "\n",
    "model = ModifiedResNet()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d206d8fd-8c1d-48ab-91d3-febdfba7eb28",
   "metadata": {},
   "source": [
    "#### Freeze the pre-trained models' weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f0664a-f12f-4ea5-8374-f6c943e27c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in model.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c6bff1-6488-4de9-89c2-445e3dbc1592",
   "metadata": {},
   "source": [
    "#### Step 4: Data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d59b0db0-eb77-45a0-9e3b-f858bbd5c038",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms.functional import pad\n",
    "\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomRotation(10),\n",
    "    transforms.Grayscale(num_output_channels=3),  # Convert to RGB format\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5], std=[0.5])  # Normalize as before\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.Grayscale(num_output_channels=3),  # Convert to RGB format\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5], std=[0.5])\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c953d30-954a-4adb-a929-27afd3236b2d",
   "metadata": {},
   "source": [
    "#### Step 5: Data augmentation (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b6757a6-e540-4b17-88c4-9a524565a443",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = MNIST(root='./data', train=True, download=True, transform=transform_train)\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7e39cc-ec38-44a9-b8e9-b7de6a6a15d6",
   "metadata": {},
   "source": [
    "#### Step 6: Define the model architecture\n",
    "The code underneath is the same that was used in the creation of a ResNet model for transfer learning,here’s a breakdown:\n",
    "\n",
    "1. `super(CustomResNet, self).__init__()`: Particularly in this particular line, there is meticulous care taken where the constructor of the parent class (the parent class here is `nn.Module`) is invoked to iteritialize the `CustomResNet` class.\n",
    "\n",
    "2. `self.resnet = torch.hub.load(‘pytorch/vision’, ‘resnet50’, pretrained=True)`: This implies that we incorporate the model through the `torch.hub.load(‘pytorch/vision:v0.1’, ‘resnet50’)`. The `pretrained` argument depicts that the mode is initialized with the precondition that we must give the files loaded with the pre-trained weights.\n",
    "\n",
    "3. `self.features = nn.Sequential(*list(pretrained_model.children())[:((fn.seq( [nn.Sequential( io.layer.(nn.select(-1)))] ))`: Meanwhile, the layer incl. the last layer (partially except last densely connected layer of the ResNet) is sequentialized. It is achieved by developing the framework from all the kids (which are filters of pretrained_models) and using them as parameters for nn.Sequential category of objects.\n",
    "4. `self.classifier = nn.Linear(pretrained_model.fc.in_features, 10)`: As a result, it generates a new fully connected layer (`nn.Linear`) which is titled as `classifier` that has the inputs being equal to the number of features in output of the final fully connected layer (`pretrained_model.fc.in_features`) of the ResNet model and the outputs that are ten (assuming the model takes ten classes to classify).\n",
    "\n",
    "In general, an architecture model is constructed in the design of ResNet transfer learning where the initial layer is pre-trained and having been transferred to the ResNet feature layer the new layer is added for doing the specific classification task at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a2d5e9-b687-4f4c-9508-c50da730e394",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomResNet(nn.Module):\n",
    "    def __init__(self, pretrained_model):\n",
    "        super(CustomResNet, self).__init__()\n",
    "        self.resnet = torch.hub.load('pytorch/vision', 'resnet50', pretrained=True)\n",
    "        self.features = nn.Sequential(*list(pretrained_model.children())[:-1])\n",
    "        self.classifier = nn.Linear(pretrained_model.fc.in_features, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "model = CustomResNet(pretrained_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2085c668-d100-4563-8f65-763f07b95f98",
   "metadata": {},
   "source": [
    "#### Step 7: Compile the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3325dc6a-7abf-4bf9-bcff-94313784415e",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "summary(model,(3,32,32))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f653ee14-75ea-454d-a7fe-a521ce2e60f9",
   "metadata": {},
   "source": [
    "#### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47ff4a46-02a2-4971-80e5-e8f2e973125e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable gradient computation for the last few layers\n",
    "for param in model.resnet.layer4.parameters():\n",
    "    param.requires_grad = True\n",
    "\n",
    "# Train the model\n",
    "model.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5b17b08-e9c7-43bb-8e5b-0cfee2fb8aff",
   "metadata": {},
   "source": [
    "#### Fine-tune the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "316a4bf0-12e6-41db-8267-37654feee493",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fine-tuning\n",
    "num_epochs = 10\n",
    "train_losses = []\n",
    "train_correct = 0\n",
    "train_total = 0\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    for inputs, labels in train_loader:\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        train_total += labels.size(0)\n",
    "        train_correct += (predicted == labels).sum().item()\n",
    "\n",
    "    train_losses.append(running_loss / len(train_loader.dataset))\n",
    "    train_accuracy = train_correct / train_total\n",
    "    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}')\n",
    "\n",
    "print(f'Finished fine-tuning with {train_accuracy} accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e34dae7-54c3-4812-b93e-99397e05d4b7",
   "metadata": {},
   "source": [
    "#### Evaluate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d952f730-015c-4f89-9fd2-cdd3b8b6dfc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = MNIST(root='./data', train=False, download=True, transform=transform_test)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)  # Don't shuffle test data\n",
    "\n",
    "# Initialize variables for tracking performance\n",
    "test_losses = []\n",
    "correct = 0\n",
    "total = 0\n",
    "\n",
    "# Loop through epochs for testing\n",
    "for epoch in range(num_epochs):\n",
    "    with torch.no_grad():\n",
    "        running_loss = 0.0\n",
    "\n",
    "        # Evaluate the model on the test set\n",
    "        for images, labels in test_loader:\n",
    "            # Forward pass (no need for gradients during testing)\n",
    "            outputs = model(images)\n",
    "\n",
    "            # Calculate loss (assuming your loss function is defined)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            # Update running loss\n",
    "            running_loss += loss.item()\n",
    "\n",
    "            # Calculate accuracy\n",
    "            _, predicted = torch.max(outputs.data, 1)  # Get the index of the maximum value\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "        # Calculate average loss for the epoch\n",
    "        test_loss = running_loss / len(test_loader.dataset)\n",
    "        test_losses.append(test_loss)\n",
    "\n",
    "    # Print epoch-wise performance (optional)\n",
    "    print(f'Epoch {epoch+1} - Test Loss: {test_loss:.4f}')\n",
    "\n",
    "# Calculate and print overall test accuracy\n",
    "test_accuracy = correct / total\n",
    "print(f'Accuracy of the model on the test set: {test_accuracy:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd1b68e4-7ae9-4c5c-98d3-0c72826efbe0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0b73f75-7bf4-4c7f-80f7-b365bad3ce1e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d5d7199-7b1d-4b06-b9ae-45c6525d990b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "062b8208-bc5b-45fa-921a-59295da44b33",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}