train_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b2974a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "# ^^^ pyforest auto-imports - don't write above this line\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from scipy.io import loadmat, savemat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cf0f0f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wfdb\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "612bbde4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.scoring_metrics import (\n",
    "    RefInfo, load_ans,\n",
    "    score, ue_calculate, ur_calculate,\n",
    "    compute_challenge_metric, gen_endpoint_score_mask,\n",
    ")\n",
    "from utils.scoring_metrics_test import _load_af_episodes\n",
    "\n",
    "# from database_reader.cpsc_databases import CPSC2021\n",
    "# from data_reader import CINC2021Reader\n",
    "from data_reader import CPSC2021Reader\n",
    "from dataset import CPSC2021\n",
    "from cfg import TrainCfg\n",
    "from utils.aux_metrics import compute_main_task_metric, compute_rr_metric\n",
    "from trainer import train, evaluate, _set_task\n",
    "\n",
    "from utils.misc import list_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f447d9f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3de65a",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_dir = \"/home/wenhao/Jupyter/wenhao/data/CPSC2021/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa1241f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "883bf36b",
   "metadata": {},
   "source": [
    "# data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67acc33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import CPSC2021\n",
    "from cfg import TrainCfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9592e4a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train = CPSC2021(TrainCfg, task=\"rr_lstm\", training=True)\n",
    "ds_val = CPSC2021(TrainCfg, task=\"rr_lstm\", training=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4520fe5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a7f49d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(ds_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77db2989",
   "metadata": {},
   "outputs": [],
   "source": [
    "err_list = []\n",
    "for idx, seg in enumerate(ds_train.segments):\n",
    "    sig, lb = ds_train[idx]\n",
    "    if sig.shape != (2,6000) or lb.shape != (750, 1):\n",
    "        print(\"\\n\"+f\"segment {seg} has sig.shape = {sig.shape}, lb.shape = {lb.shape}\"+\"\\n\")\n",
    "        err_list.append(seg)\n",
    "    print(f\"{idx+1}/{len(ds_train)}\", end=\"\\r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b116d0d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for idx, seg in enumerate(ds_val.segments):\n",
    "    sig, lb = ds_val[idx]\n",
    "    if sig.shape != (2,6000) or lb.shape != (750, 1):\n",
    "        print(\"\\n\"+f\"segment {seg} has sig.shape = {sig.shape}, lb.shape = {lb.shape}\"+\"\\n\")\n",
    "        err_list.append(seg)\n",
    "    print(f\"{idx+1}/{len(ds_val)}\", end=\"\\r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3a7c368",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95deface",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(err_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4072ce24",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadmat(ds_train._get_seg_data_path(err_list[-1]))[\"ecg\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2afe1780",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for idx, seg in enumerate(err_list):\n",
    "    path = ds_train._get_seg_data_path(seg)\n",
    "    os.remove(path)\n",
    "    path = ds_train._get_seg_ann_path(seg)\n",
    "    os.remove(path)\n",
    "    print(f\"{idx+1}/{len(err_list)}\", end=\"\\r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f989de7f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb5623b0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f7b8790",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "62304052",
   "metadata": {},
   "source": [
    "# Plan\n",
    "\n",
    "## R peak detection\n",
    "\n",
    "## rr-lstm\n",
    "\n",
    "## U-net\n",
    "\n",
    "## sequence labelling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c83b765e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "75c5355a",
   "metadata": {},
   "source": [
    "## R peak detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544bb4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import (\n",
    "    ECG_SEQ_LAB_NET_CPSC2021,\n",
    "    ECG_UNET_CPSC2021,\n",
    "    ECG_SUBTRACT_UNET_CPSC2021,\n",
    "    RR_LSTM_CPSC2021,\n",
    "    _qrs_detection_post_process,\n",
    ")\n",
    "from trainer import train\n",
    "from utils.misc import init_logger, dict_to_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a269936a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg import ModelCfg, TrainCfg\n",
    "from copy import deepcopy\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP, DataParallel as DP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d75abe7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg import ModelCfg\n",
    "task = \"qrs_detection\"  # or \"main\"\n",
    "model_cfg = deepcopy(ModelCfg[task])\n",
    "model_cfg.model_name = \"seq_lab\"\n",
    "model = ECG_SEQ_LAB_NET_CPSC2021(model_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2765651",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d757b40",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DP(model)\n",
    "model.to(torch.device(\"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "575e7604",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "896f3d01",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config = deepcopy(TrainCfg)\n",
    "# train_config.task = \"qrs_detection\"\n",
    "_set_task(\"qrs_detection\", train_config)\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ffbfb50",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config.main.reduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf6a3d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = init_logger(log_dir=train_config.log_dir, verbose=2)\n",
    "logger.info(f\"\\n{'*'*20}   Start Training   {'*'*20}\\n\")\n",
    "logger.info(f\"Using device {device}\")\n",
    "logger.info(f\"Using torch of version {torch.__version__}\")\n",
    "logger.info(f\"with configuration\\n{dict_to_str(train_config)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "019b792f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8af3291",
   "metadata": {},
   "outputs": [],
   "source": [
    "train(\n",
    "    model=model,\n",
    "    model_config=model_cfg,\n",
    "    config=train_config,\n",
    "    device=device,\n",
    "    logger=logger,\n",
    "    debug=train_config.debug,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "215b624d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61916914",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0be7eb4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "756a7a0c",
   "metadata": {},
   "source": [
    "## rr-lstm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df74ca40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import (\n",
    "    ECG_SEQ_LAB_NET_CPSC2021,\n",
    "    ECG_UNET_CPSC2021,\n",
    "    ECG_SUBTRACT_UNET_CPSC2021,\n",
    "    RR_LSTM_CPSC2021,\n",
    "    _qrs_detection_post_process,\n",
    ")\n",
    "from trainer import train, evaluate\n",
    "from utils.misc import init_logger, dict_to_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70a814a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg import ModelCfg, TrainCfg\n",
    "from copy import deepcopy\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP, DataParallel as DP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48e142e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "task = \"rr_lstm\"  # or \"main\"\n",
    "model_cfg = deepcopy(ModelCfg[task])\n",
    "model = RR_LSTM_CPSC2021(model_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a86394fa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80fbcc95",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config = deepcopy(TrainCfg)\n",
    "_set_task(\"rr_lstm\", train_config)\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b36f5f9f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# model = DP(model)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ddd7c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = init_logger(log_dir=train_config.log_dir, verbose=2)\n",
    "logger.info(f\"\\n{'*'*20}   Start Training   {'*'*20}\\n\")\n",
    "logger.info(f\"Using device {device}\")\n",
    "logger.info(f\"Using torch of version {torch.__version__}\")\n",
    "logger.info(f\"with configuration\\n{dict_to_str(train_config)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85aac5ae",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fadc3246",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = train(\n",
    "    model=model,\n",
    "    model_config=model_cfg,\n",
    "    config=train_config,\n",
    "    device=device,\n",
    "    logger=logger,\n",
    "    debug=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c024a3a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5a9435f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9caec945",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e351cf5f",
   "metadata": {},
   "source": [
    "## main_task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f8988f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import (\n",
    "    ECG_SEQ_LAB_NET_CPSC2021,\n",
    "    ECG_UNET_CPSC2021,\n",
    "    ECG_SUBTRACT_UNET_CPSC2021,\n",
    "    RR_LSTM_CPSC2021,\n",
    "    _qrs_detection_post_process,\n",
    "    _main_task_post_process\n",
    ")\n",
    "from trainer import train\n",
    "from utils.misc import init_logger, dict_to_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "183efb92",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg import ModelCfg, TrainCfg\n",
    "from copy import deepcopy\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP, DataParallel as DP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7f5d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cfg import ModelCfg\n",
    "task = \"main\"  # or \"main\"\n",
    "model_cfg = deepcopy(ModelCfg[task])\n",
    "# model_cfg.model_name = \"seq_lab\"\n",
    "# model = ECG_SEQ_LAB_NET_CPSC2021(model_cfg)\n",
    "model_cfg.model_name = \"unet\"\n",
    "model = ECG_UNET_CPSC2021(model_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dcf0142",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d90cec0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13f2f992",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6646ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DP(model)\n",
    "model.to(torch.device(\"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02d00217",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "866138ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config = deepcopy(TrainCfg)\n",
    "# train_config.task = \"qrs_detection\"\n",
    "_set_task(\"main\", train_config)\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201356ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config.main.model_name = \"unet\"\n",
    "train_config.main.reduction = 1\n",
    "train_config.main.cnn_name = None\n",
    "train_config.main.rnn_name = None\n",
    "train_config.main.attn_name = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc57fc7b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "344bb205",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = init_logger(log_dir=train_config.log_dir, verbose=2)\n",
    "logger.info(f\"\\n{'*'*20}   Start Training   {'*'*20}\\n\")\n",
    "logger.info(f\"Using device {device}\")\n",
    "logger.info(f\"Using torch of version {torch.__version__}\")\n",
    "logger.info(f\"with configuration\\n{dict_to_str(train_config)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6a6a397",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9519c292",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd102da",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "best_model = train(\n",
    "    model=model,\n",
    "    model_config=model_cfg,\n",
    "    config=train_config,\n",
    "    device=device,\n",
    "    logger=logger,\n",
    "    debug=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af3a56fa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d05b9aa5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0fc1d69",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547672ee",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "970653da",
   "metadata": {},
   "source": [
    "## Misc "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5344a46",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d316d95a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from entry_2021 import *\n",
    "from test_entry import run_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa7ab59",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f182a110",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_path = \"./working_dir/sample_data/data_98_1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296fce4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = challenge_entry(sample_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52f1371a",
   "metadata": {},
   "outputs": [],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a25e868b",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(out['predict_endpoints'][0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0ce9d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_val.reader.load_data(\"data_98_1\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf79edf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_val.reader.load_af_episodes(\"data_98_1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1445344",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca3680b1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a91b3d8c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "956518b6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}