05c - Transfer Learning (PyTorch).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning\n",
    "\n",
    "A Convolutional Neural Network (CNN) for image classification is made up of multiple layers that extract features, such as edges, corners, etc; and then use a final fully-connected layer to classify objects based on these features. You can visualize this like this:\n",
    "\n",
    "<table>\n",
    "    <tr><td rowspan=2 style='border: 1px solid black;'>&#x21d2;</td><td style='border: 1px solid black;'>Convolutional Layer</td><td style='border: 1px solid black;'>Pooling Layer</td><td style='border: 1px solid black;'>Convolutional Layer</td><td style='border: 1px solid black;'>Pooling Layer</td><td style='border: 1px solid black;'>Fully Connected Layer</td><td rowspan=2 style='border: 1px solid black;'>&#x21d2;</td></tr>\n",
    "    <tr><td colspan=4 style='border: 1px solid black; text-align:center;'>Feature Extraction</td><td style='border: 1px solid black; text-align:center;'>Classification</td></tr>\n",
    "</table>\n",
    "\n",
    "*Transfer Learning* is a technique where you can take an existing trained model and re-use its feature extraction layers, replacing its final classification layer with a fully-connected layer trained on your own custom images. With this technique, your model benefits from the feature extraction training that was performed on the base model (which may have been based on a larger training dataset than you have access to) to build a classification model for your own specific set of object classes.\n",
    "\n",
    "How does this help? Well, think of it this way. Suppose you take a professional tennis player and a complete beginner, and try to teach them both how to play raquetball. It's reasonable to assume that the professional tennis player will be easier to train, because many of the underlying skills involved in raquetball are already learned. Similarly, a pre-trained CNN model may be easier to train to classify specific set of objects because it's already learned how to identify the features of common objects, such as edges and corners. Fundamentally, a pre-trained model can be a great way to produce an effective classifier even when you have limited data with which to train it.\n",
    "\n",
    "In this notebook, we'll see how to implement transfer learning for a classification model using PyTorch.\n",
    "\n",
    "## Install and import libraries\n",
    "\n",
    "First, let's install and import the PyTorch libraries we're going to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import PyTorch libraries\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Other libraries we'll use\n",
    "import numpy as np\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "print(\"Libraries imported - ready to use PyTorch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the base model\n",
    "\n",
    "To use transfer learning, we need a base model from which we can use the trained feature extraction layers. The ***resnet*** model is an CNN-based image classifier that has been pre-trained using a huge dataset containing a large number of images of 1000 classes of object, so let's download it and take a look at its layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the model (download if not already present)\n",
    "model = torchvision.models.resnet34(pretrained=True)\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the image data\n",
    "\n",
    "The pretrained model has many layers, starting with a convolutional layer that starts the feature extraction process from image data, and ending with a fully-connected linear layer that maps the extracted features to 1000 class labels.\n",
    "\n",
    "For feature extraction to work with our own images, we  need to ensure that the image data we use the train our prediction layer has the same number of features (pixel values) as the images originally used to train the feaure extraction layers. The model does not explicitly give this size, but the first convolutional layer applies by a 7x7 kernel with a stride of 2x2 and results in 64 feature values, so the original size must be 64 x (7 &div; 2), which is 224.\n",
    "\n",
    "PyTorch includes functions for loading and transforming data. We'll use these to create an iterative loader for training data, and a second iterative loader for test data (which we'll use to validate the trained model). The loaders will transform the image data to match the format used to train the original resnet CNN model, convert the image data into *tensors* (which are the core data structure used in PyTorch), and normalize them.\n",
    "\n",
    "Run the following cell to define the data loaders and list the classes for our images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Function to ingest data using training and test loaders\n",
    "def load_dataset(data_path):\n",
    "    \n",
    "    # Resize to 256 x 256, then center-crop to 224x224 (to match the resnet image size)\n",
    "    transformation = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
    "    ])\n",
    "\n",
    "    # Load all of the images, transforming them\n",
    "    full_dataset = torchvision.datasets.ImageFolder(\n",
    "        root=data_path,\n",
    "        transform=transformation\n",
    "    )\n",
    "    \n",
    "    # Split into training (70%) and testing (30%) datasets)\n",
    "    train_size = int(0.7 * len(full_dataset))\n",
    "    test_size = len(full_dataset) - train_size\n",
    "    train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])\n",
    "    \n",
    "    # define a loader for the training data we can iterate through in 30-image batches\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=30,\n",
    "        num_workers=0,\n",
    "        shuffle=False\n",
    "    )\n",
    "    \n",
    "    # define a loader for the testing data we can iterate through in 30-image batches\n",
    "    test_loader = torch.utils.data.DataLoader(\n",
    "        test_dataset,\n",
    "        batch_size=30,\n",
    "        num_workers=0,\n",
    "        shuffle=False\n",
    "    )\n",
    "        \n",
    "    return train_loader, test_loader\n",
    "\n",
    "\n",
    "# Now load the images from the shapes folder\n",
    "import os  \n",
    "data_path = 'data/shapes/'\n",
    "\n",
    "# Get the iterative dataloaders for test and training data\n",
    "train_loader, test_loader = load_dataset(data_path)\n",
    "\n",
    "# Get the class names\n",
    "classes = os.listdir(data_path)\n",
    "classes.sort()\n",
    "print('class names:', classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a prediction layer\n",
    "\n",
    "We downloaded the complete *resnet* model including its final **fc** linear layer. This fully-connected linear layer takes 512 inputs (the extracted features) and produces 1000 outputs (class predictions based on the original training image classes). We need to replace this layer with one that takes the same number of inputs (so we can use the same number of extracted features), but produces a prediction for each of our image classes.\n",
    "\n",
    "We also need to freeze the feature extraction layers to retain the trained weights. Then when we train the model using our images, only the final prediction layer will learn new weight and bias values - the pre-trained weights already learned for feature extraction will remain the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the existing feature extraction layers to read-only\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# Replace the prediction layer\n",
    "num_ftrs = model.fc.in_features\n",
    "model.fc = nn.Linear(num_ftrs, len(classes))\n",
    "\n",
    "# Now print the full model, which will include the feature extraction layers of the base model and our preediction layer\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "\n",
    "With the layers of the CNN defined, we're ready to train it using our image data. The weights used in the feature extraction layers from the base resnet model will not be changed by training, only the final linear layer that maps the features to our shape classes will be trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    # Set the model to training mode\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    print(\"Epoch:\", epoch)\n",
    "    # Process the images in batches\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        # Use the CPU or GPU as appropriate\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        \n",
    "        # Reset the optimizer\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Push the data forward through the model layers\n",
    "        output = model(data)\n",
    "        \n",
    "        # Get the loss\n",
    "        loss = loss_criteria(output, target)\n",
    "        \n",
    "        # Keep a running total\n",
    "        train_loss += loss.item()\n",
    "        \n",
    "        # Backpropagate\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Print metrics for every 10 batches so we see some progress\n",
    "        if batch_idx % 10 == 0:\n",
    "            print('Training set [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(\n",
    "                batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "            \n",
    "    # return average loss for the epoch\n",
    "    avg_loss = train_loss / (batch_idx+1)\n",
    "    print('Training set: Average loss: {:.6f}'.format(avg_loss))\n",
    "    return avg_loss\n",
    "            \n",
    "            \n",
    "def test(model, device, test_loader):\n",
    "    # Switch the model to evaluation mode (so we don't backpropagate or drop)\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        batch_count = 0\n",
    "        for data, target in test_loader:\n",
    "            batch_count += 1\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            \n",
    "            # Get the predicted classes for this batch\n",
    "            output = model(data)\n",
    "            \n",
    "            # Calculate the loss for this batch\n",
    "            test_loss += loss_criteria(output, target).item()\n",
    "            \n",
    "            # Calculate the accuracy for this batch\n",
    "            _, predicted = torch.max(output.data, 1)\n",
    "            correct += torch.sum(target==predicted).item()\n",
    "\n",
    "    # Calculate the average loss and total accuracy for this epoch\n",
    "    avg_loss = test_loss/batch_count\n",
    "    print('Validation set: Average loss: {:.6f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        avg_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))\n",
    "    \n",
    "    # return average loss for the epoch\n",
    "    return avg_loss\n",
    "    \n",
    "    \n",
    "# Now use the train and test functions to train and test the model    \n",
    "\n",
    "device = \"cpu\"\n",
    "if (torch.cuda.is_available()):\n",
    "    # if GPU available, use cuda (on a cpu, training will take a considerable length of time!)\n",
    "    device = \"cuda\"\n",
    "print('Training on', device)\n",
    "\n",
    "# Create an instance of the model class and allocate it to the device\n",
    "model = model.to(device)\n",
    "\n",
    "# Use an \"Adam\" optimizer to adjust weights\n",
    "# (see https://pytorch.org/docs/stable/optim.html#algorithms for details of supported algorithms)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Specify the loss criteria\n",
    "loss_criteria = nn.CrossEntropyLoss()\n",
    "\n",
    "# Track metrics in these arrays\n",
    "epoch_nums = []\n",
    "training_loss = []\n",
    "validation_loss = []\n",
    "\n",
    "# Train over 3 epochs (in a real scenario, you'd likely use many more)\n",
    "epochs = 3\n",
    "for epoch in range(1, epochs + 1):\n",
    "        train_loss = train(model, device, train_loader, optimizer, epoch)\n",
    "        test_loss = test(model, device, test_loader)\n",
    "        epoch_nums.append(epoch)\n",
    "        training_loss.append(train_loss)\n",
    "        validation_loss.append(test_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View the loss history\n",
    "\n",
    "We tracked average training and validation loss for each epoch. We can plot these to verify that the loss reduced over the training process and to detect *over-fitting* (which is indicated by a continued drop in training loss after validation loss has levelled out or started to increase)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "plt.plot(epoch_nums, training_loss)\n",
    "plt.plot(epoch_nums, validation_loss)\n",
    "plt.xlabel('epoch')\n",
    "plt.ylabel('loss')\n",
    "plt.legend(['training', 'validation'], loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate model performance\n",
    "\n",
    "We can see the final accuracy based on the test data, but typically we'll want to explore performance metrics in a little more depth. Let's plot a confusion matrix to see how well the model is predicting each class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Pytorch doesn't have a built-in confusion matrix metric, so we'll use SciKit-Learn\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# Set the model to evaluate mode\n",
    "model.eval()\n",
    "\n",
    "# Get predictions for the test data and convert to numpy arrays for use with SciKit-Learn\n",
    "print(\"Getting predictions from test set...\")\n",
    "truelabels = []\n",
    "predictions = []\n",
    "for data, target in test_loader:\n",
    "    for label in target.cpu().data.numpy():\n",
    "        truelabels.append(label)\n",
    "    for prediction in model.cpu()(data).data.numpy().argmax(1):\n",
    "        predictions.append(prediction) \n",
    "\n",
    "# Plot the confusion matrix\n",
    "cm = confusion_matrix(truelabels, predictions)\n",
    "plt.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
    "plt.colorbar()\n",
    "tick_marks = np.arange(len(classes))\n",
    "plt.xticks(tick_marks, classes, rotation=45)\n",
    "plt.yticks(tick_marks, classes)\n",
    "plt.xlabel(\"Actual Shape\")\n",
    "plt.ylabel(\"Predicted Shape\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the trained model\n",
    "\n",
    "Now that we've trained the model, we can use it to predict the class of an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Function to create a random image (of a square, circle, or triangle)\n",
    "def create_image (size, shape):\n",
    "    from random import randint\n",
    "    import numpy as np\n",
    "    from PIL import Image, ImageDraw\n",
    "    \n",
    "    xy1 = randint(10,40)\n",
    "    xy2 = randint(60,100)\n",
    "    col = (randint(0,200), randint(0,200), randint(0,200))\n",
    "\n",
    "    img = Image.new(\"RGB\", size, (255, 255, 255))\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    \n",
    "    if shape == 'circle':\n",
    "        draw.ellipse([(xy1,xy1), (xy2,xy2)], fill=col)\n",
    "    elif shape == 'triangle':\n",
    "        draw.polygon([(xy1,xy1), (xy2,xy2), (xy2,xy1)], fill=col)\n",
    "    else: # square\n",
    "        draw.rectangle([(xy1,xy1), (xy2,xy2)], fill=col)\n",
    "    del draw\n",
    "    \n",
    "    return img\n",
    "    \n",
    "# Function to predict the class of an image\n",
    "def predict_image(classifier, image):\n",
    "    import numpy\n",
    "    \n",
    "    # Set the classifer model to evaluation mode\n",
    "    classifier.eval()\n",
    "    \n",
    "    # Apply the same transformations as we did for the training images\n",
    "    transformation = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
    "    ])\n",
    "\n",
    "    # Preprocess the image\n",
    "    image_tensor = transformation(image).float()\n",
    "\n",
    "    # Add an extra batch dimension since pytorch treats all inputs as batches\n",
    "    image_tensor = image_tensor.unsqueeze_(0)\n",
    "\n",
    "    # Turn the input into a Variable\n",
    "    input_features = Variable(image_tensor)\n",
    "\n",
    "    # Predict the class of the image\n",
    "    output = classifier(input_features)\n",
    "    index = output.data.numpy().argmax()\n",
    "    return index\n",
    "\n",
    "\n",
    "# Now let's try it with a new image\n",
    "from random import randint\n",
    "from PIL import Image\n",
    "import os, shutil\n",
    "\n",
    "# Create a random test image\n",
    "shape = classes[randint(0, len(classes)-1)]\n",
    "img = create_image ((128,128), shape)\n",
    "\n",
    "# Display the image\n",
    "plt.imshow(img)\n",
    "\n",
    "\n",
    "index = predict_image(model, img)\n",
    "print(classes[index])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learn more\n",
    "\n",
    "* [PyTorch Documentation](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html)"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}