Self_Attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62e5a503",
   "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": "3bb66c13",
   "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": "7f0debac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch size\n",
    "train_B = 16\n",
    "val_B = 1\n",
    "\n",
    "p_dropout = 0.\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": "82c9ccba",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_patients = g1 + g2 + g3 + g5\n",
    "val_patients = g4\n",
    "#del g1, g2, g3, g4, g5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a56b1292",
   "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": "fc78fcc3",
   "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": "582de498",
   "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",
    "    \n",
    "    return group_segments, group_labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b077a1d6",
   "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": "2b3122e3",
   "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": "61afc9dc",
   "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": "f37530b3",
   "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": "d1ce51a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "del files_names, train_segments, train_labels, train_patients, val_patients, val_segments, val_labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "87d2335d",
   "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",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1df36ebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskedMultiHeadSelfAttention(torch.nn.Module):\n",
    "    def __init__(self, no_input_features, H, d_qk, d_v, bias_att=True, bias_out=True):\n",
    "        super().__init__()\n",
    "        self.H = H\n",
    "        self.d_qk = d_qk\n",
    "        self.d_v = d_v\n",
    "        self.att_lin = torch.nn.Linear(in_features=no_input_features, out_features=(2 * H * d_qk) + (H * d_v), bias=bias_att)\n",
    "        self.out_lin = torch.nn.Linear(in_features=H*d_v, out_features=no_input_features, bias=bias_out)\n",
    "    \n",
    "    \n",
    "    \n",
    "    @staticmethod\n",
    "    def qkv_func(q, k, v, mask=None):\n",
    "        # q: (B, H, T, d_qk)\n",
    "        # k: (B, H, T, d_qk)\n",
    "        # v: (B, H, T, d_v)\n",
    "        \n",
    "        B, H, T, d_qk = q.shape\n",
    "        d_v = v.shape[-1]\n",
    "        \n",
    "        # logits: (B, H, T, T)\n",
    "        logits = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_qk)\n",
    "        \n",
    "        if mask is not None:\n",
    "            # mask MUST have shape (B, H, T, T).\n",
    "            # For example, to construct a mask that doesn't depend on future time steps\n",
    "            # mask = (torch.ones(size=(B, H, T, T)).tril() == 1)\n",
    "            # tril: https://pytorch.org/docs/stable/generated/torch.tril.html\n",
    "            logits[mask == False] = -9e15\n",
    "        \n",
    "        \n",
    "        s_m = logits.softmax(dim=-1)     # s_m: (B, H, T, T)\n",
    "        \n",
    "        # (B, H, T, T) * (B, H, T, d_v) = (B, H, T, d_v) ===permute===> (B, T, H, d_v) ===reshape===> (B, T, H * d_v)\n",
    "        return torch.matmul(s_m, v).permute((0, 2, 1, 3)).reshape((B, T, H * d_v))\n",
    "    \n",
    "    \n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        # x: (B, T, C=no_input_features) ====self.att_lin====> ( B, T, (H * d_qk) + (H * d_qk) + (H * d_v) )\n",
    "        \n",
    "        B, T, _ = x.shape\n",
    "        \n",
    "        q, k, v = self.att_lin(x).split(split_size=[self.H * self.d_qk, self.H * self.d_qk, self.H * self.d_v], dim=-1)\n",
    "        \n",
    "        q = q.reshape((B, T, self.H, self.d_qk)).permute((0, 2, 1, 3))     # q: (B, self.H, T, self.d_qk)\n",
    "        k = k.reshape((B, T, self.H, self.d_qk)).permute((0, 2, 1, 3))     # k: (B, self.H, T, self.d_qk)\n",
    "        v = v.reshape((B, T, self.H, self.d_v)).permute((0, 2, 1, 3))     # v: (B, self.H, T, self.d_v)\n",
    "        \n",
    "        # returned tensor: (B, T, C=no_input_features)\n",
    "        return self.out_lin( self.qkv_func(q, k, v, mask=mask) )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2110c65f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(torch.nn.Module):\n",
    "    def __init__(self, no_input_features, no_mid_features, H, d_qk, d_v, bias_att=True, bias_out=True, p_dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.att = MaskedMultiHeadSelfAttention(no_input_features, H, d_qk, d_v, bias_att=bias_att, bias_out=bias_out)\n",
    "        self.norm_att = torch.nn.LayerNorm(no_input_features)\n",
    "        \n",
    "        self.seq_feed_forward = torch.nn.Sequential(\n",
    "            torch.nn.Linear(no_input_features, no_mid_features),\n",
    "            torch.nn.Dropout(p_dropout),\n",
    "            torch.nn.LeakyReLU(inplace=True),\n",
    "            torch.nn.Linear(no_mid_features, no_input_features)\n",
    "        )\n",
    "        self.norm_feed_forward = torch.nn.LayerNorm(no_input_features)\n",
    "        \n",
    "        self.dropout = torch.nn.Dropout(p_dropout)\n",
    "    \n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        # x: (B, C=no_input_features, T)\n",
    "        \n",
    "        x.transpose_(-1,-2)     # x: (B, T, C=no_input_features)\n",
    "        \n",
    "        x = self.norm_att( x + self.dropout( self.att(x, mask) ) )\n",
    "        \n",
    "        x = self.norm_feed_forward( x + self.dropout( self.seq_feed_forward(x) ) )\n",
    "        \n",
    "        x.transpose_(-1,-2)     # x: (B, C=no_input_features, T)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b80fb8",
   "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": "00a2614a",
   "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": 10,
   "id": "76cc5b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the pytorch lightning module\n",
    "class CustomLightningModule(pl.LightningModule):\n",
    "    def __init__(self, p_dropout):\n",
    "        super().__init__()\n",
    "        self.encoders_seq = torch.nn.Sequential(\n",
    "            # i/p: (B, 19, 300) =====> o/p: (B, 19, 300)\n",
    "            PositionalEncoding(100),\n",
    "            \n",
    "            # i/p: (B, 19, 300) =====> o/p: (B, 19, 300)\n",
    "            #EncoderBlock(19, 10, 8, 2, 2, p_dropout=p_dropout),\n",
    "            EncoderBlock(19, 19, 7, 7, 7, p_dropout=p_dropout),\n",
    "            \n",
    "            # i/p: (B, 19, 300) ==/6==> o/p: (B, 19, 50)\n",
    "            torch.nn.MaxPool1d(6),\n",
    "            \n",
    "            \n",
    "            PositionalEncoding(100),\n",
    "            \n",
    "            # i/p: (B, 19, 50) =====> o/p: (B, 19, 50)\n",
    "            #EncoderBlock(19, 10, 8, 2, 2, p_dropout=p_dropout),\n",
    "            EncoderBlock(19, 19, 6, 6, 6, p_dropout=p_dropout),\n",
    "            \n",
    "            # i/p: (B, 19, 50) ==/5==> o/p: (B, 19, 10)\n",
    "            torch.nn.MaxPool1d(5),\n",
    "            \n",
    "            \n",
    "            PositionalEncoding(100),\n",
    "            \n",
    "            # i/p: (B, 19, 10) ====> o/p: (B, 19, 10)\n",
    "            EncoderBlock(19, 19, 4, 4, 4, p_dropout=p_dropout),\n",
    "            \n",
    "            # i/p: (B, 19, 10) ==/2==> o/p: (B, 19, 5)\n",
    "            torch.nn.MaxPool1d(2),\n",
    "            \n",
    "            \n",
    "            \n",
    "            # i/p: (B, 19, 5) =====> o/p: (B, 1, 5)\n",
    "            torch.nn.Conv1d(in_channels=19, out_channels=1, kernel_size=1, padding='valid'),\n",
    "            torch.nn.LeakyReLU(),\n",
    "            #torch.nn.Dropout(p=0.3),\n",
    "            \n",
    "            # i/p: (B, 1, 5) =====> o/p: (B, 1, 1)\n",
    "            torch.nn.Conv1d(in_channels=1, out_channels=1, kernel_size=5, padding='valid')\n",
    "        )\n",
    "        self.criterion = torch.nn.BCEWithLogitsLoss()\n",
    "        self.automatic_optimization = False\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        # x: (B, 19, 300) ===> Need to permute the last 2 dimensions\n",
    "        \n",
    "        # returned tensor: (B,)\n",
    "        return self.encoders_seq(x).reshape((x.shape[0],))\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",
    "    def train_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            train_dataset, batch_size=train_B, shuffle=True, num_workers=0, drop_last=True\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",
    "        # Note that this is NOT a pytorch optimizer. I think it's a wrapper\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",
    "        # I think since I am doing manual optimization, there is no need to return anything\n",
    "        return {'loss_before_optimizer_step': loss_before_optimizer_step.detach()}\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",
    "    def val_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            val_dataset, batch_size=val_B, shuffle=False, num_workers=0, drop_last=False\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",
    "    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": "3b4277b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_lightning_module = CustomLightningModule(p_dropout)\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=1000)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd8e953",
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = f'models/g4_self_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": "3ed1490c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54df3b60",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdfbe96",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "20e1ba45",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "model = CustomLightningModule(p_dropout)\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": 12,
   "id": "b497b43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "lis = [item for item in os.listdir('models/') if 'g1_self_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_self_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_self_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_self_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_self_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",
    "#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": 13,
   "id": "b258fb36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RC_24_ns.mat      0.005338\n",
       "FEsg_38_ns.mat    0.019621\n",
       "HB_2_ns.mat       0.000733\n",
       "FEsg_26_ns.mat    0.017761\n",
       "FigSa_9_ns.mat    0.004138\n",
       "                    ...   \n",
       "LM_3_ns.mat       0.040386\n",
       "RC_6_ns.mat       0.004489\n",
       "FEsg_19_ns.mat    0.019402\n",
       "NKra_10_ns.mat    0.017798\n",
       "PAT48_ns_6.mat    0.000081\n",
       "Length: 593, dtype: float64"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aaecdd96",
   "metadata": {},
   "outputs": [],
   "source": [
    "#with pd.ExcelWriter('IEDs_results.xlsx', mode='w') as writer:\n",
    "with pd.ExcelWriter('NIEDs_results.xlsx', mode='w') as writer:\n",
    "    s.to_excel(writer, sheet_name='self_attention', header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73cfd1d4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5142048c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f947911",
   "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
}