Transfer-Learning-on-Dogs-vs-Cats.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pragma cling add_library_path(\"/Users/krshrimali/Downloads/libtorch/lib/\")\n",
    "#pragma cling add_include_path(\"/Users/krshrimali/Downloads/libtorch/include/\")\n",
    "#pragma cling add_include_path(\"/Users/krshrimali/Downloads/libtorch/include/torch/csrc/api/include/\")\n",
    "#pragma cling add_library_path(\"/usr/local/Cellar/opencv/4.1.0_2/lib\")\n",
    "#pragma cling add_include_path(\"/usr/local/Cellar/opencv/4.1.0_2/include/opencv4\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libiomp5.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libmklml.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libc10.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libtorch.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libcaffe2_detectron_ops.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libcaffe2_module_test_dynamic.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libcaffe2_observers.dylib\")\n",
    "#pragma cling load(\"/Users/krshrimali/Downloads/libtorch/lib/libshm.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_datasets.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_aruco.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_bgsegm.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_bioinspired.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_calib3d.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_ccalib.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_core.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_dnn_objdetect.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_dnn.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_dpm.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_face.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_features2d.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_flann.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_freetype.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_fuzzy.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_gapi.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_hfs.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_highgui.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_img_hash.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_imgcodecs.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_imgproc.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_line_descriptor.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_ml.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_objdetect.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_optflow.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_phase_unwrapping.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_photo.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_plot.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_quality.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_reg.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_rgbd.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_saliency.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_sfm.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_shape.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_stereo.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_stitching.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_structured_light.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_superres.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_surface_matching.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_text.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_tracking.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_video.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_videoio.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_videostab.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_xfeatures2d.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_ximgproc.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_xobjdetect.4.1.0.dylib\")\n",
    "#pragma cling load(\"/usr/local/Cellar/opencv/4.1.0_2/lib/libopencv_xphoto.4.1.0.dylib\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#include <torch/torch.h>\n",
    "#include <torch/script.h>\n",
    "#include <iostream>\n",
    "#include <dirent.h>\n",
    "#include <opencv2/opencv.hpp>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer Learning\n",
    "\n",
    "Before we go ahead and discuss the **Why** question of Transfer Learning, let's have a look at **What is Transfer Learning?** Let's have a look at the <a href=\"http://cs231n.github.io/transfer-learning/\">Notes</a> from CS231n on Transfer Learning:\n",
    "\n",
    "> In practice, very few people train an entire Convolutional Network from scratch (with random initialization), because it is relatively rare to have a dataset of sufficient size. Instead, it is common to pretrain a ConvNet on a very large dataset (e.g. ImageNet, which contains 1.2 million images with 1000 categories), and then use the ConvNet either as an initialization or a fixed feature extractor for the task of interest.\n",
    "\n",
    "There are 3 scenarios possible:\n",
    "\n",
    "1. When the data you have is similar (but not enough) to data trained on pre-trained model: Take an example of a pre-trained model trained on ImageNet dataset (containing 1000 classes). And the data we have has Dogs and Cats classes. Fortunate enough, ImageNet has some of the classes of Dog and Cat breeds and thus the model must have learned important features from the data. Let's say we don't have enough data but since the data is similar to the breeds in the ImageNet data set, we can simply use the ConvNet (except the last FC layer) to extract features from our dataset and train only the last Linear (FC) layer. We do this by the following code snippet in `Python`:\n",
    "\n",
    "```python\n",
    "from torchvision import models\n",
    "# Download and load the pre-trained model\n",
    "model = models.resnet18(pretrained=True)\n",
    "\n",
    "# Set upgrading the gradients to False\n",
    "for param in model.parameters():\n",
    "\tparam.requires_grad = False\n",
    "\n",
    "# Change the output features to the FC Layer and set it to upgrade gradients as True\n",
    "resnet18.fc = torch.nn.Linear(512, 2)\n",
    "for param in resnet18.fc.parameters():\n",
    "\tparam.requires_grad = True\n",
    "```\n",
    "\n",
    "2. When you have enough data (and is similar to the data trained with pre-trained model): Then you might go for fine tuning the weights of all the layers in the network. This is largely due to the reason that we know we won't overfit because we have enough data.\n",
    "3. Using pre-trained models might just be enough if you have the data which matches the classes in the original data set. \n",
    "\n",
    "Transfer Learning came into existence (the answer of **Why Transfer Learning?**) because of some major reasons, which include:\n",
    "\n",
    "1. Lack of resources or data set to train a CNN. At times, we either don't have enough data or we don't have enough resources to train a CNN from scratch.\n",
    "2. Random Initialization of weights vs Initialization of weights from the pre-trained model. Sometimes, it's just better to initialize weights from the pre-trained model (as it must have learned the generic features from it's data set) instead of randomly initializing the weights.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the data with PyTorch C++ API\n",
    "\n",
    "At every stage, we will compare the Python and C++ codes to do the same thing, to make the analogy easier and understandable. Starting with setting up the data we have. Note that we do have enough data and it is also similar to the original data set of ImageNet, but since I don't have enough resources to fine tune through the whole network, we perform Transfer Learning on the final FC layer only.\n",
    "\n",
    "Starting with loading the dataset, as discussed in the blogs before, I will just post a flow chart of procedure.\n",
    "\n",
    "<img src='images/Steps-Loading-Data-PyTorch.png'></img>\n",
    "\n",
    "Let's go ahead and define the required utility functions to define Custom Dataset class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility Function - 1: `read_data` and `read_label`\n",
    "\n",
    "**Documentation** of `read_data` function.\n",
    "\n",
    "```\n",
    "torch::Tensor read_data(std::string location)\n",
    "\n",
    "Function to return image read at location given as type torch::Tensor\n",
    "     Resizes image to (224, 224, 3)\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. location (std::string type) - required to load image from the location\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     torch::Tensor type - image read as tensor\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch::Tensor read_data(std::string location) {\n",
    "    /*\n",
    "     Function to return image read at location given as type torch::Tensor\n",
    "     Resizes image to (224, 224, 3)\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. location (std::string type) - required to load image from the location\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     torch::Tensor type - image read as tensor\n",
    "    */\n",
    "    cv::Mat img = cv::imread(location, 1);\n",
    "    cv::resize(img, img, cv::Size(224, 224), cv::INTER_CUBIC);\n",
    "    torch::Tensor img_tensor = torch::from_blob(img.data, {img.rows, img.cols, 3}, torch::kByte);\n",
    "    img_tensor = img_tensor.permute({2, 0, 1});\n",
    "    return img_tensor.clone();\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Documentation** of `read_label` function.\n",
    "\n",
    "```\n",
    "torch::Tensor read_label(int label)\n",
    "\n",
    "Function to return label from int (0, 1 for binary and 0, 1, ..., n-1 for n-class classification) as type torch::Tensor\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. label (int type) - required to convert int to tensor\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     torch::Tensor type - label read as tensor\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch::Tensor read_label(int label) {\n",
    "    /*\n",
    "     Function to return label from int (0, 1 for binary and 0, 1, ..., n-1 for n-class classification) as type torch::Tensor\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. label (int type) - required to convert int to tensor\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     torch::Tensor type - label read as tensor\n",
    "    */\n",
    "    torch::Tensor label_tensor = torch::full({1}, label);\n",
    "    return label_tensor.clone();\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility Function 2: `process_images` and `process_labels`\n",
    "\n",
    "**Documentation** of `process_images` function.\n",
    "\n",
    "```\n",
    "Function returns vector of tensors (images) read from the list of images in a folder\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. list_images (std::vector<std::string> type) - list of image paths in a folder to be read\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<torch::Tensor> type - Images read as tensors\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<torch::Tensor> process_images(std::vector<std::string> list_images) {\n",
    "    /*\n",
    "     Function returns vector of tensors (images) read from the list of images in a folder\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. list_images (std::vector<std::string> type) - list of image paths in a folder to be read\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<torch::Tensor> type - Images read as tensors\n",
    "     */\n",
    "    std::vector<torch::Tensor> states;\n",
    "    for(std::vector<std::string>::iterator it = list_images.begin(); it != list_images.end(); ++it) {\n",
    "        torch::Tensor img = read_data(*it);\n",
    "        states.push_back(img);\n",
    "    }\n",
    "    return states;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Documentation** of `process_labels` function.\n",
    "\n",
    "```\n",
    "Function returns vector of tensors (labels) read from the list of labels\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. list_labels (std::vector<int> list_labels) -\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<torch::Tensor> type - returns vector of tensors (labels)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<torch::Tensor> process_labels(std::vector<int> list_labels) {\n",
    "    /*\n",
    "     Function returns vector of tensors (labels) read from the list of labels\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. list_labels (std::vector<int> type) - required to convert int to tensor labels\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<torch::Tensor> type - returns vector of tensors (labels)\n",
    "     */\n",
    "    std::vector<torch::Tensor> labels;\n",
    "    for(std::vector<int>::iterator it = list_labels.begin(); it != list_labels.end(); ++it) {\n",
    "        torch::Tensor label = read_label(*it);\n",
    "        labels.push_back(label);\n",
    "    }\n",
    "    return labels;\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility Function 3: `load_images` and `load_labels`:\n",
    "\n",
    "**Documentation** of `load_images` function.\n",
    "\n",
    "```\n",
    "Function returns vector of strings (image paths) read from the folder name given\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. folder_name (std::string type) - name of folder containing images\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<std::string> type - returns vector of image paths\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<std::string> load_images(std::string folder_name) {\n",
    "    /*\n",
    "    Function returns vector of strings (image paths) read from the folder name given\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. folder_name (std::string type) - name of folder containing images\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<std::string> type - returns vector of image paths\n",
    "    */\n",
    "    std::vector<std::string> list_images;\n",
    "    \n",
    "    std::string base_name = folder_name;\n",
    "    \n",
    "    DIR* dir;\n",
    "    struct dirent *ent;\n",
    "    \n",
    "    if((dir = opendir(base_name.c_str())) != NULL) {\n",
    "        while((ent = readdir(dir)) != NULL) {\n",
    "            std::string filename = ent->d_name;\n",
    "            if(filename.length() > 4 && filename.substr(filename.length() - 3) == \"jpg\") {\n",
    "                std::string newf = base_name + filename;\n",
    "                list_images.push_back(newf);\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    return list_images;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Documentation** of `load_labels` function.\n",
    "\n",
    "```\n",
    "Function returns vector of int (labels) to each image in the folder (folder_name)\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. folder_name (std::string type) - name of folder containing images\n",
    "     2. label (int type) - label of the class (0 or 1 in case of binary, 0 1 ... n-1 in case of n-class classification)\n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<std::string> type - returns vector of labels assigned to each image of each class\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<int> load_labels(std::string folder_name, int label) {\n",
    "    /*\n",
    "    Function returns vector of int (labels) to each image in the folder (folder_name)\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. folder_name (std::string type) - name of folder containing images\n",
    "     2. label (int type) - label of the class (0 or 1 in case of binary, 0 1 ... n-1 in case of n-class classification)\n",
    "     Returns\n",
    "     ===========\n",
    "     std::vector<std::string> type - returns vector of labels assigned to each image of each class\n",
    "    */\n",
    "    std::vector<int> list_labels;\n",
    "    DIR* dir;\n",
    "    \n",
    "    std::string base_name = folder_name;\n",
    "    \n",
    "    struct dirent *ent;\n",
    "    \n",
    "    if((dir = opendir(base_name.c_str())) != NULL) {\n",
    "        while((ent = readdir(dir)) != NULL) {\n",
    "            std::string filename = ent->d_name;\n",
    "            if(filename.length() > 4 && filename.substr(filename.length() - 3) == \"jpg\") {\n",
    "                list_labels.push_back(label);\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    return list_labels;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are done with all the utility functions, we can go ahead and define the `CustomDataset` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset : public torch::data::Dataset<CustomDataset> {\n",
    "private:\n",
    "    /* data */\n",
    "    // Should be 2 tensors\n",
    "    std::vector<torch::Tensor> states, labels;\n",
    "    size_t ds_size;\n",
    "public:\n",
    "    CustomDataset(std::vector<std::string> list_images, std::vector<int> list_labels) {\n",
    "        states = process_images(list_images);\n",
    "        labels = process_labels(list_labels);\n",
    "        ds_size = states.size();\n",
    "    };\n",
    "    \n",
    "    torch::data::Example<> get(size_t index) override {\n",
    "        /* This should return {torch::Tensor, torch::Tensor} */\n",
    "        torch::Tensor sample_img = states.at(index);\n",
    "        torch::Tensor sample_label = labels.at(index);\n",
    "        return {sample_img.clone(), sample_label.clone()};\n",
    "    };\n",
    "    \n",
    "    torch::optional<size_t> size() const override {\n",
    "        return ds_size;\n",
    "    };\n",
    "};"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainining the FC Layer\n",
    "\n",
    "Let's first have a look at ResNet18 Network Architecture\n",
    "\n",
    "<img src=\"images/ResNet18-Architecture.png\">Reference: https://www.researchgate.net/figure/ResNet-18-Architecture_tbl1_322476121</img>\n",
    "\n",
    "The next step is to train the Fully Connected layer that we inserted at the end of the network (`linear_layer`). This one should be pretty straight forward, let's see how to do it.\n",
    "\n",
    "**Documentation** of `train()` function.\n",
    "\n",
    "```\n",
    "This function trains the network on our data loader using optimizer.\n",
    "     \n",
    "     Also saves the model as model.pt after every epoch.\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. net (torch::jit::script::Module type) - Pre-trained model without last FC layer\n",
    "     2. lin (torch::nn::Linear type) - last FC layer with revised out_features depending on the no. of classes\n",
    "     3. data_loader (DataLoader& type) - Training data loader\n",
    "     4. optimizer (torch::optim::Optimizer& type) - Optimizer like Adam, SGD etc.\n",
    "     5. size_t (dataset_size type) - Size of training dataset\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     Nothing (void)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "template<typename Dataloader>\n",
    "void train(torch::jit::script::Module net, torch::nn::Linear lin, Dataloader& data_loader, torch::optim::Optimizer& optimizer, size_t dataset_size) {\n",
    "    /*\n",
    "     This function trains the network on our data loader using optimizer.\n",
    "     \n",
    "     Also saves the model as model.pt after every epoch.\n",
    "     Parameters\n",
    "     ===========\n",
    "     1. net (torch::jit::script::Module type) - Pre-trained model without last FC layer\n",
    "     2. lin (torch::nn::Linear type) - last FC layer with revised out_features depending on the number of classes\n",
    "     3. data_loader (DataLoader& type) - Training data loader\n",
    "     4. optimizer (torch::optim::Optimizer& type) - Optimizer like Adam, SGD etc.\n",
    "     5. size_t (dataset_size type) - Size of training dataset\n",
    "     \n",
    "     Returns\n",
    "     ===========\n",
    "     Nothing (void)\n",
    "     */\n",
    "    \n",
    "    float batch_index = 0;\n",
    "    \n",
    "    for(int i=0; i<15; i++) {\n",
    "        float mse = 0;\n",
    "        float Acc = 0.0;\n",
    "        \n",
    "        for(auto& batch: *data_loader) {\n",
    "            auto data = batch.data;\n",
    "            auto target = batch.target.squeeze();\n",
    "            \n",
    "            // Should be of length: batch_size\n",
    "            data = data.to(torch::kF32);\n",
    "            target = target.to(torch::kInt64);\n",
    "            \n",
    "            std::vector<torch::jit::IValue> input;\n",
    "            input.push_back(data);\n",
    "            optimizer.zero_grad();\n",
    "            \n",
    "            auto output = net.forward(input).toTensor();\n",
    "            // For transfer learning\n",
    "            output = output.view({output.size(0), -1});\n",
    "            output = lin(output);\n",
    "            \n",
    "            auto loss = torch::nll_loss(torch::log_softmax(output, 1), target);\n",
    "            \n",
    "            loss.backward();\n",
    "            optimizer.step();\n",
    "            \n",
    "            auto acc = output.argmax(1).eq(target).sum();\n",
    "            \n",
    "            Acc += acc.template item<float>();\n",
    "            mse += loss.template item<float>();\n",
    "            \n",
    "            batch_index += 1;\n",
    "        }\n",
    "        \n",
    "        mse = mse/float(batch_index); // Take mean of loss\n",
    "        std::cout << \"Epoch: \" << i  << \", \" << \"Accuracy: \" << Acc/dataset_size << \", \" << \"MSE: \" << mse << std::endl;\n",
    "        net.save(\"model.pt\");\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Dataset for training\n",
    "\n",
    "Let's go ahead and load our dataset into `DataLoader` class. \n",
    "\n",
    "The distribution of the dataset is: `Cat Images`: 200 and `Dog Images`: 200."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "// Set folder names for cat and dog images\n",
    "std::string name_cats = \"/Users/krshrimali/Documents/krshrimali-blogs/dataset/train/cat_test/\";\n",
    "std::string name_dogs = \"/Users/krshrimali/Documents/krshrimali-blogs/dataset/train/dog_test/\";\n",
    "\n",
    "std::vector<std::string> images_cats = load_images(name_cats);\n",
    "std::vector<int> labels_cats = load_labels(name_cats, 0);\n",
    "\n",
    "std::vector<std::string> images_dogs = load_images(name_dogs);\n",
    "std::vector<int> labels_dogs = load_labels(name_dogs, 1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<std::string> images_total;\n",
    "\n",
    "for(auto const& value: images_cats) {\n",
    "    images_total.push_back(value);\n",
    "}\n",
    "\n",
    "for(auto const& value: images_dogs) {\n",
    "    images_total.push_back(value);\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "std::vector<int> labels_total;\n",
    "\n",
    "for(auto const& value: labels_cats) {\n",
    "    labels_total.push_back(value);\n",
    "}\n",
    "\n",
    "for(auto const& value: labels_dogs) {\n",
    "    labels_total.push_back(value);\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto custom_dataset = CustomDataset(images_total, labels_total).map(torch::data::transforms::Stack<>());\n",
    "auto data_loader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(custom_dataset), 4);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the pre-trained model\n",
    "\n",
    "The steps to load the pre-trained model and perform Transfer Learning are listed below:\n",
    "\n",
    "1. Download the pre-trained model of ResNet18.\n",
    "2. Load pre-trained model.\n",
    "3. Change output features of the final FC layer of the model loaded. (Number of classes would change from 1000 - ImageNet to 2 - Dogs vs Cats).\n",
    "4. Define optimizer on parameters from the final FC layer to be trained.\n",
    "5. Train the FC layer on Dogs vs Cats dataset\n",
    "5. Save the model (#TODO)\n",
    "\n",
    "Let's go step by step.\n",
    "\n",
    "**Step-1**: Download the pre-trained model of ResNet18\n",
    "\n",
    "Thanks to the developers, we do have C++ models available in torchvision (https://github.com/pytorch/vision/pull/728) but for this tutorial, transferring the pre-trained model from Python to C++ using `torch.jit` is a good idea, as most PyTorch models in the wild are written in Python right now, and people can use this tutorial to learn how to trace their Python model and transfer it to C++.) \n",
    "\n",
    "First we download the pre-trained model and save it in the form of `torch.jit.trace` format to our local drive. \n",
    "\n",
    "```python\n",
    "# Reference: #TODO- Add Link\n",
    "from torchvision import models\n",
    "# Download and load the pre-trained model\n",
    "model = models.resnet18(pretrained=True)\n",
    "\n",
    "# Set upgrading the gradients to False\n",
    "for param in model.parameters():\n",
    "\tparam.requires_grad = False\n",
    "\n",
    "# Save the model except the final FC Layer\n",
    "resnet18 = torch.nn.Sequential(*list(resnet18.children())[:-1])\n",
    "\n",
    "example_input = torch.rand(1, 3, 224, 224)\n",
    "script_module = torch.jit.trace(resnet18, example_input)\n",
    "script_module.save('resnet18_without_last_layer.pt')\n",
    "```\n",
    "\n",
    "We will be using `resnet18_without_last_layer.pt` model file as our pre-trained model for transfer learning. \n",
    "\n",
    "**Step-2**: Load the pre-trained model\n",
    "\n",
    "Let's go ahead and load the pre-trained model using `torch::jit` module. Note that the reason we have converted `torch.nn.Module` to `torch.jit.ScriptModule` type, is because C++ API currently does not support loading Python `torch.nn.Module` models directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch::jit::script::Module module;\n",
    "module = torch::jit::load(\"/Users/krshrimali/Documents/krshrimali-blogs/codes/transfer-learning/transfer-learning/build/resnet18_without_lastlayer.pt\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experimentation\n",
    "\n",
    "Since we are almost done with defining required functions, let's go ahead and define the optimizer on our last FC layer and train the FC layer on our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch::nn::Linear lin(512, 2); // the last layer of resnet, which we want to replace, has dimensions 512x1000\n",
    "torch::optim::Adam opt(lin->parameters(), torch::optim::AdamOptions(1e-3 /*learning rate*/));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Accuracy: 0.745, MSE: 0.491316\n",
      "Epoch: 1, Accuracy: 0.8825, MSE: 0.151102\n",
      "Epoch: 2, Accuracy: 0.845, MSE: 0.111356\n",
      "Epoch: 3, Accuracy: 0.8025, MSE: 0.106991\n",
      "Epoch: 4, Accuracy: 0.885, MSE: 0.0578496\n",
      "Epoch: 5, Accuracy: 0.865, MSE: 0.0572935\n",
      "Epoch: 6, Accuracy: 0.9, MSE: 0.0337119\n",
      "Epoch: 7, Accuracy: 0.855, MSE: 0.0399212\n",
      "Epoch: 8, Accuracy: 0.865, MSE: 0.0347004\n",
      "Epoch: 9, Accuracy: 0.8425, MSE: 0.0341781\n",
      "Epoch: 10, Accuracy: 0.8825, MSE: 0.0239106\n",
      "Epoch: 11, Accuracy: 0.86, MSE: 0.026371\n",
      "Epoch: 12, Accuracy: 0.86, MSE: 0.0269321\n",
      "Epoch: 13, Accuracy: 0.88, MSE: 0.0224421\n",
      "Epoch: 14, Accuracy: 0.8875, MSE: 0.0172629\n"
     ]
    }
   ],
   "source": [
    "train(module, lin, data_loader, opt, custom_dataset.size().value());"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "C++17",
   "language": "C++17",
   "name": "xcpp17"
  },
  "language_info": {
   "codemirror_mode": "text/x-c++src",
   "file_extension": ".cpp",
   "mimetype": "text/x-c++src",
   "name": "c++",
   "version": "-std=c++17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}