06- Neural Networks/Lab/neural_network_lab.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Machine Learning for Neuroscience, <br>Department of Brain Sciences, Faculty of Medicine, <br> Imperial College London\n",
    "### Contributors: Francesca Palermo, Nan Fletcher-Lloyd, Alex Capstick, Yu Chen\n",
    "**Winter 2022**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will focus on Neural Networks, and will rely on [scikit-learn's](https://scikit-learn.org/stable/index.html) and [pytorch's](https://pytorch.org) python packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing packages\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.inspection import DecisionBoundaryDisplay\n",
    "import sklearn.datasets as datasets\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "\n",
    "# import plotting functions\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors\n",
    "from cycler import cycler\n",
    "binary_cmap = matplotlib.colors.LinearSegmentedColormap.from_list(\"\", ['#332288', 'white', '#AA4499'])\n",
    "plt.rcParams[\"axes.prop_cycle\"] = cycler(\n",
    "    color=['#332288','#88CCEE','#44AA99','#117733','#999933','#DDCC77','#CC6677','#882255','#AA4499']\n",
    "    )\n",
    "\n",
    "# class for holding the random state throughout the notebook.\n",
    "# this keeps results consistent\n",
    "class RandomState(object):\n",
    "    def __init__(self, random_state=None):\n",
    "        self.random_state = random_state\n",
    "    def next(self, n=1):\n",
    "        assert type(n) == int, \"Ensure n is an integer\"\n",
    "        if n == 1:\n",
    "            self.random_state,\\\n",
    "                out_state = np.random.default_rng(\n",
    "                    self.random_state\n",
    "                    ).integers(0, 1e9, size=(2,))\n",
    "        else:\n",
    "            self.random_state,\\\n",
    "                *out_state = np.random.default_rng(\n",
    "                    self.random_state\n",
    "                    ).integers(0, 1e9, size=(n+1,))\n",
    "        \n",
    "        return out_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_state = RandomState(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(random_state.next())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is `torch`?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pytorch is a python package that allows you to build machine learning models. It is highly flexible and allows you to perform auto-differentiation ✨!\n",
    "\n",
    "All weights, inputs and targets in pytorch are stored as `torch.Tensor`s, ([documentation](https://pytorch.org/docs/stable/tensors.html)) which are very similar to numpy arrays, in the sense that they hold values but they hold the magic of being differentiable automatically. This makes training neural networks possible.\n",
    "\n",
    "All neural network layers and loss functions are built using `nn.Module`s ([documentation](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)) which allow for differentiable manipulations of  `torch.Tensor`s."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to Train a Model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The basic pipeline for training a model is as follows:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src= \"_dependents/pipeline.png\" width=\"600\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Building Blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, we will discuss the building blocks that go into a deep learning network. This will include how to access single layers, as well as the different types of activation functions that can be used to create non-linearity in models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Single Neuron"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The single neuron underpins much of deep learning, and is used in almost all deep learning models. In its most basic form, it is a vector of values that trasforms an incoming vector of data to a scalar value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Basics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create this single neuron easily with pytorch's `torch.nn.Linear` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# neuron that takes in 100 features and outputs 1 feature\n",
    "sn = nn.Linear(10, 1, bias=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is of the following form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn.weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given an input, this class will produce the output by multiplying each of the feature values with the values in the weights. Let us see what this means:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = torch.rand(\n",
    "    size=(1,10), # 1 data point with 10 features\n",
    "    requires_grad=False, # data point should not be updated\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we transform the input using our single neuron:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that this output contains a `grad_fn`. This is pytorch's auto-grad system working in the background, ready to calculate the gradients to update weights in a model. We will see how this is used later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's say that we want to have 5 neurons stacked on top of each other to transform a feature vector of size 10, to a feature vectore of size 5. This can easily be done in the following way:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5 neurons acting together on the same input\n",
    "fiven = nn.Linear(10, 5, bias=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# given an input\n",
    "input = torch.rand(\n",
    "    size=(1,10), # 1 data point with 10 features\n",
    "    requires_grad=False, # data point should not be updated\n",
    "    )\n",
    "\n",
    "# we can get the output\n",
    "fiven(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the result of 5 neurons, each outputting a scalar value!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pytorch allows you to stack these neurons vertically (acting on the same input) and horizontally (acting on the outputs from previous neurons). This will be discussed in the section [here](#a-feed-forward-network)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Can We Train a Single Neuron?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can we attempt to train this single neuron to replicate the equation: `y=mx+c`, for some `m` and `c`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = torch.randn((1, 2), requires_grad=False) # a random m of size (1,2)\n",
    "c = torch.randn(1, requires_grad=False) # a random c of size (1,)\n",
    "\n",
    "print(f'The equation we are trying to predict is:\\n y={np.round(m,2)}@x+{np.round(c,2)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10000 train random inputs between 0 and 1 with 5 features \n",
    "X_train = torch.rand(size=(10000, 2), requires_grad=False)\n",
    "# 10000 train random inputs between 0 and 1 with 5 features \n",
    "X_val = torch.rand(size=(1000, 2), requires_grad=False)\n",
    "# normalise the data\n",
    "scaler = StandardScaler()\n",
    "X_train = torch.tensor(scaler.fit_transform(X_train)).float()\n",
    "X_val = torch.tensor(scaler.transform(X_train)).float()\n",
    "\n",
    "# with our set equation, this gives the y values\n",
    "y_train = X_train@m.T + c\n",
    "y_val = X_val@m.T + c\n",
    "\n",
    "print(f'The y outputs are of shape {y_train.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now need to build our single neuron model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# since here bias=True, this single neuron will take care of the m and c at the same time\n",
    "sn = nn.Linear(2,1, bias=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"To start with, the single neuron has weights:\\n {sn.weight} \"\\\n",
    "    f\"\\n \\n and bias:\\n {sn.bias}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train this model, we will pass the `X` through the model and see how close it was to the true `y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The first 5 predictions are:\\n {sn(X_train)[:5]}\"\\\n",
    "    f\"\\n The first 5 true values are:\\n {y_train[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At the moment this doesn't replicate the original function very well. We will therefore train this model using gradient descent to get it to perform better."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can train this model (and by extension any model) using the following code:\n",
    "\n",
    "For ease of training later in this notebook, the following code snippet will be added to the file `model_trainer.py` in this directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define training function that can be used \n",
    "# with any model, loss function, data and optimiser\n",
    "def train(\n",
    "    model, train_loader, n_epochs, optimiser, criterion, val_loader=None\n",
    "    ):\n",
    "    '''\n",
    "    A function to train any model with a given dataset, optimiser, and \n",
    "    criterion (loss function).\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    \n",
    "    - model: pytorch nn object:\n",
    "        The model to train\n",
    "    - train_loader: pytorch data loader:\n",
    "        The data to train with.\n",
    "    - n_epochs: integer:\n",
    "        The number of epochs to train for.\n",
    "    - optimiser: pytorch optimiser:\n",
    "        The optimiser to make the model updates.\n",
    "    - criterion: pytorch nn object:\n",
    "        The loss function to calculate the loss with.\n",
    "    - val_loader: pytorch data loader:\n",
    "        The data to calculate the validation loss with.\n",
    "    \n",
    "    Returns\n",
    "    ---------\n",
    "    \n",
    "    - model: pytorch nn object:\n",
    "        Trained version of the model given\n",
    "        as an input.\n",
    "    - tuple of dictionaries:\n",
    "        - train_loss_dict: dictionary:\n",
    "            Dictionary containing the training loss\n",
    "            with keys: `steps` and `loss`.\n",
    "        - val_loss_dict\n",
    "            Dictionary containing the validation loss\n",
    "            with keys: `steps` and `loss`.\n",
    "\n",
    "    '''\n",
    "    # check if GPU is available and use that if so\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    # ==== push model to GPU if available ====\n",
    "    model.to(device)\n",
    "\n",
    "    # since all functions in this function rely on the inputs above, \n",
    "    # they won't work outside of the train function\n",
    "\n",
    "    # pass a single batch of data through the model and get loss\n",
    "    def batch_loss(inputs, targets):\n",
    "        # ==== push data to GPU if available ====\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        # ==== forward pass ====\n",
    "        outputs = model(inputs)\n",
    "        # ==== calc and save loss ====\n",
    "        loss = criterion(outputs, targets)\n",
    "        # ==== return loss ====\n",
    "        return loss\n",
    "    \n",
    "    # train for an epoch\n",
    "    def train_epoch(train_loader):\n",
    "        model.train() # set model option to train - important if using dropout\n",
    "        batch_loss_list = [] # we will store all losses in a list\n",
    "        # for each batch in the train loader\n",
    "        for nb, (inputs, targets) in enumerate(train_loader):\n",
    "            # ==== set gradient to zero ====\n",
    "            optimiser.zero_grad() # really important! Common mistake to not do this!\n",
    "            # run data through batch to get loss\n",
    "            loss = batch_loss(inputs=inputs, targets=targets)\n",
    "            # ==== calc backprop gradients ====\n",
    "            loss.backward()\n",
    "            # ==== perform update step ====\n",
    "            optimiser.step()\n",
    "            # ==== store loss for later ====\n",
    "            batch_loss_list.append(loss.item())\n",
    "        # ==== calculate average loss ====\n",
    "        return batch_loss_list\n",
    "\n",
    "    # perform an epoch over the validation data to get loss\n",
    "    # dont want gradients in validation since we're not training\n",
    "    @torch.no_grad()\n",
    "    def val_epoch(val_loader):\n",
    "        model.eval() # set model option to eval - important if using dropout\n",
    "        batch_loss_list = [] # we will store all losses in a list\n",
    "        # for each batch in the val loader\n",
    "        for nb, (inputs, targets) in enumerate(val_loader):   \n",
    "            # ==== set gradient to zero ====\n",
    "            optimiser.zero_grad() # gradients shouldnt be calculated but good practise \n",
    "            # run data through batch to get loss\n",
    "            loss = batch_loss(inputs=inputs, targets=targets)\n",
    "            # ==== store loss for later ====\n",
    "            batch_loss_list.append(loss.item())\n",
    "        # ==== calculate average loss ====\n",
    "        return batch_loss_list\n",
    "\n",
    "\n",
    "    pbar = tqdm.tqdm(desc='Training', total=n_epochs) # progress bar\n",
    "    # loss stats\n",
    "    train_loss_dict = {'step': [], 'loss': []}\n",
    "    val_loss_dict = {'step': [], 'loss': []}\n",
    "\n",
    "    # train for the given n_epochs\n",
    "    for ne in range(n_epochs):\n",
    "        # ==== train for an epoch ====\n",
    "        n_batches = len(train_loader)\n",
    "        batch_lost_list_train = train_epoch(train_loader=train_loader)\n",
    "        # ==== get loss stats ====\n",
    "        train_loss_dict['loss'].extend(batch_lost_list_train) # adding loss\n",
    "        # adding step values. These are the number of steps from the beginning\n",
    "        train_loss_dict['step'].extend(\n",
    "            list(np.arange(ne*n_batches, (ne+1)*n_batches)+1) \n",
    "            )\n",
    "        avg_loss_train = np.mean(batch_lost_list_train)\n",
    "\n",
    "        # if a validation loader is passed\n",
    "        if val_loader is not None:  \n",
    "            # ==== epoch over validation ====\n",
    "            batch_lost_list_val = val_epoch(val_loader=val_loader)\n",
    "            # ==== get loss stats ====\n",
    "            avg_loss_val = np.mean(batch_lost_list_val)\n",
    "            val_loss_dict['loss'].append(avg_loss_val)\n",
    "            val_loss_dict['step'].append(\n",
    "                (ne+1)*n_batches+1 # the number of new steps is as many as the train loader\n",
    "                )\n",
    "        else:\n",
    "            avg_loss_val = np.nan\n",
    "\n",
    "        # ==== set pbar info and update ====\n",
    "        pbar.set_postfix(\n",
    "            {\n",
    "                'Train Loss': f\"{avg_loss_train:.3f}\",\n",
    "                'Val Loss': f\"{avg_loss_val:.3f}\"\n",
    "                }\n",
    "            )\n",
    "        pbar.update(1)\n",
    "        pbar.refresh()\n",
    "    \n",
    "    # put the model back on the cpu\n",
    "    model.to('cpu')\n",
    "\n",
    "    return model, (train_loss_dict, val_loss_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then set the settings for the training, including the number of epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the number of times to run through the whole dataset\n",
    "n_epochs = 10\n",
    "# the batch size for the data\n",
    "batch_size = 512\n",
    "\n",
    "# wrapper data in dataset\n",
    "train_dataset = torch.utils.data.TensorDataset(X_train, y_train)\n",
    "val_dataset = torch.utils.data.TensorDataset(X_val, y_val)\n",
    "\n",
    "# put dataset in train loader\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To do this, we need a loss function and an optimiser:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pytorch has lots of optimisers available for use. Here we use Adam:\n",
    "optimiser = torch.optim.SGD(\n",
    "    params=sn.parameters(), # need to pass the optimiser the model parameters\n",
    "    lr=0.01 # and the learning rate\n",
    "    )\n",
    "\n",
    "# pytorch also has lots of loss functions available. We can also easily make our own.\n",
    "# to start, we will use the mean squarred error:\n",
    "criterion = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And calling the training function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn, (train_loss_dict, val_loss_dict) = train(\n",
    "    model=sn,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    n_epochs=n_epochs,\n",
    "    optimiser=optimiser,\n",
    "    criterion=criterion,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the loss looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,4)) # plotting area\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    train_loss_dict['step'], train_loss_dict['loss'],\n",
    "    label='Train'\n",
    "    )\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    val_loss_dict['step'], val_loss_dict['loss'],\n",
    "    label='Val'\n",
    "    )\n",
    "\n",
    "# formatting\n",
    "ax.set_title('Training Loss')\n",
    "ax.set_ylabel('Loss')\n",
    "ax.set_xlabel('Steps')\n",
    "ax.legend()\n",
    "# show plot\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Did the values look like the true equation? Let's see:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"After training, the single neuron has weights:\\n {sn.weight} \"\\\n",
    "    f\"\\n \\n and bias:\\n {sn.bias}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The true equation has weights:\\n {m} \"\\\n",
    "    f\"\\n \\n and bias:\\n {c}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model approximated the function well, and we can see that the validation loss in the graph above went to 0. This tells us that the model didn't overfit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you will have seen from the linear models lab, this is essentially a linear regression model and training it in this way is not a good idea, since closed solutions exist! This is simply an example of one of the basic building blocks of a Neural Network. We will now turn our attention to some of the other building blocks of neural networks, and how we can use them to create complex machine learning models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Activation Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Activation functions are non-linear functions that are placed after a layer of neurons to produce non-linearity in outputs. To visualise this problem, let's try to solve the following classification problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import sklearn.datasets as datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading some synthetic data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y = datasets.make_moons(1000, noise=0.15, random_state=random_state.next())\n",
    "\n",
    "# train-val splits:\n",
    "X_train, X_val, y_train, y_val = train_test_split(\n",
    "    X, y, train_size=0.75, random_state=random_state.next()\n",
    "    )\n",
    "\n",
    "# scaling the data\n",
    "scaler = StandardScaler()\n",
    "X_train = torch.tensor(scaler.fit_transform(X_train)).float()\n",
    "X_val = torch.tensor(scaler.transform(X_val)).float()\n",
    "\n",
    "y_train = torch.tensor(y_train)\n",
    "y_val = torch.tensor(y_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,2,figsize=(8,4))\n",
    "ax1, ax2 = axes\n",
    "\n",
    "ax1.scatter(x=X_train[:,0], y=X_train[:,1], c=y_train, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "ax2.scatter(x=X_val[:,0], y=X_val[:,1], c=y_val, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "\n",
    "ax1.set_title('Train')\n",
    "ax2.set_title('Val')\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us try and use our linear neuron model to try and classify this data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# since we won't be using activation functions,\n",
    "# we need the result to be a float.\n",
    "# we also need to reshape the data to (n_data_points, 1)\n",
    "# we will see why later.\n",
    "y_train_reg = y_train.float().unsqueeze(1)\n",
    "y_val_reg = y_val.float().unsqueeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The shape of the input array is {X_train.shape} \"\\\n",
    "    f\"and the shape of the targets is {y_train_reg.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn = nn.Sequential(\n",
    "    nn.Linear(X_train.shape[1], 100*X_train.shape[1]),\n",
    "    nn.Linear(100*X_train.shape[1], 100*X_train.shape[1]),\n",
    "    nn.Linear(100*X_train.shape[1], 1),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will only work in the same directory as the file `model_trainer`\n",
    "from model_trainer import train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the number of times to run through the whole dataset\n",
    "n_epochs = 100\n",
    "# the batch size for the data\n",
    "batch_size = 128\n",
    "\n",
    "# wrapper data in dataset\n",
    "train_dataset = torch.utils.data.TensorDataset(X_train, y_train_reg)\n",
    "val_dataset = torch.utils.data.TensorDataset(X_val, y_val_reg)\n",
    "\n",
    "# put dataset in train loader\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pytorch has lots of optimisers available for use. Here we use Adam:\n",
    "optimiser = torch.optim.Adam(\n",
    "    params=sn.parameters(), # need to pass the optimiser the model parameters\n",
    "    lr=0.001, # and the learning rate\n",
    "    weight_decay=0.0001\n",
    "    )\n",
    "\n",
    "# mean squarred error is used since we are doing regression\n",
    "criterion = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn, (train_loss_dict, val_loss_dict) = train(\n",
    "    model=sn,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    n_epochs=n_epochs,\n",
    "    optimiser=optimiser,\n",
    "    criterion=criterion,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,4)) # plotting area\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    train_loss_dict['step'], train_loss_dict['loss'],\n",
    "    label='Train'\n",
    "    )\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    val_loss_dict['step'], val_loss_dict['loss'],\n",
    "    label='Val'\n",
    "    )\n",
    "\n",
    "# formatting\n",
    "ax.set_title('Training Loss')\n",
    "ax.set_ylabel('Loss')\n",
    "ax.set_xlabel('Steps')\n",
    "ax.legend()\n",
    "# show plot\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"A single prediction from the model looks as follows: {sn(X_train[0])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above isn't really a prediction of label, and instead is a regression prediction of the label. Let's see what this decision boundary looks like on a graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the code in this cell is not important and is written to enable us to \n",
    "# plot a decision boundary with a pytorch model\n",
    "from decision_plotter import pytorch_decision_boundary\n",
    "\n",
    "class PTtoDB(object):\n",
    "    def __init__(self, model):\n",
    "        self.model=model\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def predict(self, inputs):\n",
    "        if type(inputs) != torch.Tensor:\n",
    "            inputs = torch.tensor(inputs).float()\n",
    "        self.model.eval()\n",
    "        return self.model(inputs).reshape(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us look at the decision boundary. We will take the prediction as being the larger of of the"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax = pytorch_decision_boundary(PTtoDB(sn), X=X_val, ax=ax, cmap=binary_cmap)\n",
    "ax.scatter(x=X_val[:,0], y=X_val[:,1], c=y_val, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "ax.set_title('Boundaries on the Test Set')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's try to add a non-linear decision boundary to the code and predict the labels instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl = nn.Sequential(\n",
    "    nn.Linear(X_train.shape[1], 100*X_train.shape[1]),\n",
    "    nn.ReLU(), # non linear\n",
    "    nn.Linear(100*X_train.shape[1], 100*X_train.shape[1]),\n",
    "    nn.ReLU(), # non linear\n",
    "    nn.Linear(100*X_train.shape[1], 2),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The shape of the input array is {X_train.shape} \"\\\n",
    "    f\"and the shape of the targets is {y_train.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will only work in the same directory as the file `model_trainer`\n",
    "from model_trainer import train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the number of times to run through the whole dataset\n",
    "n_epochs = 100\n",
    "# the batch size for the data\n",
    "batch_size = 128\n",
    "\n",
    "# wrapper data in dataset\n",
    "train_dataset = torch.utils.data.TensorDataset(X_train, y_train)\n",
    "val_dataset = torch.utils.data.TensorDataset(X_val, y_val)\n",
    "\n",
    "# put dataset in train loader\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset, batch_size=batch_size, shuffle=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pytorch has lots of optimisers available for use. Here we use Adam:\n",
    "optimiser = torch.optim.Adam(\n",
    "    params=nl.parameters(), # need to pass the optimiser the model parameters\n",
    "    lr=0.01, # and the learning rate\n",
    "    weight_decay=0.0001\n",
    "    )\n",
    "\n",
    "# This time we will use cross entropy loss, \n",
    "# which includes a non-linear activation function\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl, (train_loss_dict, val_loss_dict) = train(\n",
    "    model=nl,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    n_epochs=n_epochs,\n",
    "    optimiser=optimiser,\n",
    "    criterion=criterion,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,4)) # plotting area\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    train_loss_dict['step'], train_loss_dict['loss'],\n",
    "    label='Train'\n",
    "    )\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    val_loss_dict['step'], val_loss_dict['loss'],\n",
    "    label='Val'\n",
    "    )\n",
    "\n",
    "# formatting\n",
    "ax.set_title('Training Loss')\n",
    "ax.set_ylabel('Loss')\n",
    "ax.set_xlabel('Steps')\n",
    "ax.legend()\n",
    "# show plot\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"A single prediction from the model \"\\\n",
    "    f\"looks as follows: {nl(X_train[0].unsqueeze(0))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"These values can easily be turned to probabilities \"\\\n",
    "    f\"as by using the softmax function:\\n {F.softmax(nl(X_train[0].unsqueeze(0)), dim=1, )}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above isn't really a prediction of label, and instead is a regression prediction of the label. Let's see what this decision boundary looks like on a graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the code in this cell is not important and is written to enable us to \n",
    "# plot a decision boundary with a pytorch model\n",
    "from decision_plotter import pytorch_decision_boundary\n",
    "\n",
    "class PTtoDB(object):\n",
    "    def __init__(self, model):\n",
    "        self.model=model\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def predict(self, inputs):\n",
    "        if type(inputs) != torch.Tensor:\n",
    "            inputs = torch.tensor(inputs).float()\n",
    "        self.model.eval()\n",
    "        output = self.model(inputs)\n",
    "        probabilities =  F.softmax(output, dim=1)\n",
    "        return probabilities[:,1].reshape(-1) # return probability of class 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us look at the decision boundary. We will take the prediction as being the larger of of the"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax = pytorch_decision_boundary(PTtoDB(nl), X=X_val, ax=ax, cmap=binary_cmap)\n",
    "ax.scatter(x=X_val[:,0], y=X_val[:,1], c=y_val, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "ax.set_title('Boundaries on the Test Set')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the decision bounday fits the data much closer! All we did was add a non-linear layer to the model and include loss function that contains a non linear layer, as well as using a non-linear function to turn log-probabilities to probabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we understand the motivation behind using non linear activation functions, let's take a look at some of the most common functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-1,1,100)\n",
    "y = F.relu(x)\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax.plot(x, y)\n",
    "ax.grid()\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_title('ReLU Function')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-1,1,100)\n",
    "y = F.leaky_relu(x, negative_slope=0.1)\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax.plot(x, y)\n",
    "ax.grid()\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_title('Leaky ReLU Function')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-3,3,100)\n",
    "y = F.tanh(x)\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax.plot(x, y)\n",
    "ax.grid()\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_title('Tanh Function')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-8,8,100)\n",
    "y = F.sigmoid(x)\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "ax.plot(x, y)\n",
    "ax.grid()\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_title('Sigmoid Function')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Feed Forward Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have seen how to access a single layer of neurons and some of the most common activation functions, as well as studying how these models are trained we will now see how more complex networks can be built."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basics of Building a Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pytorch allows for the user to build complex, feed forward networks easily. Let us see how you might create a multi-layer perceptron:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `nn.Sequential` function (documentation [here](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)) allows you to write a sequence with any number of `nn.Module` ([here](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)) objects in order of execution. For example, a model with 3 layers of neurons in order can be written as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_size = 10 # example input size\n",
    "output_size = 2 # example output size\n",
    "\n",
    "# a 3 layer network mapping 10 features -> 2 features\n",
    "ffn = nn.Sequential(\n",
    "    nn.Linear(input_size, 20), # (layer_input_size, layer_output_size)\n",
    "    nn.Linear(20, 20),\n",
    "    nn.Linear(20, output_size),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can view the network by printing it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ffn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding non-linear layers is also simple, we simply add those to the `nn.Sequential` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_size = 10 # example input size\n",
    "output_size = 2 # example output size\n",
    "\n",
    "# a 3 layer network mapping 10 features -> 2 features\n",
    "nn.Sequential(\n",
    "    nn.Linear(input_size, 20),\n",
    "    nn.ReLU(), # relu function added\n",
    "    nn.Linear(20, 20),\n",
    "    nn.Tanh(), # tanh function added\n",
    "    nn.Linear(20, output_size),\n",
    "    nn.LeakyReLU(), # leaky relu added\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, to create more complicated models, like transformers, this can be as simple as:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transformers are the backbone of many incredible deep learning models such as [ChatGPT](https://openai.com/blog/chatgpt/) which is worth testing out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example transformer model \n",
    "nn.Sequential(\n",
    "    nn.Linear(1000, 512),\n",
    "    # transformers need a positional encoding layer here, see:\n",
    "    #https://pytorch.org/tutorials/beginner/transformer_tutorial.html\n",
    "    # to learn how to add one.\n",
    "    nn.Transformer(512, batch_first=True), \n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, this is a much more complicated model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to build custom models or blocks, we can do the following using python classes and sub-classing the `nn.Module` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FFBlock(nn.Module):\n",
    "    def __init__(self, input_size, output_size):\n",
    "        # always needs to be the first thing done when building a custom model\n",
    "        super(FFBlock, self).__init__() \n",
    "        # defining a custom block\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(input_size, output_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(output_size), # can even add batch norm like this\n",
    "            nn.Dropout(p=0.2) # dropout with rate = 0.2\n",
    "            )\n",
    "    def forward(self, inputs):\n",
    "        return self.net(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use this class in other models too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "larger_ff_model = nn.Sequential(\n",
    "    FFBlock(10, 100),\n",
    "    FFBlock(100, 100),\n",
    "    FFBlock(100, 100),\n",
    "    nn.Linear(100, 2)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This produces the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "larger_ff_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of these models can be used in regression or classification tasks. In general:\n",
    "\n",
    "- **Regression**: \n",
    "    - Ensure that the feature output size is the same as the output size of the targets.\n",
    "    - Use a loss function designed for regression such as `nn.MSELoss`.\n",
    "    - **Example**: Predicting the temperature tomorrow based on today's data.\n",
    "\n",
    "- **Classification**:\n",
    "    - Ensure that the feature output size is the same as the number of classes being predicted. For example, if you're performning binary classification output two features, etc. The outputs can be interpretted as the log-probability of the input being from the coressponding class. To get the actual probabilities, you can use `torch.nn.functional.softmax(outputs, dim=1)`. \n",
    "    - Use a loss function designed for classification such as `nn.CrossEntropyLoss`.\n",
    "    - For multi-label classification (each data point can have multiple labels), you may want to use `nn.BCEWithLogitsLoss()`.\n",
    "    - **Example**: Predicting whether it will rain or not tomorrow, or if doing multi-label classification, predicting whether it will rain tomorrow and/or the professor wears a tie to lectures tomorrow (since these are likely to be mutually exclusive)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following are examples of classification and regression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# regression - 3 linear layers mapping 10 input features to 1 output feature\n",
    "r_model = nn.Sequential(nn.Linear(10,10), nn.Linear(10,10), nn.Linear(10,1))\n",
    "criterion = nn.MSELoss()\n",
    "optimiser = torch.optim.Adam(r_model.parameters(), lr=0.01)\n",
    "\n",
    "# classification - 3 linear layers mapping 10 input features to 5 classes\n",
    "c_model = nn.Sequential(nn.Linear(10,10), nn.Linear(10,10), nn.Linear(10,5))\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimiser = torch.optim.Adam(c_model.parameters(), lr=0.01)\n",
    "\n",
    "# multi-label classification - 3 linear layers mapping 10 input features to 5 attributes\n",
    "mc_model = nn.Sequential(nn.Linear(10,10), nn.Linear(10,10), nn.Linear(10,5))\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimiser = torch.optim.Adam(mc_model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Maybe an Easier Way"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If all you need is a multilayer perceptron (linear layers combined with drop out and non linear functions), then sklearn can provide this. See [MLPClassifier](https://scikit-learn.org/stable/modules/generated/sklearn.neural_network.MLPClassifier.html) and [MLPRegressor](https://scikit-learn.org/stable/modules/generated/sklearn.neural_network.MLPRegressor.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neural_network import MLPClassifier, MLPRegressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLPClassifier()\n",
    "model.get_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLPRegressor()\n",
    "model.get_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These models are great if you don't need anything custom and just want to have a simple network perform a classification or regression task. For example, let's look at the task we gave to our neural networks earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import sklearn.datasets as datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading some synthetic data. These do not need to be tensors, since sklearn works on numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y = datasets.make_moons(1000, noise=0.15, random_state=random_state.next())\n",
    "\n",
    "# train-val splits:\n",
    "X_train, X_val, y_train, y_val = train_test_split(\n",
    "    X, y, train_size=0.75, random_state=random_state.next()\n",
    "    )\n",
    "\n",
    "# scaling the data\n",
    "scaler = StandardScaler()\n",
    "X_train = scaler.fit_transform(X_train)\n",
    "X_val = scaler.transform(X_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,2,figsize=(8,4))\n",
    "ax1, ax2 = axes\n",
    "\n",
    "ax1.scatter(x=X_train[:,0], y=X_train[:,1], c=y_train, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "ax2.scatter(x=X_val[:,0], y=X_val[:,1], c=y_val, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "\n",
    "ax1.set_title('Train')\n",
    "ax2.set_title('Val')\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# same as our pytorch model, but comes with extra performance options as default\n",
    "mlp = MLPClassifier(\n",
    "    hidden_layer_sizes=[\n",
    "        100*X_train.shape[1],\n",
    "        100*X_train.shape[1],\n",
    "        ],\n",
    "    activation='relu',\n",
    "    solver='adam',\n",
    "    batch_size=128,\n",
    "    learning_rate='constant',\n",
    "    learning_rate_init=0.01,\n",
    "    random_state=random_state.next()\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,6))\n",
    "\n",
    "dbd = DecisionBoundaryDisplay.from_estimator(\n",
    "    estimator=mlp,\n",
    "    X=X_val,\n",
    "    grid_resolution=200,\n",
    "    plot_method='contourf',\n",
    "    response_method='predict_proba',\n",
    "    ax=ax,\n",
    "    cmap=binary_cmap,\n",
    "    alpha=0.5,\n",
    "    eps=0.3,\n",
    "    levels=100,\n",
    "    )\n",
    "\n",
    "ax.scatter(x=X_val[:,0], y=X_val[:,1], c=y_val, alpha=0.5, cmap=binary_cmap, edgecolor='black')\n",
    "ax.set_title('Sklearn MLP')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, this was considerably easier to implement, and the `MLPClassifier` and `MLPRegressor` classes give you quick access to a wide variety of interesting options."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how a feed forward network can be used to predict the [fashion MNIST dataset](https://github.com/zalandoresearch/fashion-mnist). This is a simple dataset containing images of clothes. We will flatten these images so that it is possible to make predictions using a feed forward MLP. We will then compare the pytorch and sklearn implementations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision # to get data\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from dataset import MemoryDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_images = transforms.Compose([\n",
    "                        transforms.PILToTensor(),\n",
    "                        transforms.ConvertImageDtype(torch.float),\n",
    "                        transforms.Normalize(mean=0, std=1),\n",
    "                        nn.Flatten(start_dim=0),\n",
    "                        ])\n",
    "\n",
    "train_dataset = torchvision.datasets.FashionMNIST(\n",
    "    root='./data/',\n",
    "    train=True,\n",
    "    download=True, \n",
    "    transform=transform_images\n",
    "    )\n",
    "\n",
    "test_dataset = torchvision.datasets.FashionMNIST(\n",
    "    root='./data/',\n",
    "    train=False,\n",
    "    download=True, \n",
    "    transform=transform_images\n",
    "    )\n",
    "\n",
    "# keeps data in memory after loading\n",
    "train_dataset = MemoryDataset(train_dataset, now=False,)\n",
    "test_dataset = MemoryDataset(test_dataset, now=False,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few example images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_to_show = [train_dataset[i][0].reshape(28,28) for i in range(5)]\n",
    "\n",
    "fig, axes = plt.subplots(1,5,figsize=(15,3))\n",
    "\n",
    "for nax, ax in enumerate(axes):\n",
    "    ax.imshow(\n",
    "        images_to_show[nax].numpy(),\n",
    "        cmap='Greys'\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us try and train a pytorch model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the training function we made\n",
    "from model_trainer import train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting up the data loaders\n",
    "# splitting the data into train and val\n",
    "n_train_set, n_val_set = [\n",
    "    int(0.75*len(train_dataset)), len(train_dataset)-int(0.75*len(train_dataset))\n",
    "    ]\n",
    "\n",
    "train_dataset_split, val_dataset_split = torch.utils.data.random_split(\n",
    "    train_dataset, lengths=[n_train_set, n_val_set],\n",
    "    )\n",
    "\n",
    "# the number of times to run through the whole dataset\n",
    "n_epochs = 10\n",
    "# the batch size for the data\n",
    "batch_size = 512\n",
    "\n",
    "# put datasets in train loader\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset_split, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset_split, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    test_dataset, batch_size=batch_size, shuffle=False,\n",
    "    )\n",
    "\n",
    "# model\n",
    "mlp = nn.Sequential(\n",
    "    nn.Linear(784, 16),\n",
    "    #nn.Dropout(0.2), # dropout can be easily added like this\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(16, 16), \n",
    "    #nn.Dropout(0.2), # dropout can be easily added like this\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(16,10),\n",
    "    )\n",
    "\n",
    "optimiser = torch.optim.Adam(mlp.parameters(), lr=0.01, weight_decay=0.0001)\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, we will use our function to train the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the model\n",
    "model, (train_loss_dict, val_loss_dict) = train(\n",
    "    model=mlp,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    n_epochs=n_epochs,\n",
    "    optimiser=optimiser,\n",
    "    criterion=criterion\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives us the loss plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,4)) # plotting area\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    train_loss_dict['step'], train_loss_dict['loss'],\n",
    "    label='Train'\n",
    "    )\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    val_loss_dict['step'], val_loss_dict['loss'],\n",
    "    label='Val'\n",
    "    )\n",
    "\n",
    "# formatting\n",
    "ax.set_title('Training Loss')\n",
    "ax.set_ylabel('Loss')\n",
    "ax.set_xlabel('Steps')\n",
    "ax.legend()\n",
    "# show plot\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculate the accuracy, we need to count the number of correct instances over the test loader:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly to our training function, we can define a predictions function that can return the predictions, true labels, and the raw outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad() # no gradients to calculate\n",
    "def predict(\n",
    "    model:nn.Module, \n",
    "    test_loader:torch.utils.data.DataLoader,\n",
    "    ):\n",
    "    '''\n",
    "    This function will iterate over the :code:`test_loader`\n",
    "    and return the outputs of the model\n",
    "    applied to the data.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "\n",
    "    - model: nn.Module:\n",
    "        The model used to make predictions.\n",
    "\n",
    "    - test_loader: torch.utils.data.DataLoader:\n",
    "        The data to predict on.\n",
    "    \n",
    "    Returns\n",
    "    ---------\n",
    "\n",
    "    - outputs: torch.Tensor:\n",
    "        The outputs of the model\n",
    "    \n",
    "    - labels: torch.Tensor:\n",
    "        The ground truth labels collected\n",
    "        from :code:`test_loader`.\n",
    "\n",
    "    '''\n",
    "    \n",
    "    # lists to contain the output data\n",
    "    targets = []\n",
    "    outputs = []\n",
    "\n",
    "    # adding model to GPU if available\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    # iterating over the test_loader with a progress bar\n",
    "    for input, target in tqdm.tqdm(test_loader, desc='Predicting'):\n",
    "        # ==== push data to GPU if available ====\n",
    "        input = input.to(device)\n",
    "        # ==== forward pass ====\n",
    "        output = model(input)\n",
    "        # ==== saving outputs and labels ====\n",
    "        outputs.append(output.cpu())\n",
    "        targets.append(target) # target was never pushed to GPU so remains on cpu\n",
    "    \n",
    "    # turning outputs into torch tensors instead of lists\n",
    "    outputs = torch.cat(outputs)\n",
    "    targets = torch.cat(targets)\n",
    "\n",
    "    model.to('cpu') # return the model to the CPU\n",
    "\n",
    "    return outputs, targets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will save the above function in the file `model_predictor.py` so that we can import it easily later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mlp_predictions, test_labels = predict(model, test_loader=test_loader)\n",
    "# finding the class of max probability\n",
    "_, test_mlp_predictions = torch.max(test_mlp_predictions, dim=1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mlp_accuracy = accuracy_score(test_dataset.targets.numpy(), test_mlp_predictions.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The accuracy of the pytorch MLP is {test_mlp_accuracy*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's see how well an sklearn MLP does:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neural_network import MLPClassifier\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to make numpy arrays from the same data used in the pytorch training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# since the data originated from pytorch dataloaders, we need to turn these into numpy \n",
    "# arrays for our sklearn model\n",
    "\n",
    "X_train, y_train = [], []\n",
    "for inputs, targets in train_loader:\n",
    "    X_train.append(inputs.numpy())\n",
    "    y_train.append(targets.numpy())\n",
    "\n",
    "X_val, y_val = [], []\n",
    "for inputs, targets in val_loader:\n",
    "    X_val.append(inputs.numpy())\n",
    "    y_val.append(targets.numpy())\n",
    "\n",
    "X_test, y_test = [], []\n",
    "for inputs, targets in test_loader:\n",
    "    X_test.append(inputs.numpy())\n",
    "    y_test.append(targets.numpy())\n",
    "\n",
    "X_train = np.vstack(X_train)\n",
    "y_train = np.concatenate(y_train)\n",
    "X_val = np.vstack(X_val)\n",
    "y_val = np.concatenate(y_val)\n",
    "X_test = np.vstack(X_test)\n",
    "y_test = np.concatenate(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# same as our pytorch model, but comes with extra performance options as default\n",
    "# also MLPClassifier does not support dropout\n",
    "mlp = MLPClassifier(\n",
    "    hidden_layer_sizes=[\n",
    "        16,\n",
    "        16,        \n",
    "        ],\n",
    "    activation='relu',\n",
    "    solver='adam',\n",
    "    batch_size=512,\n",
    "    learning_rate='constant',\n",
    "    learning_rate_init=0.01,\n",
    "    random_state=random_state.next()\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mlp_sk_predictions = mlp.predict(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mlp_sk_accuracy = accuracy_score(y_test, test_mlp_sk_predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The accuracy of the pytorch MLP is {test_mlp_sk_accuracy*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, Sklearn's model was arguably easier to implement, and performed better. However, when more complex datasets or models are being experimented with, Pytorch is required.\n",
    "\n",
    "Also, pytorch allows you to keep all datasets on disk and ensures that large datasets do not cause memory problems. Sklearn prefers data to be in memory. This is why it ran faster here, since converting the data from pytorch datasets to numpy arrays loaded the data into memory."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More Complex Neural Networks"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The below will take a long time to run, and you might find it more helpful to view the version of this notebook that is already run**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now look at an example of training a complex model on ECG data, located online. This will also allow us to provide an example on how to build custom datasets from online data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These more complicated examples require pytorch models and their flexibility:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will only work in the same directory as the file `model_trainer`\n",
    "from model_trainer import train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Within the `dataset.py` file in this directory, we have provided an example of loading a complex ECG dataset from online ([here, PTB-XL](https://physionet.org/content/ptb-xl/1.0.0/)) and loading it in a pytorch dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import PTB_XL, MemoryDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE:** The following code will download the ECG dataset and unpack it. This is quite a large dataset and might not work on google colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = PTB_XL(\n",
    "    data_path='./data/',\n",
    "    train=True,\n",
    "    sampling_rate=100,\n",
    "    binary=True,\n",
    "    )\n",
    "\n",
    "test_dataset = PTB_XL(\n",
    "    data_path='./data/',\n",
    "    train=False,\n",
    "    sampling_rate=100,\n",
    "    binary=True,\n",
    "    )\n",
    "\n",
    "feature_names = train_dataset.feature_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What does this data look like?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx_normal = np.argmax(train_dataset.targets == 0) # first normal ecg\n",
    "idx_abnormal = np.argmax(train_dataset.targets == 1) # first abnormal ecg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, _ = train_dataset[idx_normal]\n",
    "\n",
    "fig, axes = plt.subplots(3, inputs.shape[0]//3, figsize=(12,8))\n",
    "\n",
    "# getting the colours to make the plot look nicer!\n",
    "colors = plt.rcParams[\"axes.prop_cycle\"]()\n",
    "\n",
    "for nax, ax in enumerate(np.ravel(axes)):\n",
    "    ax.plot(inputs[nax,:], color=next(colors)[\"color\"])\n",
    "    ax.set_title(f'Feature {feature_names[nax]}')\n",
    "\n",
    "fig.suptitle('Features for Normal ECG Result')\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.3)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, _ = train_dataset[idx_abnormal]\n",
    "\n",
    "fig, axes = plt.subplots(3, inputs.shape[0]//3, figsize=(12,8))\n",
    "\n",
    "# getting the colours to make the plot look nicer!\n",
    "colors = plt.rcParams[\"axes.prop_cycle\"]()\n",
    "\n",
    "for nax, ax in enumerate(np.ravel(axes)):\n",
    "    ax.plot(inputs[nax,:], color=next(colors)[\"color\"])\n",
    "    ax.set_title(f'Feature {feature_names[nax]}')\n",
    "\n",
    "fig.suptitle('Features for Abnormal ECG Result')\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.3)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Within the `dataset.py` file, we have also given an example of a wrapper that can be built for a dataset to do interesting things! We have included a wrapper that loads all of the dataset into memory after the first epoch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = MemoryDataset(train_dataset, now=False)\n",
    "test_dataset = MemoryDataset(test_dataset, now=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now load a 1D ResNet model that will be used to train a model on this dataset. This is a complicated model, with skip connections and many layers. It is an example of how flexible pytorch can be:\n",
    "\n",
    "To view the code for this model, feel free to look through the file `resnet.py`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing a ResNet model from the file resnet.py\n",
    "from resnet import ResNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's train this ResNet model to make predictions on an ECG dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting up the data loaders\n",
    "# splitting the data into train and val\n",
    "n_train_set, n_val_set = [\n",
    "    int(0.75*len(train_dataset)), len(train_dataset)-int(0.75*len(train_dataset))\n",
    "    ]\n",
    "\n",
    "train_dataset_split, val_dataset_split = torch.utils.data.random_split(\n",
    "    train_dataset, lengths=[n_train_set, n_val_set],\n",
    "    )\n",
    "\n",
    "# the number of times to run through the whole dataset\n",
    "n_epochs = 40\n",
    "# the batch size for the data\n",
    "batch_size = 4196\n",
    "\n",
    "# put datasets in train loader\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset_split, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset_split, batch_size=batch_size, shuffle=True,\n",
    "    )\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    test_dataset, batch_size=batch_size, shuffle=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet = ResNet(\n",
    "    input_dim=1000,\n",
    "    input_channels=12,\n",
    "    kernel_size=15,\n",
    "    n_output=2,\n",
    "    dropout_rate=0.2,\n",
    "    )\n",
    "\n",
    "optimiser = torch.optim.Adam(\n",
    "    params=resnet.parameters(), \n",
    "    lr=0.001, \n",
    "    weight_decay=0.0001, \n",
    "    betas=(0.99, 0.999),\n",
    "    )\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the model\n",
    "resnet, (train_loss_dict, val_loss_dict) = train(\n",
    "    model=resnet,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    n_epochs=n_epochs,\n",
    "    optimiser=optimiser,\n",
    "    criterion=criterion\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives us the loss plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,4)) # plotting area\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    train_loss_dict['step'], train_loss_dict['loss'],\n",
    "    label='Train'\n",
    "    )\n",
    "\n",
    "# plot training loss\n",
    "ax.plot(\n",
    "    val_loss_dict['step'], val_loss_dict['loss'],\n",
    "    label='Val'\n",
    "    )\n",
    "\n",
    "# formatting\n",
    "ax.set_title('Training Loss')\n",
    "ax.set_ylabel('Loss')\n",
    "ax.set_xlabel('Steps')\n",
    "ax.legend()\n",
    "# show plot\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the function we defined earlier for \n",
    "# making predictions on a test data loader\n",
    "from model_predictor import predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_probabilities, test_targets = predict(resnet, test_loader=test_loader)\n",
    "# finding the class of max probability\n",
    "_, test_predictions = torch.max(test_probabilities, dim=1) \n",
    "test_probabilities = test_probabilities[:,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, roc_auc_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_resnet_accuracy = accuracy_score(\n",
    "    test_targets.numpy(), test_predictions.numpy()\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The accuracy of the pytorch ResNet is {test_resnet_accuracy*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_resnet_auc_roc = roc_auc_score(\n",
    "    test_targets.numpy(), test_probabilities.numpy()\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The area under the ROC of the pytorch ResNet is {test_resnet_auc_roc*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we trained a model to classify ECG results as normal or abnormal, and were able to get good accuracy and area under the ROC."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.8 (default, Apr 13 2021, 19:58:26) \n[GCC 7.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9e357fd710d72f0aa74bf0d2d79b507a37851bb5a96ddca272fe10421f421865"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}