AIOSA_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import re\n",
    "import sys \n",
    "import pickle\n",
    "from tqdm.notebook import tqdm as tqdm\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TPU imports\n",
    "\n",
    "import warnings\n",
    "import torch_xla\n",
    "import torch_xla.debug.metrics as met\n",
    "import torch_xla.distributed.data_parallel as dp\n",
    "import torch_xla.distributed.parallel_loader as pl\n",
    "import torch_xla.utils.utils as xu\n",
    "import torch_xla.core.xla_model as xm\n",
    "import torch_xla.distributed.xla_multiprocessing as xmp\n",
    "import torch_xla.test.test_utils as test_utils\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Other imports\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Defining the neural network (CNN + LSTM) \n",
    "# !pip install torchsummary\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchsummary import summary\n",
    "import random\n",
    "\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils.rnn import PackedSequence\n",
    "from typing import *\n",
    "\n",
    "\n",
    "\n",
    "class VariationalDropout(nn.Module):\n",
    "    \"\"\"\n",
    "    Applies the same dropout mask across the temporal dimension\n",
    "    See https://arxiv.org/abs/1512.05287 for more details.\n",
    "    Note that this is not applied to the recurrent activations in the LSTM like the above paper.\n",
    "    Instead, it is applied to the inputs and outputs of the recurrent layer.\n",
    "    \"\"\"\n",
    "    def __init__(self, dropout: float, batch_first: Optional[bool]=False):\n",
    "        super().__init__()\n",
    "        self.dropout = dropout\n",
    "        self.batch_first = batch_first\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        if not self.training or self.dropout <= 0.:\n",
    "            return x\n",
    "\n",
    "        is_packed = isinstance(x, PackedSequence)\n",
    "        if is_packed:\n",
    "            x, batch_sizes = x\n",
    "            max_batch_size = int(batch_sizes[0])\n",
    "        else:\n",
    "            batch_sizes = None\n",
    "            max_batch_size = x.size(0)\n",
    "\n",
    "        # Drop same mask across entire sequence\n",
    "        if self.batch_first:\n",
    "            m = x.new_empty(max_batch_size, 1, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)\n",
    "        else:\n",
    "            m = x.new_empty(1, max_batch_size, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)\n",
    "        x = x.masked_fill(m == 0, 0) / (1 - self.dropout)\n",
    "\n",
    "        if is_packed:\n",
    "            return PackedSequence(x, batch_sizes)\n",
    "        else:\n",
    "            return x\n",
    "\n",
    "class LSTM(nn.LSTM):\n",
    "    def __init__(self, *args, dropouti: float=0.,\n",
    "                 dropoutw: float=0., dropouto: float=0.,\n",
    "                 batch_first=True, unit_forget_bias=True, **kwargs):\n",
    "        super().__init__(*args, **kwargs, batch_first=batch_first)\n",
    "        self.unit_forget_bias = unit_forget_bias\n",
    "        self.dropoutw = dropoutw\n",
    "        self.input_drop = VariationalDropout(dropouti,\n",
    "                                             batch_first=batch_first)\n",
    "        self.output_drop = VariationalDropout(dropouto,\n",
    "                                              batch_first=batch_first)\n",
    "        self._init_weights()\n",
    "\n",
    "    def _init_weights(self):\n",
    "        \"\"\"\n",
    "        Use orthogonal init for recurrent layers, xavier uniform for input layers\n",
    "        Bias is 0 except for forget gate\n",
    "        \"\"\"\n",
    "        for name, param in self.named_parameters():\n",
    "            if \"weight_hh\" in name:\n",
    "                nn.init.orthogonal_(param.data)\n",
    "            elif \"weight_ih\" in name:\n",
    "                nn.init.xavier_uniform_(param.data)\n",
    "            elif \"bias\" in name and self.unit_forget_bias:\n",
    "                nn.init.zeros_(param.data)\n",
    "                param.data[self.hidden_size:2 * self.hidden_size] = 1\n",
    "\n",
    "    def _drop_weights(self):\n",
    "        for name, param in self.named_parameters():\n",
    "            if \"weight_hh\" in name:\n",
    "                getattr(self, name).data = \\\n",
    "                    torch.nn.functional.dropout(param.data, p=self.dropoutw,\n",
    "                                                training=self.training).contiguous()\n",
    "\n",
    "    def forward(self, input, hx=None):\n",
    "        self._drop_weights()\n",
    "        input = self.input_drop(input)\n",
    "        seq, state = super().forward(input, hx=hx)\n",
    "        return self.output_drop(seq), state\n",
    "    \n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Separable Convs in Pytorch\n",
    "# https://gist.github.com/iiSeymour/85a5285e00cbed60537241da7c3b8525\n",
    "\n",
    "class TCSConv1d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, dilation, padding):\n",
    "        super(TCSConv1d, self).__init__()\n",
    "        self.depthwise = nn.Conv1d(in_channels=in_channels, out_channels=in_channels,\n",
    "                                   kernel_size=kernel_size, dilation=dilation, padding=padding,\n",
    "                                   groups=in_channels, bias=False)\n",
    "        self.pointwise = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.depthwise(x)\n",
    "        x = self.pointwise(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "      \n",
    "    \n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        \n",
    "        filter_size = 16 \n",
    "        kernel_size_var = 3  \n",
    "        \n",
    "        \n",
    "        # W:input volume size\n",
    "        # F:kernel size\n",
    "        # S:stride\n",
    "        # P:amount of padding\n",
    "        # size of output volume = (W-F+2P)/S+1\n",
    "        \n",
    "        # to keep the same size, padding = dilation * (kernel - 1) / 2\n",
    "        \n",
    "\n",
    "        self.skip = TCSConv1d(in_channels=1, out_channels=filter_size, kernel_size=1,\n",
    "                              dilation=1, padding=int((1-1)/2))\n",
    "        \n",
    "        \n",
    "        # Drop 0.2\n",
    "\n",
    "        self.conv_1 = TCSConv1d(in_channels=1, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=1,\n",
    "                                     padding=int((kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_1 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_1 = nn.Dropout2d(0.2)\n",
    "        \n",
    "        \n",
    "        self.conv_2 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=2,\n",
    "                                     padding=int(2*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_2 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_2 = nn.Dropout2d(0.2)\n",
    "        \n",
    "        \n",
    "        self.conv_3 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=4,\n",
    "                                     padding=int(4*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_3 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_3 = nn.Dropout2d(0.2)  \n",
    "        \n",
    "        \n",
    "        self.conv_4 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=8,\n",
    "                                     padding=int(8*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_4 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_4 = nn.Dropout2d(0.2)\n",
    "        \n",
    "        self.conv_4b = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=16,\n",
    "                                     padding=int(16*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_4b = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_4b = nn.Dropout2d(0.2)\n",
    "        \n",
    "        \n",
    "        self.avgPool_a = nn.AvgPool1d(kernel_size=4)\n",
    "        \n",
    "        \n",
    "        \n",
    "        # Drop 0.1\n",
    "        \n",
    "        \n",
    "        self.conv_5 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=1,\n",
    "                                     padding=int((kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_5 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_5 = nn.Dropout2d(0.1)\n",
    "        \n",
    "        \n",
    "        self.conv_6 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=3,\n",
    "                                     padding=int(3*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_6 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_6 = nn.Dropout2d(0.1)\n",
    "        \n",
    "        \n",
    "        self.conv_7 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=5,\n",
    "                                     padding=int(5*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_7 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_7 = nn.Dropout2d(0.1)  \n",
    "        \n",
    "        \n",
    "        self.conv_8 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=9,\n",
    "                                     padding=int(9*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_8 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_8 = nn.Dropout2d(0.1)\n",
    "        \n",
    "        \n",
    "        self.conv_8b = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=16,\n",
    "                                     padding=int(16*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_8b = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_8b = nn.Dropout2d(0.1)\n",
    "        \n",
    "        self.avgPool_b = nn.AvgPool1d(kernel_size=4)\n",
    "        \n",
    "        \n",
    "        #0.3\n",
    "        self.conv_9 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=1,\n",
    "                                     padding=int((kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_9 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_9 = nn.Dropout2d(0.3)\n",
    "        \n",
    "        \n",
    "        self.conv_10 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=3,\n",
    "                                     padding=int(3*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_10 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_10 = nn.Dropout2d(0.3)\n",
    "        \n",
    "        \n",
    "        self.conv_11 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=5,\n",
    "                                     padding=int(5*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_11 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_11 = nn.Dropout2d(0.3)  \n",
    "        \n",
    "        \n",
    "        self.conv_12 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=9,\n",
    "                                     padding=int(9*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_12 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_12 = nn.Dropout2d(0.3)  \n",
    "        \n",
    "        self.conv_12b = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=16,\n",
    "                                     padding=int(16*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_12b = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_12b = nn.Dropout2d(0.3)\n",
    "        \n",
    "        self.avgPool_c = nn.AvgPool1d(kernel_size=5)\n",
    "        \n",
    "        \n",
    "        #0.5\n",
    "        self.conv_13 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=1,\n",
    "                                     padding=int((kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_13 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_13 = nn.Dropout2d(0.3)\n",
    "        \n",
    "        \n",
    "        self.conv_14 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=2,\n",
    "                                     padding=int(2*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_14 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_14 = nn.Dropout2d(0.3)\n",
    "        \n",
    "        \n",
    "        self.conv_15 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=4,\n",
    "                                     padding=int(4*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_15 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_15 = nn.Dropout2d(0.3)  \n",
    "        \n",
    "        \n",
    "        self.conv_16 = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=8,\n",
    "                                     padding=int(8*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_16 = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_16 = nn.Dropout2d(0.3) \n",
    "        \n",
    "        \n",
    "        self.conv_16b = TCSConv1d(in_channels=filter_size, out_channels=filter_size,\n",
    "                                     kernel_size=kernel_size_var, dilation=16,\n",
    "                                     padding=int(16*(kernel_size_var-1)/2))\n",
    "        \n",
    "        self.bn_16b = nn.BatchNorm1d(filter_size) \n",
    "        \n",
    "        self.drop_16b = nn.Dropout2d(0.3)\n",
    "        \n",
    "\n",
    "\n",
    "        self.conv_17 = TCSConv1d(in_channels=filter_size, out_channels=1,\n",
    "                                     kernel_size=1, dilation=1,\n",
    "                                     padding=int((1-1)/2))\n",
    "        \n",
    "        self.bn_17 = nn.BatchNorm1d(1) \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        self.lstm_1 = LSTM(input_size=filter_size, hidden_size=128, num_layers=1, bidirectional=True, batch_first=True, dropouti=0.1)\n",
    "        \n",
    "        \n",
    "        \n",
    "        self.fc_1 = nn.Linear(768, 1024) \n",
    "        \n",
    "        self.bn_fc_1 = nn.BatchNorm1d(1024) \n",
    "                \n",
    "        self.drop_fc_1 = nn.Dropout(0.5)  \n",
    "        \n",
    "        \n",
    "        self.fc_2 = nn.Linear(1024, 512)\n",
    "        \n",
    "        self.bn_fc_2 = nn.BatchNorm1d(512) \n",
    "        \n",
    "        self.drop_fc_2 = nn.Dropout(0.5)   \n",
    "        \n",
    "        \n",
    "        self.fc_3 = nn.Linear(256, 256)\n",
    "        \n",
    "        self.bn_fc_3 = nn.BatchNorm1d(256) \n",
    "        \n",
    "        self.drop_fc_3 = nn.Dropout(0.5)   \n",
    "        \n",
    "        \n",
    "        self.fc_4 = nn.Linear(256, 256)\n",
    "        \n",
    "        self.bn_fc_4 = nn.BatchNorm1d(256) \n",
    "        \n",
    "        self.drop_fc_4 = nn.Dropout(0.5)   \n",
    "        \n",
    "\n",
    "        self.lstm_spo2 = LSTM(input_size=1, hidden_size=256, num_layers=1, bidirectional=True, batch_first=True, dropouto=0.1, dropouti=0.1)\n",
    "        \n",
    "        self.ln = nn.LayerNorm(1024)\n",
    "        \n",
    "        \n",
    "        self.fc_5 = nn.Linear(1024, 60)\n",
    "\n",
    "        \n",
    "        \n",
    "        \n",
    "    def forward(self, x, y):\n",
    "        skip_conn = self.skip(x)\n",
    "        \n",
    "        \n",
    "        x = self.drop_1(F.relu(self.bn_1(self.conv_1(x))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "                \n",
    "        x = self.drop_2(F.relu(self.bn_2(self.conv_2(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_3(F.relu(self.bn_3(self.conv_3(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_4(F.relu(self.bn_4(self.conv_4(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "\n",
    "        skip_conn = self.avgPool_a(skip_conn)\n",
    "        \n",
    "\n",
    "        x = self.drop_5(F.relu(self.bn_5(self.conv_5(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_6(F.relu(self.bn_6(self.conv_6(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_7(F.relu(self.bn_7(self.conv_7(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_8(F.relu(self.bn_8(self.conv_8(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "       \n",
    "\n",
    "        skip_conn = self.avgPool_b(skip_conn)\n",
    "        \n",
    "        \n",
    "        x = self.drop_9(F.relu(self.bn_9(self.conv_9(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_10(F.relu(self.bn_10(self.conv_10(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_11(F.relu(self.bn_11(self.conv_11(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_12(F.relu(self.bn_12(self.conv_12(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "       \n",
    "\n",
    "        skip_conn = self.avgPool_c(skip_conn)    \n",
    "        \n",
    "        \n",
    "        x = self.drop_13(F.relu(self.bn_13(self.conv_13(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_14(F.relu(self.bn_14(self.conv_14(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_15(F.relu(self.bn_15(self.conv_15(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)\n",
    "        \n",
    "        x = self.drop_16(F.relu(self.bn_16(self.conv_16(skip_conn))))\n",
    "        skip_conn = skip_conn.add(x)  \n",
    "        \n",
    "        \n",
    "        x = x.permute(0, 2, 1)\n",
    "                \n",
    "        \n",
    "        x, states = self.lstm_1(x)        \n",
    "        x = x[:, -1, :]\n",
    "        \n",
    "        y, states = self.lstm_spo2(y)\n",
    "        y = y[:, -1, :]\n",
    "        \n",
    "        \n",
    "        \n",
    "        x = torch.cat((x, y), 1)\n",
    "        \n",
    "        x = x.view(-1, 768)\n",
    "        \n",
    "        \n",
    "        x = self.drop_fc_1(F.relu(self.ln(self.fc_1(x))))\n",
    "        \n",
    "        x = self.fc_5(x) \n",
    "        \n",
    "        \n",
    "        return x\n",
    "\n",
    "  \n",
    "    \n",
    "# Helper function that is used to initialize the weights of the model\n",
    "def init_weights(m):\n",
    "    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv1d):\n",
    "        if \"fc_5\" in str(m):\n",
    "            nn.init.xavier_uniform_(m.weight)\n",
    "        else:\n",
    "            nn.init.kaiming_normal_(m.weight)\n",
    "        nn.init.constant_(m.bias, 0.)\n",
    "    \n",
    "    \n",
    "# Helper function that is used to (re)define the model and the optimizer from scratch\n",
    "# https://pytorch.org/docs/stable/notes/randomness.html\n",
    "def reinit_model():\n",
    "    torch.manual_seed(42)\n",
    "    np.random.seed(42)\n",
    "    random.seed(42)\n",
    "    net = Net()\n",
    "    #net.apply(init_weights)\n",
    "    return net\n",
    "\n",
    "\n",
    "\n",
    "model = reinit_model()\n",
    "\n",
    "print(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weighted hinge loss function\n",
    "# input: specify the weight of class 1, with respect to class 0\n",
    "# the idea is that the weight is 1 for instances of class 0, and w_1_class for instances of class 1\n",
    "# inspired from: https://stackoverflow.com/questions/55754976/weighted-hinge-loss-function\n",
    "def weightedSquaredHingeLoss(inp, tar, device, w_1_class=1):\n",
    "    return torch.sum(torch.mean((torch.max(tar, torch.zeros(inp.shape[1], dtype=torch.float32).to(device))*(w_1_class-1)+1) * torch.max(1. - tar * inp, torch.zeros(inp.shape[1], dtype=torch.float32).to(device))**2, dim=-1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining some helper functions for the model that guide the training\n",
    "import math\n",
    "\n",
    "\n",
    "def reduce_fn_avg(vals):\n",
    "    # take average\n",
    "    return sum(vals) / len(vals)\n",
    "\n",
    "\n",
    "def reduce_fn_sum(vals):\n",
    "    # take average\n",
    "    return sum(vals)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def _run(model, EPOCHS, param, training_data_in, validation_data_in=None):\n",
    "    \n",
    "    xm.set_rng_state(42)\n",
    "    xm.save(xm.get_rng_state(), 'xm_seed')\n",
    "      \n",
    "\n",
    "    def train_fn(train_dataloader, model, optimizer, criterion, device, lr_scheduler=None):\n",
    "        \n",
    "        xm.set_rng_state(torch.load('xm_seed'), device=device)\n",
    "        xm.master_print(xm.get_rng_state())\n",
    "\n",
    "        running_loss = 0.\n",
    "        running_tp = 0.\n",
    "        running_tn = 0.\n",
    "        running_fp = 0.\n",
    "        running_fn = 0.\n",
    "        running_instances = 0.\n",
    "\n",
    "        # training() is a kind of switch for some specific layers/parts of the model that behave\n",
    "        # differently during training and inference (evaluating) time\n",
    "        # For example, Dropouts Layers, BatchNorm Layers etc. \n",
    "        model.train()\n",
    "\n",
    "        for batch_idx, (ecg, spo2, labels) in enumerate(train_dataloader, 1):\n",
    "\n",
    "            optimizer.zero_grad() # need to zero out the gradients every time, otherwise they accumulate\n",
    "            ecg = ecg.to(device) # transfer the data to the computing device\n",
    "            spo2 = spo2.to(device) # transfer the data to the computing device\n",
    "            labels = labels.to(device)# transfer the labels to the computing device     \n",
    "                \n",
    "            outputs = model(ecg, spo2)\n",
    "            \n",
    "            loss = criterion(outputs, labels, device, imbalanced_ratio)\n",
    "            \n",
    "            #xm.master_print(f'Batch: {batch_idx}, loss: {loss.item()}')\n",
    "                        \n",
    "            loss.backward() # calculate the gradients\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            xm.optimizer_step(optimizer) # update the network weights\n",
    "                                                \n",
    "            running_loss += loss.item()*len(labels)\n",
    "            running_instances += len(labels)\n",
    "            \n",
    "            predicted = torch.where(outputs.data > 0, torch.from_numpy(np.asarray([1])).to(device), torch.from_numpy(np.asarray([0])).to(device))\n",
    "            labels = torch.where(labels > 0, torch.from_numpy(np.asarray([1])).to(device), torch.from_numpy(np.asarray([0])).to(device))\n",
    "            \n",
    "            \n",
    "            fp = ((predicted - labels) == 1.).sum().item() \n",
    "            fn = ((predicted - labels) == -1.).sum().item()\n",
    "            tp = ((predicted + labels) == 2.).sum().item()\n",
    "            tn = ((predicted + labels) == 0.).sum().item()\n",
    "            fp_reduced = xm.mesh_reduce('fp_reduce', fp, reduce_fn_sum) \n",
    "            fn_reduced = xm.mesh_reduce('fn_reduce', fn, reduce_fn_sum) \n",
    "            tp_reduced = xm.mesh_reduce('tp_reduce', tp, reduce_fn_sum) \n",
    "            tn_reduced = xm.mesh_reduce('tn_reduce', tn, reduce_fn_sum) \n",
    "            \n",
    "            running_tp += tp_reduced\n",
    "            running_fp += fp_reduced\n",
    "            running_tn += tn_reduced\n",
    "            running_fn += fn_reduced\n",
    "            \n",
    "            if lr_scheduler != None:\n",
    "                lr_scheduler.step()\n",
    "                \n",
    "                \n",
    "        running_loss /= running_instances\n",
    "        loss_reduced = xm.mesh_reduce('loss_reduced', running_loss, reduce_fn_avg)\n",
    "        retval = {'loss':  loss_reduced, \n",
    "                  'tp':running_tp,\n",
    "                  'tn':running_tn,\n",
    "                  'fp':running_fp,\n",
    "                  'fn':running_fn\n",
    "                 }\n",
    "        \n",
    "        xm.save(xm.get_rng_state(), 'xm_seed')\n",
    "            \n",
    "        return retval\n",
    "            \n",
    "\n",
    "        \n",
    "    def valid_fn(valid_dataloader, model, criterion, device):\n",
    "                        \n",
    "        running_loss = 0.\n",
    "        running_tp = 0.\n",
    "        running_tn = 0.\n",
    "        running_fp = 0.\n",
    "        running_fn = 0.\n",
    "        running_instances = 0.\n",
    "         \n",
    "        # eval() is a kind of switch for some specific layers/parts of the model that behave\n",
    "        # differently during training and inference (evaluating) time\n",
    "        # For example, Dropouts Layers, BatchNorm Layers etc. \n",
    "        model.eval()\n",
    "        \n",
    "        for batch_idx, (ecg, spo2, labels) in enumerate(valid_dataloader, 1):\n",
    "\n",
    "            ecg = ecg.to(device)\n",
    "            spo2 = spo2.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            outputs = model(ecg, spo2)\n",
    "            \n",
    "            loss = criterion(outputs, labels, device, imbalanced_ratio)\n",
    "            \n",
    "            #xm.master_print(f'Batch: {batch_idx}, loss: {loss.item()}')\n",
    "\n",
    "            running_loss += loss.item()*len(labels)\n",
    "            running_instances += len(labels)\n",
    "           \n",
    "            predicted = torch.where(outputs.data > 0, torch.from_numpy(np.asarray([1])).to(device), torch.from_numpy(np.asarray([0])).to(device))\n",
    "            labels = torch.where(labels > 0, torch.from_numpy(np.asarray([1])).to(device), torch.from_numpy(np.asarray([0])).to(device))\n",
    "            \n",
    "            \n",
    "            fp = ((predicted - labels) == 1.).sum().item()\n",
    "            fn = ((predicted - labels) == -1.).sum().item()\n",
    "            tp = ((predicted + labels) == 2.).sum().item()\n",
    "            tn = ((predicted + labels) == 0.).sum().item()\n",
    "            fp_reduced = xm.mesh_reduce('val_fp_reduce', fp, reduce_fn_sum) \n",
    "            fn_reduced = xm.mesh_reduce('val_fn_reduce', fn, reduce_fn_sum) \n",
    "            tp_reduced = xm.mesh_reduce('val_tp_reduce', tp, reduce_fn_sum) \n",
    "            tn_reduced = xm.mesh_reduce('val_tn_reduce', tn, reduce_fn_sum) \n",
    "            \n",
    "            running_tp += tp_reduced\n",
    "            running_fp += fp_reduced\n",
    "            running_tn += tn_reduced\n",
    "            running_fn += fn_reduced\n",
    "        \n",
    "        running_loss /= running_instances\n",
    "        loss_reduced = xm.mesh_reduce('loss_reduced', running_loss, reduce_fn_avg)    \n",
    "        retval = {'loss': loss_reduced, \n",
    "                  'tp':running_tp,\n",
    "                  'tn':running_tn,\n",
    "                  'fp':running_fp,\n",
    "                  'fn':running_fn\n",
    "                 }\n",
    "                    \n",
    "        return retval\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Defining distributed samplers and data loaders\n",
    "    train_sampler = torch.utils.data.distributed.DistributedSampler(training_data_in,\n",
    "                                                                    num_replicas=xm.xrt_world_size(), #numcores\n",
    "                                                                    rank=xm.get_ordinal(),\n",
    "                                                                    shuffle=True)\n",
    "    \n",
    "    train_dataloader = torch.utils.data.DataLoader(training_data_in, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers=1) # only for GPUs num_workers=2, pin_memory=True)\n",
    "\n",
    "\n",
    "    \n",
    "    if validation_data_in != None:\n",
    "        validation_sampler = torch.utils.data.distributed.DistributedSampler(validation_data_in,\n",
    "                                                                             num_replicas=xm.xrt_world_size(),\n",
    "                                                                             rank=xm.get_ordinal(),\n",
    "                                                                             shuffle=True)\n",
    "\n",
    "        validation_dataloader = torch.utils.data.DataLoader(validation_data_in, batch_size=BATCH_SIZE, sampler=validation_sampler, num_workers=1)\n",
    "\n",
    "\n",
    "\n",
    "    # Defining the handle to the TPU device\n",
    "    device = xm.xla_device()\n",
    "    # Transferring the model to the computing device\n",
    "    model.to(device) \n",
    "    # Defining the loss function\n",
    "    criterion = weightedSquaredHingeLoss\n",
    "    \n",
    "    #torch.autograd.set_detect_anomaly(True)\n",
    "    \n",
    "    # Defining optimizer\n",
    "    import torch.optim as optim\n",
    "    optimizer = optim.Adam(model.parameters(), lr=3e-4, amsgrad=False, eps=1e-07) \n",
    "    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-2, epochs=200, steps_per_epoch=math.ceil(len(trainset) / BATCH_SIZE), base_momentum=0.78, max_momentum=0.99)\n",
    "    \n",
    "    # Training code\n",
    "    \n",
    "    metrics_history = {\"loss\":[], \"accuracy\":[], \"precision\":[], \"recall\":[], \"f1\":[], \"specificity\":[], \"tp\":[], \"tn\":[], \"fp\":[], \"fn\":[],\n",
    "                       \"val_loss\":[], \"val_accuracy\":[], \"val_precision\":[], \"val_recall\":[], \"val_f1\":[], \"val_specificity\":[], \"val_tp\":[], \"val_tn\":[], \"val_fp\":[], \"val_fn\":[]}\n",
    "    \n",
    "    train_begin = time.time()\n",
    "    for epoch in range(EPOCHS):\n",
    "        start = time.time()\n",
    "        para_loader = pl.ParallelLoader(train_dataloader, [device], fixed_batch_size=True) # needed for parallel training\n",
    "\n",
    "        xm.master_print(\"EPOCH:\", epoch+1)\n",
    "\n",
    "        train_metrics = train_fn(train_dataloader=para_loader.per_device_loader(device), \n",
    "                                 model=model,\n",
    "                                 optimizer=optimizer, \n",
    "                                 criterion=criterion,\n",
    "                                 device=device,\n",
    "                                 lr_scheduler=lr_scheduler)\n",
    "        \n",
    "        metrics_history[\"loss\"].append(train_metrics[\"loss\"])\n",
    "        tr_acc = (train_metrics[\"tp\"] + train_metrics[\"tn\"]) / (train_metrics[\"tp\"] + train_metrics[\"tn\"] + train_metrics[\"fp\"] + train_metrics[\"fn\"])\n",
    "        metrics_history[\"accuracy\"].append(tr_acc)\n",
    "        metrics_history[\"tp\"].append(train_metrics[\"tp\"])\n",
    "        metrics_history[\"tn\"].append(train_metrics[\"tn\"])\n",
    "        metrics_history[\"fp\"].append(train_metrics[\"fp\"])\n",
    "        metrics_history[\"fn\"].append(train_metrics[\"fn\"])\n",
    "        \n",
    "        precision = train_metrics[\"tp\"] / (train_metrics[\"tp\"] + train_metrics[\"fp\"]) if train_metrics[\"tp\"] > 0 else 0\n",
    "        recall = train_metrics[\"tp\"] / (train_metrics[\"tp\"] + train_metrics[\"fn\"]) if train_metrics[\"tp\"] > 0 else 0\n",
    "        specificity = train_metrics[\"tn\"] / (train_metrics[\"tn\"] + train_metrics[\"fp\"]) if train_metrics[\"tn\"] > 0 else 0\n",
    "        f1 = 2*precision*recall / (precision + recall) if precision*recall > 0 else 0\n",
    "        metrics_history[\"precision\"].append(precision)\n",
    "        metrics_history[\"recall\"].append(recall)\n",
    "        metrics_history[\"f1\"].append(f1)\n",
    "        metrics_history[\"specificity\"].append(specificity)\n",
    "        \n",
    "        #assert train_metrics[\"tp\"] + train_metrics[\"tn\"] + train_metrics[\"fp\"] + train_metrics[\"fn\"] == len(training_data_in)*granularita_apnee  # vale solo se input divisibile per 8\n",
    "        \n",
    "        #optimizer_sch.step() ######### <- LR SCHEDULER\n",
    "        \n",
    "        \n",
    "        if validation_data_in != None:    \n",
    "            # Calculate the metrics on the validation data, in the same way as done for training\n",
    "            with torch.no_grad(): # don't keep track of the info necessary to calculate the gradients\n",
    "                para_loader = pl.ParallelLoader(validation_dataloader, [device], fixed_batch_size=True)\n",
    "\n",
    "                val_metrics = valid_fn(valid_dataloader=para_loader.per_device_loader(device), \n",
    "                                       model=model,\n",
    "                                       criterion=criterion, \n",
    "                                       device=device)\n",
    "\n",
    "                metrics_history[\"val_loss\"].append(val_metrics[\"loss\"])\n",
    "                val_acc = (val_metrics[\"tp\"] + val_metrics[\"tn\"]) / (val_metrics[\"tp\"] + val_metrics[\"tn\"] + val_metrics[\"fp\"] + val_metrics[\"fn\"])\n",
    "                metrics_history[\"val_accuracy\"].append(val_acc)\n",
    "                metrics_history[\"val_tp\"].append(val_metrics[\"tp\"])\n",
    "                metrics_history[\"val_tn\"].append(val_metrics[\"tn\"])\n",
    "                metrics_history[\"val_fp\"].append(val_metrics[\"fp\"])\n",
    "                metrics_history[\"val_fn\"].append(val_metrics[\"fn\"])\n",
    "\n",
    "                val_precision = val_metrics[\"tp\"] / (val_metrics[\"tp\"] + val_metrics[\"fp\"]) if val_metrics[\"tp\"] > 0 else 0\n",
    "                val_recall = val_metrics[\"tp\"] / (val_metrics[\"tp\"] + val_metrics[\"fn\"]) if val_metrics[\"tp\"] > 0 else 0\n",
    "                val_specificity = val_metrics[\"tn\"] / (val_metrics[\"tn\"] + val_metrics[\"fp\"]) if val_metrics[\"tn\"] > 0 else 0\n",
    "                val_f1 = 2*val_precision*val_recall / (val_precision + val_recall) if val_precision*val_recall > 0 else 0\n",
    "                metrics_history[\"val_precision\"].append(val_precision)\n",
    "                metrics_history[\"val_recall\"].append(val_recall)\n",
    "                metrics_history[\"val_f1\"].append(val_f1)\n",
    "                metrics_history[\"val_specificity\"].append(val_specificity)\n",
    "                \n",
    "            #assert val_metrics[\"tp\"] + val_metrics[\"tn\"] + val_metrics[\"fp\"] + val_metrics[\"fn\"] == len(validation_data_in)*granularita_apnee  # vale solo se input divisibile per 8\n",
    "\n",
    "\n",
    "            xm.master_print(\"  > Training/validation loss:\", round(train_metrics['loss'], 4), round(val_metrics['loss'], 4))\n",
    "            xm.master_print(\"  > Training/validation accuracy:\", round(tr_acc, 4), round(val_acc, 4))\n",
    "            xm.master_print(\"  > Training/validation precision:\", round(precision, 4), round(val_precision, 4))\n",
    "            xm.master_print(\"  > Training/validation recall:\", round(recall, 4), round(val_recall, 4))\n",
    "            xm.master_print(\"  > Training/validation f1:\", round(f1, 4), round(val_f1, 4))\n",
    "            xm.master_print(\"  > Training/validation specificity:\", round(specificity, 4), round(val_specificity, 4))\n",
    "            xm.master_print(\"  > TRAIN tp tn fp fn :\", train_metrics[\"tp\"], train_metrics[\"tn\"], train_metrics[\"fp\"], train_metrics[\"fn\"])\n",
    "            xm.master_print(\"  > VAL tp tn fp fn :\", val_metrics[\"tp\"], val_metrics[\"tn\"], val_metrics[\"fp\"], val_metrics[\"fn\"])\n",
    "        else:\n",
    "            xm.master_print(\"  > Training loss:\", round(train_metrics['loss'], 4))\n",
    "            xm.master_print(\"  > Training accuracy:\", round(tr_acc, 4))\n",
    "            xm.master_print(\"  > Training precision:\", round(precision, 4))\n",
    "            xm.master_print(\"  > Training recall:\", round(recall, 4))\n",
    "            xm.master_print(\"  > Training f1:\", round(f1, 4))\n",
    "            xm.master_print(\"  > Training specificity:\", round(specificity, 4))\n",
    "            xm.master_print(\"  > TRAIN tp tn fp fn :\", train_metrics[\"tp\"], train_metrics[\"tn\"], train_metrics[\"fp\"], train_metrics[\"fn\"])\n",
    "\n",
    "\n",
    "        xm.master_print(\"Completed in:\", round(time.time() - start, 1), \"seconds \\n\")\n",
    "\n",
    "    xm.master_print(\"Training completed in:\", round((time.time()- train_begin)/60, 1), \"minutes\")    \n",
    "\n",
    "    \n",
    "    \n",
    "    # Save the model weights\n",
    "    xm.save(\n",
    "        model.state_dict(), './nnet_model_physio.pt'\n",
    "    )\n",
    "    \n",
    "    # Save the metrics history\n",
    "    xm.save(metrics_history, 'training_history')\n",
    "    \n",
    "    \n",
    "    \n",
    "\n",
    "# Prediction function, it runs on the CPU\n",
    "def predict_(model, data_in): \n",
    "    import math\n",
    "    \n",
    "    data_dataloader = torch.utils.data.DataLoader(data_in, batch_size=BATCH_SIZE, shuffle=False, num_workers=1)\n",
    "    \n",
    "    model.load_state_dict(torch.load('./nnet_model_physio.pt'))\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    predictions = []\n",
    "    with torch.no_grad(): # don't keep track of the info necessary to calculate the gradients \n",
    "        pbar = tqdm(desc=\"Minibatches: \", total=math.ceil(len(data_in)/BATCH_SIZE))\n",
    "        \n",
    "        for batch_idx, (ecg, spo2, labels) in enumerate(data_dataloader, 1):\n",
    "            outputs = model(ecg, spo2)\n",
    "            predicted = outputs.data\n",
    "            predictions.extend(predicted.numpy())\n",
    "            pbar.update(1)\n",
    "        pbar.close()\n",
    "    \n",
    "    return np.asarray(predictions)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}