Lab 10 - Data prep & transfer learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 10:  Data Preparation & CNN Transfer Learning\n",
    "\n",
    "[**Haiping Lu**](http://staffwww.dcs.shef.ac.uk/people/H.Lu/) -  [COM4509/6509 MLAI2020](https://github.com/maalvarezl/MLAI) @ The University of Sheffield\n",
    "\n",
    "**Sources**: This notebook is based on Lab 4 of my [SimplyDeep](https://github.com/haipinglu/SimplyDeep/) notebooks, with sources from PyTorch tutorials on [*Writing Custom Datasets, DataLoaders and Transforms*](https://pytorch.org/tutorials/beginner/data_loading_tutorial.html) and [*Transfer Learning for Computer Vision Tutorial*](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html) by [Sasank Chilamkurthy](https://chsasank.github.io).\n",
    "\n",
    "No question has been set. For more advanced examples, you may refer to our [PyKale examples](https://github.com/pykale/pykale/tree/master/examples) or other GitHub repositories of interest to you."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective\n",
    "\n",
    "* To load and preprocess/augment data from a non trivial dataset.\n",
    "* To make use of pretrained models for image classification\n",
    "\n",
    "**Suggested reading**: \n",
    "* [PyTorch tutorial on *Writing Custom Datasets, DataLoaders and Transforms*](https://pytorch.org/tutorials/beginner/data_loading_tutorial.html)\n",
    "* [PyTorch tutorial on *Transfer Learning for Computer Vision Tutorial*](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html)\n",
    "* [Data preparation - Wikipedia](https://en.wikipedia.org/wiki/Data_preparation)\n",
    "* [Transfer learning - Stanford CNN course](https://cs231n.github.io/transfer-learning/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why\n",
    "\n",
    "**Data preparation** is important when applying machine learning to real-world problems. Mostly, data are not in a format that can be directly analysed by machine learning algorithms, and we need to preprocess the data to such a format. Often, it is also helpful to inspect data to know it (well) before machine learning, to avoid unpleasant surprises.\n",
    "\n",
    "**Transfer learning** can help improve the performance on a particular learning problem by leveraging knowledge gained from a different but related problem. This can also alleviate the demanded resources and accelerate development. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Data Preparation\n",
    "\n",
    "PyTorch provides many tools to make data loading easy and hopefully, to make your code more readable. With conda, we should have ``scikit-image`` (for image io and transforms) and ``pandas`` (for easier csv parsing) already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, utils\n",
    "\n",
    "# Ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "plt.ion()   # interactive mode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 The facial pose dataset\n",
    "The dataset we are going to deal with is that of facial pose, with 68 different landmark points are annotated for each face.\n",
    "\n",
    "**Download** the dataset from [here](https://download.pytorch.org/tutorial/faces.zip) and unzip the images to a directory named `data/faces/`.\n",
    "    \n",
    "**Landmarks**: the dataset comes with a csv file `face_landmarks.csv` with each row containing the image filename followed by the $x$ and $y$ coordinates of the 68 landmarks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "landmarks_frame = pd.read_csv('data/faces/face_landmarks.csv')\n",
    "landmarks_frame.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Inspection**: let us take a look at the third image (index $n=2$) `10comm-decarlo.jpg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 2 #The index \n",
    "img_name = landmarks_frame.iloc[n, 0]\n",
    "landmarks = landmarks_frame.iloc[n, 1:].as_matrix()\n",
    "landmarks = landmarks.astype('float').reshape(-1, 2)\n",
    "\n",
    "print('Image name: {}'.format(img_name))\n",
    "print('Landmarks shape: {}'.format(landmarks.shape))\n",
    "print('First 4 Landmarks: {}'.format(landmarks[:4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write a simple helper function to show an image and its landmarks\n",
    "and use it to show a sample.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_landmarks(image, landmarks):\n",
    "    \"\"\"Show image with landmarks\"\"\"\n",
    "    plt.imshow(image)\n",
    "    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')\n",
    "    plt.pause(0.001)  # pause a bit so that plots are updated\n",
    "\n",
    "plt.figure()\n",
    "show_landmarks(io.imread(os.path.join('data/faces/', img_name)),landmarks)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a custom torch Dataset \n",
    "\n",
    "``torch.utils.data.Dataset`` is an abstract class representing a dataset. Our custom dataset should inherit ``Dataset`` and **override** the following methods:\n",
    "\n",
    "-  ``__len__`` so that ``len(dataset)`` returns the size of the dataset.\n",
    "-  ``__getitem__`` to support the indexing such that ``dataset[i]`` can\n",
    "   be used to get $i$\\ th sample\n",
    "\n",
    "Let's create a dataset class for our face landmarks dataset. We will read the csv in ``__init__`` but leave the reading of images to ``__getitem__``. This is memory efficient because all the images are not stored in the memory at once but read as required.\n",
    "\n",
    "Sample of our dataset will be a dict ``{'image': image, 'landmarks': landmarks}``. Our dataset will take an\n",
    "optional argument ``transform`` so that any required processing can be applied on the sample. We will see the usefulness of ``transform`` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FaceLandmarksDataset(Dataset):\n",
    "    \"\"\"Face Landmarks dataset.\"\"\"\n",
    "\n",
    "    def __init__(self, csv_file, root_dir, transform=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            csv_file (string): Path to the csv file with annotations.\n",
    "            root_dir (string): Directory with all the images.\n",
    "            transform (callable, optional): Optional transform to be applied on a sample.\n",
    "        \"\"\"\n",
    "        self.landmarks_frame = pd.read_csv(csv_file)\n",
    "        self.root_dir = root_dir\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.landmarks_frame)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        img_name = os.path.join(self.root_dir, self.landmarks_frame.iloc[idx, 0])\n",
    "        image = io.imread(img_name)\n",
    "        landmarks = self.landmarks_frame.iloc[idx, 1:]\n",
    "        landmarks = np.array([landmarks])\n",
    "        landmarks = landmarks.astype('float').reshape(-1, 2)\n",
    "        sample = {'image': image, 'landmarks': landmarks}\n",
    "\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Test our custom Dataset\n",
    "Let's instantiate this class and iterate through the data samples. We will print the sizes of first $N$ samples and show their landmarks.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv', root_dir='data/faces/')\n",
    "\n",
    "fig = plt.figure()\n",
    "N=3  #number to show\n",
    "for i in range(len(face_dataset)):\n",
    "    sample = face_dataset[i]\n",
    "\n",
    "    print(i, sample['image'].shape, sample['landmarks'].shape)\n",
    "\n",
    "    ax = plt.subplot(1, N, i + 1)\n",
    "    plt.tight_layout()\n",
    "    ax.set_title('Sample #{}'.format(i))\n",
    "    ax.axis('off')\n",
    "    show_landmarks(**sample)\n",
    "\n",
    "    if i == N-1:\n",
    "        #plt.show()\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Transforms\n",
    "\n",
    "#### Define transforms\n",
    "One issue we can see from the above is that the samples are not of the\n",
    "same size. Most neural networks expect the images of a fixed size.\n",
    "Therefore, we will need to write some prepocessing code.\n",
    "Let's create three transforms:\n",
    "\n",
    "-  ``Rescale``: to scale the image\n",
    "-  ``RandomCrop``: to crop from image randomly. This is data\n",
    "   augmentation.\n",
    "-  ``ToTensor``: to convert the numpy images to torch images (we need to\n",
    "   swap axes).\n",
    "\n",
    "We will write them as **callable classes** instead of simple functions so\n",
    "that parameters of the transform need not be passed everytime it's\n",
    "called. For this, we just need to implement ``__call__`` method and\n",
    "if required, ``__init__`` method. We can then use a transform like this:\n",
    "\n",
    "```python\n",
    "    tsfm = Transform(params)\n",
    "    transformed_sample = tsfm(sample)\n",
    "```\n",
    "Observe below how these transforms had to be applied **both** on the image and\n",
    "landmarks.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Rescale(object):\n",
    "    \"\"\"Rescale the image in a sample to a given size.\n",
    "\n",
    "    Args:\n",
    "        output_size (tuple or int): Desired output size. If tuple, output is matched to output_size. \n",
    "        If int, smaller of image edges is matched to output_size keeping aspect ratio the same.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, output_size):\n",
    "        assert isinstance(output_size, (int, tuple))\n",
    "        self.output_size = output_size\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, landmarks = sample['image'], sample['landmarks']\n",
    "\n",
    "        h, w = image.shape[:2]\n",
    "        if isinstance(self.output_size, int):\n",
    "            if h > w:\n",
    "                new_h, new_w = self.output_size * h / w, self.output_size\n",
    "            else:\n",
    "                new_h, new_w = self.output_size, self.output_size * w / h\n",
    "        else:\n",
    "            new_h, new_w = self.output_size\n",
    "\n",
    "        new_h, new_w = int(new_h), int(new_w)\n",
    "\n",
    "        img = transform.resize(image, (new_h, new_w))\n",
    "\n",
    "        # h and w are swapped for landmarks because for images,\n",
    "        # x and y axes are axis 1 and 0 respectively\n",
    "        landmarks = landmarks * [new_w / w, new_h / h]\n",
    "\n",
    "        return {'image': img, 'landmarks': landmarks}\n",
    "\n",
    "\n",
    "class RandomCrop(object):\n",
    "    \"\"\"Crop randomly the image in a sample.\n",
    "\n",
    "    Args:\n",
    "        output_size (tuple or int): Desired output size. If int, square crop is made.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, output_size):\n",
    "        assert isinstance(output_size, (int, tuple))\n",
    "        if isinstance(output_size, int):\n",
    "            self.output_size = (output_size, output_size)\n",
    "        else:\n",
    "            assert len(output_size) == 2\n",
    "            self.output_size = output_size\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, landmarks = sample['image'], sample['landmarks']\n",
    "\n",
    "        h, w = image.shape[:2]\n",
    "        new_h, new_w = self.output_size\n",
    "\n",
    "        top = np.random.randint(0, h - new_h)\n",
    "        left = np.random.randint(0, w - new_w)\n",
    "\n",
    "        image = image[top: top + new_h, \n",
    "                      left: left + new_w]\n",
    "\n",
    "        landmarks = landmarks - [left, top]\n",
    "\n",
    "        return {'image': image, 'landmarks': landmarks}\n",
    "\n",
    "\n",
    "class ToTensor(object):\n",
    "    \"\"\"Convert ndarrays in sample to Tensors.\"\"\"\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, landmarks = sample['image'], sample['landmarks']\n",
    "\n",
    "        # swap color axis because\n",
    "        # numpy image: H x W x C\n",
    "        # torch image: C X H X W\n",
    "        image = image.transpose((2, 0, 1))\n",
    "        return {'image': torch.from_numpy(image),\n",
    "                'landmarks': torch.from_numpy(landmarks)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compose transforms and apply them to a sample\n",
    "\n",
    "Now, we apply the transforms on a sample.\n",
    "\n",
    "Let's say we want to rescale the shorter side of the image to 256 and\n",
    "then randomly crop a square of size 224 from it. i.e, we want to compose\n",
    "``Rescale`` and ``RandomCrop`` transforms.\n",
    "``torchvision.transforms.Compose`` is a simple callable class which allows us\n",
    "to do this.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = Rescale(256)\n",
    "crop = RandomCrop(128)\n",
    "composed = transforms.Compose([Rescale(256),\n",
    "                               RandomCrop(224)])\n",
    "\n",
    "n=65  #Choose image 65 to apply the transform\n",
    "# Apply each of the above transforms on sample.\n",
    "#fig = plt.figure()\n",
    "sample = face_dataset[n]\n",
    "for i, tsfrm in enumerate([scale, crop, composed]):\n",
    "    transformed_sample = tsfrm(sample)\n",
    "    print(i)\n",
    "\n",
    "    ax = plt.subplot(1, 3, i + 1)\n",
    "    plt.tight_layout()\n",
    "    ax.set_title(type(tsfrm).__name__)\n",
    "    show_landmarks(**transformed_sample)\n",
    "\n",
    "#plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  1.4 Iterating through the dataset\n",
    "Let's put this all together to create a dataset with composed\n",
    "transforms. To summarize, every time this dataset is sampled:\n",
    "\n",
    "-  An image is read from the file on the fly\n",
    "-  Transforms are applied on the read image\n",
    "-  Since one of the transforms is random, data is augmentated on\n",
    "   sampling\n",
    "\n",
    "We can iterate over the created dataset with a ``for i in range``\n",
    "loop as before. Here, we rescale all images to the same size of 256x256 and then crop a 224x224 subimage from each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',\n",
    "                                           root_dir='data/faces/',\n",
    "                                           transform=transforms.Compose([\n",
    "                                               Rescale(256),\n",
    "                                               RandomCrop(224),\n",
    "                                               ToTensor()\n",
    "                                           ]))\n",
    "\n",
    "for i in range(len(transformed_dataset)):\n",
    "    sample = transformed_dataset[i]\n",
    "\n",
    "    print(i, sample['image'].size(), sample['landmarks'].size())\n",
    "\n",
    "    if i == 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, we are losing a lot of features by using a simple ``for`` loop to\n",
    "iterate over the data. In particular, we are missing out on:\n",
    "\n",
    "-  Batching the data\n",
    "-  Shuffling the data\n",
    "-  Load the data in parallel using ``multiprocessing`` workers.\n",
    "\n",
    "``torch.utils.data.DataLoader`` is an iterator which provides all these\n",
    "features. Parameters used below should be clear. One parameter of\n",
    "interest is ``collate_fn``. You can specify how exactly the samples need\n",
    "to be batched using ``collate_fn``. However, default collate should work\n",
    "fine for most use cases. Check out the [DataLoader API here](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader). \n",
    "\n",
    "**`num_workers`** specifies the number of workers for [multi-process data loading](https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading). On (my) Windows, if `num_workers` is set to a positive number, there will be a `BrokenPipeError`. On (my) Linux, there is no such probem, e.g., setting `num_workers=4` works fine. I am not sure about Mac."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(transformed_dataset, batch_size=4,\n",
    "                        shuffle=True, num_workers=0)\n",
    "\n",
    "\n",
    "# Helper function to show a batch\n",
    "def show_landmarks_batch(sample_batched):\n",
    "    \"\"\"Show image with landmarks for a batch of samples.\"\"\"\n",
    "    images_batch, landmarks_batch = \\\n",
    "            sample_batched['image'], sample_batched['landmarks']\n",
    "    batch_size = len(images_batch)\n",
    "    im_size = images_batch.size(2)\n",
    "    grid_border_size = 2\n",
    "\n",
    "    grid = utils.make_grid(images_batch)\n",
    "    plt.imshow(grid.numpy().transpose((1, 2, 0)))\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size + (i + 1) * grid_border_size,\n",
    "                    landmarks_batch[i, :, 1].numpy() + grid_border_size,\n",
    "                    s=10, marker='.', c='r')\n",
    "\n",
    "        plt.title('Batch from dataloader')\n",
    "\n",
    "for i_batch, sample_batched in enumerate(dataloader):\n",
    "    print(i_batch, sample_batched['image'].size(),\n",
    "          sample_batched['landmarks'].size())\n",
    "\n",
    "    # observe 4th batch and stop.\n",
    "    if i_batch == 3:\n",
    "        plt.figure()\n",
    "        show_landmarks_batch(sample_batched)\n",
    "        plt.axis('off')\n",
    "        plt.ioff()\n",
    "        plt.show()\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## 2. CNN Transfer Learning in Computer Vision \n",
    "Next, we will train a convolutional neural network (CNN) for image classification using transfer learning. As pointed out in the [cs231n notes](https://cs231n.github.io/transfer-learning)\n",
    ">  In practice, very few people train an entire Convolutional Network\n",
    "   from scratch (with random initialization), because it is relatively\n",
    "    rare to have a dataset of sufficient size. Instead, it is common to\n",
    "    pretrain a ConvNet on a very large dataset (e.g. ImageNet, which\n",
    "    contains 1.2 million images with 1000 categories), and then use the\n",
    "    ConvNet either as an initialization or a fixed feature extractor for\n",
    "    the task of interest.\n",
    "\n",
    "These two major transfer learning scenarios look as follows:\n",
    "\n",
    "-  **Finetuning the convnet**: Instead of random initializaion, we\n",
    "   initialize the network with a pretrained network, like the one that is\n",
    "   trained on imagenet 1000 dataset. Rest of the training looks as\n",
    "   usual.\n",
    "-  **ConvNet as fixed feature extractor**: Here, we will freeze the weights\n",
    "   for all of the network except that of the final fully connected\n",
    "   layer. This last fully connected layer is replaced with a new one\n",
    "   with random weights and only this layer is trained.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import torchvision\n",
    "from torchvision import datasets, models\n",
    "import time\n",
    "import copy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Load the `ants` & `bees` Data\n",
    "\n",
    "We will train a model to classify **ants** and **bees**. We have about 120 training images each for ants and bees. There are 75 validation images for each class. Usually, this is a very small dataset to generalize upon, if trained from scratch. Since we are using transfer learning, we should be able to generalize reasonably well. This dataset is a very small subset of [imagenet](https://en.wikipedia.org/wiki/ImageNet). We will use torchvision and torch.utils.data packages for loading the data.\n",
    "\n",
    "**Download** the data from [here](https://download.pytorch.org/tutorial/hymenoptera_data.zip) to our `data` directory and extract it there.\n",
    "\n",
    "#### Compose the transforms\n",
    "\n",
    "Check out the [transforms API here](https://pytorch.org/docs/stable/torchvision/transforms.html) for avaiable transforms in PyTorch.\n",
    "\n",
    "**Note**: the numbers `[0.485, 0.456, 0.406]` and `[0.229, 0.224, 0.225]` in `Normalize` are from the [imagenet](https://en.wikipedia.org/wiki/ImageNet). See discussions [here](https://stackoverflow.com/questions/58151507/why-pytorch-officially-use-mean-0-485-0-456-0-406-and-std-0-229-0-224-0-2) and also [here ](https://github.com/jacobgil/pytorch-grad-cam/issues/6).\n",
    "\n",
    "**ImageFolder** is a method of `torchvision.datasets`. See the [`ImageFolder` API here](https://pytorch.org/docs/stable/torchvision/datasets.html#imagefolder).\n",
    "It assumes that images are organized in the following way: ::\n",
    "\n",
    "    root/ants/xxx.png\n",
    "    root/ants/xxy.jpeg\n",
    "    root/ants/xxz.png\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    root/bees/123.jpg\n",
    "    root/bees/nsdf3.png\n",
    "    root/bees/asd932_.png\n",
    "\n",
    "where 'ants', 'bees' etc. are class labels. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data augmentation and normalization for training\n",
    "# Just normalization for validation\n",
    "data_transforms = {\n",
    "    'train': transforms.Compose([\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'val': transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "}\n",
    "\n",
    "data_dir = 'data/hymenoptera_data'\n",
    "image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])\n",
    "                  for x in ['train', 'val']}\n",
    "\n",
    "dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,\n",
    "                                             shuffle=True, num_workers=4)\n",
    "              for x in ['train', 'val']}\n",
    "dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}\n",
    "class_names = image_datasets['train'].classes\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inspect data by visualizing a few images\n",
    "\n",
    "Let's visualize a few training images so as to understand the data augmentations.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(inp, title=None):\n",
    "    \"\"\"Imshow for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0)) #Tensor to numpy format conversion\n",
    "    mean = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.pause(0.001)  # pause a bit so that plots are updated\n",
    "\n",
    "\n",
    "# Get a batch of training data\n",
    "inputs, classes = next(iter(dataloaders['train']))\n",
    "\n",
    "# Make a grid from batch\n",
    "out = torchvision.utils.make_grid(inputs)\n",
    "\n",
    "imshow(out, title=[class_names[x] for x in classes])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Training the model\n",
    "\n",
    "Now, let's write a general function to train a model. Here, we will\n",
    "illustrate:\n",
    "\n",
    "-  Scheduling the learning rate\n",
    "-  Saving the best model\n",
    "-  Recording the computational time.\n",
    "\n",
    "In the following, parameter ``scheduler`` is an LR scheduler object from\n",
    "``torch.optim.lr_scheduler``. Check out [its API here](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # Set model to training mode\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "\n",
    "            # Iterate over data.\n",
    "            for inputs, labels in dataloaders[phase]:\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\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",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\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('{} Loss: {:.4f} Acc: {:.4f}'.format(\n",
    "                phase, epoch_loss, epoch_acc))\n",
    "\n",
    "            # deep copy the model\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('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "    print('Best val Acc: {:4f}'.format(best_acc))\n",
    "\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualizing the model predictions\n",
    "\n",
    "Generic function to display predictions for a few images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_model(model, num_images=6):\n",
    "    was_training = model.training\n",
    "    model.eval()\n",
    "    images_so_far = 0\n",
    "    fig = plt.figure()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i, (inputs, labels) in enumerate(dataloaders['val']):\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "\n",
    "            for j in range(inputs.size()[0]):\n",
    "                images_so_far += 1\n",
    "                ax = plt.subplot(num_images//2, 2, images_so_far)\n",
    "                ax.axis('off')\n",
    "                ax.set_title('predicted: {}'.format(class_names[preds[j]]))\n",
    "                imshow(inputs.cpu().data[j])\n",
    "\n",
    "                if images_so_far == num_images:\n",
    "                    model.train(mode=was_training)\n",
    "                    return\n",
    "        model.train(mode=was_training)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Finetuning the convnet\n",
    "\n",
    "Load a pretrained model and reset final fully connected layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft = models.resnet18(pretrained=True)\n",
    "num_ftrs = model_ft.fc.in_features\n",
    "# Here the size of each output sample is set to 2.\n",
    "# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).\n",
    "model_ft.fc = nn.Linear(num_ftrs, 2)\n",
    "\n",
    "model_ft = model_ft.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Observe that all parameters are being optimized\n",
    "optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "# Decay LR by a factor of 0.1 every 7 epochs\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train and evaluate\n",
    "\n",
    "We set the number of epochs small to save time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epochs=5;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=max_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_model(model_ft)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 ConvNet as fixed feature extractor\n",
    "\n",
    "Here, we need to freeze all the network except the final layer. We need to set ``requires_grad == False`` to freeze the parameters so that the gradients are not computed in ``backward()``.\n",
    "\n",
    "You can read more about this in the [documentation](https://pytorch.org/docs/master/notes/autograd.html#excluding-subgraphs-from-backward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_conv = torchvision.models.resnet18(pretrained=True)\n",
    "for param in model_conv.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# Parameters of newly constructed modules have requires_grad=True by default\n",
    "num_ftrs = model_conv.fc.in_features\n",
    "model_conv.fc = nn.Linear(num_ftrs, 2)\n",
    "\n",
    "model_conv = model_conv.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Observe that only parameters of final layer are being optimized as\n",
    "# opposed to before.\n",
    "optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "# Decay LR by a factor of 0.1 every 7 epochs\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train and evaluate\n",
    "\n",
    "On CPU this will take about half the time compared to previous scenario. This is expected as gradients don't need to be computed for most of the network. However, forward does need to be computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_conv = train_model(model_conv, criterion, optimizer_conv, \n",
    "                         exp_lr_scheduler, num_epochs=max_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_model(model_conv)\n",
    "\n",
    "plt.ioff()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional ideas to explore\n",
    "\n",
    "\n",
    "* Construct image transforms not yet available in `torchvision.transforms` and use them to compose various manipulations of images.\n",
    "* Replace the `resnet18` model with some other [pretrained models](https://pytorch.org/docs/stable/torchvision/models.html) and compare their performance.\n",
    "* Take some other adventures, e.g., how about freezing all the network except the **last two** layers?\n",
    "* Explore other `torchvision` APIs."
   ]
  }
 ],
 "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}