CIFAR10_TransferLearning_FineTuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torchvision import models, transforms\n",
    "import torch.optim as optim\n",
    "from torch.utils import data\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import PosixPath\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Hyperparameters\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "NUM_WORKER = 4\n",
    "\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "train_data_transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomVerticalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "train_set = torchvision.datasets.CIFAR10(root='./data',\n",
    "                                         train=True,\n",
    "                                         download=True,\n",
    "                                         transform=train_data_transform)\n",
    "\n",
    "val_data_transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "\n",
    "val_set = torchvision.datasets.CIFAR10(root='./data',\n",
    "                                       train=False,\n",
    "                                       download=True,\n",
    "                                       transform=val_data_transform)\n",
    "\n",
    "\n",
    "train_loader = data.DataLoader(train_set,\n",
    "                               batch_size=BATCH_SIZE,\n",
    "                               shuffle=True,\n",
    "                               num_workers= NUM_WORKER)\n",
    "\n",
    "val_loader = data.DataLoader(val_set,\n",
    "                            batch_size=BATCH_SIZE,\n",
    "                            shuffle=False,\n",
    "                            num_workers= NUM_WORKER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y= next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([32, 3, 224, 224]), torch.Size([32]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, loss_function, optimaizer, data_loader):\n",
    "    model.train()\n",
    "    \n",
    "    current_loss = 0.0\n",
    "    current_acc = 0\n",
    "    \n",
    "    \n",
    "    for i, (inputs, labels) in enumerate(data_loader):\n",
    "        # send the input/labels to the GPU\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        optimaizer.zero_grad()\n",
    "        \n",
    "        with torch.set_grad_enabled(True):\n",
    "            outputs = model(inputs)\n",
    "            _, predictions = torch.max(outputs, 1)\n",
    "            loss = loss_function(outputs, labels)\n",
    "            \n",
    "            # backward\n",
    "            loss.backward()\n",
    "            optimaizer.step()\n",
    "            \n",
    "        # statistics\n",
    "        current_loss += loss.item() * inputs.size(0)\n",
    "        current_acc += torch.sum(predictions == labels.data)\n",
    "        \n",
    "    total_loss = current_loss / len(data_loader.dataset)\n",
    "    total_acc = current_acc.double() / len(data_loader.dataset)\n",
    "    print('Train Loss: {:.4f}; Accuracy: {:.4f}'.format(total_loss, total_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, loss_function, data_loader):\n",
    "    # set model in evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    current_loss = 0.0\n",
    "    current_acc = 0\n",
    "\n",
    "    # iterate over  the validation data\n",
    "    for i, (inputs, labels) in enumerate(data_loader):\n",
    "        # send the input/labels to the GPU\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # forward\n",
    "        with torch.set_grad_enabled(False):\n",
    "            outputs = model(inputs)\n",
    "            _, predictions = torch.max(outputs, 1)\n",
    "            loss = loss_function(outputs, labels)\n",
    "\n",
    "        # statistics\n",
    "        current_loss += loss.item() * inputs.size(0)\n",
    "        current_acc += torch.sum(predictions == labels.data)\n",
    "\n",
    "    total_loss = current_loss / len(data_loader.dataset)\n",
    "    total_acc = current_acc.double() / len(data_loader.dataset)\n",
    "\n",
    "    print('Test Loss: {:.4f}; Accuracy: {:.4f}'.format(total_loss, total_acc))\n",
    "\n",
    "    return total_loss, total_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_accuracy(accuracy: list):\n",
    "    \"\"\"Plot accuracy\"\"\"\n",
    "    plt.figure()\n",
    "    plt.plot(accuracy)\n",
    "    plt.xticks(\n",
    "        [i for i in range(0, len(accuracy))],\n",
    "        [i + 1 for i in range(0, len(accuracy))])\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer Learning Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transfer_learning(epochs = 5):\n",
    "    \n",
    "    model = torchvision.models.resnet18(pretrained=True)\n",
    "    \n",
    "    for param in model.parameters():\n",
    "        param.requires_grad = False\n",
    "    \n",
    "    # newly constructed layers have requires_grad=True by default\n",
    "    num_features = model.fc.in_features\n",
    "    model.fc = nn.Linear(num_features, 10)\n",
    "    \n",
    "    # transfer to GPU (if available)\n",
    "    model = model.to(device)\n",
    "    \n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "    \n",
    "    # only parameters of the final layer are being optimized\n",
    "    optimizer = optim.Adam(model.fc.parameters())\n",
    "        \n",
    "    \n",
    "    # train\n",
    "    test_acc = list()  # collect accuracy for plotting\n",
    "    for epoch in range(epochs):\n",
    "        print('Epoch {}/{}'.format(epoch + 1, epochs))\n",
    "\n",
    "        train_model(model, loss_function, optimizer, train_loader)\n",
    "        _, acc = test_model(model, loss_function, val_loader)\n",
    "        test_acc.append(acc)\n",
    "\n",
    "    plot_accuracy(test_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine Tuning Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fine_tuning(epochs=5):\n",
    "    # load the pre-trained model\n",
    "    model = models.resnet18(pretrained=True)\n",
    "\n",
    "    # replace the last layer\n",
    "    num_features = model.fc.in_features\n",
    "    model.fc = nn.Linear(num_features, 10)\n",
    "\n",
    "    # transfer the model to the GPU\n",
    "    model = model.to(device)\n",
    "\n",
    "    # loss function\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "    # We'll optimize all parameters\n",
    "    optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "    # train\n",
    "    test_acc = list()  # collect accuracy for plotting\n",
    "    for epoch in range(epochs):\n",
    "        print('Epoch {}/{}'.format(epoch + 1, epochs))\n",
    "\n",
    "        train_model(model, loss_function, optimizer, train_loader)\n",
    "        _, acc = test_model(model, loss_function, val_loader)\n",
    "        test_acc.append(acc)\n",
    "\n",
    "    plot_accuracy(test_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "transfer_learning(epochs = 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tryit",
   "language": "python",
   "name": "tryit"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}