W3_MLPs.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W3_MLPs",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernel": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/arashash/Deep-Learning-Course-Tutorials/blob/main/W3_MLPs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "JnoYx68VQcfU"
      },
      "source": [
        "# Week 3: Multi-layer Perceptrons (MLPs)\n",
        "\n",
        "**By Arash Ash, adapted from Neuromatch Academy [contents](https://github.com/NeuromatchAcademy/course-content-dl)**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "Wjc4c7axQcfY"
      },
      "source": [
        "---\n",
        "# Tutorial objectives\n",
        "\n",
        "In this tutorial, we will explore the Multi-layer Perceptrons (MLPs). MLPs are arguably one of the most tractable models (due to their flexibility) that we can use to study deep learning fundamentals. Here we will learn why they are: \n",
        "\n",
        "* similar to biological networks\n",
        "* good at function approximation\n",
        "* implemented the way they are in PyTorch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "VAdzxJAOQcfZ"
      },
      "source": [
        "---\n",
        "# Setup\n",
        "\n",
        "This is a GPU free notebook!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "4D-6_gBUQcfa"
      },
      "source": [
        "# Imports\n",
        "import random\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from tqdm.auto import tqdm\n",
        "from IPython.display import display\n",
        "from torch.utils.data import DataLoader, TensorDataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "execution": {},
        "id": "kPXGeaHaQcfa"
      },
      "source": [
        "# @title Figure settings\n",
        "import ipywidgets as widgets       # interactive display\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/content-creation/main/nma.mplstyle\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "execution": {},
        "id": "0vntx2KXQcfb"
      },
      "source": [
        "# @title Plotting functions\n",
        "\n",
        "def imshow(img):\n",
        "  img = img / 2 + 0.5     # unnormalize\n",
        "  npimg = img.numpy()\n",
        "  plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
        "  plt.axis(False)\n",
        "  plt.show()\n",
        "\n",
        "\n",
        "def plot_function_approximation(x, relu_acts, y_hat):\n",
        "  fig, axes = plt.subplots(2, 1)\n",
        "\n",
        "  # Plot ReLU Activations\n",
        "  axes[0].plot(x, relu_acts.T);\n",
        "  axes[0].set(xlabel='x',\n",
        "              ylabel='Activation',\n",
        "              title='ReLU Activations - Basis Functions')\n",
        "  labels = [f\"ReLU {i + 1}\" for i in range(relu_acts.shape[0])]\n",
        "  axes[0].legend(labels, ncol = 2)\n",
        "\n",
        "  # Plot function approximation\n",
        "  axes[1].plot(x, torch.sin(x), label='truth')\n",
        "  axes[1].plot(x, y_hat, label='estimated')\n",
        "  axes[1].legend()\n",
        "  axes[1].set(xlabel='x',\n",
        "              ylabel='y(x)',\n",
        "              title='Function Approximation')\n",
        "\n",
        "  plt.tight_layout()\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "execution": {},
        "id": "n_rnXdw3Qcfc"
      },
      "source": [
        "# @title Set random seed\n",
        "\n",
        "# @markdown Executing `set_seed(seed=seed)` you are setting the seed\n",
        "\n",
        "# for DL its critical to set the random seed so that students can have a\n",
        "# baseline to compare their results to expected results.\n",
        "# Read more here: https://pytorch.org/docs/stable/notes/randomness.html\n",
        "\n",
        "# Call `set_seed` function in the exercises to ensure reproducibility.\n",
        "import random\n",
        "import torch\n",
        "\n",
        "def set_seed(seed=None, seed_torch=True):\n",
        "  if seed is None:\n",
        "    seed = np.random.choice(2 ** 32)\n",
        "  random.seed(seed)\n",
        "  np.random.seed(seed)\n",
        "  if seed_torch:\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.backends.cudnn.benchmark = False\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "\n",
        "  print(f'Random seed {seed} has been set.')\n",
        "\n",
        "\n",
        "# In case that `DataLoader` is used\n",
        "def seed_worker(worker_id):\n",
        "  worker_seed = torch.initial_seed() % 2**32\n",
        "  np.random.seed(worker_seed)\n",
        "  random.seed(worker_seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "execution": {},
        "id": "mvwTdFogQcfc"
      },
      "source": [
        "# @title Set device (GPU or CPU). Execute `set_device()`\n",
        "# especially if torch modules used.\n",
        "\n",
        "# inform the user if the notebook uses GPU or CPU.\n",
        "\n",
        "def set_device():\n",
        "  device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "  if device != \"cuda\":\n",
        "    print(\"GPU is not enabled in this notebook. \\n\"\n",
        "          \"If you want to enable it, in the menu under `Runtime` -> \\n\"\n",
        "          \"`Hardware accelerator.` and select `GPU` from the dropdown menu\")\n",
        "  else:\n",
        "    print(\"GPU is enabled in this notebook. \\n\"\n",
        "          \"If you want to disable it, in the menu under `Runtime` -> \\n\"\n",
        "          \"`Hardware accelerator.` and select `None` from the dropdown menu\")\n",
        "\n",
        "  return device"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "xuZbmxq9Qcfd"
      },
      "source": [
        "SEED = 2021\n",
        "set_seed(seed=SEED)\n",
        "DEVICE = set_device()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "vXN33V3lQcfd"
      },
      "source": [
        "---\n",
        "# Section 1: MLPs and Function Approximation Theorem"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "hwdvLFImQcfd",
        "cellView": "form"
      },
      "source": [
        "# @title Video: What a single Hidden Layer MLP can do?\n",
        "from ipywidgets import widgets\n",
        "\n",
        "out1 = widgets.Output()\n",
        "with out1:\n",
        "  from IPython.display import YouTubeVideo\n",
        "  video = YouTubeVideo(id=f\"YfJMyiwKvGQ\", width=854, height=480, fs=1, rel=0)\n",
        "  print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "  display(video)\n",
        "\n",
        "display(out1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "dfA1gnaVQcfe"
      },
      "source": [
        "## Coding Exercise 1: Function approximation with ReLU\n",
        "We learned that one hidden layer MLPs are enough to approximate any smooth function! Now let's manually fit a sine function using ReLU activation. \n",
        "\n",
        "We will approximate the sine function using a linear combination (a weighted sum) of ReLUs with slope 1. We need to determine the bias terms (which determines where the ReLU inflection point from 0 to linear occurs) and how to weight each ReLU. The idea is to set the weights iteratively so that the slope changes in the new sample's direction.\n",
        "\n",
        "First, we generate our \"training data\" from a sine function using `torch.sine` function.\n",
        "\n",
        "```python\n",
        ">>> import torch\n",
        ">>> torch.manual_seed(2021)\n",
        "<torch._C.Generator object at 0x7f8734c83830>\n",
        ">>> a = torch.randn(5)\n",
        ">>> print(a)\n",
        "tensor([ 2.2871,  0.6413, -0.8615, -0.3649, -0.6931])\n",
        ">>> torch.sin(a)\n",
        "tensor([ 0.7542,  0.5983, -0.7588, -0.3569, -0.6389])\n",
        "```\n",
        "\n",
        "These are the points we will use to learn how to approximate the function. We have 10 training data points so we will have 9 ReLUs (we don't need a ReLU for the last data point as we don't have anything to the right of it to model). \n",
        "\n",
        "We first need to figure out the bias term for each ReLU and compute the activation of each ReLU where:\n",
        "\n",
        "\\begin{equation}\n",
        "y(x) = \\text{max}(0,x+b)\n",
        "\\end{equation}\n",
        "\n",
        "We then need to figure out the correct weights on each ReLU so the linear combination approximates the desired function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "8HNURE-aQcfe"
      },
      "source": [
        "def approximate_function(x_train, y_train):\n",
        "  # Number of relus\n",
        "  n_relus = x_train.shape[0] - 1\n",
        "\n",
        "  # x axis points (more than x train)\n",
        "  x = torch.linspace(torch.min(x_train), torch.max(x_train), 1000)\n",
        "\n",
        "  ## COMPUTE RELU ACTIVATIONS\n",
        "  # First determine what bias terms should be for each of `n_relus` ReLUs\n",
        "  b = - x_train[:-1]\n",
        "\n",
        "  # Compute ReLU activations for each point along the x axis (x)\n",
        "  relu_acts = torch.zeros((n_relus, x.shape[0]))\n",
        "\n",
        "  for i_relu in range(n_relus):\n",
        "    relu_acts[i_relu, :] = torch.relu(x + b[i_relu])\n",
        "\n",
        "  ## COMBINE RELU ACTIVATIONS\n",
        "  # Set up weights for weighted sum of ReLUs\n",
        "  combination_weights = torch.zeros((n_relus, ))\n",
        "\n",
        "  ####################################################################\n",
        "  # Fill in missing code below (...),\n",
        "  # then remove or comment the line below to test your function\n",
        "  raise NotImplementedError(\"Complete approximate_function!\")\n",
        "  ####################################################################\n",
        "  prev_slope = 0\n",
        "  for i in range(n_relus):\n",
        "    # Figure out weights on each ReLU\n",
        "    delta_x = x_train[i+1] - x_train[i]\n",
        "    slope = ...\n",
        "    combination_weights[i] = slope - prev_slope\n",
        "    prev_slope = slope\n",
        "\n",
        "  # Get output of weighted sum of ReLU activations for every point along x axis\n",
        "  y_hat = ...\n",
        "\n",
        "  return y_hat, relu_acts, x\n",
        "\n",
        "# Make training data from sine function\n",
        "N_train = 10\n",
        "x_train = torch.linspace(0, 2*np.pi, N_train).view(-1, 1)\n",
        "y_train = torch.sin(x_train)\n",
        "\n",
        "## Uncomment the lines below to test your function approximation\n",
        "# y_hat, relu_acts, x = approximate_function(x_train, y_train)\n",
        "# plot_function_approximation(x, relu_acts, y_hat)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "T65X1-EfQcff"
      },
      "source": [
        "---\n",
        "# Section 2: MLPs in Pytorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "jpHSTv6UQcff",
        "cellView": "form"
      },
      "source": [
        "# @title Video: Math to Pytorch Implementation\n",
        "from ipywidgets import widgets\n",
        "\n",
        "out1 = widgets.Output()\n",
        "with out1:\n",
        "  from IPython.display import YouTubeVideo\n",
        "  video = YouTubeVideo(id=f\"Dd-CdIOJWnk\", width=854, height=480, fs=1, rel=0)\n",
        "  print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "  display(video)\n",
        "\n",
        "display(out1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "4CW5YsnEQcfg"
      },
      "source": [
        "In the previous segment, we implemented a function to approximate any smooth function using MLPs. We saw that using Lipschitz continuity; we can prove that our approximation is mathematically correct. MLPs are fascinating, but before we get into the details on designing them, let's familiarize ourselves with some basic terminology of MLPs- layer, neuron, depth, width, weight, bias, and activation function. Armed with these ideas, we can now design an MLP given its input, hidden layers, and output size."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "vslYnMk0Qcfg"
      },
      "source": [
        "## Coding Exercise 2: Implement a general-purpose MLP in Pytorch\n",
        "The objective is to design an MLP with these properties:\n",
        "* works with any input (1D, 2D, etc.)\n",
        "* construct any number of given hidden layers using `nn.Sequential()` and `add_module()` function\n",
        "* use the same given activation function (i.e., [Leaky ReLU](https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html)) in all hidden layers.\n",
        "\n",
        "**Leaky ReLU** is described by the following mathematical formula:\n",
        "\n",
        "\\begin{equation}\n",
        "\\text{LeakyReLU}(x) = \\text{max}(0,x) + \\text{negative_slope} \\cdot \\text{min}(0, x) =\n",
        "\\left\\{\n",
        "  \\begin{array}{ll}\n",
        "  x & ,\\; \\text{if} \\; x \\ge 0 \\\\\n",
        "  \\text{negative_slope} \\cdot x & ,\\;  \\text{otherwise}\n",
        "  \\end{array}\n",
        "\\right.\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "DqQxSsLbQcfg"
      },
      "source": [
        "class Net(nn.Module):\n",
        "  def __init__(self, actv, input_feature_num, hidden_unit_nums, output_feature_num):\n",
        "    super(Net, self).__init__()\n",
        "    self.input_feature_num = input_feature_num # save the input size for reshapinng later\n",
        "    self.mlp = nn.Sequential() # Initialize layers of MLP\n",
        "\n",
        "    in_num = input_feature_num # initialize the temporary input feature to each layer\n",
        "    for i in range(len(hidden_unit_nums)): # Loop over layers and create each one\n",
        "\n",
        "      ####################################################################\n",
        "      # Fill in missing code below (...),\n",
        "      # then remove or comment the line below to test your function\n",
        "      raise NotImplementedError(\"Create MLP Layer\")\n",
        "      ####################################################################\n",
        "\n",
        "      out_num = hidden_unit_nums[i] # assign the current layer hidden unit from list\n",
        "      layer = ... # use nn.Linear to define the layer\n",
        "      in_num = out_num # assign next layer input using current layer output\n",
        "      self.mlp.add_module('Linear_%d'%i, layer) # append layer to the model with a name\n",
        "\n",
        "      actv_layer = eval('nn.%s'%actv) # Assign activation function (eval allows us to instantiate object from string)\n",
        "      self.mlp.add_module('Activation_%d'%i, actv_layer) # append activation to the model with a name\n",
        "\n",
        "    out_layer = nn.Linear(in_num, output_feature_num) # Create final layer\n",
        "    self.mlp.add_module('Output_Linear', out_layer) # append the final layer\n",
        "\n",
        "  def forward(self, x):\n",
        "    # reshape inputs to (batch_size, input_feature_num)\n",
        "    # just in case the input vector is not 2D, like an image!\n",
        "    x = x.view(-1, self.input_feature_num)\n",
        "\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Run MLP model\")\n",
        "    ####################################################################\n",
        "\n",
        "    logits = ... # forward pass of MLP\n",
        "    return logits\n",
        "\n",
        "\n",
        "input = torch.zeros((100, 2))\n",
        "## Uncomment below to create network and test it on input\n",
        "# net = Net(actv='LeakyReLU(0.1)', input_feature_num=2, hidden_unit_nums=[100, 10, 5], output_feature_num=1).to(DEVICE)\n",
        "# y = net(input.to(DEVICE))\n",
        "# print(f'The output shape is {y.shape} for an input of shape {input.shape}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "r2DURIYBQcfh"
      },
      "source": [
        "```\n",
        "The output shape is torch.Size([100, 1]) for an input of shape torch.Size([100, 2])\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "4JgZ_Tn_Qcfl"
      },
      "source": [
        "---\n",
        "# Section 3: Neuron Physiology and Motivation to Deep Learning\n",
        "\n",
        "This section will motivate one of the most popular nonlinearities in deep learning, the ReLU nonlinearity, by starting from the biophysics of neurons and obtaining the ReLU nonlinearity through a sequence of approximations. We will also show that neuronal biophysics sets a time scale for signal propagation speed through the brain. This time scale implies that neural circuits underlying fast perceptual and motor processing in the brain may not be excessively deep. \n",
        "\n",
        "This biological motivation for deep learning is not strictly necessary to follow the rest of this course, so this section can be safely skipped and **is a bonus section**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "H5tgF2LIQcfl",
        "cellView": "form"
      },
      "source": [
        "# @title Video: Artificial and Biological Neuron Model and importance of the ReLU activation\n",
        "from ipywidgets import widgets\n",
        "\n",
        "out1 = widgets.Output()\n",
        "with out1:\n",
        "  from IPython.display import YouTubeVideo\n",
        "  video = YouTubeVideo(id=f\"FRcCtlqtK7Q\", width=854, height=480, fs=1, rel=0)\n",
        "  print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "  display(video)\n",
        "\n",
        "display(out1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "uyueo5NvQcfl"
      },
      "source": [
        "## Leaky Integrate-and-fire (LIF)\n",
        "The basic idea of LIF neuron was proposed in 1907 by Louis Édouard Lapicque, long before we understood the electrophysiology of a neuron (see a translation of [Lapicque's paper](https://pubmed.ncbi.nlm.nih.gov/17968583/) ). More details of the model can be found in the book [**Theoretical neuroscience**](http://www.gatsby.ucl.ac.uk/~dayan/book/) by Peter Dayan and Laurence F. Abbott.\n",
        "\n",
        "The model dynamics is defined with the following formula,\n",
        "\n",
        "\n",
        "\\begin{equation}\n",
        "\\frac{d V}{d t}=\\left\\{\\begin{array}{cc}\n",
        "\\frac{1}{C}\\left(-\\frac{V}{R}+I \\right) & t>t_{r e s t} \\\\\n",
        "0 & \\text { otherwise }\n",
        "\\end{array}\\right.\n",
        "\\end{equation}\n",
        "\n",
        "\n",
        "Note that $V$, $C$, and $R$ are the membrane voltage, capacitance, and resitance of the neuron respectively and $-\\frac{V}{R}$ is the leakage current. When $I$ is sufficiently strong such that $V$ reaches a certain threshold value $V_{\\rm th}$, it momentarily spikes and then $V$ is reset to $V_{\\rm reset}< V_{\\rm th}$, and voltage stays at $V_{\\rm reset}$ for $\\tau_{\\rm ref}$ ms, mimicking the refractoriness of the neuron during an action potential (note that $V_{\\rm reset}$ and $\\tau_{\\rm ref}$ is assumed to be zero in the lecture):\n",
        "\n",
        "\n",
        "\\begin{eqnarray}\n",
        "V(t)=V_{\\rm reset} \\text{  for } t\\in(t_{\\text{sp}}, t_{\\text{sp}} + \\tau_{\\text{ref}}]\n",
        "\\end{eqnarray}\n",
        "\n",
        "\n",
        "where $t_{\\rm sp}$ is the spike time when $V(t)$ just exceeded $V_{\\rm th}$.\n",
        "\n",
        "Thus, the LIF model captures the facts that a neuron:\n",
        "- performs spatial and temporal integration of synaptic inputs \n",
        "- generates a spike when the voltage reaches a certain threshold\n",
        "- goes refractory during the action potential\n",
        "- has a leaky membrane \n",
        "\n",
        "For in-depth content on computational models of neurons, follow the [NMA](https://www.neuromatchacademy.org/) tutorial 1 of *Biological Neuron Models*. Specifically, for NMA-CN 2021 follow this [Tutorial](https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/W2D3_BiologicalNeuronModels/W2D3_Tutorial1.ipynb).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "e916QP9zQcfm"
      },
      "source": [
        "## Simulating an LIF Neuron\n",
        "\n",
        "In the cell below is given a function for LIF neuron model with it's arguments described.\n",
        "\n",
        "Note that we will use Euler's method to make a numerical approximation to a derivative. Hence we will use the following implementation of the model dynamics,\n",
        "\n",
        "\\begin{equation}\n",
        "V_n=\\left\\{\\begin{array}{cc}\n",
        "V_{n-1} + \\frac{1}{C}\\left(-\\frac{V}{R}+I \\right) \\Delta t & t>t_{r e s t} \\\\\n",
        "0 & \\text { otherwise }\n",
        "\\end{array}\\right.\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "9t9HhHRyQcfm"
      },
      "source": [
        "def run_LIF(I, T=50, dt=0.1, t_ref=10,\n",
        "            Rm=1, Cm=10, Vth=1, V_spike=0.5):\n",
        "  \"\"\"\n",
        "  Simulate the LIF dynamics with external input current\n",
        "\n",
        "  Args:\n",
        "    I          : input current (mA)\n",
        "    T          : total time to simulate (msec)\n",
        "    dt         : simulation time step (msec)\n",
        "    t_ref      : refractory period (msec)\n",
        "    Rm         : resistance (kOhm)\n",
        "    Cm         : capacitance (uF)\n",
        "    Vth        : spike threshold (V)\n",
        "    V_spike    : spike delta (V)\n",
        "\n",
        "  Returns:\n",
        "    time       : time points\n",
        "    Vm         : membrane potentials\n",
        "  \"\"\"\n",
        "\n",
        "  # Set up array of time steps\n",
        "  time = torch.arange(0, T+dt, dt)\n",
        "\n",
        "  # Set up array for tracking Vm\n",
        "  Vm = torch.zeros(len(time))\n",
        "\n",
        "  # Iterate over each time step\n",
        "  t_rest = 0\n",
        "  for i, t in enumerate(time):\n",
        "\n",
        "    # If t is after refractory period\n",
        "    if t > t_rest:\n",
        "      Vm[i] = Vm[i-1] + 1/Cm*(-Vm[i-1]/Rm + I)  * dt\n",
        "\n",
        "    # If Vm is over the threshold\n",
        "    if Vm[i] >= Vth:\n",
        "\n",
        "      # Increase volatage by change due to spike\n",
        "      Vm[i] += V_spike\n",
        "\n",
        "      # Set up new refactory period\n",
        "      t_rest = t + t_ref\n",
        "\n",
        "  return time, Vm\n",
        "\n",
        "\n",
        "sim_time, Vm = run_LIF(1.5)\n",
        "# Plot the membrane voltage across time\n",
        "plt.plot(sim_time, Vm)\n",
        "plt.title('LIF Neuron Output')\n",
        "plt.ylabel('Membrane Potential (V)')\n",
        "plt.xlabel('Time (msec)')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "UPWOTMRpQcfm"
      },
      "source": [
        "### Interactive Demo: Neuron's transfer function explorer for different $R_m$ and $t_{ref}$\n",
        "We know that real neurons communicate by modulating the spike count meaning that more input current causes a neuron to spike more often. Therefore, to find an input-output relationship, it makes sense to characterize their spike count as a function of input current. This is called the neuron's input-output transfer function. Let's plot the neuron's transfer function and see how it changes with respect to the **membrane resistance** and **refractory time**? "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {},
        "id": "Fy075zCrQcfm",
        "cellView": "form"
      },
      "source": [
        "# @title\n",
        "\n",
        "# @markdown Make sure you execute this cell to enable the widget!\n",
        "my_layout = widgets.Layout()\n",
        "\n",
        "@widgets.interact(Rm=widgets.FloatSlider(1., min=1, max=100.,\n",
        "                                         step=0.1, layout=my_layout),\n",
        "                  t_ref=widgets.FloatSlider(1., min=1, max=100.,\n",
        "                                            step=0.1, layout=my_layout)\n",
        "                  )\n",
        "\n",
        "\n",
        "def plot_IF_curve(Rm, t_ref):\n",
        "  T = 1000 # total time to simulate (msec)\n",
        "  dt = 1 # simulation time step (msec)\n",
        "  Vth = 1 # spike threshold (V)\n",
        "  Is_max = 2\n",
        "  Is = torch.linspace(0, Is_max, 10)\n",
        "  spike_counts = []\n",
        "  for I in Is:\n",
        "    _, Vm = run_LIF(I, T=T, dt=dt, Vth=Vth, Rm=Rm, t_ref=t_ref)\n",
        "    spike_counts += [torch.sum(Vm > Vth)]\n",
        "\n",
        "  plt.plot(Is, spike_counts)\n",
        "  plt.title('LIF Neuron: Transfer Function')\n",
        "  plt.ylabel('Spike count')\n",
        "  plt.xlabel('I (mA)')\n",
        "  plt.xlim(0, Is_max)\n",
        "  plt.ylim(0, 80)\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "execution": {},
        "id": "QLecvjRsQcfm"
      },
      "source": [
        "### Think!: Real and Artificial neuron similarities\n",
        "\n",
        "What happens at infinite membrane resistance ($R_m$) and small refactory time ($t_{ref}$)? Why?\n",
        "\n",
        "Take 10 mins to discuss the similarity between a real neuron and an artificial one with your pod."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_twSTgsWYF9O"
      },
      "source": [
        "---\n",
        "# Summary\n",
        "\n",
        "In this tutorial we have explored the Multi-leayer Perceptrons (MLPs). More specifically, we have discuss the similarities of artificial and biological neural networks (for more information see the Bonus section), we have learned the Universal Approximation Theorem, and we have implemented MLPs in PyTorch."
      ]
    }
  ]
}