assignment_ecg_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/uu-sml/wasp-assigninmen-af-classification/blob/main/assignment_ecg_classification.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z00pmDU9qUai"
   },
   "source": [
    "# WASP Course: Artificial Intelligence and Machine Learning\n",
    "\n",
    "Lecturer: Dave Zachariah\n",
    "\n",
    "Assignment responsible: Philipp Pilar, Jingwei Hu, Paul Häusner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Student and Group Information\n",
    "\n",
    "Fill this out for the submission of the assignment (you submit this notebook with your solution)\n",
    "\n",
    "- **Student names:** <font color='red'>Fill in</font>\n",
    "\n",
    "- **Team ID:** <font color='red'>Fill in</font>\n",
    "\n",
    "Make sure that the team id is the same as the one with which you submit your model predictions (see coding task 7) such that we can check your performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z_X8-WK_Htgr"
   },
   "source": [
    "---\n",
    "# Module 3 - Assignment Overview: ECG classification\n",
    "\n",
    "The [electrocardiogram (ECG)](https://www.mayoclinic.org/tests-procedures/ekg/about/pac-20384983) records the electrical signals in the heart. It is a common  test used to quickly detect heart problems and to monitor the heart's health. \n",
    "In this assignment you will implement and evaluate a model to classify whether the person has [atrial fibrillation (AF)](https://www.mayoclinic.org/diseases-conditions/atrial-fibrillation/symptoms-causes/syc-20350624.) or not based on measurements from the ECG exam. \n",
    "\n",
    "\n",
    "**Submission:** You submit the deliverables (see below) at https://canvas.kth.se/courses/48799/assignments\n",
    "\n",
    "**Due Date:** August 23, 2024.\n",
    "\n",
    "---\n",
    "## Basic Tasks\n",
    "Your task is to implement a classification model, train this model on training data, and evaluate its performance on validation data. We provide skeleton code for the implementation of a simple convolution neural network model.\n",
    "\n",
    "The steps required to implement this model are presented as numbered tasks below. In total there are seven (7) coding tasks and five (5) explanation tasks. \n",
    "\n",
    "## Competitive setting\n",
    "\n",
    "You have to compute the predictions for the test data (you do not have the labels for it) and submit your predictions to be evaluated to a leaderboard. These predictions will be scored and your submission will be ranked according to the F1 score and compared with your colleagues. In the end a winning team will be determined.\n",
    "\n",
    "### Deliverables\n",
    "There are two deliverables:\n",
    "1. You have to submit this Jupyter notebook on the course web-page (Canvas) together with your code and explanations (where asked for it) that describe your implementation and your experimental results. The notebook should run as a standalone in google colab.\n",
    "2. You have to have at least **three (3)** submissions (for instructions on how to submit, see coding task 7) where you try to improve the model architecture, the training procedure or the problem formulation. In the submission of this notebook you have to provide a short explanation of what changed between each submission and justify why you decided to make these changes.\n",
    "\n",
    "### Grading\n",
    "To pass the assignment, you must submit a complete and working implementation of a model and a well-motivated description and evaluation of it. Your model should reach an Area under the ROC curve (AUROC) on the test data of at least 0.97 and an Average Precision (AP) score of 0.95. Note that the leaderboard to is sorted by F1 score and not AUROC, hence you would want to balance all three metrics.\n",
    "\n",
    "### GPU Acceleration\n",
    "To be able to use the GPUs provided by colab in order to speed up your computations, you want to check that the `Hardware accelerator` is set to `GPU` under `Runtime > change runtime type`. Note that notebooks run by connecting to virtual machines that have maximum lifetimes that can be as much as 12 hours. Notebooks will also disconnect from VMs when left idle for too long. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# helper function\n",
    "def exists(path):\n",
    "    val = os.path.exists(path)\n",
    "    if val:\n",
    "        print(f'{path} already exits. Using cached. Delete it manually to recieve it again!')\n",
    "    return val\n",
    "\n",
    "# clone requirements.txt if not yet available\n",
    "if not exists('requirements.txt'):\n",
    "    !git clone https://gist.github.com/dgedon/8a7b91714568dc35d0527233e9ceada4.git req\n",
    "    !mv req/requirements.txt .\n",
    "    !yes | rm -r req"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RZ_o0wKcm7WM"
   },
   "outputs": [],
   "source": [
    "# Install packages (python>=3.9 is required)\n",
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JkjZfvuKHd6q"
   },
   "outputs": [],
   "source": [
    "# Import\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from tqdm.notebook import trange, tqdm\n",
    "import h5py\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TTfxCz1YIpGP"
   },
   "source": [
    "---\n",
    "## The data set\n",
    "\n",
    "The dataset is a subset of the [*CODE dataset*](https://scilifelab.figshare.com/articles/dataset/CODE_dataset/15169716): an anotated database of ECGs. The ECG exams were recorded in Brazil by the Telehealth Network of the state Minas Gerais between 2010 and 2016. The dataset and its usage for the development of deep learning methods was described in [\"Automatic diagnosis of the 12-lead ECG using a deep neural network\"](https://www.nature.com/articles/s41467-020-15432-4).\n",
    "The full dataset is available for research upon request.\n",
    "\n",
    "\n",
    "For the training dataset you have labels. \n",
    "For the test dataset you only have the ECG exams but no labels. Evaluation is done by submitting to the leaderboard.\n",
    "\n",
    "Download the dataset from the given dropbox link and unzip the folder containing the files. The downloaded files are in WFDB format (see [here](https://www.physionet.org/content/wfdb-python/3.4.1/) for details)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PeNTb95FM6R1"
   },
   "outputs": [],
   "source": [
    "# 1. Download dataset\n",
    "if not exists('codesubset.tar.gz'):\n",
    "    !wget https://www.dropbox.com/s/9zkqa5y5jqakdil/codesubset.tar.gz?dl=0 -O codesubset.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rQR2EI49OVP0"
   },
   "outputs": [],
   "source": [
    "# 1. unzip the downloaded data set folder\n",
    "if not exists('codesubset'):\n",
    "    !tar -xf codesubset.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the extraced folder 'codesubset' contains\n",
    "1. subfolders with the ECG exam traces. These have to be further preprocessed which we do in the next steps.\n",
    "2. a csv file which contain the labels and other features for the training data set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1H8tsO-QABmw"
   },
   "source": [
    "\n",
    "### Preprocessing\n",
    "\n",
    "Run the cells below to  Clone the GitHub repository which we use for [data preprocessing](https://github.com/antonior92/ecg-preprocessing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7yz7VW7nQEO-",
    "outputId": "9edf7844-03f0-4833-ecb6-2c53cda99f03"
   },
   "outputs": [],
   "source": [
    "# 2. clone the code files for data preprocessing\n",
    "if not exists('ecg-preprocessing'):\n",
    "    !git clone https://github.com/paulhausner/ecg-preprocessing.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KlvD4bEtQUhc"
   },
   "source": [
    "Let us plot an ECG sample. We can plot ECGs using the `ecg_plot` library for example by using the following code snippet where `ecg_sample` is an array of size `(number of leads * sequence length)`. Now we can view an ECG before preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 987
    },
    "id": "QOpX0vlEQVUU",
    "outputId": "4f57693a-b62e-4881-ab61-daace2bdeb31"
   },
   "outputs": [],
   "source": [
    "import ecg_plot\n",
    "runfile(\"ecg-preprocessing/read_ecg.py\")\n",
    "\n",
    "PATH_TO_WFDB = 'codesubset/train/TNMG100046'\n",
    "ecg_sample, sample_rate, _ = read_ecg(PATH_TO_WFDB)\n",
    "\n",
    "# ECG plot\n",
    "plt.figure()\n",
    "lead = ['I', 'II', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n",
    "ecg_plot.plot(ecg_sample, sample_rate=sample_rate, style='bw', row_height=8, lead_index=lead, columns=1, title='Sample ECG before pre-processing')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i_icbSxgDowE"
   },
   "source": [
    "\n",
    "The preprocessing consist of:\n",
    "- resampling all ECG traces to the sample sampling period (400 Hz). Option: ``--new_freq 400``\n",
    "- zero padding if necessary such that all ECG have the same number of samples (4096). Option: ``--new_len 4096``.\n",
    "- removing trends in the ECG signal. Option: ``--remove_baseline``\n",
    "- remove possible power line noise. Option: ``--powerline 60``\n",
    "\n",
    "You can run the script bellow to plot the same ECG after the preprocessing.  The script also use the  `ecg_plot` library (as you did above).  You can try also with different command line options to see how the preprocessing affects the signal that will be used by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "lDQfCECcDoGN",
    "outputId": "33526310-7605-4a9d-b055-97a0366ed6ba"
   },
   "outputs": [],
   "source": [
    "%run ecg-preprocessing/plot_from_ecg.py codesubset/train/TNMG100046 --new_freq 400 --new_len 4096 --remove_baseline --powerline 60"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w2Kh8RkYD8bE"
   },
   "source": [
    "\n",
    "Next we perform the preprocessing in all exams and convert them into one single h5 file (see [here](https://www.h5py.org/#:~:text=The%20h5py%20package%20is%20a,they%20were%20real%20NumPy%20arrays.) for details about the format). The resulting h5 files contains the traces as arrays with the shape `(number of traces * sequence length * number of leads)` where sequence length is 4096 and number of leads is 8. \n",
    "The files `train.h5` and `test.h5` will be saved inside the folder `codesubset/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cyS0WXTzQU5c",
    "outputId": "d56a623b-48a0-4284-9711-d92fbec1a0f6"
   },
   "outputs": [],
   "source": [
    "# 3. Generate train\n",
    "if not exists('codesubset/train.h5'):\n",
    "    !python ecg-preprocessing/generate_h5.py --new_freq 400 --new_len 4096 --remove_baseline --powerline 60 codesubset/train/RECORDS.txt codesubset/train.h5\n",
    "# 3. Generate test\n",
    "if not exists('codesubset/test.h5'):\n",
    "    !python ecg-preprocessing/generate_h5.py --new_freq 400 --new_len 4096 --remove_baseline --powerline 60 codesubset/test/RECORDS.txt codesubset/test.h5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IvelAs8NJExH"
   },
   "source": [
    "### Coding Task 1: Data Analysis\n",
    "\n",
    "Before starting to model you have to analyse the dataset. You can be creative in your way of *getting a feeling* for the data. What you have to do is:\n",
    "- plot an ECG after proprocessing saved in the hdf5 file. For this use the `ecg_plot()` example above and see below for how to access the preprocessed data in h5 format.\n",
    "\n",
    "Some further ideas to explore are:\n",
    "- check the balance of the data set,\n",
    "- evaluate the distribution of age and sex of the patients,\n",
    "- think about the performance that a best naive classifier would achieve, e.g. by random guessing or always predicting one class.\n",
    "\n",
    "<br />\n",
    "\n",
    "**How to access the data?**\n",
    "\n",
    "You can acces the data in the h5 file in the following way\n",
    "```\n",
    "import h5py\n",
    "\n",
    "PATH_TO_H5_FILE = 'codesubset/train.h5'\n",
    "f = h5py.File(PATH_TO_H5_FILE, 'r')\n",
    "data = f['tracings']\n",
    "```\n",
    "Then, `data[i]` is an numpy array of the $i$th ECG exam (including all time points and leads).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "ja8j1xOYJdqg",
    "outputId": "498aacf4-f018-44ee-cb38-aaccfa943177"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "TASK: Insert your code here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uu4tpfc3STcD"
   },
   "source": [
    "### Explanation task 1: Data Analysis\n",
    "\n",
    "Please explain your main findings of the data analysis task in a few bullet points. Explain also what the preprocessing does and why it is necessary.\n",
    "\n",
    "<br />\n",
    "\n",
    "\n",
    "**<font color='red'>Your explanation here:</font>**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yO_E8qGUJ2Db"
   },
   "source": [
    "---\n",
    "## Model\n",
    "\n",
    "The model class consists of two methods: \n",
    "- `__init__(self, args)`: This methods initializes the class, e.g. by using `mymodel=ModelBaseline(args)`.\n",
    "- `forward(self,input_data)`: This method is called when we run `model_output=mymodel(input_data)`.\n",
    "\n",
    "The dimension of the input data is  `(batch size * sequence length * number of leads)`. Where **batch size** is a hyperparameter, **sequence length** is the number of ECG time samples (=4096) and **number of leads** (=8).\n",
    "\n",
    "The `ModelBaseline` (provided below) is a 2 layer model with one convolutional layers and one linear layer. Some explanations: \n",
    "- The conv layer downsamples the input traces from 4096 samples to 128 samples and increases the number of channels from 8 (=number of leads) to 32. Here we use a kernel size of 3.\n",
    "- The linear layer uses the flattened output from the conv and outputs one prediction. Since we have a binary problem, a single prediction is sufficient.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pONc-F25K-Z5"
   },
   "outputs": [],
   "source": [
    "class ModelBaseline(nn.Module):\n",
    "    def __init__(self,):\n",
    "        super(ModelBaseline, self).__init__()\n",
    "        self.kernel_size = 3\n",
    "\n",
    "        # conv layer\n",
    "        downsample = self._downsample(4096, 128)\n",
    "        self.conv1 = nn.Conv1d(in_channels=8, \n",
    "                               out_channels=32, \n",
    "                               kernel_size=self.kernel_size, \n",
    "                               stride=downsample,\n",
    "                               padding=self._padding(downsample),\n",
    "                               bias=False)\n",
    "        \n",
    "        # linear layer\n",
    "        self.lin = nn.Linear(in_features=32*128,\n",
    "                             out_features=1)\n",
    "        \n",
    "        # ReLU\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def _padding(self, downsample):\n",
    "        return max(0, int(np.floor((self.kernel_size - downsample + 1) / 2)))\n",
    "\n",
    "    def _downsample(self, seq_len_in, seq_len_out):\n",
    "        return int(seq_len_in // seq_len_out)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x= x.transpose(2,1)\n",
    "\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x_flat= x.view(x.size(0), -1)\n",
    "        x = self.lin(x_flat)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2Q5oKHb1Ls87"
   },
   "source": [
    "### Coding Task 2: Define your model\n",
    "\n",
    "In the cell below you have to define your model. You can be inspired by the baseline model above but you can also define any other kind of neural network architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1BHovxZZLvkd"
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self,):\n",
    "        super(Model, self).__init__()\n",
    "        \n",
    "    def forward(x):\n",
    "        return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v7besXJ1Qjax"
   },
   "source": [
    "### Explanation Task 2: Final Model\n",
    "Please explain and motivate in short sentences or bullet points the choice of your final model.\n",
    "\n",
    "<br />\n",
    "\n",
    "\n",
    "**<font color='red'>Your explanation here:</font>**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PeWVnhT1L72u"
   },
   "source": [
    "---\n",
    "## Train function\n",
    "\n",
    "The function `train(...)` is called to in every epoch to train the model. The function loads the training data, makes predictions, compares predictions with true labels in the loss function and adapting the model parameters using stochastic gradient descent.\n",
    "\n",
    "In the code cell below there is the basic structure to load data from the data loader and to log your loss. The arguments of the function are explained by the use in the `main(...)` function below.\n",
    "\n",
    "If you are unfamiliar with PyTorch training loops, then this official [tutorial](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html) might help (especially section \"4. Train your Network\")."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nHcflxJ1Wprw"
   },
   "source": [
    "### Coding Task 3: Fill training loop\n",
    "\n",
    "Fill the code cell below such that the model is training when `train(...)` is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QMLCx9Cahr7f"
   },
   "outputs": [],
   "source": [
    "def train_loop(epoch, dataloader, model, optimizer, loss_function, device):\n",
    "    # model to training mode (important to correctly handle dropout or batchnorm layers)\n",
    "    model.train()\n",
    "    # allocation\n",
    "    total_loss = 0  # accumulated loss\n",
    "    n_entries = 0   # accumulated number of data points\n",
    "    # progress bar def\n",
    "    train_pbar = tqdm(dataloader, desc=\"Training Epoch {epoch:2d}\".format(epoch=epoch), leave=True)\n",
    "    # training loop\n",
    "    for traces, diagnoses in train_pbar:\n",
    "        # data to device (CPU or GPU if available)\n",
    "        traces, diagnoses = traces.to(device), diagnoses.to(device)\n",
    "\n",
    "        \"\"\"\n",
    "        TASK: Insert your code here. This task can be done in 5 lines of code.\n",
    "        \"\"\"\n",
    "\n",
    "        # Update accumulated values\n",
    "        total_loss += loss.detach().cpu().numpy()\n",
    "        n_entries += len(traces)\n",
    "\n",
    "        # Update progress bar\n",
    "        train_pbar.set_postfix({'loss': total_loss / n_entries})\n",
    "    train_pbar.close()\n",
    "    return total_loss / n_entries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "djQnstiLiWB1"
   },
   "source": [
    "---\n",
    "## Eval function\n",
    "\n",
    "The `eval(...)` function is similar to the `train(...)` function but is used to evaluate the model on validation data without adapting the model parameters. You can prohibit computing gradients by using a `with torch.no_grad():` statement.\n",
    "\n",
    "Currenlty only the loss is logged here. Additionally you have to collect all your predictions and the true values in order to compute more metrics such as AUROC."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "22G-f_ooWunl"
   },
   "source": [
    "### Coding Task 4: Fill evaluation loop\n",
    "Fill the code cell below such we obtain model predictions to evaluate the validation loss and collect the predictoin in order to compute other validation metrics in the `main(...)` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xWOvM9e5ijqk"
   },
   "outputs": [],
   "source": [
    "def eval_loop(epoch, dataloader, model, loss_function, device):\n",
    "    # model to evaluation mode (important to correctly handle dropout or batchnorm layers)\n",
    "    model.eval()\n",
    "    # allocation\n",
    "    total_loss = 0  # accumulated loss\n",
    "    n_entries = 0   # accumulated number of data points\n",
    "    valid_probs = []  # accumulated predicted probabilities\n",
    "    valid_true = [] # accumulated true labels\n",
    "    \n",
    "    # progress bar def\n",
    "    eval_pbar = tqdm(dataloader, desc=\"Evaluation Epoch {epoch:2d}\".format(epoch=epoch), leave=True)\n",
    "    # evaluation loop\n",
    "    for traces_cpu, diagnoses_cpu in eval_pbar:\n",
    "        # data to device (CPU or GPU if available)\n",
    "        traces, diagnoses = traces_cpu.to(device), diagnoses_cpu.to(device)\n",
    "\n",
    "        \"\"\"\n",
    "        TASK: Insert your code here. This task can be done in 6 lines of code.\n",
    "        \"\"\"\n",
    "\n",
    "        # Update accumulated values\n",
    "        total_loss += loss.detach().cpu().numpy()\n",
    "        n_entries += len(traces)\n",
    "\n",
    "        # Update progress bar\n",
    "        eval_pbar.set_postfix({'loss': total_loss / n_entries})\n",
    "    eval_pbar.close()\n",
    "    return total_loss / n_entries, np.vstack(valid_probs), np.vstack(valid_true)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dtBEPHo7jEZP"
   },
   "source": [
    "---\n",
    "## Run Training\n",
    "\n",
    "In the code cell below there are some initial (non-optimal!) training hyperparameters. Further, we combine everything from above into training code. That means that we build the dataloaders, define the model/loss/optimizer and then train/validate the model over multiple epochs. Here, we save the model with the lowest validation loss as the best model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1Eh5IsvQWy0y"
   },
   "source": [
    "### Coding Task 5: Combine everything to train/validate the model\n",
    "\n",
    "The following tasks are necessary in the code below\n",
    "- split the data into training and validation data\n",
    "- define the loss function\n",
    "- decide and implement validation metric(s) to evaluate and compare the model on\n",
    "\n",
    "Optional task:\n",
    "- include learning rate scheduler\n",
    "- take specific care about possible data inbalance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wOz48rnxdGfp"
   },
   "source": [
    "### Coding Task 6: Run your model and adapt hyperparameters\n",
    "\n",
    "After you combined everything in task 5, now you run the code to evaluate the model. Based on the resulting validation metrics you tune\n",
    "- the training hyperparameters\n",
    "- the model architecture\n",
    "- the model hyperparameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pjGcqXpOdSbA"
   },
   "source": [
    "### Explanation Task 3: Hyperparameter\n",
    "Please explain and motivate in short sentences or bullet points the final choice of hyperparamer and how you developed them.\n",
    "\n",
    "<br />\n",
    "\n",
    "\n",
    "**<font color='red'>Your explanation here:</font>**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3VNfdwn8IqQQ"
   },
   "outputs": [],
   "source": [
    "# set seed\n",
    "seed = 42\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# choose variables\n",
    "\"\"\"\n",
    "TASK: Adapt the following hyperparameters if necessary\n",
    "\"\"\"\n",
    "learning_rate = 1e-2  \n",
    "weight_decay = 1e-1  \n",
    "num_epochs = 15\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CQvRyaQcyvcM",
    "outputId": "6f0a25f4-da7d-44ef-ac3b-34ac3087db09"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, random_split, DataLoader\n",
    "\n",
    "# Set device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "tqdm.write(\"Use device: {device:}\\n\".format(device=device))\n",
    "\n",
    "# =============== Build data loaders ======================================#\n",
    "tqdm.write(\"Building data loaders...\")\n",
    "\n",
    "path_to_h5_train, path_to_csv_train, path_to_records = 'codesubset/train.h5', 'codesubset/train.csv', 'codesubset/train/RECORDS.txt'\n",
    "# load traces\n",
    "traces = torch.tensor(h5py.File(path_to_h5_train, 'r')['tracings'][()], dtype=torch.float32)\n",
    "# load labels\n",
    "ids_traces = [int(x.split('TNMG')[1]) for x in list(pd.read_csv(path_to_records, header=None)[0])] # Get order of ids in traces\n",
    "df = pd.read_csv(path_to_csv_train)\n",
    "df.set_index('id_exam', inplace=True)\n",
    "df = df.reindex(ids_traces) # make sure the order is the same\n",
    "labels = torch.tensor(np.array(df['AF']), dtype=torch.float32).reshape(-1,1)\n",
    "# load dataset\n",
    "dataset = TensorDataset(traces, labels)\n",
    "len_dataset = len(dataset)\n",
    "n_classes = len(torch.unique(labels))\n",
    "# split data\n",
    "\"\"\"\n",
    "TASK: Split the dataset in train and validation; Insert your code here.\n",
    "This can be done in <=4 line of code\n",
    "\"\"\"\n",
    "\n",
    "# build data loaders\n",
    "train_dataloader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)\n",
    "valid_dataloader = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False)                         \n",
    "tqdm.write(\"Done!\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GMOKrUn9jfca"
   },
   "outputs": [],
   "source": [
    "# =============== Define model ============================================#\n",
    "tqdm.write(\"Define model...\")\n",
    "\"\"\"\n",
    "TASK: Replace the baseline model with your model; Insert your code here\n",
    "\"\"\"\n",
    "model = ModelBaseline()\n",
    "model.to(device=device)\n",
    "tqdm.write(\"Done!\\n\")\n",
    "\n",
    "# =============== Define loss function ====================================#\n",
    "\"\"\"\n",
    "TASK: define the loss; Insert your code here. This can be done in 1 line of code\n",
    "\"\"\"\n",
    "\n",
    "# =============== Define optimizer ========================================#\n",
    "tqdm.write(\"Define optimiser...\")\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n",
    "tqdm.write(\"Done!\\n\")\n",
    "\n",
    "# =============== Define lr scheduler =====================================#\n",
    "# TODO advanced students (non mandatory)\n",
    "\"\"\"\n",
    "OPTIONAL: define a learning rate scheduler; Insert your code here\n",
    "\"\"\"\n",
    "lr_scheduler = None\n",
    "\n",
    "# =============== Train model =============================================#\n",
    "tqdm.write(\"Training...\")\n",
    "best_loss = np.Inf\n",
    "# allocation\n",
    "train_loss_all, valid_loss_all = [], []\n",
    "\n",
    "# loop over epochs\n",
    "for epoch in trange(1, num_epochs + 1):\n",
    "    # training loop\n",
    "    train_loss = train_loop(epoch, train_dataloader, model, optimizer, loss_function, device)\n",
    "    # validation loop\n",
    "    valid_loss, y_pred, y_true = eval_loop(epoch, valid_dataloader, model, loss_function, device)\n",
    "\n",
    "    # collect losses\n",
    "    train_loss_all.append(train_loss)\n",
    "    valid_loss_all.append(valid_loss)\n",
    "\n",
    "    # compute validation metrics for performance evaluation\n",
    "    \"\"\"\n",
    "    TASK: compute validation metrics (e.g. AUROC); Insert your code here\n",
    "    This can be done e.g. in 5 lines of code\n",
    "    \"\"\"\n",
    "\n",
    "    # save best model: here we save the model only for the lowest validation loss\n",
    "    if valid_loss < best_loss:\n",
    "        # Save model parameters\n",
    "        torch.save({'model': model.state_dict()}, 'model.pth') \n",
    "        # Update best validation loss\n",
    "        best_loss = valid_loss\n",
    "        # statement\n",
    "        model_save_state = \"Best model -> saved\"\n",
    "    else:\n",
    "        model_save_state = \"\"\n",
    "\n",
    "    # Print message\n",
    "    tqdm.write('Epoch {epoch:2d}: \\t'\n",
    "                'Train Loss {train_loss:.6f} \\t'\n",
    "                'Valid Loss {valid_loss:.6f} \\t'\n",
    "                '{model_save}'\n",
    "                .format(epoch=epoch,\n",
    "                        train_loss=train_loss,\n",
    "                        valid_loss=valid_loss,\n",
    "                        model_save=model_save_state)\n",
    "                    )\n",
    "\n",
    "    # Update learning rate with lr-scheduler\n",
    "    if lr_scheduler:\n",
    "        lr_scheduler.step()\n",
    "\n",
    "\"\"\"\n",
    "TASK: Here it can make sense to plot your learning curve; Insert your code here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ouhd4vgt4jlS"
   },
   "source": [
    "---\n",
    "## Model Testing\n",
    "\n",
    "Since we saved our best model, we can now load the trained model and make predictions on the test data set. We save the predictions in a csv file which will be uploaded as part of the deliverables. Note that we take a `Sigmoid()` function on the model prediction in order to obtain soft predictions (probabilities) instead of hard predictions (0s or 1s)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aXJRnGTpW7Qv"
   },
   "source": [
    "### Coding Task 7: Make prediction for test data\n",
    "\n",
    "Here you do not really need to code but you have to:\n",
    "- replace the baseline model with your model. If you do not use colab then change the path to the model location to load the trained model)\n",
    "- run the script. The predictions are saved in the variable `soft_pred`.\n",
    "- upload your predictions to the leaderboard online (see instruction details below). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build the dataloader once and re-use when running the cell below possibly multiple times.\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# =============== Build data loaders ==========================================#\n",
    "tqdm.write(\"Building data loaders...\")\n",
    "# load data\n",
    "path_to_h5_test, path_to_csv_test = 'codesubset/test.h5', 'codesubset/test.csv'\n",
    "traces = torch.tensor(h5py.File(path_to_h5_test, 'r')['tracings'][()], dtype=torch.float32)\n",
    "dataset = TensorDataset(traces)\n",
    "len_dataset = len(dataset)\n",
    "# build data loaders\n",
    "test_dataloader = DataLoader(dataset, batch_size=32, shuffle=False)\n",
    "tqdm.write(\"Done!\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0-m4CYeW_hvO"
   },
   "outputs": [],
   "source": [
    "# =============== Define model ================================================#\n",
    "tqdm.write(\"Define model...\")\n",
    "\"\"\"\n",
    "TASK: Replace the baseline model with your model; Insert your code here\n",
    "\"\"\"\n",
    "model = ModelBaseline()\n",
    "\n",
    "# load stored model parameters\n",
    "ckpt = torch.load('model.pth', map_location=lambda storage, loc: storage)\n",
    "model.load_state_dict(ckpt['model'])\n",
    "# put model on device\n",
    "model.to(device=device)\n",
    "tqdm.write(\"Done!\\n\")\n",
    "\n",
    "# =============== Evaluate model ==============================================#\n",
    "model.eval()\n",
    "# allocation\n",
    "test_pred = torch.zeros(len_dataset,1)\n",
    "# progress bar def\n",
    "test_pbar = tqdm(test_dataloader, desc=\"Testing\")\n",
    "# evaluation loop\n",
    "end=0\n",
    "for traces in test_pbar:\n",
    "    # data to device\n",
    "    traces = traces[0].to(device)\n",
    "    start = end\n",
    "    with torch.no_grad():\n",
    "        # Forward pass\n",
    "        model_output = model(traces)\n",
    "\n",
    "        # store output\n",
    "        end = min(start + len(model_output), test_pred.shape[0])\n",
    "        test_pred[start:end] = torch.nn.Sigmoid()(model_output).detach().cpu()\n",
    "\n",
    "test_pbar.close()\n",
    "\n",
    "# =============== Save predictions ============================================#\n",
    "soft_pred = np.stack((1-test_pred.numpy(), test_pred.numpy()),axis=1).squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qMxVJxHsgGdT"
   },
   "source": [
    "To upload your predictions to the leaderboard, use the following code. There are the following steps to follow:\n",
    "1. Download the GitHub repository for the leaderboard submission system.\n",
    "2. Register your team with a **team id** and **password**. The password ensures that only your team can upload to your team id. Do only run the registration once.\n",
    "3. Upload you predictions as a new submission. There are some things to obey here:\n",
    "    - For each submission you have to attach a note for you to keep track of the submission in the leaderboard and for us to know which submission you refer to in your explanation. Choose something meaningful such as \"submission A\" or \"model B\".\n",
    "    - You can only get one prediction evaluated per day and you get the score the following day. If you do multiple submissions on the same day, the initial submission will be overwritten and thus only the final submission will be evaluated.\n",
    "    - Only a maximum of ***FIVE*** submissions will be evaluated. So make them count! (If you update an submission before it is evaluated it doesn't count)\n",
    "    - The evaluation score is published with you team_id and note at http://hyperion.it.uu.se:5050/leaderboard\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-BPWcP20RHwu",
    "outputId": "9effbfc3-14fe-440c-a0e5-d64566754784"
   },
   "outputs": [],
   "source": [
    "# 1. Download repository for leaderboard submission system\n",
    "if not exists('leaderboard'):\n",
    "    !git clone https://gist.github.com/3ff6c4c867331c0bf334301842d753c7.git leaderboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DRaCnCLsgHPg"
   },
   "outputs": [],
   "source": [
    "# 2. Registration of your team\n",
    "host = \"http://hyperion.it.uu.se:5050/\"\n",
    "runfile(\"leaderboard/leaderboard_helpers.py\")\n",
    "\n",
    "\"\"\"\n",
    "TASK: Decide for a team_id (max 20 chars) and password. \n",
    "Do not change this after you have registered your team\n",
    "\"\"\"\n",
    "team_id = '' #Fill in a string\n",
    "password = '' #Fill in a string\n",
    "\n",
    "# run the registration\n",
    "r = register_team(team_id, password)\n",
    "if (r.status_code == 201):\n",
    "    print(\"Team registered successfully! Good luck\")\n",
    "elif not (r.status_code == 200):\n",
    "    raise Exception(\"You can not change your password once created. If you need help, please contact the teachers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "I03QEBiegF7_",
    "outputId": "778464c4-f0c4-41b9-a40a-c0b1cc734c0f"
   },
   "outputs": [],
   "source": [
    "# 3. Upload the prediction as submission\n",
    "\n",
    "# Write a note about the training procedure so you can identify it in the leaderboard. e.g. 5 epochs, or First  (Max 20 characters)\n",
    "\"\"\"\n",
    "TASK: Add a note for you submission\n",
    "\"\"\"\n",
    "note = '' #Fill in a string\n",
    "\n",
    "# Submit the predictions to the leaderboard. Note, this also saves your submissions in your colab folder\n",
    "r = submit(team_id, password, soft_pred.tolist(), note)\n",
    "if r.status_code == 201:\n",
    "    print(\"Submission successful!\")\n",
    "elif r.status_code == 200:\n",
    "    print(\"Submission updated!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jTKe9X-zTt7f"
   },
   "source": [
    "### Explanation Task 4: Submissions\n",
    "One of the grading criteria are three submissions to the leaderboard. List the three main submissions in the table below and explain the main changes in your code for each submission.\n",
    "\n",
    "<br />\n",
    "\n",
    "\n",
    "**<font color='red'>Your explanation here:</font>**\n",
    "\n",
    "Your team id: **<font color='red'>Fill in</font>**\n",
    "\n",
    "| Submission note | Accuracy | F1 | AUC | AP | Submission description |\n",
    "| --------------- | -------- | -- | --  | -- | ---------------------- |\n",
    "|xxx              | 0        | 0  | 0   | 0  | desc                   |\n",
    "|xxx              | 0        | 0  | 0   | 0  | desc                   |\n",
    "|xxx              | 0        | 0  | 0   | 0  | desc                   |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lw-GKZ4fVDhD"
   },
   "source": [
    "### Explanation Task 5: Reflection on Metrics\n",
    "Your were asked to reach a certain value in AUC and AP while maximising F1 for the leaderboard position. Explain in bullet points what aspect each of the metrics covers and why it is important not to just focus on one metric. What can happen if you only focus on AUC for example?\n",
    "\n",
    "<br />\n",
    "\n",
    "**<font color='red'>Your explanation here:</font>**"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "assignment_ecg_classification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}