COMP5329 Assignment 2 - Code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "512OAXAvdyiA"
   },
   "source": [
    "# Initial Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iSOjenw7ebQN"
   },
   "source": [
    "## Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7witdSwueH4u"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import csv\n",
    "import string\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import statistics\n",
    "import itertools\n",
    "from PIL import Image\n",
    "import time\n",
    "from io import StringIO\n",
    "from skimage import io\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "import torchvision\n",
    "from torchvision import models\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from torchvision.transforms import ToTensor, Lambda, Resize, Compose, ToPILImage, Normalize, RandomCrop, RandomHorizontalFlip, RandomVerticalFlip\n",
    "import matplotlib.pyplot as plt\n",
    "import csv\n",
    "from sklearn.metrics import f1_score, confusion_matrix, precision_score, recall_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o4PwnFxReYis"
   },
   "source": [
    "## Drive/file/device setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iQ8VroMdeWzV"
   },
   "outputs": [],
   "source": [
    "#Set up google drive directory\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "\n",
    "#ZIP_PATH = #Comment path here\n",
    "\n",
    "ZIP_FILENAME = '2021s1comp5329assignment2.zip'\n",
    "\n",
    "!cp '{ZIP_PATH}' .\n",
    "!unzip -q '{ZIP_FILENAME}'\n",
    "!rm '{ZIP_FILENAME}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XTJ9ZhpAeeq0"
   },
   "outputs": [],
   "source": [
    "DIR = '/content/COMP5329S1A2Dataset'\n",
    "\n",
    "TRAIN_CSV = os.path.join(DIR, \"train.csv\")\n",
    "TEST_CSV = os.path.join(DIR, \"test.csv\")\n",
    "IMAGES_DIR = os.path.join(DIR, \"data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_wRCIqkEehDb"
   },
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using {device} device')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BZnhUTm0ejJU"
   },
   "source": [
    "# Data Read"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GoFQuMTvekP0"
   },
   "source": [
    "## AssignmentDataset Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-tB4F0dnem-V"
   },
   "outputs": [],
   "source": [
    "class AssignmentDataset(Dataset):\n",
    "    '''\n",
    "    The AssignmentDataset Class, child of torch.utils.data.Dataset.\n",
    "\n",
    "    Attributes:\n",
    "    csv_file (str): A string representation of the file directory to the csv data file.\n",
    "    image_dir (str): A string representation of the file directory to the images data file.\n",
    "    transform (torchvision.transforms.Compose): A torchvision.transforms.Compose object consisting of desired transforms to be made to the image data.\n",
    "    target_transform (function): A function to be passed to the label data for any desired transformations. Used to apply one-hot encoding.\n",
    "    has_labels (Bool): A boolean value to flag whether the dataset has labels or not, i.e. if it is training or testing.\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, csv_file, image_dir, transform=None, target_transform=None, has_labels = True):\n",
    "        self.image_dir = image_dir\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "        self.has_labels = has_labels\n",
    "        \n",
    "        with open(csv_file) as file:\n",
    "            lines = [re.sub(r'([^,])\"(\\s*[^\\n])', r'\\1/\"\\2', line) for line in file]\n",
    "            self.dataframe = pd.read_csv(StringIO(''.join(lines)), escapechar=\"/\")\n",
    "                \n",
    "\n",
    "    def __len__(self):\n",
    "        '''\n",
    "        Returns the length of the dataframe representation of the csv component of the dataset.\n",
    "\n",
    "        Parameters:\n",
    "        None\n",
    "\n",
    "        Returns:\n",
    "        self.dataframe.shape[0] (int): The length of the datasets dataframe representation.\n",
    "\n",
    "        '''\n",
    "        return self.dataframe.shape[0]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        '''\n",
    "        Reads the data from the dataframe then outputs the observation's image tensor, OHE label tensor (if there are labels present), ImageID string and caption string in a tuple form.\n",
    "                \n",
    "        Parameters:\n",
    "        idx (int): The index for which observation to return.\n",
    "\n",
    "        Returns:\n",
    "        sample (tuple): The tuple containing the observation's image tensor, OHE label tensor (if present), ImageID string and caption string.\n",
    "        '''\n",
    "        \n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "            \n",
    "        img_path = os.path.join(self.image_dir, \n",
    "                                self.dataframe.iloc[idx, self.dataframe.columns.get_loc('ImageID')])\n",
    "        \n",
    "        img = io.imread(img_path)\n",
    "\n",
    "        img_id = self.dataframe.iloc[idx, self.dataframe.columns.get_loc('ImageID')]\n",
    "        caption = self.dataframe.iloc[idx, self.dataframe.columns.get_loc('Caption')]\n",
    "\n",
    "        if self.has_labels:\n",
    "          labels = self.dataframe.iloc[idx, self.dataframe.columns.get_loc('Labels')]\n",
    "          labels = labels.split(' ') # converts the string into an iterable Python list\n",
    "          labels = [int(x) for x in labels] # convert the string of numbers into integer using Pytorch for computation speed\n",
    "                     \n",
    "          if self.target_transform:\n",
    "              labels = self.target_transform(labels)\n",
    "          if self.transform:\n",
    "              img = self.transform(img)\n",
    "\n",
    "          sample = (img, labels, img_id, caption)\n",
    "          \n",
    "        else:\n",
    "          if self.transform:\n",
    "              img = self.transform(img)          \n",
    "          sample = (img, img_id, caption)\n",
    "            \n",
    "        return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oTCwSZI9g7M8"
   },
   "source": [
    "## Finding mean/standard deviation for normalization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_GIWqdwmigsd"
   },
   "source": [
    "Define image_stats function to return means and standard deviations per channel in a dataloader."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TFhr2cz6klfH"
   },
   "source": [
    "## image_stats function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F6rn-NyThXgd"
   },
   "outputs": [],
   "source": [
    "def image_stats(checking_dataloader):\n",
    "  '''\n",
    "  A function to calculate the channel means and standard deviations of the image input in a training dataset.\n",
    "  Sums the channel means & squared means per image, then calculates the overall mean and standard deviation by taking the first moment (mean) and square root of the second moment (variance).\n",
    "\n",
    "  Parameters:\n",
    "  checking_dataloader (torch.utils.data.Dataloader): The Dataloader object with a batch size of 1 and containing images, image labels, image IDs and captions.\n",
    "\n",
    "  Returns:\n",
    "    means (torch.Tensor): A tensor object containing the channel means as floats. \n",
    "    stdevs (torch.Tensor): A tensor object containing the channel standard deviations as floats.\n",
    "  '''\n",
    "\n",
    "  sum_channels, sumsq_channels, n_batches = 0, 0, 0\n",
    "\n",
    "  for step, (x, _, _, _ ) in enumerate(checking_dataloader):\n",
    "    sum_channels += torch.mean(x, dim = [0, 2, 3])\n",
    "    sumsq_channels += torch.mean(x**2 , dim = [0, 2, 3])\n",
    "    n_batches += 1\n",
    "\n",
    "  means = sum_channels/n_batches\n",
    "  stdevs = (sumsq_channels/n_batches - means**2)**0.5\n",
    "\n",
    "  return means, stdevs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ky77ggwqimc2"
   },
   "source": [
    "##Checking dataset/loader\n",
    "\n",
    "<br> Set up a \"checking\" dataset of the data to use in the image_stats function - note the batch size of 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hLF-va9Cg6Ca"
   },
   "outputs": [],
   "source": [
    "##Finding Mean/std\n",
    "NUM_LABELS = 19\n",
    "BATCH_SIZE = 1\n",
    "\n",
    "transforms = Compose([\n",
    "    ToTensor(),\n",
    "])\n",
    "\n",
    "train_dataset_check = AssignmentDataset(csv_file = TRAIN_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                 transform = transforms,\n",
    "                                 target_transform = Lambda(lambda y: torch.zeros(NUM_LABELS, dtype=torch.uint8).scatter_(dim=0, index=torch.sub(torch.tensor(y), 1), value=1)),\n",
    "                                  has_labels = True,\n",
    "                                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zyaXyaDGkO7S"
   },
   "outputs": [],
   "source": [
    "train_dataloader_check = DataLoader(train_dataset_check, batch_size=BATCH_SIZE, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4Sr6XhsRiuEd"
   },
   "source": [
    "## Calculate & print mean/standard deviation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lYXXftXpi8Bs"
   },
   "outputs": [],
   "source": [
    "means, stdevs = image_stats(train_dataloader_check)\n",
    "print(f'Mean is: {means} \\n Standard Deviation is: {stdevs}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sJIhhyLHjM6t"
   },
   "source": [
    "## Sanity check image sizes - create get_image_size function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q9C4HABhito7"
   },
   "outputs": [],
   "source": [
    "def get_image_sizes(checking_dataloader):\n",
    "    '''\n",
    "    Obtains the heights and widths of the images within the supplied dataset.\n",
    "\n",
    "    Parameters:\n",
    "    checking_dataloader (torch.utils.data.Dataloader): The Dataloader object with a batch size of 1 and containing images, image labels, image IDs and captions.\n",
    "\n",
    "    Returns:\n",
    "      heights (list(int)): A list of the image heights as integers.\n",
    "      widths (list(int)): A list of the image widths as integers.\n",
    "    '''\n",
    "    \n",
    "    heights, widths = [], []\n",
    "\n",
    "    for step, (img, _, _, _) in enumerate(checking_dataloader):\n",
    "        heights.append(img.size()[2])\n",
    "        widths.append(img.size()[3])\n",
    "\n",
    "    return heights, widths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ndY0IZJgkUlt"
   },
   "outputs": [],
   "source": [
    "heights, widths = get_image_sizes(train_dataloader_check)\n",
    "\n",
    "print(f'Median height is: {statistics.median(heights)}')\n",
    "print(f'Mean height is: {statistics.mean(heights)}')\n",
    "print(f'Median width is: {statistics.median(widths)}')\n",
    "print(f'Mean width is: {statistics.mean(widths)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v-Pn1aWPlXAL"
   },
   "source": [
    "## Creating Train/Test/Validation Datasets and Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kYWYkFBWkge_"
   },
   "outputs": [],
   "source": [
    "# Concluded 19 labels - one-hot encoded. Leaving label 12 in there as it'll just have zeros for this index\n",
    "NUM_LABELS = 19\n",
    "BATCH_SIZE = 50\n",
    "RESIZE_SIZE = (299, 299)\n",
    "CROP_SIZE = 224 # Seems to be the standard image size used (224)\n",
    "TRAIN_VAL_PROP = 0.75\n",
    "SEED = 5329\n",
    "\n",
    "transforms = Compose([\n",
    "    ToTensor(),\n",
    "    Normalize(mean = means, std = stdevs),\n",
    "    Resize(RESIZE_SIZE),\n",
    "    RandomHorizontalFlip(),\n",
    "    RandomVerticalFlip(),\n",
    "    RandomCrop(CROP_SIZE)\n",
    "])\n",
    "\n",
    "# The main dataset available with labels\n",
    "main_dataset = AssignmentDataset(csv_file = TRAIN_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                 transform = transforms,\n",
    "                                 target_transform = Lambda(lambda y: torch.zeros(NUM_LABELS, dtype=torch.uint8).scatter_(dim=0, index=torch.sub(torch.tensor(y), 1), value=1)),\n",
    "                                 has_labels = True\n",
    "                                 )\n",
    "\n",
    "# Need to further split it up\n",
    "train_dataset, val_dataset = random_split(main_dataset, \n",
    "                                          [int(round(TRAIN_VAL_PROP * len(main_dataset))), int(round((1 - TRAIN_VAL_PROP) * len(main_dataset)))], \n",
    "                                          generator=torch.Generator().manual_seed(SEED)) # Setting seed to ensure consistency\n",
    "\n",
    "# The final test dataset with no labels\n",
    "test_dataset = AssignmentDataset(csv_file = TEST_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                transform = transforms,\n",
    "                                 has_labels = False)\n",
    "\n",
    "main_dataloader = DataLoader(main_dataset, batch_size = BATCH_SIZE, shuffle=True)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R2TEdYtMmRBx"
   },
   "source": [
    "# Miscellaneous Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DJaKVxFvyWrg"
   },
   "source": [
    "## Encoding/Decoding label functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cheI1vtHyWTg"
   },
   "outputs": [],
   "source": [
    "def encode_outputs(output, threshold):\n",
    "  ''' \n",
    "  A function to transform sigmoid outputs to one-hot encoded representations, based on a threshold value which deems the sigmoid output to be positive.\n",
    "  Creates a numpy array filled with zeros per observation, transforms this to a list and appends an index with 1s if the sigmoid output of the index is greater than the threshold then appends this observation to the encoded_outputs list containing all observations.\n",
    "  \n",
    "  Parameters:\n",
    "  output (torch.Tensor): A torch tensor of model output/predictions, calculated on a dataset with a batch size of 1.\n",
    "  threshold (float): The desired probability threshold to consider a sigmoid output as a positive classification.\n",
    "\n",
    "  Returns:\n",
    "  encoded_outputs (list(list(int))): A list per observation containing a list of one-hot encoded predictions as integers with length 19.\n",
    "  '''\n",
    "\n",
    "  encoded_outputs = []\n",
    "  for i in range(len(output)):\n",
    "    encoded_output = np.zeros(19).tolist()\n",
    "    for j in range(len(output[i][0])):\n",
    "      encoded_output[j] = 1 if output[i][0][j].item() >= threshold else 0\n",
    "    encoded_outputs.append(encoded_output)\n",
    "  return encoded_outputs\n",
    "\n",
    "def decode_labels(labels, threshold):\n",
    "  '''\n",
    "  A function to transform sigmoid outputs or OHE predictions to a list of numeric labels.\n",
    "  Creates an empty list for an individual observation, appends with the value of the index + 1 which has the prediction greater than the threshold and aggregates in the decoded_labels list for all observations.\n",
    "\n",
    "  Parameters:\n",
    "  labels (list(torch.Tensor)): A list per observation of the one-hot encoded labels in a torch.Tensor.\n",
    "  threshold (float): The desired probability threshold to consider a sigmoid output as a positive classification.\n",
    "\n",
    "  Returns:\n",
    "  decoded_labels (list(list(str))): A list for each observation containing a list of the labels as strings.\n",
    "  '''\n",
    "\n",
    "  decoded_labels = []\n",
    "  for i in range(len(labels)):\n",
    "    decoded_label = []\n",
    "    for j in range(len(labels[i][0])):  \n",
    "      if labels[i][0][j].item() >= threshold:\n",
    "        decoded_label.append(j + 1)\n",
    "    decoded_labels.append(decoded_label)\n",
    "  return decoded_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o_mUa8z2n6_t"
   },
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tk1tapy4oReF"
   },
   "source": [
    "show img function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o6Eu2VO5oPrw"
   },
   "outputs": [],
   "source": [
    "def show_img(dataset):\n",
    "  '''\n",
    "  A function to show a 2 x 3 plot of 6 random images and display the one-hot encoded labels above it.\n",
    "\n",
    "  Parameters:\n",
    "  dataset (AssignmentDataset): The dataset containing the desired set of images for which samples are to be plotted.\n",
    "\n",
    "  Returns:\n",
    "  None \n",
    "  '''\n",
    "\n",
    "  loader = DataLoader(train_dataset, batch_size = 6, shuffle = True)\n",
    "  batch = next(iter(loader))\n",
    "  images, labels,_,_ = batch\n",
    "\n",
    "  grid = torchvision.utils.make_grid(images, nrow = 3)\n",
    "  plt.figure(figsize = (12, 12))\n",
    "  plt.imshow(np.transpose(grid, (1,2,0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h-_n8H2eohpq"
   },
   "outputs": [],
   "source": [
    "show_img(train_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "75OgEj01rYZA"
   },
   "source": [
    "## Weights setup - not used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kbs8wtvhskp8"
   },
   "outputs": [],
   "source": [
    "label_count = np.zeros(19, dtype = int)\n",
    "\n",
    "for label in main_dataset.dataframe.Labels:\n",
    "  individual_label = label.split(' ')  \n",
    "  individual_label = [int(i) for i in individual_label]\n",
    "  for j in individual_label:\n",
    "    label_count[j-1] += 1\n",
    "    \n",
    "print(label_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_OgGmDUutIwT"
   },
   "outputs": [],
   "source": [
    "unique_labels = np.zeros(19).tolist()\n",
    "for i in range(0, 19):\n",
    "  unique_labels[i] = i+1\n",
    "unique_label_occurence_df = pd.DataFrame(data = {\"Count\": label_count, \"Labels\": unique_labels})\n",
    "unique_label_occurence_df['Weight'] = [(len(main_dataset)/label_count[i] - 1) if label_count[i] != 0 else 0 for i in range(len(label_count))]\n",
    "\n",
    "unique_label_occurence_df = unique_label_occurence_df.sort_values(by = ['Count'], ascending = False, ignore_index = True)\n",
    "label_weights = torch.tensor(unique_label_occurence_df['Weight']).to(device)\n",
    "\n",
    "unique_label_occurence_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v-n43Dkettbv"
   },
   "source": [
    "## Training Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "scZFF6fYnymV"
   },
   "source": [
    "### Training loop funciton"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YLO5JAjJlsbs"
   },
   "outputs": [],
   "source": [
    "def train_loop(dataloader, val_dataloader, model, loss_fn, optimizer, test_threshold=None, with_captions=False):\n",
    "    '''\n",
    "    The function to run a training loop on a selected model and returns training losses per batch as well as validation loss & validation f1 if applicable.\n",
    "\n",
    "    Sets the model to train model, iterates through the dataloader and passes the data forward through the network then calculates loss. Removes accumulated gradients in the optimizer and runs the backpropagation process, finally updating the parameters.\n",
    "    Prints training updates as well as validation scores based on the test_threshold value.\n",
    "    \n",
    "    Parameters:\n",
    "    dataloader (torch.utils.data.Dataloader): The dataloader object for the training data.\n",
    "    val_dataloader(torch.utils.data.Dataloader): The dataloader object for the validation data.\n",
    "    Model (child of torch.nn.Module): The chosen Model to train.\n",
    "    loss_fn: The chosen loss function for learning & evaluation.\n",
    "    optimizer: The chosen optimizer for parameter updates.\n",
    "    test_threshold (float or None): The threshold on which to convert sigmoid outputs to positive labels in validation, or None to ensure validation is not conducted in training. Set to None by default.\n",
    "    with_captions (bool): The boolean value indicating whether image captions are to incorporated into the model.\n",
    "    \n",
    "    Returns:\n",
    "      train_batch_losses (list): A list of the loss values per training batch.\n",
    "      val_batch_losses (list): A list of the loss values per validation epoch.\n",
    "      val_f1_scores (list): A list of the mean F1 scores per validation epoch.\n",
    "    '''\n",
    "\n",
    "    train_batch_losses = []\n",
    "    val_batch_losses = []\n",
    "    val_f1_scores = []\n",
    "    \n",
    "    size = len(dataloader.dataset)\n",
    "    # Set the model to training mode\n",
    "    model.train()\n",
    "\n",
    "    for batch, (X, y, _, captions) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        if with_captions:\n",
    "            # Perform encoding for the captions\n",
    "            word_encoding_array = encode_captions(captions)\n",
    "            # Compute prediction and loss\n",
    "            pred = model(X, word_encoding_array)\n",
    "        else:\n",
    "            pred = model(X)\n",
    "        \n",
    "        if type(model) == torchvision.models.inception.Inception3:\n",
    "          loss = loss_fn(pred.logits, y.type(torch.float))\n",
    "        else:\n",
    "          loss = loss_fn(pred, y.type(torch.float))\n",
    "        \n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "        if batch % EVAL_EVERY == 0:\n",
    "            train_batch_losses.append(loss.item())\n",
    "            train_loss, current = loss.item(), batch * len(X)\n",
    "            print(f'Current observation is {current}')\n",
    "\n",
    "            # setting test_loss to dummy value \n",
    "            test_loss = 0\n",
    "            test_f1 = 0\n",
    "\n",
    "            if test_threshold is not None:\n",
    "                sig = nn.Sigmoid()\n",
    "                pred_list = []\n",
    "                true_list = []\n",
    "                # Switch to eval mode for evaluating test set\n",
    "                model.eval()\n",
    "                with torch.no_grad():\n",
    "                    for X, y, _, captions in val_dataloader:\n",
    "                        X, y = X.to(device), y.to(device)\n",
    "                        true_list.append(y.tolist()[0])\n",
    "\n",
    "                        if with_captions:\n",
    "                            word_encoding_array = encode_captions(captions)\n",
    "                            pred = model.forward(X, word_encoding_array)\n",
    "                        else:\n",
    "                            pred = model.forward(X)\n",
    "\n",
    "                        sig_pred = sig(pred.cpu())\n",
    "                        pred_list.append(sig_pred)\n",
    "                        test_loss = loss_fn(pred, y.type(torch.float))\n",
    "\n",
    "                encoded_val_outputs = encode_outputs(pred_list, test_threshold)\n",
    "\n",
    "                test_f1 = f1_score(y_true = true_list, y_pred = encoded_val_outputs, average = \"weighted\", zero_division = 0)\n",
    "                \n",
    "                val_batch_losses.append(test_loss.item())\n",
    "                val_f1_scores.append(test_f1)\n",
    "\n",
    "                # Switch back to train mode to resume training\n",
    "                model.train()\n",
    "\n",
    "            if test_threshold is None:\n",
    "              print(f\"train loss: {train_loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "            else:\n",
    "              print(f\"train loss: {train_loss:>7f}  [{current:>5d}/{size:>5d}], validation loss: {test_loss:>7f}, validation mean f1: {test_f1:>7f}\")\n",
    "    \n",
    "    return train_batch_losses, val_batch_losses, val_f1_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Oy-YC2osn0k_"
   },
   "source": [
    "### Output to submission function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7MCkBRVLmKdk"
   },
   "outputs": [],
   "source": [
    "def output_to_submission(test_dataloader, model, threshold, with_captions=False):\n",
    "    '''\n",
    "    The function to run a model on testing data and provide a dataframe as well as a space delimited list for submission.\n",
    "\n",
    "    First sets up relevant empty lists and headings, then feeds the test data through to model and converts the output to decoded predictions, finally organizing these decoded preidctions in the dataframe & list with image IDs.\n",
    "\n",
    "    Parameters:\n",
    "    test_dataloader (torch.utils.data.Dataloader): The dataloader containing the testing data.\n",
    "    model (child of torch.nn.Module): The model to feed the training data through.\n",
    "    threshold (float)\n",
    "\n",
    "    Returns:\n",
    "    text_output(list(list(str))): The space delimited list of test data ImageIDs and respective model predictions.\n",
    "    dataframe_output (pandas DataFrame): The Dataframe with the test data ImageIDs and respective model predictions.\n",
    "    '''\n",
    "\n",
    "    sig = torch.nn.Sigmoid()\n",
    "\n",
    "    predictions = []\n",
    "    text_predictions = []\n",
    "    img_ids = test_dataloader.dataset.dataframe['ImageID']\n",
    "\n",
    "    headings = ['ImageID', 'Labels']\n",
    "    all_labels = [headings]\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "      for X, _, captions in test_dataloader:\n",
    "        X = X.to(device)\n",
    "\n",
    "        if with_captions:\n",
    "            word_encoding_array = encode_captions(captions)\n",
    "            output = model.forward(X, word_encoding_array).cpu()\n",
    "        else:\n",
    "            output = model.forward(X).cpu()\n",
    "            \n",
    "        output = sig(output)\n",
    "        predictions.append(output)\n",
    "\n",
    "    decoded_predictions = decode_labels(predictions, threshold)\n",
    "    for i in range(len(decoded_predictions)):\n",
    "      single_text = [img_ids[i], \" \".join(map(str, decoded_predictions[i]))]\n",
    "      text_predictions.append(\" \".join(map(str, decoded_predictions[i])))\n",
    "      \n",
    "      all_labels.append(single_text)\n",
    "\n",
    "    text_output = all_labels\n",
    "    dataframe_output = pd.DataFrame({\"ImageID\": test_dataloader.dataset.dataframe['ImageID'], \"Labels\": text_predictions})\n",
    "\n",
    "    return text_output, dataframe_output            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oidAXBAHoGVq"
   },
   "source": [
    "### Caption Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X6nFZqduoA9F"
   },
   "outputs": [],
   "source": [
    "# Generate mapping dictionary between alphabet combination to Numpy array index\n",
    "alphabet_combinations = {'1': string.ascii_lowercase, '2':string.ascii_lowercase, '3': string.ascii_lowercase}\n",
    "alphabet_mapping = {}\n",
    "array_index = 0\n",
    "\n",
    "for combo in itertools.product(*[alphabet_combinations[k] for k in sorted(alphabet_combinations.keys())]):\n",
    "    letter_combinations = ''.join(combo)\n",
    "    alphabet_mapping[letter_combinations] = array_index\n",
    "    array_index += 1\n",
    "\n",
    "\n",
    "def encode_captions(captions):\n",
    "    '''\n",
    "    Transforms a list/tuple of N captions, into a N x NUM_LABELS size Torch tensor.\n",
    "\n",
    "    Parameters:\n",
    "    captions (list or tuple): A list or data structure containing strings of captions.\n",
    "\n",
    "    Returns:\n",
    "    word_encoding_array (torch.Tensor): A torch.Tensor containing 3-letter encoded captions for each image.\n",
    "    '''\n",
    "\n",
    "    # Create an empty Numpy 2D array of size (N x NUM_LABELS)\n",
    "    word_encoding_array = np.zeros((len(captions), len(alphabet_mapping)))\n",
    "\n",
    "    # Iterate over the captions\n",
    "    for caption_num, caption in enumerate(captions):\n",
    "\n",
    "        # Extract the alphabet characters\n",
    "        split_text = [x for x in caption.lower() if x in string.ascii_lowercase]\n",
    "\n",
    "        for idx in range(len(split_text)):\n",
    "            if idx+3 > len(split_text):\n",
    "                break\n",
    "\n",
    "            combination = (''.join(split_text[idx:idx+3]))\n",
    "            # Perform a lookup on which index it should do a +1\n",
    "            array_idx = alphabet_mapping[combination]\n",
    "            # Increase the combination count for the array\n",
    "            word_encoding_array[caption_num][array_idx] += 1\n",
    "\n",
    "    # Convert numpy array to Torch tensor\n",
    "    word_encoding_array = torch.from_numpy(word_encoding_array).to(device)\n",
    "\n",
    "    return word_encoding_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LR-4b_jdvs7_"
   },
   "source": [
    "# Basic CNN from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mk2lQfU47MUO"
   },
   "source": [
    "## BasicCNN Class Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TIEeybQBwOoQ"
   },
   "outputs": [],
   "source": [
    "class BasicCNN(nn.Module):\n",
    "  '''\n",
    "  The BasicCNN class for a basic Convolutional Neural Network, child of torch.nn.Module.\n",
    "  Contains 3 convolutional layers, 3 fully connected layers & 3 instances of dropout.\n",
    "  '''\n",
    "\n",
    "  def __init__(self):\n",
    "    super(BasicCNN, self).__init__()\n",
    "\n",
    "    self.conv1 = nn.Conv2d(3, 32, 10, stride = 1, padding = 1)\n",
    "    self.conv2 = nn.Conv2d(32, 64, 5, stride = 2, padding = 1)\n",
    "    self.conv3 = nn.Conv2d(64, 128, 3, stride = 1, padding = 1)\n",
    "\n",
    "    self.fc1 = nn.Linear(21632, 128)\n",
    "    self.fc2 = nn.Linear(128, 256)\n",
    "    self.fc3 = nn.Linear(256, 19)\n",
    "\n",
    "    self.dropout1 = nn.Dropout2d(0.25)\n",
    "    self.dropout2 = nn.Dropout2d(1/3)\n",
    "    self.dropout3 = nn.Dropout2d(0.5)\n",
    "\n",
    "  def forward(self, x):\n",
    "    '''\n",
    "    The forward pass of the Basic CNN.\n",
    "    Takes the input and passes it sequentially through the convolutional layers, leaky relu fucntions, max pooling layers then applies dropout to pass to the fully connected layer.\n",
    "\n",
    "    Parameters:\n",
    "    x (torch.Tensor): A torch Tensor object containing the input.\n",
    "\n",
    "    Returns:\n",
    "    logits (torch.Tensor): A torch Tensor object containing logit representations of the output.\n",
    "    '''\n",
    "\n",
    "    x = self.conv1(x)\n",
    "    x = F.leaky_relu(x)\n",
    "    x = F.max_pool2d(x, 2)\n",
    "    \n",
    "    x = self.conv2(x)\n",
    "    x = F.leaky_relu(x)\n",
    "    x = F.max_pool2d(x, 2)\n",
    "    \n",
    "    x = self.conv3(x)\n",
    "    x = F.leaky_relu(x)\n",
    "    x = F.max_pool2d(x, 2)\n",
    "\n",
    "    x = self.dropout1(x)\n",
    "\n",
    "    x = torch.flatten(x, 1)\n",
    "    x = self.fc1(x)\n",
    "    x = F.relu(x)\n",
    "\n",
    "    x = self.dropout2(x)\n",
    "    x = self.fc2(x)\n",
    "    x = F.relu(x)\n",
    "\n",
    "    x = self.dropout3(x)\n",
    "    x = self.fc3(x)\n",
    "    logits = x\n",
    "    \n",
    "    return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZMLSPaBjwTSn"
   },
   "source": [
    "## BasicCNN Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GFSKLrXtrGcR"
   },
   "source": [
    "### Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SHwugF8owRwg"
   },
   "outputs": [],
   "source": [
    "BasicCNN_model = BasicCNN().to(device)\n",
    "print(\"Model Initialized\")\n",
    "print(BasicCNN_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hPHCdDLBx_3O"
   },
   "source": [
    "### Training BasicCNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sd8aPQInwe6u"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(BasicCNN_model.parameters(), lr=LR)\n",
    "\n",
    "BasicCNN_train_loss = []\n",
    "BasicCNN_val_loss = []\n",
    "BasicCNN_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, BasicCNN_model, loss_fn, optimizer, test_threshold=0.5, with_captions=False)\n",
    "    BasicCNN_train_loss += train_loss\n",
    "    BasicCNN_val_loss += val_loss\n",
    "    BasicCNN_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tm2HTrYaIjJ3"
   },
   "source": [
    "# CNN2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wff9NV_2IrNM"
   },
   "source": [
    "## CNN2 Class Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0TOJWc6AIxdc"
   },
   "outputs": [],
   "source": [
    "class CNN2(nn.Module):\n",
    "    '''\n",
    "    The (more advanced) CNN2 Class, child of torch.nn.Module, that incorporates the ability to return feature maps as well as utilise captions.\n",
    "    Contains 4 convolutional layers, 2 fully connected layers & 1 instances of dropout.\n",
    "\n",
    "    Attributes:\n",
    "    return_fmals (bool): The boolean value to dictate whether feature maps are returned or not. Set to False by default.\n",
    "    with_captions (bool): The boolean value to dictate whether the model is to incorporate captions when training. Set to True by default.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, return_fmaps=False, with_captions=False):\n",
    "        super(CNN2, self).__init__()\n",
    "        # things to note, we use Conv2d because the convolution slides in 2 dimension, not 3\n",
    "        # also the outer channel represents the number of filters, since 1 filter flattens a 3d image\n",
    "        self.conv1 = nn.Conv2d(3, 64, 3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv2d(64, 128, 3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(128, 128, 3, stride=1, padding=1)\n",
    "        self.conv4 = nn.Conv2d(128, 128, 3, stride=1, padding=1)\n",
    "        self.dropout1 = nn.Dropout(0.5)\n",
    "        if with_captions:\n",
    "            self.fc1 = nn.Linear((128*14*14) + (26**3), 1024)\n",
    "        else:\n",
    "            self.fc1 = nn.Linear((128*14*14), 1024)\n",
    "        self.fc2 = nn.Linear(1024, NUM_LABELS)\n",
    "        #self.fc3 = nn.Linear(1024, NUM_LABELS)\n",
    "\n",
    "        self.return_fmaps = return_fmaps\n",
    "\n",
    "    \n",
    "    def set_return_fmaps(self, val=True):\n",
    "        '''\n",
    "        The setter function to toggle returning feature maps on/off.\n",
    "\n",
    "        Parameters:\n",
    "        val (bool): The boolean value for which to set the return_fmaps attribute to. Set to True by default.\n",
    "\n",
    "        Returns:\n",
    "        None\n",
    "        '''\n",
    "\n",
    "        self.return_fmaps = val\n",
    "\n",
    "\n",
    "    def forward(self, x, caption = None):\n",
    "        '''\n",
    "        The forward pass of the CNN2.\n",
    "        Takes the input and passes it sequentially through the convolutional layers, relu fucntions, max pooling layers and dropout (if applicable) to pass to the fully connected layer.\n",
    "\n",
    "        Parameters:\n",
    "        x (torch.Tensor): A torch Tensor object containing the input.\n",
    "        caption (tuple): A tuple containing the captions in the data.\n",
    "\n",
    "        Returns:\n",
    "        x (torch.Tensor): A torch Tensor object containing logit representations of the output.\n",
    "        fmaps (list): A list containing the feature maps values. Returned only if self.return_fmaps == True.\n",
    "        '''\n",
    "\n",
    "        fmaps = []\n",
    "\n",
    "        x = self.conv1(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.conv2(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        x = self.conv3(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.conv4(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        \n",
    "        if caption == None:\n",
    "            x = self.fc1(x)\n",
    "        else:\n",
    "            # Mixing in the caption data here, concat in dimension 1\n",
    "            x = self.fc1(torch.cat((x, caption.type(torch.float)), 1)) \n",
    "\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout1(x)\n",
    "        x = self.fc2(x)\n",
    "\n",
    "        if self.return_fmaps:\n",
    "            return x, fmaps\n",
    "        else:\n",
    "            return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1CAp77jvsbTm"
   },
   "source": [
    "## CNN2 Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PTXGTuQfI3MB"
   },
   "outputs": [],
   "source": [
    "CNN2_model = CNN2(with_captions=True).to(device)\n",
    "print('Model initialised')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4uP4kO69sgAf"
   },
   "source": [
    "## CNN2 Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bvuE2JZFI9MB"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(CNN2_model.parameters(), lr=LR)\n",
    "\n",
    "CNN2_train_loss = []\n",
    "CNN2_val_loss = []\n",
    "CNN2_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, CNN2_model, loss_fn, optimizer, test_threshold=0.5, with_captions=True)\n",
    "    CNN2_train_loss += train_loss\n",
    "    CNN2_val_loss += val_loss\n",
    "    CNN2_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VhYi9rqBydMB"
   },
   "source": [
    "# AlexNet from Scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qV2C5-F17dPp"
   },
   "source": [
    "### AlexNetScratch Class Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2tvdUBjw3X-o"
   },
   "outputs": [],
   "source": [
    "class AlexNetScratch(nn.Module):\n",
    "  '''\n",
    "  The AlexNetScratch Class, child of torch.nn.Module, that is a CNN that follows the AlexNet Structure.\n",
    "  Contains 5 convolutional layers, 2 fully connected layers & 2 instances of dropout. Convolutional & classifier layers are fed through torch.nn.Sequential.\n",
    "\n",
    "  Attributes:\n",
    "  num_classes (int): The number of classes in the data which is set as the number of output layers in the \n",
    "  stem_stride (int): The stride value used in the first convolutional layer.\n",
    "  '''\n",
    "\n",
    "  def __init__(self, num_classes = 19, stem_stride = 4):\n",
    "    super(AlexNetScratch, self).__init__()\n",
    "\n",
    "    self.features = nn.Sequential(\n",
    "        nn.Conv2d(3, 64, kernel_size = 11, stride = stem_stride, padding = 2),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "\n",
    "        nn.Conv2d(64, 192, kernel_size = 5, padding = 2),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "\n",
    "        nn.Conv2d(192, 384, kernel_size = 3, stride = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Conv2d(384, 256, kernel_size = 2, padding = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Conv2d(256, 256, kernel_size = 3, padding = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "    )\n",
    "\n",
    "    self.avgpool = nn.AdaptiveAvgPool2d((6, 6))\n",
    "\n",
    "    self.classifier = nn.Sequential(\n",
    "        nn.Dropout(),\n",
    "        nn.Linear(256 * 6 * 6, 4096),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Dropout(),\n",
    "        nn.Linear(4096, 4096),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Linear(4096, num_classes),\n",
    "    )\n",
    "\n",
    "  def forward(self, x):\n",
    "    '''\n",
    "    The forward pass of the AlexNetScratch model.\n",
    "    Takes the input and passes it sequentially through the Sequential feature layers, applies average pooling & then passes through to the classifier layers.\n",
    "\n",
    "    Parameters:\n",
    "    x (torch.Tensor): A torch Tensor object containing the input.\n",
    "\n",
    "    Returns:\n",
    "    x (torch.tensor): A torch Tensor containing the output of the network feedforward process.\n",
    "    '''\n",
    "    \n",
    "    x = self.features(x)\n",
    "    x = self.avgpool(x)\n",
    "    x = torch.flatten(x, 1)\n",
    "    x = self.classifier(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I-cp5V_p7kx7"
   },
   "source": [
    "### AlexNet Scratch Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V-wyFoHI4tdN"
   },
   "outputs": [],
   "source": [
    "#Create network\n",
    "ANS_model = AlexNetScratch(num_classes = 19, stem_stride = 2).to(device)\n",
    "print(\"AlexNet from Scratch Initialized\")\n",
    "print(ANS_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sXAk7mRq7vPo"
   },
   "source": [
    "### AlexNet Scratch Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_5nDWiT4uTN"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(ANS_model.parameters(), lr=LR)\n",
    "\n",
    "ANS_train_loss = []\n",
    "ANS_val_loss = []\n",
    "ANS_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, ANS_model, loss_fn, optimizer, test_threshold=0.5, with_captions=True)\n",
    "    ANS_train_loss += train_loss\n",
    "    ANS_val_loss += val_loss\n",
    "    ANS_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3VDhBYvI5db1"
   },
   "source": [
    "# Pre-trained AlexNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iXc2eqky8ITr"
   },
   "source": [
    "## Pre-trained AlexNet Model Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g0BxQCTO8J3q"
   },
   "outputs": [],
   "source": [
    "from torchvision.models import alexnet\n",
    "ANP_model = alexnet(pretrained = True)\n",
    "print(ANP_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GJbG8iY-8fEj"
   },
   "source": [
    "## Override the Pre-trained AlexNet FC/classification layer and view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1DHA3H8e8fUr"
   },
   "outputs": [],
   "source": [
    "ANP_model.classifier[6] = nn.Linear(4096, 19)\n",
    "print(ANP_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L2BKZVB88nPz"
   },
   "source": [
    "## Pre-trained AlexNet training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RT_-Xkc48u48"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(ANP_model.parameters(), lr=LR)\n",
    "\n",
    "ANP_train_loss = []\n",
    "ANP_val_loss = []\n",
    "ANP_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, ANP_model, loss_fn, optimizer, test_threshold=0.5, with_captions=False)\n",
    "    ANP_train_loss += train_loss\n",
    "    ANP_val_loss += val_loss\n",
    "    ANP_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Aj4e5_pWWUoZ"
   },
   "source": [
    "# Custom AlexNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HF9YiQ74vDEs"
   },
   "source": [
    "## Custom AlexNet Class Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "59BRNT6YWRCQ"
   },
   "outputs": [],
   "source": [
    "class CustomAlexNet(nn.Module):\n",
    "  '''\n",
    "  The CustomAlexNet Class, child of torch.nn.Module, that incorporates the ability to utilise captions (which requires an extra fully connected layer).\n",
    "  Contains 5 convolutional layers, 3 fully connected layers & 2 instances of dropout.\n",
    "\n",
    "  Attributes:\n",
    "  num_classes (int): The number of classes in the data which is set as the number of output layers in the \n",
    "  stem_stride (int): The stride value used in the first convolutional layer.\n",
    "  with_captions (bool): The boolean value to dictate whether the model is to incorporate captions when training. Set to False by default.\n",
    "  '''\n",
    "\n",
    "  def __init__(self, num_classes = 19, stem_stride = 4, with_captions = False):\n",
    "    super(CustomAlexNet, self).__init__()\n",
    "\n",
    "    self.dropout1 = nn.Dropout()\n",
    "\n",
    "    if with_captions:\n",
    "        self.fc1 = nn.Linear((256 * 6 * 6) + (26**3), 4096)\n",
    "    else:\n",
    "        self.fc1 = nn.Linear(256 * 6 * 6, 4096)\n",
    "\n",
    "    self.leaky_relu1 = nn.LeakyReLU(inplace = True)\n",
    "    self.dropout2 = nn.Dropout()\n",
    "    self.fc2 = nn.Linear(4096, 4096)\n",
    "    self.leaky_relu2 = nn.LeakyReLU(inplace = True)\n",
    "    self.fc3 = nn.Linear(4096, num_classes)\n",
    "\n",
    "\n",
    "\n",
    "    self.features = nn.Sequential(\n",
    "        nn.Conv2d(3, 64, kernel_size = 11, stride = stem_stride, padding = 2),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "\n",
    "        nn.Conv2d(64, 192, kernel_size = 5, padding = 2),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "\n",
    "        nn.Conv2d(192, 384, kernel_size = 3, stride = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Conv2d(384, 256, kernel_size = 2, padding = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.Conv2d(256, 256, kernel_size = 3, padding = 1),\n",
    "        nn.LeakyReLU(inplace = True),\n",
    "        nn.MaxPool2d(kernel_size = 3, stride = 2),\n",
    "    )\n",
    "\n",
    "    self.avgpool = nn.AdaptiveAvgPool2d((6, 6))\n",
    "\n",
    "\n",
    "  def forward(self, x, caption=None):\n",
    "    '''\n",
    "    The forward pass of the CustomAlexNet model.\n",
    "    Takes the input and passes it sequentially through the Sequential feature layers, applies average pooling & then passes through to the fully connected layers.\n",
    "\n",
    "    Parameters:\n",
    "    x (torch.Tensor): A torch Tensor object containing the input.\n",
    "    caption (tuple): A tuple conatining the captions in the data.\n",
    "    Returns:\n",
    "    x (torch.tensor): A torch Tensor containing the output of the network feedforward process.\n",
    "    '''\n",
    "\n",
    "    x = self.features(x)\n",
    "    x = self.avgpool(x)\n",
    "    x = torch.flatten(x, 1)\n",
    "    x = self.dropout1(x)\n",
    "    if caption == None:\n",
    "        x = self.fc1(x)\n",
    "    else:\n",
    "        x = self.fc1(torch.cat((x, caption.type(torch.float)), 1))\n",
    "    x = self.leaky_relu1(x)\n",
    "    x = self.dropout2(x)\n",
    "    x = self.fc2(x)\n",
    "    x = self.leaky_relu2(x)\n",
    "    x = self.fc3(x)\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L7AINqUhvGzG"
   },
   "source": [
    "## Custom AlexNet Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IFJ35bXyWYLp"
   },
   "outputs": [],
   "source": [
    "#Create network\n",
    "ANC_model = CustomAlexNet(num_classes = 19, stem_stride = 2, with_captions=True).to(device)\n",
    "print(\"Custom AlexNet Initialized\")\n",
    "print(ANC_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8uq0xf3GvL2y"
   },
   "source": [
    "## Custom AlexNet Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QIm5zRSsWbev"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(ANC_model.parameters(), lr=LR)\n",
    "\n",
    "ANC_train_loss = []\n",
    "ANC_val_loss = []\n",
    "ANC_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, ANC_model, loss_fn, optimizer, test_threshold= 0.5, with_captions=True)\n",
    "    ANC_train_loss += train_loss\n",
    "    ANC_val_loss += val_loss\n",
    "    ANC_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "um3wBU-FYJOf"
   },
   "outputs": [],
   "source": [
    "sum(p.numel() for p in ANC_model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AJfq1gw9BrmJ"
   },
   "source": [
    "# Pre-trained ResNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u7kYxL5ICY0q"
   },
   "source": [
    "## Pre-trained ResNet Model Import - using ResNet 18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oJ3_s3PJCf2B"
   },
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "RNP_model = models.resnet18(pretrained = True)\n",
    "print(RNP_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xE6XbzoPCacl"
   },
   "source": [
    "## Override the Pre-trained ResNet FC/classification layer and view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XCB1Zsl-CgI7"
   },
   "outputs": [],
   "source": [
    "RNP_model.fc = nn.Sequential(\n",
    "    nn.Linear(512 + (26**3), 19)\n",
    ")\n",
    "\n",
    "print(RNP_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yf4QEXlxCcxq"
   },
   "source": [
    "## Pre-trained ResNet training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t0I2VdybCgmz"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(RNP_model.parameters(), lr=LR)\n",
    "\n",
    "RNP_train_loss = []\n",
    "RNP_val_loss = []\n",
    "RNP_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, RNP_model, loss_fn, optimizer, test_threshold=0.5, with_captions=False)\n",
    "    RNP_train_loss += train_loss\n",
    "    RNP_val_loss += val_loss\n",
    "    RNP_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mY0IlQMToLHJ"
   },
   "source": [
    "# Custom ResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dmnP9yx2oOZk"
   },
   "outputs": [],
   "source": [
    "pretrained_resnet = models.resnet18(pretrained = True)\n",
    "pretrained_resnet.fc = nn.Linear(512, 19)\n",
    "\n",
    "class CustomResNet(nn.Module):\n",
    "    '''\n",
    "    The CustomResNet Class, child of torch.nn.Module, that follows the resnet18 architecture and incorporates the ability to utilise captions.\n",
    "\n",
    "    Attributes:\n",
    "    pretrained_model (child of torch.nn.Module): The pretrained resnet model for the architecture to follow.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, pretrained_model):\n",
    "        super(CustomResNet, self).__init__()\n",
    "        self.pretrained_model = pretrained_model\n",
    "        self.my_layers = nn.Sequential(\n",
    "            nn.Linear(19 + (26**3), 2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(2048, 19)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, caption = None):\n",
    "        '''\n",
    "        The forward pass of the CustomResNet model.\n",
    "        Takes the input and passes it through the pretrained model forward pass, then includes concatenates the image input with the captions and passes this through the amended output layers.\n",
    "\n",
    "        Parameters:\n",
    "        x (torch.Tensor): A torch Tensor object containing the input.\n",
    "        caption (tuple): A tuple containing the captions in the data.\n",
    "\n",
    "        Returns:\n",
    "        x (torch.tensor): A torch Tensor containing the output of the network feedforward process.\n",
    "        '''\n",
    "\n",
    "        x = self.pretrained_model(x)\n",
    "        x = self.my_layers(torch.cat((x, caption.type(torch.float)), 1)) # the magic\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X27lcJauoSOH"
   },
   "outputs": [],
   "source": [
    "RNC_model = CustomResNet(pretrained_resnet).to(device)\n",
    "RNC_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bts5AQ5woUQB"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(RNC_model.parameters(), lr=LR)\n",
    "\n",
    "RNC_train_loss = []\n",
    "RNC_val_loss = []\n",
    "RNC_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, RNC_model, loss_fn, optimizer, test_threshold=0.5, with_captions=True)\n",
    "    RNC_train_loss += train_loss\n",
    "    RNC_val_loss += val_loss\n",
    "    RNC_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-aui-zbw65B2"
   },
   "source": [
    "# Data Re-read for Inception Models - no cropping, smaller batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gyaOM8S9670C"
   },
   "outputs": [],
   "source": [
    "# Concluded 19 labels - one-hot encoded. \n",
    "NUM_LABELS = 19\n",
    "BATCH_SIZE = 25\n",
    "RESIZE_SIZE = (299, 299)\n",
    "TRAIN_VAL_PROP = 0.75\n",
    "SEED = 5329\n",
    "\n",
    "transforms = Compose([\n",
    "    ToTensor(),\n",
    "    Normalize(mean = means, std = stdevs),\n",
    "    Resize(RESIZE_SIZE),\n",
    "    RandomHorizontalFlip(),\n",
    "    RandomVerticalFlip()\n",
    "])\n",
    "\n",
    "# The main dataset available with labels\n",
    "main_dataset = AssignmentDataset(csv_file = TRAIN_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                 transform = transforms,\n",
    "                                 target_transform = Lambda(lambda y: torch.zeros(NUM_LABELS, dtype=torch.uint8).scatter_(dim=0, index=torch.sub(torch.tensor(y), 1), value=1)),\n",
    "                                 has_labels = True\n",
    "                                 )\n",
    "\n",
    "# Need to further split it up\n",
    "train_dataset, val_dataset = random_split(main_dataset, \n",
    "                                          [int(round(TRAIN_VAL_PROP * len(main_dataset))), int(round((1 - TRAIN_VAL_PROP) * len(main_dataset)))], \n",
    "                                          generator=torch.Generator().manual_seed(SEED)) # Setting seed to ensure consistency\n",
    "\n",
    "# The final test dataset with no labels\n",
    "test_dataset = AssignmentDataset(csv_file = TEST_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                transform = transforms,\n",
    "                                 has_labels = False)\n",
    "\n",
    "main_dataloader = DataLoader(main_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EjN0xclkTVB-"
   },
   "source": [
    "# Pre-trained Inception"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N9rS0rj8xN_6"
   },
   "source": [
    "## Set Parameter Requires Grad function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hPetufIyTafF"
   },
   "outputs": [],
   "source": [
    "def set_grad(model, FE):\n",
    "  '''\n",
    "  A function switch off the torch.requires_grad switch when feature extraction is being conducted or the model is being altered (i.e. parameters does not require gradients)\n",
    "\n",
    "  Parameters:\n",
    "  model (child of nn.Module): The model in use.\n",
    "  FE (bool): A boolean to reflect if feature extraction is being conducted.\n",
    "\n",
    "  Returns:\n",
    "  None\n",
    "  '''\n",
    "  if FE == True:\n",
    "    for p in model.parameters():\n",
    "      p.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aPGkUh49xRTS"
   },
   "source": [
    "## Pre-trained Inception Model Import - using Inception v3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Miiu92a3Taid"
   },
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "inception = models.inception_v3(pretrained=True)\n",
    "\n",
    "set_grad(inception, True)\n",
    "\n",
    "        # Handle the auxilary net\n",
    "num_ftrs = inception.AuxLogits.fc.in_features\n",
    "inception.AuxLogits.fc = nn.Linear(num_ftrs, 19)\n",
    "        # Handle the primary net\n",
    "num_ftrs = inception.fc.in_features\n",
    "\n",
    "inception.fc = nn.Linear(num_ftrs,19)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KIvKRuqe4VNT"
   },
   "source": [
    "## Pre-trained Inception Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E-K2gICYTamf"
   },
   "outputs": [],
   "source": [
    "INCP_model = inception.cuda()\n",
    "\n",
    "print(INCP_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZdsZ848v4hPR"
   },
   "source": [
    "## Pre-trained Inception Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pPcDIoqQTaq8"
   },
   "outputs": [],
   "source": [
    "INCP_model.to(device)\n",
    "\n",
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(INCP_model.parameters(), lr=LR)\n",
    "\n",
    "INCP_train_loss = []\n",
    "INCP_val_loss = []\n",
    "INCP_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, INCP_model, loss_fn, optimizer, test_threshold=0.5, with_captions=False)\n",
    "    INCP_train_loss += train_loss\n",
    "    INCP_val_loss += val_loss\n",
    "    INCP_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SNKfrLIaqy5W"
   },
   "source": [
    "# Custom Inception"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rw-uPb4S1yB1"
   },
   "source": [
    "## Importing inception"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hjmEuEkDq3Za"
   },
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "pretrained_inception = models.inception_v3(pretrained=True)\n",
    "\n",
    "pretrained_inception.fc = nn.Linear(2048,19)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hw1O-nlX11LZ"
   },
   "source": [
    "## Custom Inception Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "45RI9Hvoq6hb"
   },
   "outputs": [],
   "source": [
    "class CustomInception(nn.Module):\n",
    "    '''\n",
    "    The CustomInception Class, child of torch.nn.Module, that follows the inception architecture and incorporates the ability to utilise captions.\n",
    "\n",
    "    Attributes:\n",
    "    pretrained_model (child of torch.nn.Module): The pretrained inception model for the architecture to follow.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, pretrained_model):\n",
    "        super(CustomInception, self).__init__()\n",
    "        self.pretrained_model = pretrained_model\n",
    "        self.my_layers = nn.Sequential(\n",
    "            nn.Linear(19 + (26**3), 2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(2048, 19)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, caption = None):\n",
    "        '''\n",
    "        The forward pass of the CustomInception Model.\n",
    "\n",
    "        Takes the input and passes it through the pretrained model forward pass, then includes concatenates the image input with the captions and passes this through the amended output layers.\n",
    "        As Inception returns auxiliary logits when training but not when evaluating, the self.training value has to be checked if logits have to be specified or not.\n",
    "\n",
    "        Parameters:\n",
    "        x (torch.Tensor): A torch Tensor object containing the input.\n",
    "        caption (tuple): A tuple containing the captions in the data.\n",
    "\n",
    "        Returns:\n",
    "        x (torch.tensor): A torch Tensor containing the output of the network feedforward process.\n",
    "        '''\n",
    "\n",
    "        if self.training:\n",
    "          x = self.pretrained_model(x).logits\n",
    "          \n",
    "        else:\n",
    "          x = self.pretrained_model(x)\n",
    "        x = self.my_layers(torch.cat((x, caption.type(torch.float)), 1)) # the magic\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eT4hdMaC6hXv"
   },
   "source": [
    "## Custom Inception Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "41ShFY_k6fHH"
   },
   "outputs": [],
   "source": [
    "INCC_model = CustomInception(pretrained_inception).to(device)\n",
    "print(INCC_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mpWutOH54rSD"
   },
   "source": [
    "## Custom Inception Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tPq4qtT3rBzM"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 3\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(INCC_model.parameters(), lr=LR)\n",
    "\n",
    "INCC_train_loss = []\n",
    "INCC_val_loss = []\n",
    "INCC_val_f1 = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(train_dataloader, val_dataloader, INCC_model, loss_fn, optimizer, test_threshold=0.5, with_captions=True)\n",
    "    INCC_train_loss += train_loss\n",
    "    INCC_val_loss += val_loss\n",
    "    INCC_val_f1 += val_f1\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dzCZtNuT4zqu"
   },
   "source": [
    "# Chosen Model Submission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fTJzIaewJ_i5"
   },
   "source": [
    "## Redefine & initialize chosen model for convenience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FEEMqG33KLK6"
   },
   "outputs": [],
   "source": [
    "class CNN2(nn.Module):\n",
    "    '''\n",
    "    The (more advanced) CNN2 Class, child of torch.nn.Module, that incorporates the ability to return feature maps as well as utilise captions.\n",
    "    Contains 4 convolutional layers, 2 fully connected layers & 1 instances of dropout.\n",
    "\n",
    "    Attributes:\n",
    "    return_fmals (bool): The boolean value to dictate whether feature maps are returned or not. Set to False by default.\n",
    "    with_captions (bool): The boolean value to dictate whether the model is to incorporate captions when training. Set to True by default.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, return_fmaps=False, with_captions=False):\n",
    "        super(CNN2, self).__init__()\n",
    "        # things to note, we use Conv2d because the convolution slides in 2 dimension, not 3\n",
    "        # also the outer channel represents the number of filters, since 1 filter flattens a 3d image\n",
    "        self.conv1 = nn.Conv2d(3, 64, 3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv2d(64, 128, 3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(128, 128, 3, stride=1, padding=1)\n",
    "        self.conv4 = nn.Conv2d(128, 128, 3, stride=1, padding=1)\n",
    "        self.dropout1 = nn.Dropout(0.5)\n",
    "        if with_captions:\n",
    "            self.fc1 = nn.Linear((128*14*14) + (26**3), 1024)\n",
    "        else:\n",
    "            self.fc1 = nn.Linear((128*14*14), 1024)\n",
    "        self.fc2 = nn.Linear(1024, NUM_LABELS)\n",
    "        #self.fc3 = nn.Linear(1024, NUM_LABELS)\n",
    "\n",
    "        self.return_fmaps = return_fmaps\n",
    "\n",
    "    \n",
    "    def set_return_fmaps(self, val=True):\n",
    "        '''\n",
    "        The setter function to toggle returning feature maps on/off.\n",
    "\n",
    "        Parameters:\n",
    "        val (bool): The boolean value for which to set the return_fmaps attribute to. Set to True by default.\n",
    "\n",
    "        Returns:\n",
    "        None\n",
    "        '''\n",
    "\n",
    "        self.return_fmaps = val\n",
    "\n",
    "\n",
    "    def forward(self, x, caption = None):\n",
    "        '''\n",
    "        The forward pass of the CNN2.\n",
    "        Takes the input and passes it sequentially through the convolutional layers, relu fucntions, max pooling layers and dropout (if applicable) to pass to the fully connected layer.\n",
    "\n",
    "        Parameters:\n",
    "        x (torch.Tensor): A torch Tensor object containing the input.\n",
    "        caption (tuple): A tuple containing the captions in the data.\n",
    "\n",
    "        Returns:\n",
    "        x (torch.Tensor): A torch Tensor object containing logit representations of the output.\n",
    "        fmaps (list): A list containing the feature maps values. Returned only if self.return_fmaps == True.\n",
    "        '''\n",
    "\n",
    "        fmaps = []\n",
    "\n",
    "        x = self.conv1(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.conv2(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        x = self.conv3(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.conv4(x)\n",
    "        fmaps.append(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        \n",
    "        if caption == None:\n",
    "            x = self.fc1(x)\n",
    "        else:\n",
    "            # Mixing in the caption data here, concat in dimension 1\n",
    "            x = self.fc1(torch.cat((x, caption.type(torch.float)), 1)) \n",
    "\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout1(x)\n",
    "        x = self.fc2(x)\n",
    "\n",
    "        if self.return_fmaps:\n",
    "            return x, fmaps\n",
    "        else:\n",
    "            return x\n",
    "            \n",
    "chosen_model = CNN2(with_captions=True).to(device)\n",
    "print('Model initialised')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9M0Fm95-8lJ9"
   },
   "source": [
    "## Re-Read Data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lNmdbSOM84o1"
   },
   "outputs": [],
   "source": [
    "# Concluded 19 labels - one-hot encoded. Leaving label 12 in there as it'll just have zeros for this index\n",
    "NUM_LABELS = 19\n",
    "BATCH_SIZE = 100\n",
    "RESIZE_SIZE = (299, 299)\n",
    "CROP_SIZE = 224 # Seems to be the standard image size used (224)\n",
    "TRAIN_VAL_PROP = 0.75\n",
    "SEED = 5329\n",
    "\n",
    "transforms = Compose([\n",
    "    ToTensor(),\n",
    "    Normalize(mean = [0.4635, 0.4492, 0.4212], std = [0.2753, 0.2716, 0.2870]),\n",
    "    Resize(RESIZE_SIZE),\n",
    "    RandomHorizontalFlip(),\n",
    "    RandomVerticalFlip(),\n",
    "    RandomCrop(CROP_SIZE)\n",
    "])\n",
    "\n",
    "# The main dataset available with labels\n",
    "main_dataset = AssignmentDataset(csv_file = TRAIN_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                 transform = transforms,\n",
    "                                 target_transform = Lambda(lambda y: torch.zeros(NUM_LABELS, dtype=torch.uint8).scatter_(dim=0, index=torch.sub(torch.tensor(y), 1), value=1)),\n",
    "                                 has_labels = True\n",
    "                                 )\n",
    "\n",
    "# Need to further split it up\n",
    "train_dataset, val_dataset = random_split(main_dataset, \n",
    "                                          [int(round(TRAIN_VAL_PROP * len(main_dataset))), int(round((1 - TRAIN_VAL_PROP) * len(main_dataset)))], \n",
    "                                          generator=torch.Generator().manual_seed(SEED)) # Setting seed to ensure consistency\n",
    "\n",
    "# The final test dataset with no labels\n",
    "test_dataset = AssignmentDataset(csv_file = TEST_CSV,\n",
    "                                 image_dir = IMAGES_DIR,\n",
    "                                transform = transforms,\n",
    "                                 has_labels = False)\n",
    "\n",
    "main_dataloader = DataLoader(main_dataset, batch_size = BATCH_SIZE, shuffle=True)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F-_h0V3hBYOj"
   },
   "source": [
    "## Train over full training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fVBetCq55hmw"
   },
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "NUM_EPOCHS = 1\n",
    "EVAL_EVERY = 50\n",
    "\n",
    "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
    "optimizer = torch.optim.Adadelta(chosen_model.parameters(), lr=LR)\n",
    "\n",
    "chosen_model_train_loss = []\n",
    "\n",
    "t0 = time.time()\n",
    "for t in range(NUM_EPOCHS):\n",
    "    print(f\"-------------------------------\\nEpoch {t+1}\")\n",
    "    train_loss, val_loss, val_f1 = train_loop(main_dataloader, val_dataloader, chosen_model, loss_fn, optimizer, test_threshold=None, with_captions=True)\n",
    "    chosen_model_train_loss += train_loss\n",
    "t1 = time.time()\n",
    "print(\"Model Trained - training time: %i:%i\"% ((t1 - t0) // 60, (t1 - t0) % 60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rQ13n7a1F8wZ"
   },
   "source": [
    "## Generate submissions & export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BANfwCn7KfPX"
   },
   "outputs": [],
   "source": [
    "threshold = 0.5\n",
    "submission = output_to_submission(test_dataloader, chosen_model, threshold, with_captions=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZeyodYuIORvz"
   },
   "outputs": [],
   "source": [
    "with open('submission_file.txt', 'w') as file:\n",
    "    file.writelines(', '.join(prediction) + '\\n' for prediction in submission[0])\n",
    "\n",
    "with open('submission_file.csv', 'w', newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "    writer.writerows(submission[0])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "512OAXAvdyiA",
    "iSOjenw7ebQN",
    "o4PwnFxReYis",
    "BZnhUTm0ejJU",
    "GoFQuMTvekP0",
    "TFhr2cz6klfH",
    "ky77ggwqimc2",
    "4Sr6XhsRiuEd",
    "sJIhhyLHjM6t",
    "v-Pn1aWPlXAL",
    "R2TEdYtMmRBx",
    "DJaKVxFvyWrg",
    "o_mUa8z2n6_t",
    "75OgEj01rYZA",
    "scZFF6fYnymV",
    "Oy-YC2osn0k_",
    "oidAXBAHoGVq",
    "mk2lQfU47MUO",
    "GFSKLrXtrGcR",
    "hPHCdDLBx_3O",
    "tm2HTrYaIjJ3",
    "wff9NV_2IrNM",
    "1CAp77jvsbTm",
    "4uP4kO69sgAf",
    "VhYi9rqBydMB",
    "qV2C5-F17dPp",
    "I-cp5V_p7kx7",
    "sXAk7mRq7vPo",
    "3VDhBYvI5db1",
    "iXc2eqky8ITr",
    "GJbG8iY-8fEj",
    "L2BKZVB88nPz",
    "HF9YiQ74vDEs",
    "L7AINqUhvGzG",
    "u7kYxL5ICY0q",
    "xE6XbzoPCacl",
    "yf4QEXlxCcxq",
    "mY0IlQMToLHJ",
    "-aui-zbw65B2",
    "EjN0xclkTVB-",
    "N9rS0rj8xN_6",
    "aPGkUh49xRTS",
    "KIvKRuqe4VNT",
    "ZdsZ848v4hPR",
    "SNKfrLIaqy5W",
    "rw-uPb4S1yB1",
    "hw1O-nlX11LZ",
    "eT4hdMaC6hXv",
    "mpWutOH54rSD",
    "dzCZtNuT4zqu",
    "fTJzIaewJ_i5",
    "9M0Fm95-8lJ9",
    "F-_h0V3hBYOj",
    "rQ13n7a1F8wZ"
   ],
   "name": "500525438_500448953_500393860 - COMP5329 2021S1A2 - Code.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}