Hierarchical_Attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6ffc9421",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For some reasons, file descriptors (FDs) do not get released.\n",
    "# This is a work around which increases the allowed limit.\n",
    "import resource\n",
    "rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)\n",
    "resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "244275c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ahmedmohammed/opt/miniconda3/envs/pytorch_env/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "import os\n",
    "import re\n",
    "import random\n",
    "import multiprocessing as mp\n",
    "import numpy as np\n",
    "from scipy.io import loadmat\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from copy import deepcopy\n",
    "from sklearn.metrics import roc_auc_score, average_precision_score\n",
    "from torch.optim.swa_utils import AveragedModel, SWALR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b8140bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch size\n",
    "train_B = 16\n",
    "val_B = 1\n",
    "\n",
    "montage = 'BH'\n",
    "p_dropout = 0.\n",
    "\n",
    "window_len = 50\n",
    "#overlap = 0.8\n",
    "overlap = 0.5\n",
    "sample_H = 3\n",
    "sample_d_qk = 5\n",
    "sample_d_v = 5\n",
    "mid_features = 15\n",
    "epoch_H = 3\n",
    "epoch_d_qk = 5\n",
    "epoch_d_v = 5\n",
    "\n",
    "SCALE = True\n",
    "data_folder = 'BH average 100-sps (0.75-38 Hz)'\n",
    "files_names = os.listdir(data_folder)\n",
    "\n",
    "g1 = ['PAT31', 'MC', 'FEsg', 'DHut', 'NKra', 'PAT2']\n",
    "g2 = ['RC', 'MBra', 'ESow', 'PAT1', 'EG', 'FigSa', 'PAT28']\n",
    "g3 = ['PAT47', 'RA', 'PAT51', 'MPi', 'GNA', 'LM', 'DG']\n",
    "g4 = ['PAT33', 'PB', 'PAT22', 'MAXJ', 'PAT19', 'LP', 'HeMod', 'PAT50', 'PJuly']\n",
    "g5 = ['PAT48', 'KS', 'PAT3', 'PAT8', 'LRio', 'HB', 'ALo', 'JAlv', 'PAT25', 'AR']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e37c78c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_patients = g2 + g3 + g4 + g5\n",
    "val_patients = g1\n",
    "del g1, g2, g3, g4, g5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae1c2c23",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, segments, labels, flip_prob, montage):\n",
    "        self.segments = segments\n",
    "        self.labels = labels\n",
    "        self.flip_prob = flip_prob\n",
    "        \n",
    "        if flip_prob > 0:\n",
    "            self.idx_swaps = self.montage_swap_func(montage)\n",
    "    \n",
    "    \n",
    "    @staticmethod\n",
    "    def montage_swap_func(montage):\n",
    "        if montage == 'BH':\n",
    "            electrodes = ['Fp1','F7','T3','T5','O1','F3','C3','P3','Fz','Cz','Pz','Fp2','F8','T4','T6','O2','F4','C4','P4']\n",
    "        elif montage == 'MCH':\n",
    "            electrodes = ['C3','C4','O1','O2','Cz','F3','F4','F7','F8','Fz','Fp1','Fp2','P3','P4','Pz','T3','T4','T5','T6']\n",
    "        else:\n",
    "            raise NameError('Unavailable Montage')\n",
    "        \n",
    "        dic = {electrode: idx for idx, electrode in enumerate(electrodes)}\n",
    "        \n",
    "        label_swaps = [('Fp1', 'Fp2'), ('F7', 'F8'), ('F3', 'F4'), ('C3', 'C4'), ('P3', 'P4'), ('T3', 'T4'), ('T5', 'T6'), ('O1', 'O2')]\n",
    "        return [(dic[e_left], dic[e_right]) for e_left, e_right in label_swaps]\n",
    "    \n",
    "    \n",
    "    def flip_lr(self, segment):\n",
    "        # flip based on self.idx_swaps\n",
    "        for e_left, e_right in self.idx_swaps:\n",
    "            cloned_channel = torch.clone(segment[e_left])\n",
    "            segment[e_left] = segment[e_right]\n",
    "            segment[e_right] = cloned_channel\n",
    "        \n",
    "        return segment\n",
    "    \n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "    \n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        if random.random() < self.flip_prob:\n",
    "            return self.flip_lr( torch.clone(self.segments[index]) ), self.labels[index]\n",
    "        else:\n",
    "            return self.segments[index], self.labels[index]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "facd5453",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prep_patient_data(patient_name):\n",
    "    \n",
    "    patient_segments = []\n",
    "    patient_labels = []\n",
    "    \n",
    "    regex = re.compile(f'{patient_name}\\D[0-9.a-z_A-Z]*')\n",
    "    \n",
    "    for file_name in files_names:\n",
    "        if regex.match(file_name):\n",
    "            segment = torch.tensor( loadmat(os.path.join(data_folder, file_name))['segment'] ).float()\n",
    "            if 'sp' in file_name:\n",
    "                label = torch.tensor(1.0).float()\n",
    "            elif 'ns' in file_name:\n",
    "                label = torch.tensor(0.0).float()\n",
    "            \n",
    "            patient_segments.append(segment)\n",
    "            patient_labels.append(label)\n",
    "    \n",
    "    return patient_segments, patient_labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b8afd0f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(patients_names):\n",
    "    group_segments = []\n",
    "    group_labels = []\n",
    "    \n",
    "    #with mp.Pool(processes=mp.cpu_count() - 1) as p:\n",
    "    #    for patient_segments, patient_labels in p.imap_unordered(prep_patient_data, patients_names):\n",
    "    #        group_segments += patient_segments\n",
    "    #        group_labels += patient_labels\n",
    "    \n",
    "    for patient_name in patients_names:\n",
    "        patient_segments, patient_labels = prep_patient_data(patient_name)\n",
    "        group_segments += patient_segments\n",
    "        group_labels += patient_labels\n",
    "    \n",
    "    return group_segments, group_labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952da550",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_segments, train_labels = prepare_dataset(train_patients)\n",
    "val_segments, val_labels = prepare_dataset(val_patients)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c844167d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def median_IQR_func(group_segments):\n",
    "    numbers = []\n",
    "    for segment in group_segments:\n",
    "        for row in range(segment.shape[0]):\n",
    "            for col in range(segment.shape[1]):\n",
    "                numbers.append(segment[row, col].item())\n",
    "    \n",
    "    quantiles = torch.tensor(numbers).quantile( torch.tensor([0.25, 0.5, 0.75]) )\n",
    "    median = quantiles[1].item()\n",
    "    IQR = quantiles[2].item() - quantiles[0].item()\n",
    "    \n",
    "    return median, IQR\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a1651ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaler_func(group_segments, median, IQR):\n",
    "    # Scales the group_segments list inplace\n",
    "    for idx in range(len(group_segments)):\n",
    "        group_segments[idx] = ((group_segments[idx] - median) / IQR)\n",
    "    \n",
    "    return group_segments\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d72820",
   "metadata": {},
   "outputs": [],
   "source": [
    "if SCALE:\n",
    "    median, IQR = median_IQR_func(train_segments)\n",
    "    train_segments = scaler_func(train_segments, median, IQR)\n",
    "    val_segments = scaler_func(val_segments, median, IQR)\n",
    "\n",
    "#train_dataset = CustomDataset(train_segments, train_labels, 0.5, 'BH')\n",
    "#val_dataset = CustomDataset(val_segments, val_labels, 0., 'BH')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19cf67bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "del files_names, train_patients, val_patients\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5a62a63e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(torch.nn.Module):\n",
    "    def __init__(self, base):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "    \n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x: (B, C=no_input_features, T)\n",
    "        _, d, T = x.shape\n",
    "        \n",
    "        pe = torch.zeros((d, T), device=x.device)\n",
    "        t = torch.arange(T, device=x.device).reshape(1, T)\n",
    "        \n",
    "        col_len = math.ceil(d / 2)\n",
    "        col = torch.arange(start=1, end=col_len+1, device=x.device).reshape(col_len, 1)\n",
    "        \n",
    "        angle = t / (self.base ** ((2 * col) / d))\n",
    "        sin = torch.sin(angle)\n",
    "        \n",
    "        if d % 2 == 0:\n",
    "            cos = torch.cos(angle)\n",
    "        else:\n",
    "            cos = torch.cos(angle[0:-1])\n",
    "        \n",
    "        #cos = torch.where(torch.tensor(d % 2 == 0, device=x.device), torch.cos(angle), torch.cos(angle[0:-1]))\n",
    "        \n",
    "        pe[0::2, :] = sin\n",
    "        pe[1::2, :] = cos\n",
    "        \n",
    "        # return shape: (B, C=no_input_features, T)\n",
    "        return pe + x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bacf0b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskedMultiHeadSelfAttention(torch.nn.Module):\n",
    "    def __init__(self, no_input_features, H, d_qk, d_v, bias_att=True):\n",
    "        super().__init__()\n",
    "        self.H = H\n",
    "        self.d_qk = d_qk\n",
    "        self.d_v = d_v\n",
    "        \n",
    "        self.q = torch.nn.Parameter(torch.empty(H, d_qk))\n",
    "        torch.nn.init.xavier_normal_(self.q)\n",
    "        \n",
    "        self.att_lin = torch.nn.Linear(\n",
    "            in_features=no_input_features,\n",
    "            out_features=(H * d_qk) + (H * d_v),\n",
    "            bias=bias_att\n",
    "        )\n",
    "    \n",
    "    \n",
    "    \n",
    "    @staticmethod\n",
    "    def qkv_func(q, k, v):\n",
    "        # num_freq_bands: F or E or T\n",
    "        # q: (B, H, 1, d_qk)\n",
    "        # k: (B, H, F or E or T, d_qk)\n",
    "        # v: (B, H, F or E or T, d_v)\n",
    "        \n",
    "        B, H, T, d_qk = q.shape\n",
    "        d_v = v.shape[-1]\n",
    "        \n",
    "        # logits: (B, H, 1, T)\n",
    "        logits = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_qk)     # logits: (B, H, 1, T)\n",
    "        s_m = logits.softmax(dim=-1)                                        # s_m: (B, H, 1, T)\n",
    "        \n",
    "        # (B, H, 1, T) * (B, H, T, d_v) = (B, H, 1, d_v) ===permute===> (B, 1, H, d_v) ===reshape===> (B, H * d_v)\n",
    "        return torch.matmul(s_m, v).transpose(-2,-3).reshape((B, H * d_v))\n",
    "    \n",
    "    \n",
    "    \n",
    "    def forward(self, x, dim):\n",
    "        if dim not in {-1, -2}:\n",
    "            raise NameError(f'Error: dim ({dim}) has to be either -1 or -2 !')\n",
    "        \n",
    "        if dim == -1:\n",
    "            # x: (B, C=no_input_features, T)\n",
    "            x.transpose_(-1, -2)\n",
    "        \n",
    "        # x: (B, T, C=no_input_features)\n",
    "        B, T, _ = x.shape\n",
    "        \n",
    "        k, v = self.att_lin(x).split(split_size=[self.H * self.d_qk, self.H * self.d_v], dim=-1)\n",
    "        \n",
    "        q = self.q.expand(B, 1, self.H, self.d_qk).transpose(-2, -3)     # q: (B, self.H, 1, self.d_qk)\n",
    "        k = k.reshape((B, T, self.H, self.d_qk)).transpose(-2, -3)       # k: (B, self.H, T, self.d_qk)\n",
    "        v = v.reshape((B, T, self.H, self.d_v)).transpose(-2, -3)        # v: (B, self.H, T, self.d_v)\n",
    "        \n",
    "        qkv = self.qkv_func(q, k, v)                                     # qkv: (B, self.H * self.d_v)\n",
    "        \n",
    "        if dim == -1:\n",
    "            x.transpose_(-1, -2)                                         # (B, C=no_input_features, T)\n",
    "        \n",
    "        return qkv\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "078f1848",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler):\n",
    "\n",
    "    def __init__(self, optimizer, warmup, max_num_iters):\n",
    "        self.warmup = warmup\n",
    "        self.max_num_iters = max_num_iters\n",
    "        super().__init__(optimizer)\n",
    "    \n",
    "    \n",
    "    def get_lr(self):\n",
    "        # returns a list of the learning rates\n",
    "        lr_factor = self.get_lr_factor(epoch=self.last_epoch)\n",
    "        return [base_lr * lr_factor for base_lr in self.base_lrs]\n",
    "    \n",
    "    \n",
    "    def get_lr_factor(self, epoch):\n",
    "        # Optional method that computes lr_factor, NOT learning rate itself.\n",
    "        if epoch <= self.warmup:\n",
    "            slope = (1 - 0) / (self.warmup - 0)\n",
    "            lr_factor = slope * epoch\n",
    "        else:\n",
    "            f = 1 / (2 * (self.max_num_iters - self.warmup))\n",
    "            lr_factor = 0.5 * ( 1 + math.cos( 2 * math.pi * f * (epoch - self.warmup) ) )\n",
    "        \n",
    "        return lr_factor\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cff7461f",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_validation_loss = math.inf\n",
    "best_epoch = 0\n",
    "best_model_state = None\n",
    "AP_at_min_val_loss = 0.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8fa61470",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the pytorch lightning module\n",
    "class CustomLightningModule(pl.LightningModule):\n",
    "    def __init__(self, p_dropout, window_len, overlap, sample_H, sample_d_qk, sample_d_v, mid_features, epoch_H, epoch_d_qk, epoch_d_v):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.window_len = window_len\n",
    "        self.increment = round((1 - overlap) * window_len)\n",
    "        \n",
    "        \n",
    "        self.sample_pe = PositionalEncoding(100)\n",
    "        \n",
    "        self.sample_attention = torch.nn.ModuleList()\n",
    "        self.sample_nonlinearity = torch.nn.ModuleList()\n",
    "        for end_idx in range(window_len - 1, 300, self.increment):\n",
    "            att = MaskedMultiHeadSelfAttention(19, sample_H, sample_d_qk, sample_d_v)\n",
    "            self.sample_attention.append(att)\n",
    "            \n",
    "            non_linear = torch.nn.Sequential(\n",
    "                torch.nn.Linear(in_features=sample_H*sample_d_v, out_features=mid_features),\n",
    "                torch.nn.LeakyReLU(negative_slope=0.01),\n",
    "                torch.nn.Dropout(p=p_dropout)\n",
    "            )\n",
    "            self.sample_nonlinearity.append(non_linear)\n",
    "            \n",
    "        \n",
    "        if end_idx != 299:\n",
    "            raise NameError(f'Error: Inspect carefully the overlap ({overlap}) argument!')\n",
    "        \n",
    "        \n",
    "        self.epoch_pe = PositionalEncoding(100)\n",
    "        self.epoch_att = MaskedMultiHeadSelfAttention(mid_features, epoch_H, epoch_d_qk, epoch_d_v)\n",
    "        \n",
    "        self.seq = torch.nn.Sequential(\n",
    "            torch.nn.Linear(in_features=epoch_H*epoch_d_v, out_features=10),\n",
    "            torch.nn.LeakyReLU(negative_slope=0.01),\n",
    "            torch.nn.Dropout(p=p_dropout),\n",
    "            torch.nn.Linear(in_features=10, out_features=5),\n",
    "            torch.nn.LeakyReLU(negative_slope=0.01),\n",
    "            torch.nn.Dropout(p=p_dropout),\n",
    "            torch.nn.Linear(in_features=5, out_features=1)\n",
    "        )\n",
    "        \n",
    "        \n",
    "        self.criterion = torch.nn.BCEWithLogitsLoss()\n",
    "        self.automatic_optimization = False\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x: (B, 19, T)\n",
    "        \n",
    "        lis = []\n",
    "        for i, (att, non_linear) in enumerate(zip(self.sample_attention, self.sample_nonlinearity)):\n",
    "            start_idx = i * self.increment\n",
    "            end_idx = start_idx + self.window_len\n",
    "            \n",
    "            encoded_epoch = self.sample_pe(x[:, :, start_idx:end_idx])\n",
    "            attention_output = att(encoded_epoch, -1)            # (B, sample_H * sample_d_v)\n",
    "            \n",
    "            non_linear_output = non_linear(attention_output)     # (B, mid_features)\n",
    "            lis.append(non_linear_output)\n",
    "        \n",
    "        \n",
    "        epochs = torch.stack(lis, -1)                     # (B, mid_features, no_epochs)\n",
    "        epochs = self.epoch_pe(epochs)                    # (B, mid_features, no_epochs)\n",
    "        attention_output = self.epoch_att(epochs, -1)     # (B, epoch_H * epoch_d_v)\n",
    "        \n",
    "        logits = self.seq(attention_output)             # (B, 1)\n",
    "        \n",
    "        return logits.squeeze(dim=-1)\n",
    "    \n",
    "    \n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), weight_decay=0)\n",
    "        cosine_scheduler = CosineWarmupScheduler(optimizer=optimizer, warmup=100, max_num_iters=self.trainer.max_epochs)\n",
    "        swa_scheduler = SWALR(optimizer, swa_lr=0.0001, anneal_epochs=100, anneal_strategy='cos')\n",
    "        return [optimizer], [cosine_scheduler, swa_scheduler]\n",
    "    \n",
    "    \n",
    "    \n",
    "    def train_dataloader(self):\n",
    "        dataset = CustomDataset(train_segments, train_labels, 0.5, montage)\n",
    "        # mp.cpu_count()-1\n",
    "        return torch.utils.data.DataLoader(dataset, batch_size=train_B, shuffle=True, num_workers=0,\n",
    "                                           pin_memory=True, drop_last=True)\n",
    "    \n",
    "    \n",
    "    \n",
    "    def val_dataloader(self):\n",
    "        dataset = CustomDataset(val_segments, val_labels, 0, montage)\n",
    "        # mp.cpu_count()-1\n",
    "        return torch.utils.data.DataLoader(dataset, batch_size=val_B, shuffle=False, num_workers=0,\n",
    "                                           pin_memory=True, drop_last=False)\n",
    "    \n",
    "    \n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        \n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            logits_before_optimizer_step = self(x)\n",
    "            loss_before_optimizer_step = self.criterion(logits_before_optimizer_step, y)\n",
    "        self.train()\n",
    "        \n",
    "        \n",
    "        opt = self.optimizers()\n",
    "\n",
    "        def closure():\n",
    "            logits = self(x)\n",
    "            loss = self.criterion(logits, y)\n",
    "            opt.zero_grad()\n",
    "            self.manual_backward(loss)\n",
    "            return loss\n",
    "\n",
    "        opt.step(closure=closure)\n",
    "        \n",
    "        return {'loss_before_optimizer_step': loss_before_optimizer_step.detach()}\n",
    "    \n",
    "    \n",
    "    \n",
    "    def training_epoch_end(self, training_step_outputs):\n",
    "        if self.global_rank == 0:\n",
    "            train_loss = torch.stack([dic['loss_before_optimizer_step'] for dic in training_step_outputs]).mean()\n",
    "            print(f'Epoch: {self.current_epoch}')\n",
    "            print(f'Training loss BEFORE optimizer step: {round(train_loss.item(), 5)}')\n",
    "            print('--------------------------------------')\n",
    "    \n",
    "    \n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        loss = self.criterion(logits, y)\n",
    "        return {'loss': loss, 'y': y, 'logits': logits}\n",
    "    \n",
    "    \n",
    "    \n",
    "    def validation_epoch_end(self, validation_step_outputs):\n",
    "        global min_validation_loss, best_epoch, best_model_state, AP_at_min_val_loss\n",
    "        \n",
    "        swa_start_epoch = 1000\n",
    "        \n",
    "        cosine_scheduler, swa_scheduler = self.lr_schedulers()\n",
    "        \n",
    "        if self.current_epoch > 0:\n",
    "            if self.current_epoch < swa_start_epoch:\n",
    "                cosine_scheduler.step()\n",
    "            else:\n",
    "                swa_scheduler.step()\n",
    "            \n",
    "        \n",
    "        \n",
    "        if self.global_rank == 0:\n",
    "            \n",
    "            if self.current_epoch > (swa_start_epoch + swa_scheduler.anneal_epochs):\n",
    "                swa_model.update_parameters(self)\n",
    "                \n",
    "            \n",
    "            val_loss = torch.stack([dic['loss'] for dic in validation_step_outputs]).mean()\n",
    "            val_loss = val_loss.item()\n",
    "            print(f'val_loss: {val_loss}')\n",
    "            \n",
    "            y = torch.cat([dic['y'] for dic in validation_step_outputs]).cpu()\n",
    "            logits = torch.cat([dic['logits'] for dic in validation_step_outputs]).cpu()\n",
    "            pred = torch.sigmoid(logits)\n",
    "            \n",
    "            y = np.array(y)\n",
    "            pred = np.array(pred)\n",
    "            \n",
    "            try:\n",
    "                validation_AP = average_precision_score(y, pred)\n",
    "                \n",
    "                if (val_loss < min_validation_loss) and (self.current_epoch > 1):\n",
    "                    min_validation_loss = val_loss\n",
    "                    best_epoch = self.current_epoch\n",
    "                    best_model_state = deepcopy(self.state_dict())\n",
    "                    AP_at_min_val_loss = validation_AP\n",
    "                \n",
    "                \n",
    "                print(f'validation_AP = {round(validation_AP, 5)} ... AP_at_min_val_loss = {round(AP_at_min_val_loss, 5)} @ epoch {best_epoch}')\n",
    "                \n",
    "            except ValueError:\n",
    "                print(f'ValueError @ epoch: {self.current_epoch} ====> y = {y}')\n",
    "            \n",
    "            print('=======================================================================================')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbf8924e",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_lightning_module = CustomLightningModule(p_dropout, window_len, overlap, sample_H, sample_d_qk,\n",
    "                                            sample_d_v, mid_features, epoch_H, epoch_d_qk, epoch_d_v)\n",
    "\n",
    "swa_model = AveragedModel(my_lightning_module)\n",
    "\n",
    "trainer = pl.Trainer(gpus=0, enable_checkpointing=False, enable_progress_bar=False, logger=False, max_epochs=400)\n",
    "\n",
    "trainer.fit(my_lightning_module)\n",
    "\n",
    "# Update bn statistics for the swa_model at the end\n",
    "torch.optim.swa_utils.update_bn(my_lightning_module.train_dataloader(), swa_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f436adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = f'models/g1_hierarchical_attention_10k_AP_{round(100 * AP_at_min_val_loss, 2)}.pt'\n",
    "torch.save(best_model_state, PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f83a21a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07da70df",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2768e1d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4ed4919a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "\n",
    "model = CustomLightningModule(p_dropout, window_len, overlap,\n",
    "                              sample_H, sample_d_qk, sample_d_v,\n",
    "                              mid_features, epoch_H, epoch_d_qk, epoch_d_v)\n",
    "\n",
    "#IEDs_names = [item for item in os.listdir(data_folder) if 'sp' in item]\n",
    "NIEDs_names = [item for item in os.listdir(data_folder) if 'ns' in item]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "945e7b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "lis = [item for item in os.listdir('models/') if 'g1_hierarchical_attention' in item]\n",
    "g1_model_name = max(lis)\n",
    "train_segments, _ = prepare_dataset(g2 + g3 + g4 + g5)\n",
    "g1_median, g1_IQR = median_IQR_func(train_segments)\n",
    "\n",
    "lis = [item for item in os.listdir('models/') if 'g2_hierarchical_attention' in item]\n",
    "g2_model_name = max(lis)\n",
    "train_segments, _ = prepare_dataset(g1 + g3 + g4 + g5)\n",
    "g2_median, g2_IQR = median_IQR_func(train_segments)\n",
    "\n",
    "\n",
    "lis = [item for item in os.listdir('models/') if 'g3_hierarchical_attention' in item]\n",
    "g3_model_name = max(lis)\n",
    "train_segments, _ = prepare_dataset(g1 + g2 + g4 + g5)\n",
    "g3_median, g3_IQR = median_IQR_func(train_segments)\n",
    "\n",
    "\n",
    "lis = [item for item in os.listdir('models/') if 'g4_hierarchical_attention' in item]\n",
    "g4_model_name = max(lis)\n",
    "train_segments, _ = prepare_dataset(g1 + g2 + g3 + g5)\n",
    "g4_median, g4_IQR = median_IQR_func(train_segments)\n",
    "\n",
    "\n",
    "lis = [item for item in os.listdir('models/') if 'g5_hierarchical_attention' in item]\n",
    "g5_model_name = max(lis)\n",
    "train_segments, _ = prepare_dataset(g1 + g2 + g3 + g4)\n",
    "g5_median, g5_IQR = median_IQR_func(train_segments)\n",
    "\n",
    "\n",
    "#s = pd.Series(index=IEDs_names, dtype='float64')\n",
    "s = pd.Series(index=NIEDs_names, dtype='float64')\n",
    "\n",
    "\n",
    "#for IED_name in IEDs_names:\n",
    "for NIED_name in NIEDs_names:\n",
    "    #segment = loadmat(os.path.join(data_folder, IED_name))['segment']\n",
    "    segment = loadmat(os.path.join(data_folder, NIED_name))['segment']\n",
    "    segment = torch.tensor(segment).float()\n",
    "    \n",
    "    #patient_name = IED_name[: IED_name.find('_')]\n",
    "    patient_name = NIED_name[: NIED_name.find('_')]\n",
    "    \n",
    "    if patient_name in g1:\n",
    "        segment = (segment - g1_median) / g1_IQR\n",
    "        model.load_state_dict(torch.load(f'models/{g1_model_name}'))\n",
    "    elif patient_name in g2:\n",
    "        segment = (segment - g2_median) / g2_IQR\n",
    "        model.load_state_dict(torch.load(f'models/{g2_model_name}'))\n",
    "    elif patient_name in g3:\n",
    "        segment = (segment - g3_median) / g3_IQR\n",
    "        model.load_state_dict(torch.load(f'models/{g3_model_name}'))\n",
    "    elif patient_name in g4:\n",
    "        segment = (segment - g4_median) / g4_IQR\n",
    "        model.load_state_dict(torch.load(f'models/{g4_model_name}'))\n",
    "    elif patient_name in g5:\n",
    "        segment = (segment - g5_median) / g5_IQR\n",
    "        model.load_state_dict(torch.load(f'models/{g5_model_name}'))\n",
    "    else:\n",
    "        raise NameError('Unknown Patient')\n",
    "    \n",
    "    \n",
    "    segment.unsqueeze_(0)\n",
    "    model.eval()\n",
    "    \n",
    "    #s.at[IED_name] = torch.sigmoid(model(segment)).item()\n",
    "    s.at[NIED_name] = torch.sigmoid(model(segment)).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7db47305",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RC_24_ns.mat      0.002088\n",
       "FEsg_38_ns.mat    0.013555\n",
       "HB_2_ns.mat       0.010820\n",
       "FEsg_26_ns.mat    0.218849\n",
       "FigSa_9_ns.mat    0.083698\n",
       "                    ...   \n",
       "LM_3_ns.mat       0.047767\n",
       "RC_6_ns.mat       0.009134\n",
       "FEsg_19_ns.mat    0.042846\n",
       "NKra_10_ns.mat    0.115910\n",
       "PAT48_ns_6.mat    0.001987\n",
       "Length: 593, dtype: float64"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "be215587",
   "metadata": {},
   "outputs": [],
   "source": [
    "#with pd.ExcelWriter('IEDs_results.xlsx', mode='a') as writer:\n",
    "with pd.ExcelWriter('NIEDs_results.xlsx', mode='a') as writer:\n",
    "    s.to_excel(writer, sheet_name='hierarchical_attention', header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c6de9d0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b0935d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19836318",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_env",
   "language": "python",
   "name": "pytorch_env"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}