deep_sleep2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CTS5vGYAruq4"
   },
   "source": [
    "# Deep Sleep 2.0 - Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M7b-e6cXruq_"
   },
   "source": [
    "## Main Switches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QB2piScSsAe3"
   },
   "outputs": [],
   "source": [
    "''' \n",
    "MODEL: \n",
    "Select an existing model configuration (see the 'models' folder and paper) or \n",
    "define a custom configuration using a JSON configuration file \n",
    "'''\n",
    "MODEL_NAME = 'model_2' # available: {'model_0'; 'model_1'; 'model_2'; 'model_3'}\n",
    "\n",
    "''' \n",
    "MODE: \n",
    "Select between model training (True) or model inference (False)\n",
    "'''\n",
    "TRAIN_MODE = True  # available options: {True; False}\n",
    "\n",
    "''' \n",
    "CHECKPOINT: \n",
    "Load the last available training checkpoint?\n",
    "'''\n",
    "LOAD_CHECKPOINT = False  # available options: {True; False}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zu4kHt-rt7UU"
   },
   "source": [
    "## Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FYBk22Erruq_"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import pickle\n",
    "import json\n",
    "import csv\n",
    "import time\n",
    "import importlib\n",
    "import matplotlib.pyplot as plt\n",
    "from itertools import zip_longest\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.data.dataset import random_split\n",
    "from torchsummary import summary\n",
    "from losses import CustomBCELoss, CustomBCEWithLogitsLoss\n",
    "from tqdm import tqdm\n",
    "from score2018 import Challenge2018Score\n",
    "from utils import preprocess, get_record, folders_to_records_txt\n",
    "\n",
    "print('GPU is', '' if torch.cuda.is_available() else 'not', 'available.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_PATH        = './'\n",
    "DATA_PATH        = os.path.join(ROOT_PATH, 'data')\n",
    "TRAIN_DATA_PATH  = os.path.join(DATA_PATH, 'training')\n",
    "TEST_DATA_PATH   = os.path.join(DATA_PATH, 'test')\n",
    "MODEL_PATH       = os.path.join(ROOT_PATH, 'models', MODEL_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "geuVZ74wt7UV"
   },
   "source": [
    "## Hyperparameters and Reproducibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GP7skJyit7UV"
   },
   "outputs": [],
   "source": [
    "assert os.path.isfile(os.path.join(MODEL_PATH, 'hyperparameters.txt')), \\\n",
    "\"File hyperparameters.txt does not exist in: \" + MODEL_PATH + os.path.sep\n",
    "\n",
    "# load the model configuration and hyperparameters from the JSON file\n",
    "with open(os.path.join(MODEL_PATH, 'hyperparameters.txt')) as f:\n",
    "    hyperparameters = json.load(f)\n",
    "print('\\033[1m',\"Hyperparameters specific to the selected model:\",'\\033[0m', \\\n",
    "      *hyperparameters.items(), sep='\\n')\n",
    "    \n",
    "ARCHITECTURE_NAME = hyperparameters['ARCHITECTURE_NAME']\n",
    "SEED              = hyperparameters['SEED']\n",
    "MAX_NUM_EPOCHS    = hyperparameters['MAX_NUM_EPOCHS']\n",
    "CHANNELS          = hyperparameters['CHANNELS']\n",
    "LEARNING_RATE     = hyperparameters['LEARNING_RATE']\n",
    "DECAY_RATE        = hyperparameters['DECAY_RATE']\n",
    "STOP_STRIP        = hyperparameters['STOP_STRIP']\n",
    "DEVICE            = hyperparameters['DEVICE']\n",
    "BATCH_SIZE        = hyperparameters['BATCH_SIZE']\n",
    "NUM_WORKERS       = hyperparameters['NUM_WORKERS']\n",
    "PIN_MEMORY        = hyperparameters['PIN_MEMORY']\n",
    "LINEAR            = hyperparameters['LINEAR']\n",
    "Z_NORM            = hyperparameters['Z_NORM']\n",
    "TRANSFORMS        = hyperparameters['TRANSFORMS']\n",
    "\n",
    "# to enable analysis/inference on CPU-only devices\n",
    "if not TRAIN_MODE: \n",
    "    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\" \n",
    "    BATCH_SIZE = 2 if torch.cuda.is_available() else 1\n",
    "    NUM_WORKERS = 2 if torch.cuda.is_available() else 0\n",
    "\n",
    "# set seed for reproducibility\n",
    "random.seed(SEED) \n",
    "np.random.seed(SEED) \n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "os.environ[\"PYTHONHASHSEED\"] = str(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K2ApC-ATrurA"
   },
   "source": [
    "## Data Availability & Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KfzhDTc5rurB"
   },
   "outputs": [],
   "source": [
    "'''\n",
    "Identify all records stored in '/data/*' (incl. subfolders) and save \n",
    "their respective paths in 'RECORDS.txt'.\n",
    "'''\n",
    "folders_to_records_txt(DATA_PATH)\n",
    "\n",
    "'''\n",
    "Uniform all recording/arousal signals to have the same 8-milion length \n",
    "(2^23 = 8,388,608) by padding zeros and centering the recording/arousal \n",
    "signals. Perform Z-score normalization, if applicable.\n",
    "'''\n",
    "preprocess(DATA_PATH, Z_NORM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BiPIA1EirurE"
   },
   "source": [
    "## Custom Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JyRSxXstrurE"
   },
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    \"\"\"Custom dataset for multichannel PSG recordings and sleep arousal annotations\"\"\"\n",
    "    \n",
    "    def __init__(self, data_path, records, channels = [*range(13)], z_norm = False, transforms = None):     \n",
    "        self.data_path = data_path\n",
    "        self.records = records\n",
    "        self.channels = channels\n",
    "        self.n_channels = len(channels)\n",
    "        self.normalization = 2 if z_norm else 1\n",
    "        self.transforms = transforms\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.records)        \n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        file_path = os.path.join(self.data_path, self.records[idx])\n",
    "        recording, arousal = get_record(file_path, self.channels, self.normalization)\n",
    "        \n",
    "        if self.transforms is not None:\n",
    "            recording, arousal = self.transforms(recording, arousal)\n",
    "        \n",
    "        return recording, arousal, idx\n",
    "\n",
    "    def get_record_name(self, idx):\n",
    "        return self.records[idx]   \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3T2FUZDFrurG"
   },
   "source": [
    "## Custom Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q2r55TlkrurG"
   },
   "outputs": [],
   "source": [
    "class MagScale(object):\n",
    "    \"\"\"Rescale the magnitude of all PSG channels with the same random scale factor\"\"\"\n",
    "    \n",
    "    def __init__(self, low = 0.8, high = 1.25):\n",
    "        self.low = low\n",
    "        self.high = high\n",
    "    \n",
    "    def __call__(self, recording, arousal):\n",
    "        scale = self.low + torch.rand(1)*(self.high - self.low)\n",
    "        recording = scale*recording\n",
    "\n",
    "        return recording, arousal   \n",
    "\n",
    "    \n",
    "class MagScaleRandCh(object):\n",
    "    \"\"\"Rescale the magnitude of a randomly selected PSG channel with a random scale factor\"\"\"\n",
    "    \n",
    "    def __init__(self, n_channels = 13, low = 0.8, high = 1.25):\n",
    "        self.n_channels = n_channels        \n",
    "        self.low = low\n",
    "        self.high = high\n",
    "    \n",
    "    def __call__(self, recording, arousal):\n",
    "        scales = self.low + torch.rand(self.n_channels).view(-1,1)*(self.high - self.low)\n",
    "        recording = scales*recording\n",
    "\n",
    "        return recording, arousal       \n",
    "\n",
    "    \n",
    "class RandShuffle(object):\n",
    "    \"\"\"Randomly reshuffle a subset of related PSG channels\"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.r2 = torch.LongTensor([6, 7])       # fixed channels (E1-M2, Chin)\n",
    "        self.r4 = torch.LongTensor([10, 11, 12]) # fixed channels (Airflow, SaO2, ECG)\n",
    "    \n",
    "    def __call__(self, recording, arousal):    \n",
    "        r1 = torch.randperm(6)     # shuffle EEG channels (F3-M2, F4-M1, C3-M2, C4-M1, O1-M2, O2-M1)\n",
    "        r3 = 8 + torch.randperm(2) # shuffle EMG channels (ABD, Chest)\n",
    "        r = torch.cat((r1, self.r2, r3, self.r4)).type(torch.long) \n",
    "        \n",
    "        return recording.index_select(0, r), arousal     \n",
    "\n",
    "    \n",
    "class AddRandGaussian2All(object):\n",
    "    \"\"\"Add zero-mean Gaussian noise to all PSG channels\"\"\"\n",
    "    \n",
    "    def __init__(self, z_norm = True):\n",
    "        self.z_norm = z_norm\n",
    "    \n",
    "    def __call__(self, recording, arousal):\n",
    "        if self.z_norm:\n",
    "            std_dev = 0.1 \n",
    "        else:\n",
    "            std_dev = 0.1*torch.std(recording, 1, keepdim = True)\n",
    "        recording = recording + std_dev*torch.randn(recording.shape)\n",
    "        \n",
    "        return recording, arousal    \n",
    "    \n",
    "    \n",
    "class InjectRandGaussian(object):\n",
    "    \"\"\"Replace a randomly selected PSG channel with a standard Gaussian noise sequence\"\"\"\n",
    "    \n",
    "    def __init__(self, n_channels = 13):\n",
    "        self.n_channels = n_channels\n",
    "    \n",
    "    def __call__(self, recording, arousal):\n",
    "        ri = torch.randint(0,self.n_channels,(1,)).type(torch.long)\n",
    "        recording[ri] = torch.normal(mean = 0, std = 1, size = (1, recording.shape[1]))\n",
    "        \n",
    "        return recording, arousal     \n",
    "    \n",
    "\n",
    "class TimeScale(object):\n",
    "    \"\"\"Stretch/shrink the recording and arousal signals with a random time scale while \n",
    "       maintaining the original lengths\"\"\"\n",
    "\n",
    "    def __init__(self, interval, n_channels = 13):\n",
    "        self.interval = interval\n",
    "        self.n_channels = n_channels        \n",
    "\n",
    "    def __call__(self, recording, arousal):\n",
    "        scale = 1 + self.interval*(torch.rand(1) - 0.5)\n",
    "        recording = F.interpolate(recording.reshape((1,self.n_channels,-1)), \\\n",
    "                                  scale_factor = scale, recompute_scale_factor = True)\n",
    "        arousal = F.interpolate(arousal.reshape((1,1,-1)), scale_factor = scale, \\\n",
    "                                recompute_scale_factor = True)\n",
    "\n",
    "        return recording, arousal\n",
    "\n",
    "    \n",
    "class ToTensor(object):\n",
    "    \"\"\"Convert the recording and arousal signals to Tensors\"\"\"\n",
    "\n",
    "    def __call__(self, recording, arousal):\n",
    "        return torch.Tensor(recording), torch.Tensor(arousal)     \n",
    "    \n",
    "    \n",
    "class Compose:\n",
    "    \"\"\"Stack multiple transforms together\"\"\"\n",
    "    def __init__(self, transforms):\n",
    "        self.transforms = transforms\n",
    "\n",
    "    def __call__(self, recording, arousal):\n",
    "        for t in self.transforms:\n",
    "            recording, arousal = t(recording, arousal)\n",
    "            \n",
    "        return recording, arousal   \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Kh1KB67HrurH"
   },
   "source": [
    "## Dataset Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pOk6Y7W9rurH"
   },
   "outputs": [],
   "source": [
    "assert os.path.isfile(os.path.join(TRAIN_DATA_PATH, 'RECORDS.txt')), \\\n",
    "\"File RECORDS.txt does not exist in: \" + TRAIN_DATA_PATH + os.path.sep\n",
    "\n",
    "with open(os.path.join(TRAIN_DATA_PATH, 'RECORDS.txt'), 'r') as f:\n",
    "    records_in_train_folder = f.read().splitlines()\n",
    "try:\n",
    "    with open(os.path.join(TEST_DATA_PATH, 'RECORDS.txt'), 'r') as f:\n",
    "        records_in_test_folder = f.read().splitlines()\n",
    "        print(\"PhysioNet test dataset available and will be used for testing (100%).\")\n",
    "except FileNotFoundError:\n",
    "        records_in_test_folder = None\n",
    "        print(\"PhysioNet test dataset not available.\")\n",
    "\n",
    "train_tot = len(records_in_train_folder)\n",
    "\n",
    "if not records_in_test_folder: \n",
    "    \n",
    "    '''\n",
    "    Default Case: complete PhysioNet test dataset not available. \n",
    "    The PhysioNet training dataset is split in training (60%), validation (15%), \n",
    "    and test (25%) sets.\n",
    "    '''\n",
    "    \n",
    "    TEST_DATA_PATH = TRAIN_DATA_PATH\n",
    "    \n",
    "    train_split = int(0.6*train_tot)\n",
    "    val_split   = int(0.15*train_tot)\n",
    "    test_split  = train_tot - train_split - val_split    \n",
    "    \n",
    "    train_records, tmp = train_test_split(records_in_train_folder, \\\n",
    "                                          test_size = val_split + test_split, \\\n",
    "                                          random_state = SEED)\n",
    "    val_records, test_records = train_test_split(tmp, test_size = test_split, \\\n",
    "                                                 random_state = SEED) \n",
    "    \n",
    "    print(\"The PhysioNet training dataset has been split in\", \n",
    "          \"training ({:2.0%}),\".format(train_split/train_tot),\n",
    "          \"validation ({:2.0%}),\".format(val_split/train_tot),\n",
    "          \"and test ({:2.0%}) sets.\".format(test_split/train_tot))\n",
    "    \n",
    "else:\n",
    "    \n",
    "    '''\n",
    "    Exceptional Case: complete PhysioNet test dataset is available (incl.labels!). \n",
    "    The PhysioNet training dataset is split in training (80%) and validation (20%) \n",
    "    sets. The complete PhysioNet test dataset is used for testing purposes.\n",
    "    '''\n",
    "    \n",
    "    train_split = int(0.8*train_tot)\n",
    "    val_split = train_tot - train_split\n",
    "    \n",
    "    train_records, val_records = train_test_split(records_in_train_folder, \\\n",
    "                                                  test_size = val_split, \\\n",
    "                                                  random_state = SEED)\n",
    "    test_records = records_in_test_folder\n",
    "    \n",
    "    print(\"The PhysioNet training dataset has been split in\",\n",
    "          \"training ({:2.0%})\".format(train_split/train_tot),\n",
    "          \"and validation ({:2.0%}) sets.\".format(val_split/train_tot))\n",
    "\n",
    "# Record the splits\n",
    "if TRAIN_MODE:\n",
    "    with open(os.path.join(MODEL_PATH, 'records.csv'),\"w+\") as f:\n",
    "        writer = csv.writer(f)\n",
    "        writer.writerow(['train_records', 'val_records', 'test_records'])\n",
    "        for values in zip_longest(*[train_records, val_records, test_records]):\n",
    "            writer.writerow(values)     \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transforms_list = [ToTensor()]\n",
    "\n",
    "for transform in TRANSFORMS:\n",
    "    if transform == 'MagScale':\n",
    "        train_transforms_list.append(MagScale())\n",
    "    if transform == 'MagScaleRandCh':\n",
    "        train_transforms_list.append(MagScaleRandCh())        \n",
    "    if transform == 'TimeScale':\n",
    "        train_transforms_list.append(TimeScale())\n",
    "    if transform == 'RandShuffle':\n",
    "        train_transforms_list.append(RandShuffle())\n",
    "    if transform == 'AddRandGaussian2All':\n",
    "        train_transforms_list.append(AddRandGaussian2All())\n",
    "    if transform == 'InjectRandGaussian':\n",
    "        train_transforms_list.append(InjectRandGaussian())\n",
    "        \n",
    "train_transforms = Compose(train_transforms_list)\n",
    "val_test_transforms = ToTensor()\n",
    "\n",
    "train_dataset = CustomDataset(TRAIN_DATA_PATH, train_records, CHANNELS, \\\n",
    "                              Z_NORM, train_transforms)\n",
    "val_dataset   = CustomDataset(TRAIN_DATA_PATH, val_records, CHANNELS, \\\n",
    "                              Z_NORM, val_test_transforms)\n",
    "test_dataset  = CustomDataset(TEST_DATA_PATH, test_records, CHANNELS, \\\n",
    "                              Z_NORM, val_test_transforms)\n",
    "\n",
    "print(\"Length of the train dataset is:\", len(train_dataset))\n",
    "print(\"Length of the val dataset is:\", len(val_dataset))\n",
    "print(\"Length of the test dataset is:\", len(test_dataset))\n",
    "\n",
    "train_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle = True, \\\n",
    "                          num_workers = NUM_WORKERS, pin_memory = PIN_MEMORY)\n",
    "val_loader   = DataLoader(val_dataset, BATCH_SIZE, shuffle = False, \\\n",
    "                          num_workers = NUM_WORKERS, pin_memory = PIN_MEMORY)\n",
    "test_loader  = DataLoader(test_dataset, BATCH_SIZE, shuffle = False, \\\n",
    "                          num_workers = NUM_WORKERS, pin_memory = PIN_MEMORY)\n",
    "\n",
    "print(\"Total number of train batches:\", len(train_loader))\n",
    "print(\"Total number of val batches:\", len(val_loader))\n",
    "print(\"Total number of test batches:\", len(test_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PouKkgXOxeYc"
   },
   "source": [
    "## Evaluation Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RrqrL-GTrurK"
   },
   "outputs": [],
   "source": [
    "def eval_fn(model, loader, DEVICE, comp_score = False):\n",
    "    \n",
    "    # prepare model for evaluation (disable batch normalization)\n",
    "    model.eval()\n",
    "    \n",
    "    # use the official PhysioNet scoring function\n",
    "    scores = Challenge2018Score() if comp_score else None\n",
    "    \n",
    "    # define loss function\n",
    "    if comp_score:\n",
    "        loss_fn = CustomBCELoss().to(DEVICE)\n",
    "    else:\n",
    "        loss_fn = CustomBCEWithLogitsLoss().to(DEVICE)\n",
    "    \n",
    "    # compute loss and, if applicable, score\n",
    "    with torch.no_grad():\n",
    "        loss_epoch_sum = 0\n",
    "        for x, y, idx in loader:          \n",
    "            x = x.to(device = DEVICE)\n",
    "            y = y.to(device = DEVICE)\n",
    "\n",
    "            with torch.cuda.amp.autocast(enabled = torch.cuda.is_available() \\\n",
    "                                         and not comp_score):\n",
    "\n",
    "                y_hat = model(x, comp_score)\n",
    "                loss = loss_fn(y_hat, y)\n",
    "            \n",
    "                # compute AUROC/AUPRC score for each record in batch\n",
    "                if comp_score:                \n",
    "                    for i, single_idx in enumerate(idx):\n",
    "                        record = loader.dataset.get_record_name(single_idx)\n",
    "                        scores.score_record(y[i].view(-1).to('cpu'), \\\n",
    "                                            y_hat[i].view(-1).to('cpu'), record)\n",
    "                        auroc = scores.record_auroc(record)\n",
    "                        auprc = scores.record_auprc(record)               \n",
    "                        print('%-11s  AUROC: %8.6f,  AUPRC: %8.6f' % \\\n",
    "                              (record, auroc, auprc))\n",
    "\n",
    "            loss_epoch_sum += float(loss.item())                 \n",
    "\n",
    "    return loss_epoch_sum/len(loader), scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YfyEP8bSrurL"
   },
   "source": [
    "## Training Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9EK38AZmrurL"
   },
   "outputs": [],
   "source": [
    "def train_fn(model, loader, optimizer, loss_fn, scaler, DEVICE):\n",
    "    \n",
    "    # prepare model for evaluation (enable batch normalization)\n",
    "    model.train()\n",
    "\n",
    "    loss_epoch_sum = 0\n",
    "    \n",
    "    with tqdm(total = len(loader)) as pbar:\n",
    "        start_time = time.time()\n",
    "        \n",
    "        for batch_idx, (x, y, _) in enumerate(loader):\n",
    "\n",
    "            \"\"\" Step 0. move the tensors to the right device \"\"\"    \n",
    "            \n",
    "            x = x.to(device = DEVICE)\n",
    "            y = y.to(device = DEVICE)  \n",
    "            \n",
    "            prepare_time = start_time - time.time()\n",
    "\n",
    "            \"\"\" Step 1. clear gradients \"\"\"\n",
    "            optimizer.zero_grad(set_to_none = True)  \n",
    "\n",
    "            with torch.cuda.amp.autocast(enabled = torch.cuda.is_available()):     \n",
    "\n",
    "                \"\"\" Step 2. Forward pass \"\"\"\n",
    "                y_hat = model(x)\n",
    "\n",
    "                \"\"\" Step 3. Loss calculation \"\"\"                \n",
    "                loss = loss_fn(y_hat, y)\n",
    "\n",
    "            \"\"\" Step 4. Backward pass \"\"\"\n",
    "            scaler.scale(loss).backward()\n",
    "\n",
    "            \"\"\" Step 5. Optimization (parameter update) \"\"\"\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "\n",
    "            \"\"\" Step 6. Timing and logging \"\"\"\n",
    "            loss_epoch_sum += float(loss.item())\n",
    "                   \n",
    "            process_time = start_time - time.time() - prepare_time\n",
    "            compute_efficiency = process_time / (process_time + prepare_time)    \n",
    "            \n",
    "            pbar.update(1)\n",
    "            pbar.set_postfix({'Running loss' : loss_epoch_sum/(batch_idx + 1), \\\n",
    "                              'Compute efficiency': compute_efficiency}) \n",
    "            start_time = time.time()\n",
    "        \n",
    "    return loss_epoch_sum/len(loader) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgeIo8bnt7Uc"
   },
   "source": [
    "## Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bx3UdMBst7Uc"
   },
   "outputs": [],
   "source": [
    "''' \n",
    "Create (and load) model and define hyperparameters \n",
    "'''\n",
    "\n",
    "DeepSleepNet = getattr(importlib.import_module('architectures.' + ARCHITECTURE_NAME), 'DeepSleepNet')\n",
    "\n",
    "model     = DeepSleepNet(in_channels = len(CHANNELS), linear = LINEAR).to(DEVICE)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE, weight_decay = DECAY_RATE)\n",
    "loss_fn   = CustomBCEWithLogitsLoss().to(DEVICE)\n",
    "scaler    = torch.cuda.amp.GradScaler(enabled = torch.cuda.is_available())\n",
    "\n",
    "if LOAD_CHECKPOINT:\n",
    "    try:\n",
    "        for chp_no in range(99,-1,-1):\n",
    "            last_checkpoint_path = os.path.join(MODEL_PATH, 'my_checkpoint_' + \\\n",
    "                                                str(chp_no)  +'.pth.tar')\n",
    "            if os.path.exists(last_checkpoint_path):\n",
    "                break\n",
    "        checkpoint = torch.load(last_checkpoint_path, \\\n",
    "                                map_location = torch.device(DEVICE))\n",
    "\n",
    "        model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        epoch0 = checkpoint['epoch'] + 1\n",
    "        train_loss_history = checkpoint['train_loss_history']\n",
    "        val_loss_history = checkpoint['val_loss_history']\n",
    "        \n",
    "        print('Model saved in my_checkpoint_'+str(chp_no)+'.pth.tar) loaded.')    \n",
    "    except: \n",
    "        print('No checkpoint file found.')\n",
    "    \n",
    "else:\n",
    "    epoch0 = 0        \n",
    "    train_loss_history = []\n",
    "    val_loss_history = []\n",
    "    \n",
    "''' \n",
    "Start the training procedure \n",
    "'''\n",
    "if TRAIN_MODE:\n",
    "    for epoch in range(epoch0, MAX_NUM_EPOCHS):\n",
    "        np.random.seed(SEED + epoch)\n",
    "\n",
    "        # evaluate early stop criterion\n",
    "        if len(val_loss_history) > STOP_STRIP:           \n",
    "            if val_loss_history[-(STOP_STRIP + 1)] < min(val_loss_history[-STOP_STRIP::]) \\\n",
    "            or np.isnan([val_loss_history, train_loss_history]).any():\n",
    "                print(\"Early stop has been triggered.\")\n",
    "                break        \n",
    "\n",
    "        # run one epoch training\n",
    "        train_loss_epoch = train_fn(model, train_loader, optimizer, loss_fn, scaler, DEVICE)\n",
    "\n",
    "        # check validation loss\n",
    "        val_loss_epoch, _ = eval_fn(model, val_loader, DEVICE) \n",
    "\n",
    "        # loss logging\n",
    "        train_loss_history.append(train_loss_epoch)\n",
    "        val_loss_history.append(val_loss_epoch)          \n",
    "\n",
    "        # loss printing\n",
    "        print(\"\\nEpoch: {} \\t Mean training loss {:.6f}\".format(epoch, train_loss_epoch))\n",
    "        print(\"Epoch: {} \\t Mean validation loss {:.6f}\".format(epoch, val_loss_epoch))    \n",
    "\n",
    "        # save current epoch model checkpoint\n",
    "        checkpoint = {\n",
    "            \"epoch\"                : epoch,\n",
    "            \"model_state_dict\"     : model.state_dict(),\n",
    "            \"optimizer_state_dict\" : optimizer.state_dict(),\n",
    "            \"train_loss_history\"   : train_loss_history,\n",
    "            \"val_loss_history\"     : val_loss_history,\n",
    "        }\n",
    "        torch.save(checkpoint, os.path.join(MODEL_PATH, \\\n",
    "                                            'my_checkpoint_'+str(epoch)+'.pth.tar'))\n",
    "        print(\">Checkpoint {} saved<\".format(epoch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EPMlNo9_t7Uc"
   },
   "source": [
    "## Training & Validation Loss Histories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3KIP2oTCg9HP"
   },
   "outputs": [],
   "source": [
    "plt.plot(range(len(train_loss_history)), train_loss_history, '-rx', \\\n",
    "         label = \"Training loss\")\n",
    "plt.plot(range(len(val_loss_history)), val_loss_history, '-bx', \\\n",
    "         label = \"Validation loss\")\n",
    "plt.xlabel(\"Epoch (#)\", fontdict = None, labelpad = None)\n",
    "plt.ylabel(\"Cross Entropy Loss\", fontdict = None, labelpad = None)\n",
    "plt.legend(loc = 'upper right', borderaxespad = 0.7, shadow = True)\n",
    "plt.grid(linestyle = '--', linewidth = 0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FJitvHyveW1h"
   },
   "source": [
    "## Identify the Smallest Cross-Validation Loss Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PwM07BpseW1h"
   },
   "outputs": [],
   "source": [
    "best_model_idx = val_loss_history.index(min(val_loss_history))\n",
    "best_model_path = os.path.join(MODEL_PATH, 'my_checkpoint_' + str(best_model_idx) + '.pth.tar')\n",
    "checkpoint = torch.load(best_model_path, map_location = torch.device(DEVICE))\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "print(\"Best model has been loaded from my_checkpoint_\" + str(best_model_idx) + \".pth.tar\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute Test Loss and Gross AUROC/AUPRC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss_path = os.path.join(MODEL_PATH, 'test_loss_' + str(best_model_idx) + '.pickle')\n",
    "test_score_path = os.path.join(MODEL_PATH, 'test_score_' + str(best_model_idx) + '.pickle')\n",
    "\n",
    "if os.path.exists(test_loss_path) and os.path.exists(test_score_path):\n",
    "    with open(test_loss_path, 'rb') as f1, open(test_score_path, 'rb') as f2:\n",
    "        test_loss = pickle.load(f1)\n",
    "        test_score = pickle.load(f2)\n",
    "else:\n",
    "    test_loss, test_score = eval_fn(model, test_loader, DEVICE, comp_score = True)\n",
    "    with open(test_loss_path, 'wb') as f1, open(test_score_path, 'wb') as f2:\n",
    "        pickle.dump(test_loss, f1)    \n",
    "        pickle.dump(test_score, f2)            \n",
    "\n",
    "print(\"Test results based on {:d} test cases.\".format(len(test_score._record_auc)))\n",
    "print(\"Cross-entropy loss = {:.6f}\".format(test_loss))\n",
    "print('Gross AUROC: %8.6f' % test_score.gross_auroc())\n",
    "print('Gross AUPRC: %8.6f' % test_score.gross_auprc())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tksGKQyht7Ue"
   },
   "source": [
    "## Model Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WPLctKg-xeYc"
   },
   "outputs": [],
   "source": [
    "summary(model, input_size=(len(CHANNELS), 2**23))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "main.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": ".venv-dlh-p",
   "language": "python",
   "name": ".venv-dlh-p"
  },
  "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.8"
  },
  "metadata": {
   "interpreter": {
    "hash": "f8b2cc128ef1477fdc89dad98627a78c44d5fcbfcbaa32b380aba07844063c75"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}