Image-Classification-in-PyTorch/image_classification_using_transfer_learning_in_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#%env CUDA_VISIBLE_DEVICES=2\n",
    "\n",
    "import torch, torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "import time\n",
    "from torchsummary import summary\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Applying Transforms to the Data\n",
    "image_transforms = { \n",
    "    'train': transforms.Compose([\n",
    "        transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),\n",
    "        transforms.RandomRotation(degrees=15),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.CenterCrop(size=224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                             [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'valid': transforms.Compose([\n",
    "        transforms.Resize(size=256),\n",
    "        transforms.CenterCrop(size=224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                             [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'test': transforms.Compose([\n",
    "        transforms.Resize(size=256),\n",
    "        transforms.CenterCrop(size=224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                             [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Load the Data\n",
    "\n",
    "# Set train and valid directory paths\n",
    "\n",
    "dataset = 'caltech_10'\n",
    "\n",
    "train_directory = os.path.join(dataset, 'train')\n",
    "valid_directory = os.path.join(dataset, 'valid')\n",
    "test_directory = os.path.join(dataset, 'test')\n",
    "\n",
    "# Batch size\n",
    "bs = 32\n",
    "\n",
    "# Number of classes\n",
    "num_classes = len(os.listdir(valid_directory))-1  #10#2#257\n",
    "print(num_classes)\n",
    "\n",
    "# Load Data from folders\n",
    "data = {\n",
    "    'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']),\n",
    "    'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid']),\n",
    "    'test': datasets.ImageFolder(root=test_directory, transform=image_transforms['test'])\n",
    "}\n",
    "\n",
    "# Get a mapping of the indices to the class names, in order to see the output classes of the test images.\n",
    "idx_to_class = {v: k for k, v in data['train'].class_to_idx.items()}\n",
    "print(idx_to_class)\n",
    "\n",
    "# Size of Data, to be used for calculating Average Loss and Accuracy\n",
    "train_data_size = len(data['train'])\n",
    "valid_data_size = len(data['valid'])\n",
    "test_data_size = len(data['test'])\n",
    "\n",
    "# Create iterators for the Data loaded using DataLoader module\n",
    "train_data_loader = DataLoader(data['train'], batch_size=bs, shuffle=True)\n",
    "valid_data_loader = DataLoader(data['valid'], batch_size=bs, shuffle=True)\n",
    "test_data_loader = DataLoader(data['test'], batch_size=bs, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_data_size, valid_data_size, test_data_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Load pretrained ResNet50 Model\n",
    "resnet50 = models.resnet50(pretrained=True)\n",
    "resnet50 = resnet50.to('cuda:0')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Freeze model parameters\n",
    "for param in resnet50.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Change the final layer of ResNet50 Model for Transfer Learning\n",
    "fc_inputs = resnet50.fc.in_features\n",
    "\n",
    "resnet50.fc = nn.Sequential(\n",
    "    nn.Linear(fc_inputs, 256),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.4),\n",
    "    nn.Linear(256, num_classes), # Since 10 possible outputs\n",
    "    nn.LogSoftmax(dim=1) # For using NLLLoss()\n",
    ")\n",
    "\n",
    "# Convert model to be used on GPU\n",
    "resnet50 = resnet50.to('cuda:0')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define Optimizer and Loss Function\n",
    "loss_func = nn.NLLLoss()\n",
    "optimizer = optim.Adam(resnet50.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_and_validate(model, loss_criterion, optimizer, epochs=25):\n",
    "    '''\n",
    "    Function to train and validate\n",
    "    Parameters\n",
    "        :param model: Model to train and validate\n",
    "        :param loss_criterion: Loss Criterion to minimize\n",
    "        :param optimizer: Optimizer for computing gradients\n",
    "        :param epochs: Number of epochs (default=25)\n",
    "  \n",
    "    Returns\n",
    "        model: Trained Model with best validation accuracy\n",
    "        history: (dict object): Having training loss, accuracy and validation loss, accuracy\n",
    "    '''\n",
    "    \n",
    "    start = time.time()\n",
    "    history = []\n",
    "    best_acc = 0.0\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        epoch_start = time.time()\n",
    "        print(\"Epoch: {}/{}\".format(epoch+1, epochs))\n",
    "        \n",
    "        # Set to training mode\n",
    "        model.train()\n",
    "        \n",
    "        # Loss and Accuracy within the epoch\n",
    "        train_loss = 0.0\n",
    "        train_acc = 0.0\n",
    "        \n",
    "        valid_loss = 0.0\n",
    "        valid_acc = 0.0\n",
    "        \n",
    "        for i, (inputs, labels) in enumerate(train_data_loader):\n",
    "\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            # Clean existing gradients\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # Forward pass - compute outputs on input data using the model\n",
    "            outputs = model(inputs)\n",
    "            \n",
    "            # Compute loss\n",
    "            loss = loss_criterion(outputs, labels)\n",
    "            \n",
    "            # Backpropagate the gradients\n",
    "            loss.backward()\n",
    "            \n",
    "            # Update the parameters\n",
    "            optimizer.step()\n",
    "            \n",
    "            # Compute the total loss for the batch and add it to train_loss\n",
    "            train_loss += loss.item() * inputs.size(0)\n",
    "            \n",
    "            # Compute the accuracy\n",
    "            ret, predictions = torch.max(outputs.data, 1)\n",
    "            correct_counts = predictions.eq(labels.data.view_as(predictions))\n",
    "            \n",
    "            # Convert correct_counts to float and then compute the mean\n",
    "            acc = torch.mean(correct_counts.type(torch.FloatTensor))\n",
    "            \n",
    "            # Compute total accuracy in the whole batch and add to train_acc\n",
    "            train_acc += acc.item() * inputs.size(0)\n",
    "            \n",
    "            #print(\"Batch number: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}\".format(i, loss.item(), acc.item()))\n",
    "\n",
    "            \n",
    "        # Validation - No gradient tracking needed\n",
    "        with torch.no_grad():\n",
    "\n",
    "            # Set to evaluation mode\n",
    "            model.eval()\n",
    "\n",
    "            # Validation loop\n",
    "            for j, (inputs, labels) in enumerate(valid_data_loader):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # Forward pass - compute outputs on input data using the model\n",
    "                outputs = model(inputs)\n",
    "\n",
    "                # Compute loss\n",
    "                loss = loss_criterion(outputs, labels)\n",
    "\n",
    "                # Compute the total loss for the batch and add it to valid_loss\n",
    "                valid_loss += loss.item() * inputs.size(0)\n",
    "\n",
    "                # Calculate validation accuracy\n",
    "                ret, predictions = torch.max(outputs.data, 1)\n",
    "                correct_counts = predictions.eq(labels.data.view_as(predictions))\n",
    "\n",
    "                # Convert correct_counts to float and then compute the mean\n",
    "                acc = torch.mean(correct_counts.type(torch.FloatTensor))\n",
    "\n",
    "                # Compute total accuracy in the whole batch and add to valid_acc\n",
    "                valid_acc += acc.item() * inputs.size(0)\n",
    "\n",
    "                #print(\"Validation Batch number: {:03d}, Validation: Loss: {:.4f}, Accuracy: {:.4f}\".format(j, loss.item(), acc.item()))\n",
    "            \n",
    "        # Find average training loss and training accuracy\n",
    "        avg_train_loss = train_loss/train_data_size \n",
    "        avg_train_acc = train_acc/train_data_size\n",
    "\n",
    "        # Find average training loss and training accuracy\n",
    "        avg_valid_loss = valid_loss/valid_data_size \n",
    "        avg_valid_acc = valid_acc/valid_data_size\n",
    "\n",
    "        history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])\n",
    "                \n",
    "        epoch_end = time.time()\n",
    "    \n",
    "        print(\"Epoch : {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \\n\\t\\tValidation : Loss : {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s\".format(epoch, avg_train_loss, avg_train_acc*100, avg_valid_loss, avg_valid_acc*100, epoch_end-epoch_start))\n",
    "        \n",
    "        # Save if the model has best accuracy till now\n",
    "        torch.save(model, dataset+'_model_'+str(epoch)+'.pt')\n",
    "            \n",
    "    return model, history\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Print the model to be trained\n",
    "#summary(resnet50, input_size=(3, 224, 224), batch_size=bs, device='cuda')\n",
    "\n",
    "# Train the model for 25 epochs\n",
    "num_epochs = 30\n",
    "trained_model, history = train_and_validate(resnet50, loss_func, optimizer, num_epochs)\n",
    "\n",
    "torch.save(history, dataset+'_history.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "history = np.array(history)\n",
    "plt.plot(history[:,0:2])\n",
    "plt.legend(['Tr Loss', 'Val Loss'])\n",
    "plt.xlabel('Epoch Number')\n",
    "plt.ylabel('Loss')\n",
    "plt.ylim(0,1)\n",
    "plt.savefig(dataset+'_loss_curve.png')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.plot(history[:,2:4])\n",
    "plt.legend(['Tr Accuracy', 'Val Accuracy'])\n",
    "plt.xlabel('Epoch Number')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.ylim(0,1)\n",
    "plt.savefig(dataset+'_accuracy_curve.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    " def computeTestSetAccuracy(model, loss_criterion):\n",
    "    '''\n",
    "    Function to compute the accuracy on the test set\n",
    "    Parameters\n",
    "        :param model: Model to test\n",
    "        :param loss_criterion: Loss Criterion to minimize\n",
    "    '''\n",
    "\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    test_acc = 0.0\n",
    "    test_loss = 0.0\n",
    "\n",
    "    # Validation - No gradient tracking needed\n",
    "    with torch.no_grad():\n",
    "\n",
    "        # Set to evaluation mode\n",
    "        model.eval()\n",
    "\n",
    "        # Validation loop\n",
    "        for j, (inputs, labels) in enumerate(test_data_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Forward pass - compute outputs on input data using the model\n",
    "            outputs = model(inputs)\n",
    "\n",
    "            # Compute loss\n",
    "            loss = loss_criterion(outputs, labels)\n",
    "\n",
    "            # Compute the total loss for the batch and add it to valid_loss\n",
    "            test_loss += loss.item() * inputs.size(0)\n",
    "\n",
    "            # Calculate validation accuracy\n",
    "            ret, predictions = torch.max(outputs.data, 1)\n",
    "            correct_counts = predictions.eq(labels.data.view_as(predictions))\n",
    "\n",
    "            # Convert correct_counts to float and then compute the mean\n",
    "            acc = torch.mean(correct_counts.type(torch.FloatTensor))\n",
    "\n",
    "            # Compute total accuracy in the whole batch and add to valid_acc\n",
    "            test_acc += acc.item() * inputs.size(0)\n",
    "\n",
    "            print(\"Test Batch number: {:03d}, Test: Loss: {:.4f}, Accuracy: {:.4f}\".format(j, loss.item(), acc.item()))\n",
    "\n",
    "    # Find average test loss and test accuracy\n",
    "    avg_test_loss = test_loss/test_data_size \n",
    "    avg_test_acc = test_acc/test_data_size\n",
    "\n",
    "    print(\"Test accuracy : \" + str(avg_test_acc))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def predict(model, test_image_name):\n",
    "    '''\n",
    "    Function to predict the class of a single test image\n",
    "    Parameters\n",
    "        :param model: Model to test\n",
    "        :param test_image_name: Test image\n",
    "\n",
    "    '''\n",
    "    \n",
    "    transform = image_transforms['test']\n",
    "\n",
    "    test_image = Image.open(test_image_name)\n",
    "    plt.imshow(test_image)\n",
    "    \n",
    "    test_image_tensor = transform(test_image)\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        test_image_tensor = test_image_tensor.view(1, 3, 224, 224).cuda()\n",
    "    else:\n",
    "        test_image_tensor = test_image_tensor.view(1, 3, 224, 224)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        # Model outputs log probabilities\n",
    "        out = model(test_image_tensor)\n",
    "        ps = torch.exp(out)\n",
    "        topk, topclass = ps.topk(3, dim=1)\n",
    "        for i in range(3):\n",
    "            print(\"Predcition\", i+1, \":\", idx_to_class[topclass.cpu().numpy()[0][i]], \", Score: \", topk.cpu().numpy()[0][i])\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Test a particular model on a test image\n",
    "\n",
    "dataset = 'caltech_10'\n",
    "model = torch.load('caltech_10_model_8.pt')\n",
    "predict(model, 'pixabay-test-animals/triceratops-954293_640.jpg')\n",
    "\n",
    "# Load Data from folders\n",
    "#computeTestSetAccuracy(model, loss_func)\n",
    "        \n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}