Transfer_Learning.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": "markdown",
      "source": [
        "# Transfer Learning\n",
        "\n",
        "Using an already trained model and only training the last layer of the neural network to speed up the training process."
      ],
      "metadata": {
        "id": "qsVlJj0Y3S0b"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 1** - Importing the required libraries"
      ],
      "metadata": {
        "id": "yqaqdZjn4T8F"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2B0js8zw3OXy"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "import numpy as np\n",
        "import torchvision\n",
        "from torchvision import datasets, models, transforms\n",
        "import matplotlib.pyplot as plt\n",
        "import time\n",
        "import os\n",
        "import copy"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 2** - Uploading the file containing the dataset\n",
        "\n",
        "Instructions for Running:\n",
        "1. Add the shared dataset folder to your Google Drive. (find it in the github repo https://github.com/DarioVajda/ML_basics)\n",
        "2. Mount your Google Drive in Colab.\n",
        "3. Access the dataset using the path where you added the folder."
      ],
      "metadata": {
        "id": "9ELMFFCu7J8V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "data_dir = '/content/drive/My Drive/google_colab/hymenoptera_data'"
      ],
      "metadata": {
        "id": "RvyZeQXi7SDj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 3** -  Loading the datasets"
      ],
      "metadata": {
        "id": "v1q-aD6n4uXT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "mean = np.array([0.485, 0.456, 0.406])\n",
        "std = np.array([0.229, 0.224, 0.225])\n",
        "\n",
        "data_transforms = {\n",
        "    'train': transforms.Compose([\n",
        "        transforms.RandomResizedCrop(224),\n",
        "        transforms.RandomHorizontalFlip(),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean, std)\n",
        "    ]),\n",
        "    'val': transforms.Compose([\n",
        "        transforms.Resize(256),\n",
        "        transforms.CenterCrop(224),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean, std)\n",
        "    ])\n",
        "}\n",
        "\n",
        "sets = [ 'train', 'val' ]\n",
        "image_datasets = { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [ 'train', 'val' ] }\n",
        "dataloaders = { x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=0) for x in [ 'train', 'val' ] }\n",
        "\n",
        "dataset_sizes = { x: len(image_datasets[x]) for x in [ 'train', 'val' ]}\n",
        "class_names = image_datasets['train'].classes\n",
        "print(class_names)"
      ],
      "metadata": {
        "id": "GH-WxBZ04-8i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 4** - Function for training the model"
      ],
      "metadata": {
        "id": "3BV_LFw4trlO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_model(model, criterion, optimizer, scheduler, num_epochs=25):\n",
        "    since = time.time()\n",
        "\n",
        "    best_model_wts = copy.deepcopy(model.state_dict())\n",
        "    best_acc = 0\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        print(f'Epoch {epoch}/{num_epochs}')\n",
        "        print('_' * 20)\n",
        "\n",
        "        for phase in [ 'train', 'val' ]:\n",
        "            if phase == 'train':\n",
        "                model.train()\n",
        "            else:\n",
        "                model.eval()\n",
        "\n",
        "            running_loss = 0\n",
        "            running_corrects = 0\n",
        "\n",
        "            for inputs, labels in dataloaders[phase]:\n",
        "                inputs, labels = inputs.to(device), labels.to(device)\n",
        "\n",
        "                with torch.set_grad_enabled(phase=='train'):\n",
        "                    outputs = model(inputs)\n",
        "                    _, preds = torch.max(outputs, 1)\n",
        "                    loss = criterion(outputs, labels)\n",
        "\n",
        "                    if phase=='train':\n",
        "                        optimizer.zero_grad()\n",
        "                        loss.backward()\n",
        "                        optimizer.step()\n",
        "\n",
        "                running_loss += loss.item() * inputs.size(0)\n",
        "                running_corrects += torch.sum(preds == labels.data)\n",
        "\n",
        "            if phase=='train':\n",
        "                scheduler.step()\n",
        "\n",
        "            epoch_loss = running_loss / dataset_sizes[phase]\n",
        "            epoch_acc = running_corrects.double() / dataset_sizes[phase]\n",
        "\n",
        "            print(f'{phase}, Loss: {epoch_loss:.4f}; Acc: {epoch_acc:.4f}')\n",
        "\n",
        "            if phase == 'val' and epoch_acc > best_acc:\n",
        "                best_acc = epoch_acc\n",
        "                best_model_wts = copy.deepcopy(model.state_dict())\n",
        "\n",
        "        print()\n",
        "\n",
        "    time_elapsed = time.time() - since\n",
        "    print(f'Training complete in {time_elapsed//60:.0f}m {time_elapsed%60:.0f}s')\n",
        "    print(f'Best val Acc: {best_acc:.4f}')\n",
        "\n",
        "    model.load_state_dict(best_model_wts)\n",
        "    return model"
      ],
      "metadata": {
        "id": "mIbHtUegtvB0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 5.1** - Fine tuning the pretrained model\n",
        "\n",
        "Fine tuning is training an entire existing model on a lot less new data. Now all of the weights will be updated in the old model."
      ],
      "metadata": {
        "id": "ifMMzV8Uzt3Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# importing the pretrained model\n",
        "model = models.resnet18(pretrained=True)\n",
        "\n",
        "# creating a new last layer for the neural network, there are 2 outputs because it will be used for classification of two classes\n",
        "num_features = model.fc.in_features\n",
        "model.fc = nn.Linear(num_features, 2)\n",
        "model.to(device)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(model.parameters(),  lr=0.001)\n",
        "\n",
        "# defining the scheduler - it will update the learning rate while training the model\n",
        "step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)\n",
        "\n",
        "# calling the training loop function\n",
        "model = train_model(model, criterion, optimizer, step_lr_scheduler, num_epochs=20)"
      ],
      "metadata": {
        "id": "SgSXuQ3WzwEl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 5.2** - Training only the last layer of the pretrained model\n",
        "\n",
        "All layers except the first one in the pretrained model will be frozen and the old weights will remain unchanged. Only the last layer of the pretrained model will be changed in our training on the new dataset."
      ],
      "metadata": {
        "id": "EaPz_2m11gYf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# importing the pretrained model\n",
        "model = models.resnet18(pretrained=True)\n",
        "\n",
        "# freezing all layers in the pretrained model\n",
        "for param in model.parameters():\n",
        "    param.requires_grad = False\n",
        "\n",
        "# creating a new last layer for the neural network, there are 2 outputs because it will be used for classification of two classes\n",
        "num_features = model.fc.in_features\n",
        "model.fc = nn.Linear(num_features, 2)\n",
        "model.to(device)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(model.parameters(),  lr=0.001)\n",
        "\n",
        "# defining the scheduler - it will update the learning rate while training the model\n",
        "step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)\n",
        "\n",
        "# calling the training loop function\n",
        "model = train_model(model, criterion, optimizer, step_lr_scheduler, num_epochs=20)"
      ],
      "metadata": {
        "id": "hV4ygmmQ2Qg6"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}