Neural_ODE_Demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Neural ODE Demo.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyMrbpnO8gchdb8T8DIKVfuv",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/GitTeaching/Predicting-using-Neural-ODE/blob/main/Neural_ODE_Demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ggAy3tbELaqq"
      },
      "source": [
        "## Neural Ordinary Differential Equations based on : https://arxiv.org/abs/1806.07366\n",
        "\n",
        "#### - Dataset used: ECG Heartbeat Categorization Dataset - https://www.kaggle.com/shayanfazeli/heartbeat\n",
        "#### - Code Sources from : \n",
        "\n",
        "https://medium.com/analytics-vidhya/intro-to-neural-odes-part-3-the-basics-9697b3bd1946\n",
        "\n",
        "https://github.com/abaietto/neural_ode_classification/blob/master/ECG_Classification.ipynb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vH0UK6Z9E3fH"
      },
      "source": [
        "## Load Data from Kaggle"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EQcHmXX6CfH-",
        "outputId": "e823be64-b418-4a60-a344-08dc7b7157ac",
        "colab": {
          "resources": {
            "http://localhost:8080/nbextensions/google.colab/files.js": {
              "data": "Ly8gQ29weXJpZ2h0IDIwMTcgR29vZ2xlIExMQwovLwovLyBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgIkxpY2Vuc2UiKTsKLy8geW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLgovLyBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXQKLy8KLy8gICAgICBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjAKLy8KLy8gVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZQovLyBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiAiQVMgSVMiIEJBU0lTLAovLyBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC4KLy8gU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZAovLyBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS4KCi8qKgogKiBAZmlsZW92ZXJ2aWV3IEhlbHBlcnMgZm9yIGdvb2dsZS5jb2xhYiBQeXRob24gbW9kdWxlLgogKi8KKGZ1bmN0aW9uKHNjb3BlKSB7CmZ1bmN0aW9uIHNwYW4odGV4dCwgc3R5bGVBdHRyaWJ1dGVzID0ge30pIHsKICBjb25zdCBlbGVtZW50ID0gZG9jdW1lbnQuY3JlYXRlRWxlbWVudCgnc3BhbicpOwogIGVsZW1lbnQudGV4dENvbnRlbnQgPSB0ZXh0OwogIGZvciAoY29uc3Qga2V5IG9mIE9iamVjdC5rZXlzKHN0eWxlQXR0cmlidXRlcykpIHsKICAgIGVsZW1lbnQuc3R5bGVba2V5XSA9IHN0eWxlQXR0cmlidXRlc1trZXldOwogIH0KICByZXR1cm4gZWxlbWVudDsKfQoKLy8gTWF4IG51bWJlciBvZiBieXRlcyB3aGljaCB3aWxsIGJlIHVwbG9hZGVkIGF0IGEgdGltZS4KY29uc3QgTUFYX1BBWUxPQURfU0laRSA9IDEwMCAqIDEwMjQ7CgpmdW5jdGlvbiBfdXBsb2FkRmlsZXMoaW5wdXRJZCwgb3V0cHV0SWQpIHsKICBjb25zdCBzdGVwcyA9IHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCk7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICAvLyBDYWNoZSBzdGVwcyBvbiB0aGUgb3V0cHV0RWxlbWVudCB0byBtYWtlIGl0IGF2YWlsYWJsZSBmb3IgdGhlIG5leHQgY2FsbAogIC8vIHRvIHVwbG9hZEZpbGVzQ29udGludWUgZnJvbSBQeXRob24uCiAgb3V0cHV0RWxlbWVudC5zdGVwcyA9IHN0ZXBzOwoKICByZXR1cm4gX3VwbG9hZEZpbGVzQ29udGludWUob3V0cHV0SWQpOwp9CgovLyBUaGlzIGlzIHJvdWdobHkgYW4gYXN5bmMgZ2VuZXJhdG9yIChub3Qgc3VwcG9ydGVkIGluIHRoZSBicm93c2VyIHlldCksCi8vIHdoZXJlIHRoZXJlIGFyZSBtdWx0aXBsZSBhc3luY2hyb25vdXMgc3RlcHMgYW5kIHRoZSBQeXRob24gc2lkZSBpcyBnb2luZwovLyB0byBwb2xsIGZvciBjb21wbGV0aW9uIG9mIGVhY2ggc3RlcC4KLy8gVGhpcyB1c2VzIGEgUHJvbWlzZSB0byBibG9jayB0aGUgcHl0aG9uIHNpZGUgb24gY29tcGxldGlvbiBvZiBlYWNoIHN0ZXAsCi8vIHRoZW4gcGFzc2VzIHRoZSByZXN1bHQgb2YgdGhlIHByZXZpb3VzIHN0ZXAgYXMgdGhlIGlucHV0IHRvIHRoZSBuZXh0IHN0ZXAuCmZ1bmN0aW9uIF91cGxvYWRGaWxlc0NvbnRpbnVlKG91dHB1dElkKSB7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICBjb25zdCBzdGVwcyA9IG91dHB1dEVsZW1lbnQuc3RlcHM7CgogIGNvbnN0IG5leHQgPSBzdGVwcy5uZXh0KG91dHB1dEVsZW1lbnQubGFzdFByb21pc2VWYWx1ZSk7CiAgcmV0dXJuIFByb21pc2UucmVzb2x2ZShuZXh0LnZhbHVlLnByb21pc2UpLnRoZW4oKHZhbHVlKSA9PiB7CiAgICAvLyBDYWNoZSB0aGUgbGFzdCBwcm9taXNlIHZhbHVlIHRvIG1ha2UgaXQgYXZhaWxhYmxlIHRvIHRoZSBuZXh0CiAgICAvLyBzdGVwIG9mIHRoZSBnZW5lcmF0b3IuCiAgICBvdXRwdXRFbGVtZW50Lmxhc3RQcm9taXNlVmFsdWUgPSB2YWx1ZTsKICAgIHJldHVybiBuZXh0LnZhbHVlLnJlc3BvbnNlOwogIH0pOwp9CgovKioKICogR2VuZXJhdG9yIGZ1bmN0aW9uIHdoaWNoIGlzIGNhbGxlZCBiZXR3ZWVuIGVhY2ggYXN5bmMgc3RlcCBvZiB0aGUgdXBsb2FkCiAqIHByb2Nlc3MuCiAqIEBwYXJhbSB7c3RyaW5nfSBpbnB1dElkIEVsZW1lbnQgSUQgb2YgdGhlIGlucHV0IGZpbGUgcGlja2VyIGVsZW1lbnQuCiAqIEBwYXJhbSB7c3RyaW5nfSBvdXRwdXRJZCBFbGVtZW50IElEIG9mIHRoZSBvdXRwdXQgZGlzcGxheS4KICogQHJldHVybiB7IUl0ZXJhYmxlPCFPYmplY3Q+fSBJdGVyYWJsZSBvZiBuZXh0IHN0ZXBzLgogKi8KZnVuY3Rpb24qIHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCkgewogIGNvbnN0IGlucHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKGlucHV0SWQpOwogIGlucHV0RWxlbWVudC5kaXNhYmxlZCA9IGZhbHNlOwoKICBjb25zdCBvdXRwdXRFbGVtZW50ID0gZG9jdW1lbnQuZ2V0RWxlbWVudEJ5SWQob3V0cHV0SWQpOwogIG91dHB1dEVsZW1lbnQuaW5uZXJIVE1MID0gJyc7CgogIGNvbnN0IHBpY2tlZFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgaW5wdXRFbGVtZW50LmFkZEV2ZW50TGlzdGVuZXIoJ2NoYW5nZScsIChlKSA9PiB7CiAgICAgIHJlc29sdmUoZS50YXJnZXQuZmlsZXMpOwogICAgfSk7CiAgfSk7CgogIGNvbnN0IGNhbmNlbCA9IGRvY3VtZW50LmNyZWF0ZUVsZW1lbnQoJ2J1dHRvbicpOwogIGlucHV0RWxlbWVudC5wYXJlbnRFbGVtZW50LmFwcGVuZENoaWxkKGNhbmNlbCk7CiAgY2FuY2VsLnRleHRDb250ZW50ID0gJ0NhbmNlbCB1cGxvYWQnOwogIGNvbnN0IGNhbmNlbFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgY2FuY2VsLm9uY2xpY2sgPSAoKSA9PiB7CiAgICAgIHJlc29sdmUobnVsbCk7CiAgICB9OwogIH0pOwoKICAvLyBXYWl0IGZvciB0aGUgdXNlciB0byBwaWNrIHRoZSBmaWxlcy4KICBjb25zdCBmaWxlcyA9IHlpZWxkIHsKICAgIHByb21pc2U6IFByb21pc2UucmFjZShbcGlja2VkUHJvbWlzZSwgY2FuY2VsUHJvbWlzZV0pLAogICAgcmVzcG9uc2U6IHsKICAgICAgYWN0aW9uOiAnc3RhcnRpbmcnLAogICAgfQogIH07CgogIGNhbmNlbC5yZW1vdmUoKTsKCiAgLy8gRGlzYWJsZSB0aGUgaW5wdXQgZWxlbWVudCBzaW5jZSBmdXJ0aGVyIHBpY2tzIGFyZSBub3QgYWxsb3dlZC4KICBpbnB1dEVsZW1lbnQuZGlzYWJsZWQgPSB0cnVlOwoKICBpZiAoIWZpbGVzKSB7CiAgICByZXR1cm4gewogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbXBsZXRlJywKICAgICAgfQogICAgfTsKICB9CgogIGZvciAoY29uc3QgZmlsZSBvZiBmaWxlcykgewogICAgY29uc3QgbGkgPSBkb2N1bWVudC5jcmVhdGVFbGVtZW50KCdsaScpOwogICAgbGkuYXBwZW5kKHNwYW4oZmlsZS5uYW1lLCB7Zm9udFdlaWdodDogJ2JvbGQnfSkpOwogICAgbGkuYXBwZW5kKHNwYW4oCiAgICAgICAgYCgke2ZpbGUudHlwZSB8fCAnbi9hJ30pIC0gJHtmaWxlLnNpemV9IGJ5dGVzLCBgICsKICAgICAgICBgbGFzdCBtb2RpZmllZDogJHsKICAgICAgICAgICAgZmlsZS5sYXN0TW9kaWZpZWREYXRlID8gZmlsZS5sYXN0TW9kaWZpZWREYXRlLnRvTG9jYWxlRGF0ZVN0cmluZygpIDoKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJ24vYSd9IC0gYCkpOwogICAgY29uc3QgcGVyY2VudCA9IHNwYW4oJzAlIGRvbmUnKTsKICAgIGxpLmFwcGVuZENoaWxkKHBlcmNlbnQpOwoKICAgIG91dHB1dEVsZW1lbnQuYXBwZW5kQ2hpbGQobGkpOwoKICAgIGNvbnN0IGZpbGVEYXRhUHJvbWlzZSA9IG5ldyBQcm9taXNlKChyZXNvbHZlKSA9PiB7CiAgICAgIGNvbnN0IHJlYWRlciA9IG5ldyBGaWxlUmVhZGVyKCk7CiAgICAgIHJlYWRlci5vbmxvYWQgPSAoZSkgPT4gewogICAgICAgIHJlc29sdmUoZS50YXJnZXQucmVzdWx0KTsKICAgICAgfTsKICAgICAgcmVhZGVyLnJlYWRBc0FycmF5QnVmZmVyKGZpbGUpOwogICAgfSk7CiAgICAvLyBXYWl0IGZvciB0aGUgZGF0YSB0byBiZSByZWFkeS4KICAgIGxldCBmaWxlRGF0YSA9IHlpZWxkIHsKICAgICAgcHJvbWlzZTogZmlsZURhdGFQcm9taXNlLAogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbnRpbnVlJywKICAgICAgfQogICAgfTsKCiAgICAvLyBVc2UgYSBjaHVua2VkIHNlbmRpbmcgdG8gYXZvaWQgbWVzc2FnZSBzaXplIGxpbWl0cy4gU2VlIGIvNjIxMTU2NjAuCiAgICBsZXQgcG9zaXRpb24gPSAwOwogICAgd2hpbGUgKHBvc2l0aW9uIDwgZmlsZURhdGEuYnl0ZUxlbmd0aCkgewogICAgICBjb25zdCBsZW5ndGggPSBNYXRoLm1pbihmaWxlRGF0YS5ieXRlTGVuZ3RoIC0gcG9zaXRpb24sIE1BWF9QQVlMT0FEX1NJWkUpOwogICAgICBjb25zdCBjaHVuayA9IG5ldyBVaW50OEFycmF5KGZpbGVEYXRhLCBwb3NpdGlvbiwgbGVuZ3RoKTsKICAgICAgcG9zaXRpb24gKz0gbGVuZ3RoOwoKICAgICAgY29uc3QgYmFzZTY0ID0gYnRvYShTdHJpbmcuZnJvbUNoYXJDb2RlLmFwcGx5KG51bGwsIGNodW5rKSk7CiAgICAgIHlpZWxkIHsKICAgICAgICByZXNwb25zZTogewogICAgICAgICAgYWN0aW9uOiAnYXBwZW5kJywKICAgICAgICAgIGZpbGU6IGZpbGUubmFtZSwKICAgICAgICAgIGRhdGE6IGJhc2U2NCwKICAgICAgICB9LAogICAgICB9OwogICAgICBwZXJjZW50LnRleHRDb250ZW50ID0KICAgICAgICAgIGAke01hdGgucm91bmQoKHBvc2l0aW9uIC8gZmlsZURhdGEuYnl0ZUxlbmd0aCkgKiAxMDApfSUgZG9uZWA7CiAgICB9CiAgfQoKICAvLyBBbGwgZG9uZS4KICB5aWVsZCB7CiAgICByZXNwb25zZTogewogICAgICBhY3Rpb246ICdjb21wbGV0ZScsCiAgICB9CiAgfTsKfQoKc2NvcGUuZ29vZ2xlID0gc2NvcGUuZ29vZ2xlIHx8IHt9OwpzY29wZS5nb29nbGUuY29sYWIgPSBzY29wZS5nb29nbGUuY29sYWIgfHwge307CnNjb3BlLmdvb2dsZS5jb2xhYi5fZmlsZXMgPSB7CiAgX3VwbG9hZEZpbGVzLAogIF91cGxvYWRGaWxlc0NvbnRpbnVlLAp9Owp9KShzZWxmKTsK",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/javascript"
                ]
              ],
              "status": 200,
              "status_text": ""
            }
          },
          "base_uri": "https://localhost:8080/",
          "height": 143
        }
      },
      "source": [
        "from google.colab import files\n",
        "files.upload()\n",
        "\n",
        "!pip install -q kaggle\n",
        "\n",
        "!mkdir -p ~/.kaggle\n",
        "!cp kaggle.json ~/.kaggle/\n",
        "\n",
        "!kaggle datasets download shayanfazeli/heartbeat"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-cb87e438-927f-497e-ae70-7f8741fc4e32\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-cb87e438-927f-497e-ae70-7f8741fc4e32\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script src=\"/nbextensions/google.colab/files.js\"></script> "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Saving kaggle.json to kaggle.json\n",
            "Warning: Your Kaggle API key is readable by other users on this system! To fix this, you can run 'chmod 600 /root/.kaggle/kaggle.json'\n",
            "Downloading heartbeat.zip to /content\n",
            " 90% 89.0M/98.8M [00:01<00:00, 75.3MB/s]\n",
            "100% 98.8M/98.8M [00:01<00:00, 92.8MB/s]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SZFgZJLUD2-n",
        "outputId": "2763099b-7a44-45c3-a488-871572d38e04",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121
        }
      },
      "source": [
        "!ls /content\n",
        "!unzip heartbeat.zip"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "heartbeat.zip  kaggle.json  sample_data\n",
            "Archive:  heartbeat.zip\n",
            "  inflating: mitbih_test.csv         \n",
            "  inflating: mitbih_train.csv        \n",
            "  inflating: ptbdb_abnormal.csv      \n",
            "  inflating: ptbdb_normal.csv        \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MDl67l0-EsC7"
      },
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "%matplotlib inline\n",
        "sns.set_style('darkgrid')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A8600BukFIUx"
      },
      "source": [
        "## Data loading and spliting"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rQ6bHlx7EEYS",
        "outputId": "a822eb2e-8d85-4b24-8ae7-25ace38d3cdb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        }
      },
      "source": [
        "mit_train = pd.read_csv('mitbih_train.csv', header=None)\n",
        "mit_test = pd.read_csv('mitbih_test.csv', header=None)\n",
        "\n",
        "print(f\"Train shape: {mit_train.shape}\")\n",
        "print(f\"Test shape: {mit_test.shape}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Train shape: (87554, 188)\n",
            "Test shape: (21892, 188)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TMG04mneEkaU"
      },
      "source": [
        "# Separate target from data\n",
        "y_train = mit_train[187]\n",
        "X_train = mit_train.loc[:, :186]\n",
        "\n",
        "y_test = mit_test[187]\n",
        "X_test = mit_test.loc[:, :186]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZvpAYa95ELdq",
        "outputId": "aad50bc9-f354-4326-8db0-31d8d675ec0c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121
        }
      },
      "source": [
        "# Proportions of each class\n",
        "round(mit_train[187].value_counts(normalize=True).sort_index(), 2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0    0.83\n",
              "1.0    0.03\n",
              "2.0    0.07\n",
              "3.0    0.01\n",
              "4.0    0.07\n",
              "Name: 187, dtype: float64"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0pAcUCP3Fvvz"
      },
      "source": [
        "## Preparing and converting data for Pytorch - tensors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LhwOWiFXFX-v"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import TensorDataset, DataLoader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C9QJr_WLFehz"
      },
      "source": [
        "X_train, y_train, X_test, y_test = map(torch.from_numpy, (X_train.values, y_train.values, X_test.values, y_test.values))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zdCu6nREFoEd"
      },
      "source": [
        "# Convert to 3D tensor\n",
        "X_train = X_train.unsqueeze(1)\n",
        "X_test = X_test.unsqueeze(1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7dwI_hpDFpiT"
      },
      "source": [
        "# Batch size\n",
        "bs = 128\n",
        "\n",
        "train_ds = TensorDataset(X_train, y_train)\n",
        "train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)\n",
        "\n",
        "test_ds = TensorDataset(X_test, y_test)\n",
        "test_dl = DataLoader(test_ds, batch_size=bs * 2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fgdhfgqmGM2H"
      },
      "source": [
        "## Import models from torchdiffeq  - ResBlock and NeuralODEBlock"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o5alZdMuFshb"
      },
      "source": [
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CVXnAml9Gvg3",
        "outputId": "c548c713-8d5c-43d7-ddb8-5c039f87d0f0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 159
        }
      },
      "source": [
        "!pip install torchdiffeq"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting torchdiffeq\n",
            "  Downloading https://files.pythonhosted.org/packages/67/af/377e42c20058f4891dedc1827c1a6b7b16772a452d18097d05c1db06b338/torchdiffeq-0.1.1-py3-none-any.whl\n",
            "Requirement already satisfied: torch>=1.3.0 in /usr/local/lib/python3.6/dist-packages (from torchdiffeq) (1.6.0+cu101)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch>=1.3.0->torchdiffeq) (1.18.5)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch>=1.3.0->torchdiffeq) (0.16.0)\n",
            "Installing collected packages: torchdiffeq\n",
            "Successfully installed torchdiffeq-0.1.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JMKWzZmzG7wF"
      },
      "source": [
        "### ResBlock class, ConcatConv1d class, ODEfunc class, ODENet class, and Flatten class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EJ5iwyNjGVCo"
      },
      "source": [
        "\"\"\"\n",
        "ResNet and ODENet classes for ECG classification.\n",
        "Code adapted from:\n",
        "https://github.com/rtqichen/torchdiffeq/blob/master/examples/odenet_mnist.py\n",
        "\"\"\"\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torchdiffeq import odeint_adjoint\n",
        "\n",
        "\n",
        "def norm(dim):\n",
        "    \"\"\"\n",
        "    Group normalization to improve model accuracy and training speed.\n",
        "    \"\"\"\n",
        "    return nn.GroupNorm(min(32, dim), dim)\n",
        "\n",
        "\n",
        "class ResBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    Simple residual block used to construct ResNet.\n",
        "    \"\"\"\n",
        "    def __init__(self, dim):\n",
        "        super(ResBlock, self).__init__()\n",
        "        self.gn1 = norm(dim)\n",
        "        self.conv1 = nn.Conv1d(dim, dim, kernel_size=3, padding=1, bias=False)\n",
        "        self.gn2 = norm(dim)\n",
        "        self.conv2 = nn.Conv1d(dim, dim, kernel_size=3, padding=1, bias=False)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Shortcut\n",
        "        identity = x\n",
        "\n",
        "        # First convolution\n",
        "        out = self.gn1(x)\n",
        "        out = self.relu(out)\n",
        "        out = self.conv1(out)\n",
        "\n",
        "        # Second convolution\n",
        "        out = self.gn2(out)\n",
        "        out = self.relu(out)\n",
        "        out = self.conv2(out)\n",
        "\n",
        "        # Add shortcut\n",
        "        out += identity\n",
        "\n",
        "        return out\n",
        "\n",
        "\n",
        "class ConcatConv1d(nn.Module):\n",
        "    \"\"\"\n",
        "    1d convolution concatenated with time for usage in ODENet.\n",
        "    \"\"\"\n",
        "    def __init__(self, dim_in, dim_out, kernel_size=3, stride=1, padding=0, bias=True, transpose=False):\n",
        "        super(ConcatConv1d, self).__init__()\n",
        "        module = nn.ConvTranspose1d if transpose else nn.Conv1d\n",
        "        self._layer = module(\n",
        "            dim_in + 1, dim_out, kernel_size=kernel_size, stride=stride, padding=padding,\n",
        "            bias=bias\n",
        "        )\n",
        "\n",
        "    def forward(self, t, x):\n",
        "        tt = torch.ones_like(x[:, :1, :]) * t\n",
        "        ttx = torch.cat([tt, x], 1)\n",
        "        return self._layer(ttx)\n",
        "\n",
        "\n",
        "class ODEfunc(nn.Module):\n",
        "    \"\"\"\n",
        "    Network architecture for ODENet.\n",
        "    \"\"\"\n",
        "    def __init__(self, dim):\n",
        "        super(ODEfunc, self).__init__()\n",
        "        self.norm1 = norm(dim)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.conv1 = ConcatConv1d(dim, dim, 3, 1, 1)\n",
        "        self.norm2 = norm(dim)\n",
        "        self.conv2 = ConcatConv1d(dim, dim, 3, 1, 1)\n",
        "        self.norm3 = norm(dim)\n",
        "        self.nfe = 0    # Number of function evaluations\n",
        "\n",
        "    def forward(self, t, x):\n",
        "        self.nfe += 1\n",
        "        out = self.norm1(x)\n",
        "        out = self.relu(out)\n",
        "        out = self.conv1(t, out)\n",
        "        out = self.norm2(out)\n",
        "        out = self.relu(out)\n",
        "        out = self.conv2(t, out)\n",
        "        out = self.norm3(out)\n",
        "        return out\n",
        "\n",
        "\n",
        "class ODENet(nn.Module):\n",
        "    \"\"\"\n",
        "    Neural ODE.\n",
        "    Uses ODE solver (dopri5 by default) to yield model output.\n",
        "    Backpropagation is done with the adjoint method as described in\n",
        "    https://arxiv.org/abs/1806.07366.\n",
        "    Parameters\n",
        "    ----------\n",
        "    odefunc : nn.Module\n",
        "        network architecture\n",
        "    rtol : float\n",
        "        relative tolerance of ODE solver\n",
        "    atol : float\n",
        "        absolute tolerance of ODE solver\n",
        "    \"\"\"\n",
        "    def __init__(self, odefunc, rtol=1e-3, atol=1e-3):\n",
        "        super(ODENet, self).__init__()\n",
        "        self.odefunc = odefunc\n",
        "        self.integration_time = torch.tensor([0, 1]).float()\n",
        "        self.rtol = rtol\n",
        "        self.atol = atol\n",
        "\n",
        "    def forward(self, x):\n",
        "        self.integration_time = self.integration_time.type_as(x)\n",
        "        out = odeint_adjoint(self.odefunc, x, self.integration_time, self.rtol, self.atol)\n",
        "        return out[1]\n",
        "\n",
        "    # Update number of function evaluations (nfe)\n",
        "    @property\n",
        "    def nfe(self):\n",
        "        return self.odefunc.nfe\n",
        "\n",
        "    @nfe.setter\n",
        "    def nfe(self, value):\n",
        "        self.odefunc.nfe = value\n",
        "\n",
        "\n",
        "class Flatten(nn.Module):\n",
        "    \"\"\"\n",
        "    Flatten feature maps for input to linear layer.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super(Flatten, self).__init__()\n",
        "\n",
        "    def forward(self, x):\n",
        "        shape = torch.prod(torch.tensor(x.shape[1:])).item()\n",
        "        return x.view(-1, shape)\n",
        "\n",
        "    \n",
        "def count_parameters(model):\n",
        "    \"\"\"\n",
        "    Count number of tunable parameters.\n",
        "    \"\"\"\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ctfmmp1HHYQk"
      },
      "source": [
        "### Helpers adapted from https://pytorch.org/tutorials/beginner/nn_tutorial.html"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6HrLW9rgHSva"
      },
      "source": [
        "# Helpers adapted from https://pytorch.org/tutorials/beginner/nn_tutorial.html\n",
        "\n",
        "def get_model(is_odenet=True, dim=64, adam=False, **kwargs):\n",
        "    \"\"\"\n",
        "    Initialize ResNet or ODENet with optimizer.\n",
        "    \"\"\"\n",
        "    downsampling_layers = [\n",
        "        nn.Conv1d(1, dim, 3, 1),\n",
        "        norm(dim),\n",
        "        nn.ReLU(inplace=True),\n",
        "        nn.Conv1d(dim, dim, 4, 2, 1),\n",
        "        norm(dim),\n",
        "        nn.ReLU(inplace=True),\n",
        "        nn.Conv1d(dim, dim, 4, 2, 1)\n",
        "    ]\n",
        "\n",
        "    feature_layers = [ODENet(ODEfunc(dim), **kwargs)] if is_odenet else [ResBlock(dim) for _ in range(6)]\n",
        "\n",
        "    fc_layers = [norm(dim), nn.ReLU(inplace=True), nn.AdaptiveAvgPool1d(1), Flatten(), nn.Linear(dim, 5)]\n",
        "\n",
        "    model = nn.Sequential(*downsampling_layers, *feature_layers, *fc_layers)\n",
        "\n",
        "    opt = optim.Adam(model.parameters()) if adam else optim.SGD(model.parameters(), lr=0.1, momentum=0.9)\n",
        "\n",
        "    return model, opt\n",
        "\n",
        "\n",
        "def loss_batch(model, loss_func, xb, yb, opt=None):\n",
        "    \"\"\"\n",
        "    Calculate loss and update weights if training.\n",
        "    \"\"\"\n",
        "    loss = loss_func(model(xb.float()), yb.long())\n",
        "\n",
        "    if opt is not None:\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "        opt.zero_grad()\n",
        "\n",
        "    return loss.item(), len(xb)\n",
        "\n",
        "\n",
        "def fit(epochs, model, loss_func, opt, train_dl, valid_dl):\n",
        "    \"\"\"\n",
        "    Train neural network model.\n",
        "    \"\"\"\n",
        "    num_batches = len(train_dl)\n",
        "    \n",
        "    for epoch in range(epochs):\n",
        "        print(f\"Training... epoch {epoch + 1}\")\n",
        "        \n",
        "        model.train()   # Set model to training mode\n",
        "        batch_count = 0\n",
        "        start = time.time()\n",
        "        for xb, yb in train_dl:\n",
        "            batch_count += 1\n",
        "            curr_time = time.time()\n",
        "            percent = round(batch_count/len(train_dl) * 100, 1)\n",
        "            elapsed = round((curr_time - start)/60, 1)\n",
        "            print(f\"    Percent trained: {percent}%  Time elapsed: {elapsed} min\", end='\\r')\n",
        "            loss_batch(model, loss_func, xb, yb, opt)\n",
        "            \n",
        "            \n",
        "\n",
        "        model.eval()    # Set model to validation mode\n",
        "        with torch.no_grad():\n",
        "            losses, nums = zip(\n",
        "                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]\n",
        "            )\n",
        "        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)\n",
        "\n",
        "        print(f\"\\n    val loss: {round(val_loss, 2)}\\n\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqI796PMHn9z"
      },
      "source": [
        "## ResNet modeling\n",
        "\n",
        "#### Its architecture : [ Downsampling layers + ResBlock / features layers + fully connected layers ]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dQqx2fEFHhS3"
      },
      "source": [
        "resnet, resopt = get_model(is_odenet=False, adam=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v1geFhICHw_y",
        "outputId": "cb32a47e-ee5b-4a6a-ebad-518e98310437",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 156
        }
      },
      "source": [
        "fit(2, resnet, F.cross_entropy, resopt, train_dl, test_dl)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Training... epoch 1\n",
            "\n",
            "    val loss: 0.25\n",
            "\n",
            "Training... epoch 2\n",
            "\n",
            "    val loss: 0.14\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDYJW9u0H6NS"
      },
      "source": [
        "## ODENet modeling\n",
        "\n",
        "#### Its architecture : [ Downsampling layers + ODENet Block / features layers + fully connected layers ]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NfhYelXoH9AG"
      },
      "source": [
        "odenet, odeopt = get_model(adam=False, rtol=1e-3, atol=1e-3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yKcKsHL5H-ZD",
        "outputId": "6af8125e-dd93-4ab9-b365-2a0ec65d2b1b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 156
        }
      },
      "source": [
        "fit(2, odenet, F.cross_entropy, odeopt, train_dl, test_dl)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Training... epoch 1\n",
            "\n",
            "    val loss: 0.21\n",
            "\n",
            "Training... epoch 2\n",
            "\n",
            "    val loss: 0.12\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gr869tLfK2za"
      },
      "source": [
        "## Accuracy and memory usage"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EO5gyycQK2JW"
      },
      "source": [
        "def accuracy(model, X_test, y_test):\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        logits = model(X_test.float())\n",
        "    preds = torch.argmax(F.softmax(logits, dim=1), axis=1).numpy()\n",
        "    return (preds == y_test.numpy()).mean()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r81kXbQ6K58G",
        "outputId": "c26c1a66-cb0c-4ede-a7b8-179b45704c6c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        }
      },
      "source": [
        "print(f\"ResNet accuracy: {round(accuracy(resnet, X_test, y_test), 3)}\")\n",
        "print(f\"ODENet accuracy: {round(accuracy(odenet, X_test, y_test), 3)}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "ResNet accuracy: 0.961\n",
            "ODENet accuracy: 0.967\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HEeanD-TLGzj",
        "outputId": "1915425c-8b9e-4597-cd62-7805883808c7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 69
        }
      },
      "source": [
        "print(\"Number of tunable parameters in...\")\n",
        "print(f\"    ResNet: {count_parameters(resnet)}\")\n",
        "print(f\"    ODENet: {count_parameters(odenet)}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Number of tunable parameters in...\n",
            "    ResNet: 182853\n",
            "    ODENet: 59333\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}