Lab_4_(part3)_PyTorch_Transfer_Learning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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.6"
    },
    "colab": {
      "name": "Lab 4 PyTorch Transfer Learning",
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zp9qGvQ7UHRy",
        "colab_type": "text"
      },
      "source": [
        "#Lab 4 (part 3): Transfer Learning with PyTorch\n",
        "**Remember to enable GPU ;-)**\n",
        "\n",
        "This tutorial is inspired by a tutorial made by [Sasank Chilamkurthy](https://chsasank.github.io). You can find the original tutorial here:\n",
        "\n",
        "https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html\n",
        "\n",
        "In this tutorial, you will learn how to train your network using\n",
        "transfer learning. You can read more about the transfer learning at [cs231n\n",
        "notes](http://cs231n.github.io/transfer-learning/).\n",
        "\n",
        "Quoting these notes,\n",
        "\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",
        "Traditionally, people have distinguished between two types of transfer learning: finetuning and feature extraction. In **finetuning**, we start with a pretrained model and update all of the model’s parameters for our new task, in essence retraining the whole model. In **feature extraction**, we start with a pretrained model and only update the final layer weights from which we derive predictions. It is called feature extraction because we use the pretrained CNN as a fixed feature-extractor, and only change the output layer. For more technical information about transfer learning see here and here.\n",
        "\n",
        "In practise, it is often best to combine the two approaches, simply because it makes training faster. There three basic steps:\n",
        "\n",
        "1. **Initialize the model:** First download a pre-trained CNN. Use only the pre-trained convolutional base (encoder) and replace the final layers (decoder) with your own that you initialize with random weights.\n",
        "\n",
        "2. **Feature extraction:** Freeze the weights of the convolutional base and train only the final layers. Train only for a few iterations. The purpose is simply to get a better starting point for finetuning.\n",
        "\n",
        "3. **Finetuning:** Unfreeze the convolutional base and train the entire network.\n",
        "\n",
        "However, before performing these steps, we need to load the data and set up dataloaders fro both training and validation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psDB2dxWZZug",
        "colab_type": "text"
      },
      "source": [
        "##1. Data loaders\n",
        "First, include what's neede:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DGjUY7J2UHRz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from __future__ import print_function, division\n",
        "\n",
        "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\n",
        "\n",
        "plt.ion()   # interactive mode"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2sy9KuxlUHR1",
        "colab_type": "text"
      },
      "source": [
        "###1.1 Download data\n",
        "We will use [torchvision](https://pytorch.org/docs/stable/torchvision/index.html) and [torch.utils.data](https://pytorch.org/docs/stable/data.html) packages for loading the data.\n",
        "\n",
        "The problem we're going to solve today is to train a model to classify\n",
        "**ants** and **bees**. This dataset is a very small subset of imagenet. 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.\n",
        "\n",
        "First, let's download the data from [here](https://download.pytorch.org/tutorial/hymenoptera_data.zip) and extract it to the current directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "05hf3M0RUSMX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pathlib import Path\n",
        "import requests\n",
        "\n",
        "DATA_PATH = Path(\"data\")\n",
        "PATH = DATA_PATH / \"hymenoptera_data\"\n",
        "\n",
        "PATH.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "URL = \"https://download.pytorch.org/tutorial/\"\n",
        "FILENAME = \"hymenoptera_data.zip\"\n",
        "\n",
        "if not (DATA_PATH / FILENAME).exists():\n",
        "        content = requests.get(URL + FILENAME).content\n",
        "        (DATA_PATH / FILENAME).open(\"wb\").write(content)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SGkPGH1PUuXN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import zipfile\n",
        "with zipfile.ZipFile((DATA_PATH / FILENAME).as_posix(), \"r\") as f:\n",
        "    f.extractall(DATA_PATH)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T61eUYSSezfP",
        "colab_type": "text"
      },
      "source": [
        "Check that the file has been downloaded and unzipped:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DvUnh9jEUeku",
        "colab_type": "code",
        "outputId": "cf80f2b3-7c1f-48a8-cda6-c771416a5f4a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "!ls data/hymenoptera_data"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "train  val\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NVgshPbFfAlX",
        "colab_type": "text"
      },
      "source": [
        "###1.2 Simple data loader\n",
        "Let's set up a simple [DataLoader](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) based on [ImageFolder](https://pytorch.org/docs/stable/torchvision/datasets.html#imagefolder) and visualize some training examples.\n",
        "\n",
        "- ImageFolder is a generic data set where the images are arranged in sub-folders corresponding to classes.\n",
        "- DataLoader combines a dataset and a sampler, and provides an iterable over the given dataset.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ak2pUPi3fSej",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# The transformer is used here to reshape images to 224x224 and convert to tensor\n",
        "simple_transformer = transforms.Compose([\n",
        "        transforms.CenterCrop(224),\n",
        "        transforms.ToTensor()])\n",
        "\n",
        "data_dir = os.path.join(str(PATH), 'train') # 'data/hymenoptera_data/train'\n",
        "\n",
        "# ImageFolder is a generic data set where the images are arranged in sub-folders corresponding to classes\n",
        "train_set = datasets.ImageFolder(data_dir,simple_transformer)\n",
        "\n",
        "# Combines a dataset and a sampler, and provides an iterable over the given dataset.\n",
        "dataloader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=4)\n",
        "\n",
        "class_names = train_set.classes"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lr4JpGVOi5UT",
        "colab_type": "text"
      },
      "source": [
        "###Questions 1:\n",
        "- What does `shuffle=True` mean?\n",
        "- What dies `num_workers=4` mean?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RAH-Lk7xm5F_",
        "colab_type": "text"
      },
      "source": [
        "Get a batch of training images and display it:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "382XCNVqUHRu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LrQwSMoZUHR5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def imshow(inp, title=None):\n",
        "    \"\"\"Imshow for Tensor.\"\"\"\n",
        "    inp = inp.numpy().transpose((1, 2, 0))\n",
        "    plt.imshow(inp)\n",
        "    if title is not None:\n",
        "        plt.title(title)\n",
        "\n",
        "# Get a random batch of training data\n",
        "inputs, classes = next(iter(dataloader))\n",
        "\n",
        "# Make a grid of images\n",
        "out = torchvision.utils.make_grid(inputs)\n",
        "\n",
        "imshow(out, title=[class_names[x] for x in classes])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Kw8eL-8Wibh",
        "colab_type": "text"
      },
      "source": [
        "###1.3 Data augmentation\n",
        "We can add more variation to the training data by using data augmentation. We can do this using PyTorch's [transformer](https://pytorch.org/docs/stable/torchvision/transforms.html).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SnzUTcvSgwvL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Data augmentation transformer\n",
        "augmentation_transformer = transforms.Compose([\n",
        "        transforms.RandomResizedCrop(224), # Crop the given PIL Image to random size and aspect ratio\n",
        "        transforms.RandomHorizontalFlip(0.5), # Horizontally flip the given PIL Image randomly with a given probability. \n",
        "        transforms.ColorJitter(brightness=0.1), # Randomly change the brightness, contrast and saturation of an image\n",
        "        transforms.RandomRotation(degrees=10), # Rotate the image by angle.\n",
        "        transforms.ToTensor()])\n",
        "\n",
        "# Same as above, expect that we are using augmentation_transformer instead of simple_transformer\n",
        "train_set = datasets.ImageFolder(data_dir, augmentation_transformer)\n",
        "dataloader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=4)\n",
        "inputs, classes = next(iter(dataloader))\n",
        "out = torchvision.utils.make_grid(inputs)\n",
        "imshow(out, title=[class_names[x] for x in classes])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oa9sR7MWjSNG",
        "colab_type": "text"
      },
      "source": [
        "###Questions 2:\n",
        "- Why do we even need data augmentation?\n",
        "- Would `transforms.RandomHorizontalFlip` make sense on MNIST? Why/why not?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wVpbYHWLj_Ll",
        "colab_type": "text"
      },
      "source": [
        "###1.4 Set up data loaders for both training and validation\n",
        "Finally, we are ready to create data loaders for both training and validation.\n",
        "\n",
        "Notice that we use dictionaries (curly brackets) here to store both train and val in the same variable."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vCfbO6b9UHR2",
        "colab_type": "code",
        "colab": {}
      },
      "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(0.5), \n",
        "        transforms.ColorJitter(brightness=0.1),\n",
        "        transforms.RandomRotation(degrees=10),\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 = str(PATH) # 'data/hymenoptera_data'\n",
        "\n",
        "image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),\n",
        "                                          data_transforms[x])\n",
        "                  for x in ['train', 'val']}\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\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eXGxFdJokIpO",
        "colab_type": "text"
      },
      "source": [
        "###Questions 3:\n",
        "What does\n",
        "\n",
        "```\n",
        "transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
        "```\n",
        "\n",
        "do, and why do you think it is needed?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h_DjMXiVUHR4",
        "colab_type": "text"
      },
      "source": [
        "Visualize a few images."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ijmZfbaTXxrr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Redefine imshow to handle normalization\n",
        "def imshow(inp, title=None):\n",
        "    \"\"\"Imshow for Tensor.\"\"\"\n",
        "    inp = inp.numpy().transpose((1, 2, 0))\n",
        "\n",
        "    #undo normalization\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",
        "    \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])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WR45Avwcaq2u",
        "colab_type": "text"
      },
      "source": [
        "##2. Initialize the model\n",
        "PyTorch has a number of pretrained models that you can choose from. See complete list [here](https://pytorch.org/docs/stable/torchvision/models.html).\n",
        "\n",
        "We will be using resnet18 here. Let's download the weights:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QEBsbhlfneOj",
        "colab_type": "code",
        "outputId": "34fb68d8-619d-45a4-e4d1-6a4e754f9d92",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        }
      },
      "source": [
        "model = models.resnet18(pretrained=True)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading: \"https://download.pytorch.org/models/resnet18-5c106cde.pth\" to /root/.cache/torch/checkpoints/resnet18-5c106cde.pth\n",
            "100%|██████████| 44.7M/44.7M [00:00<00:00, 88.0MB/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UyeC6d6lpGUC",
        "colab_type": "text"
      },
      "source": [
        "###2.1 Displaying the network architecture\n",
        "You can print the model and its layers."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "afnCuFhwoEUV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ar8w7nFSpzs_",
        "colab_type": "text"
      },
      "source": [
        "To see individual layers, simply dot your way into the model using the names in the list above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DkKEGqRbqA-k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(model.layer1) # First convolutional block\n",
        "print(model.fc) # Final fully connected layer"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZRl9_bJkiWg",
        "colab_type": "text"
      },
      "source": [
        "###Questions 4:\n",
        "As you can see, the final layer is a Linear layer (i.e., a fully connected layer).\n",
        "- What is the dimensionality of the output? Why do you think it has this particular size?\n",
        "- What dimensions should the input vector to the final layer have?\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y-TqWY8UpPJ6",
        "colab_type": "text"
      },
      "source": [
        "###2.2 Reshape final layer\n",
        "Let's change the number of ouputs to 2 instead (for bees and ants)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MCb3HaJCUHSC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Size of input vector to fully connected layer\n",
        "num_ftrs = model.fc.in_features # 512\n",
        "\n",
        "# Replace fully connected layer\n",
        "model.fc = nn.Linear(num_ftrs, 2)\n",
        "\n",
        "print(model.fc)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xhYR92nzUHR7",
        "colab_type": "text"
      },
      "source": [
        "##3. Set up general training function\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 (read more [here](https://pytorch.org/tutorials/beginner/saving_loading_models.html))\n",
        "\n",
        "In the following, parameter ``scheduler`` is a learning rate scheduler object from [``torch.optim.lr_scheduler``](https://pytorch.org/docs/stable/optim.html).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "85vmefD3UHR8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Generic function to train a model\n",
        "def train_model(model, criterion, optimizer, scheduler, num_epochs=25):\n",
        "    since = time.time()\n",
        "\n",
        "    # Copy weights\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",
        "                scheduler.step()\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 only if 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",
        "\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(\n",
        "        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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z_fUjrTMUHR-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Generic function to display predictions for a few images\n",
        "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(2, num_images//2, images_so_far)\n",
        "                ax.axis('off')\n",
        "                ax.set_title('pred/true: {}/{}'.format(class_names[preds[j]],\n",
        "                                                       class_names[labels[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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FrZvGFJavl4s",
        "colab_type": "text"
      },
      "source": [
        "##4. Feature extraction\n",
        "Freeze the weights of the convolutional base and train only the final fully connected layer. Train only for a few iterations. The purpose is simply to get a better starting point for finetuning.\n",
        "\n",
        "We want to set set the ``.requires_grad attribute`` of the parameters in the conovolutional base to ``False``. By default, when we load a pretrained model all of the parameters have ``.requires_grad=True``, which is fine if we are training from scratch or finetuning. However, if we are feature extracting and only want to compute gradients for the newly initialized layer(s) then we want all of the other parameters to not require gradients.\n",
        "\n",
        "You can read more about this in the documentation [here](http://pytorch.org/docs/notes/autograd.html#excluding-subgraphs-from-backward).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lEWPGFwTvqYz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Freeze all layers (i.e., disable training)\n",
        "for param in model.parameters():\n",
        "    param.requires_grad = False\n",
        "\n",
        "# Unfreeze final layer (named fc)\n",
        "for param in model.fc.parameters():\n",
        "    param.requires_grad = True\n",
        "\n",
        "# Put the model on the GPU\n",
        "model = model.to(device)\n",
        "\n",
        "# Loss function\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Observe that all parameters are being optimized\n",
        "optimizer = optim.SGD(model.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, step_size=7, gamma=0.1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-yuD4-C7UHSE",
        "colab_type": "text"
      },
      "source": [
        "###4.1 Train and evaluate\n",
        "It should take around 5-10 min on CPU. On GPU though, it takes less than half a\n",
        "minute."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lGMQ2qAvUHSF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=5)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1AxW-gBxUHSH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "visualize_model(model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1sdEdO-hUHSJ",
        "colab_type": "text"
      },
      "source": [
        "##5. Fine-tune model\n",
        "Finally, let's unfreeze all layers to fine-tune the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7emNRHpOUHSJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Unfreeze all layers\n",
        "for param in model.parameters():\n",
        "    param.requires_grad = True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9hmwJdw6UHSL",
        "colab_type": "text"
      },
      "source": [
        "###5.1 Train and evaluate\n",
        "On CPU this will take more time compared to the previous scenario.\n",
        "This is expected as gradients need to be computed for most of the\n",
        "network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UAIu80n8UHSM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JmjnT788UHSO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "visualize_model(model)\n",
        "\n",
        "plt.ioff()\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "deAPU394lPkz",
        "colab_type": "text"
      },
      "source": [
        "##Suggestions for optional tasks\n",
        "- Try re-running the above tutorial with another data set (see list [here](https://pytorch.org/docs/stable/torchvision/datasets.html)) or a different CNN architecture (get inspiration [here](https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html)).\n",
        "\n",
        "- [Object detection finetuning tutorial](https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html) (use [this](https://colab.research.google.com/github/pytorch/vision/blob/temp-tutorial/tutorials/torchvision_finetuning_instance_segmentation.ipynb) notebook if you are working in Google Colab).\n",
        "\n",
        "![alt text](https://pytorch.org/tutorials/_static/img/tv_tutorial/tv_image01.png)\n",
        "\n",
        "- Learn about [Spatial Transformer Networks](https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html).\n",
        "\n",
        "![alt text](https://pytorch.org/tutorials/_images/stn-arch.png)\n",
        "\n",
        "- How does [Neural Style Transfer](https://pytorch.org/tutorials/advanced/neural_style_tutorial.html) work?\n",
        "\n",
        "![alt text](https://pytorch.org/tutorials/_images/neuralstyle.png)\n",
        "\n",
        "- [Adversarial attacks](https://pytorch.org/tutorials/beginner/fgsm_tutorial.html): You may be surprised to find that adding imperceptible perturbations to an image can cause drastically different model performance.\n",
        "\n",
        "![](https://pytorch.org/tutorials/_images/fgsm_panda_image.png)\n",
        "\n",
        "- [DCGAN tutorial](https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html): Dream up fake celebrities.\n",
        "\n",
        "![](https://pytorch.org/tutorials/_images/sphx_glr_dcgan_faces_tutorial_003.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WrblJHXaPAJe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}