{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7392733,"sourceType":"datasetVersion","datasetId":4297749},{"sourceId":7465251,"sourceType":"datasetVersion","datasetId":4317718},{"sourceId":7657023,"sourceType":"datasetVersion","datasetId":4464367}],"dockerImageVersionId":30648,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# [CV 0.68 | LB 0.46] DilatedInception WaveNet in PyTorch - Training\n\n## Introduction\nAfter joining this competition, I focus on validating how far raw EEG signals can go through many experiments. Finally, I find an simple architecture mixing the concept of **dilation** and **inception**, which can be seen as an extension of [Chris' version](https://www.kaggle.com/code/cdeotte/wavenet-starter-lb-0-52). Following illustrates the model architecture,\n[![Screenshot-2024-02-19-at-1-11-40-PM.png](https://i.postimg.cc/MKp8xVVV/Screenshot-2024-02-19-at-1-11-40-PM.png)](https://postimg.cc/7bdRnCBZ)\n\nAs can be seen, I modify the original dilated convolution to **dilated inception convolution**. Concretely speaking, instead of using only one kernel size per layer in `WaveBlock`, this model tries to capture enriched temporal patterns by considering different kernel sizes (2, 3, 6, and 7 here) at the same time. Furthermore, to restrain model parameters from growing too much, the model narrows down the output channels of convolutions with different kernel sizes. Let `out_channels = 64`, each dilated convolution with a specific kernel size only outputs 16 channels, then the concatenation along channel dimension is applied to match `out_channels = 64`.\n\nThis model architecture achieves **CV 0.68** (below 0.7) and LB 0.46 (below 0.5).\n\n## About this Notebook\nIn this kernel, we'll run local cross-validation (5-fold gkf with `patient_id` as group) using the model architecture introduced above. After models are well-trained, we can run inference using [[LB 0.46] DilatedInception WaveNet - Inference](https://www.kaggle.com/abaojiang/lb-0-46-dilatedinception-wavenet-inference).\n\n## Acknowledgements\nSpecial thanks to [@cdeotte](https://www.kaggle.com/cdeotte)'s sharing, [WaveNet Starter - [LB 0.52]](https://www.kaggle.com/code/cdeotte/wavenet-starter-lb-0-52).\n\n<a id=\"toc\"></a>\n## Table of Contents\n* [1. Load Data](#load_data)\n* [2. Define Dataset](#dataset)\n* [3. Define Model Architecture](#model)\n* [4. Define Loss Criterion](#loss)\n* [5. Define Evalutor](#evaluator)\n* [6. Define Trainer](#trainer)\n* [7. Run Cross-Validation](#cv)\n* [8. Derive CV Score](#cv_score)\n\n## Import Packages","metadata":{}},{"cell_type":"code","source":"import gc\nimport logging\nimport math\nimport os\nimport random\nimport sys\nimport pickle\nimport warnings\nfrom abc import abstractmethod\nfrom datetime import datetime\nfrom pathlib import Path\nfrom tqdm.notebook import tqdm\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union\nwarnings.simplefilter(\"ignore\")\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport yaml\nfrom scipy.signal import butter, lfilter\nfrom sklearn.model_selection import GroupKFold\nfrom torch.nn.modules.loss import _Loss\nfrom torch.optim import Optimizer, lr_scheduler\nfrom torch.optim.lr_scheduler import _LRScheduler\nfrom torch import Tensor\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nfrom transformers import get_cosine_schedule_with_warmup","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:35:57.359956Z","iopub.execute_input":"2024-02-20T09:35:57.360768Z","iopub.status.idle":"2024-02-20T09:36:02.015779Z","shell.execute_reply.started":"2024-02-20T09:35:57.360710Z","shell.execute_reply":"2024-02-20T09:36:02.014821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class _Logger:\n    \"\"\"Customized logger.\n\n    Args:\n        logging_level: lowest-severity log message the logger handles\n        logging_file: file stream for logging\n            *Note: If `logging_file` isn't specified, message is only\n                logged to system standard output.\n    \"\"\"\n\n    _logger: logging.Logger = None\n\n    def __init__(\n        self,\n        logging_level: str = \"INFO\",\n        logging_file: Optional[Path] = None,\n    ):\n        self.logging_level = logging_level\n        self.logging_file = logging_file\n\n        self._build_logger()\n\n    def get_logger(self) -> logging.Logger:\n        \"\"\"Return customized logger.\"\"\"\n        return self._logger\n\n    def _build_logger(self) -> None:\n        \"\"\"Build logger.\"\"\"\n        self._logger = logging.getLogger()\n        self._logger.setLevel(self._get_level())\n        self._add_handler()\n\n    def _get_level(self) -> int:\n        \"\"\"Return lowest severity of the events the logger handles.\n\n        Returns:\n            level: severity of the events\n        \"\"\"\n        level = 0\n\n        if self.logging_level == \"DEBUG\":\n            level = logging.DEBUG\n        elif self.logging_level == \"INFO\":\n            level = logging.INFO\n        elif self.logging_level == \"WARNING\":\n            level = logging.WARNING\n        elif self.logging_level == \"ERROR\":\n            level = logging.ERROR\n        elif self.logging_level == \"CRITICAL\":\n            level = logging.CRITICAL\n\n        return level\n\n    def _add_handler(self) -> None:\n        \"\"\"Add stream and file (optional) handlers to logger.\"\"\"\n        s_handler = logging.StreamHandler(sys.stdout)\n        self._logger.addHandler(s_handler)\n\n        if self.logging_file is not None:\n            f_handler = logging.FileHandler(self.logging_file, mode=\"a\")\n            self._logger.addHandler(f_handler)\n\n            \ndef _seed_everything(seed: int) -> None:\n    \"\"\"Seed current experiment to guarantee reproducibility.\n\n    Args:\n        seed: manually specified seed number\n    \"\"\"\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    # When running with cudnn backend\n    torch.backends.cudnn.deterministic = False\n    torch.backends.cudnn.benchmark = True \n    # Set a fixed value for the hash seed\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:36:02.017373Z","iopub.execute_input":"2024-02-20T09:36:02.017797Z","iopub.status.idle":"2024-02-20T09:36:02.030877Z","shell.execute_reply.started":"2024-02-20T09:36:02.017771Z","shell.execute_reply":"2024-02-20T09:36:02.029990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Data Paths and Configuration and Metadata","metadata":{}},{"cell_type":"code","source":"DATA_PATH = Path(\"/kaggle/input/hms-harmful-brain-activity-classification\")\n\nclass CFG:\n    train_models = False\n    seed = 42\n    \n    exp_id = datetime.now().strftime(\"%m%d-%H-%M-%S\")\n    exp_dump_path = Path(\"/kaggle/working/\") / exp_id\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    \n    # == Data ==\n    gen_eegs = False\n    # Chris' 8 channels\n    feats = [\n        \"Fp1\", \"T3\", \"C3\", \"O1\",\n        \"Fp2\", \"C4\", \"T4\", \"O2\"\n    ]\n    cast_eegs = True\n    dataset = {\n        \"eeg\": {\n            \"n_feats\": 8,\n            \"apply_chris_magic_ch8\": True,\n            \"normalize\": True,\n            \"apply_butter_lowpass_filter\": True,\n            \"apply_mu_law_encoding\": False,\n            \"downsample\": 5\n        }\n    }\n\n    # == Trainer ==\n    trainer = {\n        \"epochs\": 5,\n        \"lr\": 1e-3,\n        \"dataloader\": {\n            \"batch_size\": 32,\n            \"shuffle\": True,\n            \"num_workers\": 2\n        },\n        \"use_amp\": True,\n        \"grad_accum_steps\": 1,\n        \"model_ckpt\": {\n            \"ckpt_metric\": \"kldiv\",\n            \"ckpt_mode\": \"min\",\n            \"best_ckpt_mid\": \"last\"\n        },\n        \"es\": {\"patience\": 0},\n        \"step_per_batch\": True,\n        \"one_batch_only\": False\n    }\n    \n    # == Debug ==\n    one_fold_only = True\n    \n    \nN_CLASSES = 6\nTGT_VOTE_COLS = [\n    \"seizure_vote\", \"lpd_vote\", \"gpd_vote\", \"lrda_vote\",\n    \"grda_vote\", \"other_vote\"\n]\nTGT_COL = \"target\"\nEEG_FREQ = 200  # Hz\nEEG_WLEN = 50  # sec\nEEG_PTS = int(EEG_FREQ * EEG_WLEN)","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:36:02.032123Z","iopub.execute_input":"2024-02-20T09:36:02.032404Z","iopub.status.idle":"2024-02-20T09:36:02.101319Z","shell.execute_reply.started":"2024-02-20T09:36:02.032380Z","shell.execute_reply":"2024-02-20T09:36:02.100339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not CFG.exp_dump_path.exists():\n    os.mkdir(CFG.exp_dump_path)\n    \nlogger = _Logger(logging_file=CFG.exp_dump_path / \"train_eval.log\").get_logger()\n_seed_everything(CFG.seed)","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:36:02.103132Z","iopub.execute_input":"2024-02-20T09:36:02.103443Z","iopub.status.idle":"2024-02-20T09:36:02.116823Z","shell.execute_reply.started":"2024-02-20T09:36:02.103419Z","shell.execute_reply":"2024-02-20T09:36:02.115978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"load_data\"></a>\n## 1. Load Data\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nLet's load the training data first. Following Chris' implementation, we're going to crop EEG signals and keep only the **middle 50-sec window**. To accelerate experiment iteration, we can directly load cropped EEGs from Chris' dataset [Brain-EEGs](https://www.kaggle.com/datasets/cdeotte/brain-eegs). Don't forget to upvote Chris' dataset if you use it.","metadata":{}},{"cell_type":"code","source":"def _get_eeg_window(file: Path) -> np.ndarray:\n    \"\"\"Return cropped EEG window.\n\n    Default setting is to return the middle 50-sec window.\n\n    Args:\n        file: EEG file path\n        test: if True, there's no need to truncate EEGs\n\n    Returns:\n        eeg_win: cropped EEG window \n    \"\"\"\n    eeg = pd.read_parquet(file, columns=CFG.feats)\n    n_pts = len(eeg)\n    offset = (n_pts - EEG_PTS) // 2\n    eeg = eeg.iloc[offset:offset + EEG_PTS]\n    \n    eeg_win = np.zeros((EEG_PTS, len(CFG.feats)))\n    for j, col in enumerate(CFG.feats):\n        if CFG.cast_eegs:\n            eeg_raw = eeg[col].values.astype(\"float32\")\n        else:\n            eeg_raw = eeg[col].values \n\n        # Fill missing values\n        mean = np.nanmean(eeg_raw)\n        if np.isnan(eeg_raw).mean() < 1:\n            eeg_raw = np.nan_to_num(eeg_raw, nan=mean)\n        else: \n            # All missing\n            eeg_raw[:] = 0\n        eeg_win[:, j] = eeg_raw \n        \n    return eeg_win ","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:36:08.719820Z","iopub.execute_input":"2024-02-20T09:36:08.720166Z","iopub.status.idle":"2024-02-20T09:36:08.728935Z","shell.execute_reply.started":"2024-02-20T09:36:08.720139Z","shell.execute_reply":"2024-02-20T09:36:08.727926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = pd.read_csv(DATA_PATH / \"train.csv\")\nlogger.info(f\"Train data shape | {train.shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:36:09.204373Z","iopub.execute_input":"2024-02-20T09:36:09.204971Z","iopub.status.idle":"2024-02-20T09:36:09.437708Z","shell.execute_reply.started":"2024-02-20T09:36:09.204939Z","shell.execute_reply":"2024-02-20T09:36:09.436609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"uniq_eeg_ids = train[\"eeg_id\"].unique()\nn_uniq_eeg_ids = len(uniq_eeg_ids)\n\nif CFG.gen_eegs:\n    logger.info(\"Generate cropped EEGs...\")\n    all_eegs = {}\n    for i, eeg_id in tqdm(enumerate(uniq_eeg_ids), total=n_uniq_eeg_ids):\n        eeg_win = _get_eeg_window(DATA_PATH / \"train_eegs\" / f\"{eeg_id}.parquet\")\n        all_eegs[eeg_id] = eeg_win\nelse:\n    logger.info(\"Load cropped EEGs...\")\n    all_eegs = np.load(\"/kaggle/input/brain-eegs/eegs.npy\", allow_pickle=True).item()\n    assert len(all_eegs) == n_uniq_eeg_ids\nlogger.info(f\"Demo EEG shape | {list(all_eegs.values())[0].shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:36:10.050253Z","iopub.execute_input":"2024-02-20T09:36:10.050608Z","iopub.status.idle":"2024-02-20T09:37:49.389812Z","shell.execute_reply.started":"2024-02-20T09:36:10.050580Z","shell.execute_reply":"2024-02-20T09:37:49.388860Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"logger.info(f\"Process labels...\")\ndf_tmp = train.groupby(\"eeg_id\")[[\"patient_id\"]].agg(\"first\")\nlabels_tmp = train.groupby(\"eeg_id\")[TGT_VOTE_COLS].agg(\"sum\")\nfor col in TGT_VOTE_COLS:\n    df_tmp[col] = labels_tmp[col].values\n\n# Normalize target columns\ny_data = df_tmp[TGT_VOTE_COLS].values\ny_data = y_data / y_data.sum(axis=1, keepdims=True)\ndf_tmp[TGT_VOTE_COLS] = y_data\n\ntgt = train.groupby(\"eeg_id\")[[\"expert_consensus\"]].agg(\"first\")\ndf_tmp[TGT_COL] = tgt \n\ntrain = df_tmp.reset_index()\nlogger.info(f\"Training DataFrame shape | {train.shape}\")\ntrain.head()","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.391808Z","iopub.execute_input":"2024-02-20T09:37:49.392457Z","iopub.status.idle":"2024-02-20T09:37:49.471453Z","shell.execute_reply.started":"2024-02-20T09:37:49.392422Z","shell.execute_reply":"2024-02-20T09:37:49.470567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"dataset\"></a>\n## 2. Define Dataset\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nAfter cropped EEGs are well-prepared, we then define the custom `Dataset` with an EEG data transformer for further processing.","metadata":{}},{"cell_type":"code","source":"class _EEGTransformer(object):\n    \"\"\"Data transformer for raw EEG signals.\"\"\"\n\n    FEAT2CODE = {f: i for i, f in enumerate(CFG.feats)}\n\n    def __init__(\n        self,\n        n_feats: int,\n        apply_chris_magic_ch8: bool = True,\n        normalize: bool = True,\n        apply_butter_lowpass_filter: bool = True,\n        apply_mu_law_encoding: bool = False,\n        downsample: Optional[int] = None,\n    ) -> None:\n        self.n_feats = n_feats\n        self.apply_chris_magic_ch8 = apply_chris_magic_ch8\n        self.normalize = normalize\n        self.apply_butter_lowpass_filter = apply_butter_lowpass_filter\n        self.apply_mu_law_encoding = apply_mu_law_encoding\n        self.downsample = downsample\n\n    def transform(self, x: np.ndarray) -> np.ndarray:\n        \"\"\"Apply transformation on raw EEG signals.\n        \n        Args:\n            x: raw EEG signals, with shape (L, C)\n\n        Return:\n            x_: transformed EEG signals\n        \"\"\"\n        x_ = x.copy()\n        if self.apply_chris_magic_ch8:\n            x_ = self._apply_chris_magic_ch8(x_)\n\n        if self.normalize:\n            x_ = np.clip(x_, -1024, 1024)\n            x_ = np.nan_to_num(x_, nan=0) / 32.0\n\n        if self.apply_butter_lowpass_filter:\n            x_ = self._butter_lowpass_filter(x_) \n\n        if self.apply_mu_law_encoding:\n            x_ = self._quantize_data(x_, 1)\n\n        if self.downsample is not None:\n            x_ = x_[::self.downsample, :]\n\n        return x_\n\n    def _apply_chris_magic_ch8(self, x: np.ndarray) -> np.ndarray:\n        \"\"\"Generate features based on Chris' magic formula.\"\"\" \n        x_tmp = np.zeros((EEG_PTS, self.n_feats), dtype=\"float32\")\n\n        # Generate features\n        x_tmp[:, 0] = x[:, self.FEAT2CODE[\"Fp1\"]] - x[:, self.FEAT2CODE[\"T3\"]]\n        x_tmp[:, 1] = x[:, self.FEAT2CODE[\"T3\"]] - x[:, self.FEAT2CODE[\"O1\"]]\n        \n        x_tmp[:, 2] = x[:, self.FEAT2CODE[\"Fp1\"]] - x[:, self.FEAT2CODE[\"C3\"]]\n        x_tmp[:, 3] = x[:, self.FEAT2CODE[\"C3\"]] - x[:, self.FEAT2CODE[\"O1\"]]\n        \n        x_tmp[:, 4] = x[:, self.FEAT2CODE[\"Fp2\"]] - x[:, self.FEAT2CODE[\"C4\"]]\n        x_tmp[:, 5] = x[:, self.FEAT2CODE[\"C4\"]] - x[:, self.FEAT2CODE[\"O2\"]]\n        \n        x_tmp[:, 6] = x[:, self.FEAT2CODE[\"Fp2\"]] - x[:, self.FEAT2CODE[\"T4\"]]\n        x_tmp[:, 7] = x[:, self.FEAT2CODE[\"T4\"]] - x[:, self.FEAT2CODE[\"O2\"]]\n\n        return x_tmp\n\n    def _butter_lowpass_filter(self, data, cutoff_freq=20, sampling_rate=200, order=4):\n        nyquist = 0.5 * sampling_rate\n        normal_cutoff = cutoff_freq / nyquist\n        b, a = butter(order, normal_cutoff, btype=\"low\", analog=False)\n        filtered_data = lfilter(b, a, data, axis=0)\n\n        return filtered_data\n                \n    def _quantize_data(self, data, classes):\n        mu_x = self._mu_law_encoding(data, classes)\n        \n        return mu_x\n\n    def _mu_law_encoding(self, data, mu):\n        mu_x = np.sign(data) * np.log(1 + mu * np.abs(data)) / np.log(mu + 1)\n\n        return mu_x","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:37:49.473152Z","iopub.execute_input":"2024-02-20T09:37:49.473420Z","iopub.status.idle":"2024-02-20T09:37:49.490649Z","shell.execute_reply.started":"2024-02-20T09:37:49.473399Z","shell.execute_reply":"2024-02-20T09:37:49.489686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class EEGDataset(Dataset):\n    \"\"\"Dataset for pure raw EEG signals.\n\n    Args:\n        data: processed data\n        split: data split\n\n    Attributes:\n        _n_samples: number of samples\n        _infer: if True, the dataset is constructed for inference\n            *Note: Ground truth is not provided.\n    \"\"\"\n\n    def __init__(\n        self,\n        data: Dict[str,  Any],\n        split: str,\n        **dataset_cfg: Any,\n    ) -> None:\n        self.metadata = data[\"meta\"]\n        self.all_eegs = data[\"eeg\"]\n        self.dataset_cfg = dataset_cfg\n\n        # Raw EEG data transformer\n        self.eeg_params = dataset_cfg[\"eeg\"]\n        self.eeg_trafo = _EEGTransformer(**self.eeg_params)\n\n        self._set_n_samples()\n        self._infer = True if split == \"test\" else False\n\n        self._stream_X = True if self.all_eegs is None else False\n        self._X, self._y = self._transform()\n\n    def _set_n_samples(self) -> None:\n        assert len(self.metadata) == self.metadata[\"eeg_id\"].nunique()\n        self._n_samples = len(self.metadata)\n\n    def _transform(self) -> Tuple[Optional[np.ndarray], np.ndarray]:\n        \"\"\"Transform feature and target matrices.\"\"\"\n        if self.eeg_params[\"downsample\"] is not None:\n            eeg_len = int(EEG_PTS / self.eeg_params[\"downsample\"])\n        else:\n            eeg_len = int(EEG_PTS)\n        if not self._stream_X:\n            X = np.zeros((self._n_samples, eeg_len, self.eeg_params[\"n_feats\"]), dtype=\"float32\")\n        else:\n            X = None\n        y = np.zeros((self._n_samples, N_CLASSES), dtype=\"float32\") if not self._infer else None\n\n        for i, row in tqdm(self.metadata.iterrows(), total=len(self.metadata)):\n            # Process raw EEG signals\n            if not self._stream_X:\n                # Retrieve raw EEG signals\n                eeg = self.all_eegs[row[\"eeg_id\"]]\n\n                # Apply EEG transformer\n                x = self.eeg_trafo.transform(eeg)\n\n                X[i] = x\n\n            if not self._infer:\n                y[i] = row[TGT_VOTE_COLS] \n\n        return X, y\n\n    def __len__(self) -> int:\n        return self._n_samples\n\n    def __getitem__(self, idx: int) -> Dict[str, Tensor]:\n        if self._X is None:\n            # Load data here...\n#             x = np.load(...)\n#             x = self.eeg_trafo.transform(x)\n            pass\n        else:\n            x = self._X[idx, ...]\n        data_sample = {\"x\": torch.tensor(x, dtype=torch.float32)}\n        if not self._infer:\n            data_sample[\"y\"] = torch.tensor(self._y[idx, :], dtype=torch.float32)\n\n        return data_sample","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.493413Z","iopub.execute_input":"2024-02-20T09:37:49.493696Z","iopub.status.idle":"2024-02-20T09:37:49.508551Z","shell.execute_reply.started":"2024-02-20T09:37:49.493673Z","shell.execute_reply":"2024-02-20T09:37:49.507770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"model\"></a>\n## 3. Define Model Architecture\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"class _WaveBlock(nn.Module):\n    \"\"\"WaveNet block.\n\n    Args:\n        kernel_size: kernel size, pass a list of kernel sizes for\n            inception\n    \"\"\"\n\n    def __init__(\n        self,\n        n_layers: int, \n        in_dim: int,\n        h_dim: int,\n        kernel_size: Union[int, List[int]],\n        conv_module: Optional[Type[nn.Module]] = None,\n    ) -> None:\n        super().__init__()\n\n        self.n_layers = n_layers\n        self.dilation_rates = [2**l for l in range(n_layers)]\n\n        self.in_conv = nn.Conv2d(in_dim, h_dim, kernel_size=(1, 1)) \n        self.gated_tcns = nn.ModuleList()\n        self.skip_convs = nn.ModuleList()\n        for layer in range(n_layers):\n            c_in, c_out = h_dim, h_dim\n            self.gated_tcns.append(\n                _GatedTCN(\n                    in_dim=c_in,\n                    h_dim=c_out,\n                    kernel_size=kernel_size,\n                    dilation_factor=self.dilation_rates[layer],\n                    conv_module=conv_module,\n                )\n            )\n            self.skip_convs.append(nn.Conv2d(h_dim, h_dim, kernel_size=(1, 1)))\n\n        # Initialize parameters\n        nn.init.xavier_uniform_(self.in_conv.weight, gain=nn.init.calculate_gain(\"relu\"))\n        nn.init.zeros_(self.in_conv.bias)\n        for i in range(len(self.skip_convs)):\n            nn.init.xavier_uniform_(self.skip_convs[i].weight, gain=nn.init.calculate_gain(\"relu\"))\n            nn.init.zeros_(self.skip_convs[i].bias)\n\n    def forward(self, x: Tensor) -> Tensor:\n        \"\"\"Forward pass.\n        \n        Shape:\n            x: (B, C, N, L), where C denotes in_dim\n            x_skip: (B, C', N, L), where C' denotes h_dim\n        \"\"\"\n        # Input convolution\n        x = self.in_conv(x)\n\n        x_skip = x\n        for layer in range(self.n_layers):\n            x = self.gated_tcns[layer](x)\n            x = self.skip_convs[layer](x)\n\n            # Skip-connection\n            x_skip = x_skip + x \n\n        return x_skip\n\n\nclass _GatedTCN(nn.Module):\n    \"\"\"Gated temporal convolution layer.\n\n    Parameters:\n        conv_module: customized convolution module\n    \"\"\"\n\n    def __init__(\n        self,\n        in_dim: int,\n        h_dim: int,\n        kernel_size: Union[int, List[int]],\n        dilation_factor: int,\n        dropout: Optional[float] = None,\n        conv_module: Optional[Type[nn.Module]] = None,\n    ) -> None:\n        super().__init__()\n\n        # Model blocks\n        if conv_module is None:\n            self.filt = nn.Conv2d(\n                in_channels=in_dim, out_channels=h_dim, kernel_size=(1, kernel_size), dilation=dilation_factor\n            )\n            self.gate = nn.Conv2d(\n                in_channels=in_dim, out_channels=h_dim, kernel_size=(1, kernel_size), dilation=dilation_factor\n            )\n        else:\n            self.filt = conv_module(\n                in_channels=in_dim, out_channels=h_dim, kernel_size=kernel_size, dilation=dilation_factor\n            )\n            self.gate = conv_module(\n                in_channels=in_dim, out_channels=h_dim, kernel_size=kernel_size, dilation=dilation_factor\n            )\n\n        if dropout is not None:\n            self.dropout = nn.Dropout(dropout)\n        else:\n            self.dropout = None\n\n    def forward(self, x: Tensor) -> Tensor:\n        \"\"\"Forward pass.\n\n        Parameters:\n            x: input sequence\n\n        Return:\n            h: output sequence\n\n        Shape:\n            x: (B, C, N, L), where L denotes the input sequence length\n            h: (B, h_dim, N, L')\n        \"\"\"\n        x_filt = F.tanh(self.filt(x))\n        x_gate = F.sigmoid(self.gate(x))\n        h = x_filt * x_gate\n        if self.dropout is not None:\n            h = self.dropout(h)\n\n        return h\n\n\nclass _DilatedInception(nn.Module):\n    \"\"\"Dilated inception layer.\n\n    Note that `out_channels` will be split across #kernels.\n    \"\"\"\n\n    def __init__(\n        self, \n        in_channels: int, \n        out_channels: int, \n        kernel_size: List[int], \n        dilation: int\n    ) -> None:\n        super().__init__()\n\n        # Network parameters\n        n_kernels = len(kernel_size)\n        assert out_channels % n_kernels == 0, \"`out_channels` must be divisible by #kernels.\"\n        h_dim = out_channels // n_kernels\n\n        # Model blocks\n        self.convs = nn.ModuleList()\n        for k in kernel_size:\n            self.convs.append(\n                nn.Conv2d(\n                    in_channels=in_channels, \n                    out_channels=h_dim, \n                    kernel_size=(1, k),\n                    padding=\"same\",\n                    dilation=dilation),\n            )\n\n    def forward(self, x: Tensor) -> Tensor:\n        \"\"\"Forward pass.\n\n        Parameters:\n            x: input sequence\n\n        Return:\n            h: output sequence\n\n        Shape:\n            x: (B, C, N, L), where C = in_channels\n            h: (B, C', N, L'), where C' = out_channels\n        \"\"\"\n        x_convs = []\n        for conv in self.convs:\n            x_conv = conv(x)\n            x_convs.append(x_conv)\n        h = torch.cat(x_convs, dim=1)\n\n        return h","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:37:49.509942Z","iopub.execute_input":"2024-02-20T09:37:49.510277Z","iopub.status.idle":"2024-02-20T09:37:49.534415Z","shell.execute_reply.started":"2024-02-20T09:37:49.510247Z","shell.execute_reply":"2024-02-20T09:37:49.533576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DilatedInceptionWaveNet(nn.Module):\n    \"\"\"WaveNet architecture with dilated inception conv.\"\"\"\n\n    def __init__(self,) -> None:\n        super().__init__()\n\n        kernel_size = [2, 3, 6, 7]\n\n        # Model blocks \n        self.wave_module = nn.Sequential(\n            _WaveBlock(12, 1, 16, kernel_size, _DilatedInception),\n            _WaveBlock(8, 16, 32, kernel_size, _DilatedInception),\n            _WaveBlock(4, 32, 64, kernel_size, _DilatedInception),\n            _WaveBlock(1, 64, 64, kernel_size, _DilatedInception),\n        )\n        self.output = nn.Sequential(\n            nn.Linear(64 * 4, 64),\n            nn.ReLU(),\n            nn.Linear(64, N_CLASSES)\n        ) \n\n    def forward(self, inputs: Dict[str, Tensor]) -> Tensor:\n        \"\"\"Forward pass.\n\n        Shape:\n            x: (B, L, C)\n        \"\"\"\n        x = inputs[\"x\"]\n        bs, length, in_dim = x.shape\n        x = x.transpose(1, 2).unsqueeze(dim=2)  # (B, C, N, L), N is redundant\n\n        x_ll_1 = self.wave_module(x[:, 0:1, :])\n        x_ll_2 = self.wave_module(x[:, 1:2, :])\n        x_ll = (F.adaptive_avg_pool2d(x_ll_1, (1, 1)) + F.adaptive_avg_pool2d(x_ll_2, (1, 1))) / 2\n\n        x_rl_1 = self.wave_module(x[:, 2:3, :])\n        x_rl_2 = self.wave_module(x[:, 3:4, :])\n        x_rl = (F.adaptive_avg_pool2d(x_rl_1, (1, 1)) + F.adaptive_avg_pool2d(x_rl_2, (1, 1))) / 2\n\n        x_lp_1 = self.wave_module(x[:, 4:5, :])\n        x_lp_2 = self.wave_module(x[:, 5:6, :])\n        x_lp = (F.adaptive_avg_pool2d(x_lp_1, (1, 1)) + F.adaptive_avg_pool2d(x_lp_2, (1, 1))) / 2\n\n        x_rp_1 = self.wave_module(x[:, 6:7, :])\n        x_rp_2 = self.wave_module(x[:, 7:8, :])\n        x_rp = (F.adaptive_avg_pool2d(x_rp_1, (1, 1)) + F.adaptive_avg_pool2d(x_rp_2, (1, 1))) / 2\n\n        x = torch.cat([x_ll, x_rl, x_lp, x_rp], axis=1).reshape(bs, -1)\n        output = self.output(x)\n\n        return output","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.535573Z","iopub.execute_input":"2024-02-20T09:37:49.535894Z","iopub.status.idle":"2024-02-20T09:37:49.550077Z","shell.execute_reply.started":"2024-02-20T09:37:49.535847Z","shell.execute_reply":"2024-02-20T09:37:49.549276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"loss\"></a>\n## 4. Define Loss Criterion\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nJust as many sharings, we use `KLDivLoss` as our loss criterion.","metadata":{}},{"cell_type":"code","source":"class KLDivWithLogitsLoss(nn.KLDivLoss):\n    \"\"\"Kullback-Leibler divergence loss with logits as input.\"\"\"\n\n    def __init__(self):\n        super().__init__(reduction=\"batchmean\")\n\n    def forward(self, y_pred: Tensor, y_true: Tensor) -> Tensor:\n        y_pred = F.log_softmax(y_pred,  dim=1)\n        kldiv_loss = super().forward(y_pred, y_true)\n\n        return kldiv_loss","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.551263Z","iopub.execute_input":"2024-02-20T09:37:49.551981Z","iopub.status.idle":"2024-02-20T09:37:49.563530Z","shell.execute_reply.started":"2024-02-20T09:37:49.551942Z","shell.execute_reply":"2024-02-20T09:37:49.562795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"evaluator\"></a>\n## 5. Define Evaluator\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nThen, we define the evaluator for performance tracking. Please feel free to add any evaluation metrics you're interested in.","metadata":{}},{"cell_type":"code","source":"class Evaluator(object):\n    \"\"\"Custom evaluator.\n\n    Args:\n        metric_names: evaluation metrics\n    \"\"\"\n\n    eval_metrics: Dict[str, Callable[..., float]] = {}\n    EPS: float = 1e-6\n\n    def __init__(self, metric_names: List[str]) -> None:\n        self.metric_names = metric_names\n\n        self._build()\n\n    def evaluate(\n        self,\n        y_true: Tensor,\n        y_pred: Tensor,\n        scaler: Optional[object] = None,\n    ) -> Dict[str, float]:\n        \"\"\"Run evaluation using pre-specified metrics.\n\n        Args:\n            y_true: ground truth\n            y_pred: prediction\n            scaler: scaling object\n\n        Returns:\n            eval_result: evaluation performance report\n        \"\"\"\n        if scaler is not None:\n            # Do inverse transformation to rescale y values\n            y_pred, y_true = self._rescale_y(y_pred, y_true, scaler)\n\n        eval_result = {}\n        for metric_name, metric in self.eval_metrics.items():\n            eval_result[metric_name] = metric(y_pred, y_true).item()\n\n        return eval_result\n\n    def _build(self) -> None:\n        \"\"\"Build evaluation metric instances.\"\"\"\n        for metric_name in self.metric_names:\n            if metric_name == \"kldiv\":\n                self.eval_metrics[metric_name] = KLDivWithLogitsLoss() \n            elif metric_name == \"ce\":\n                self.eval_metrics[metric_name] = nn.CrossEntropyLoss()\n\n    def _rescale_y(self, y_pred: Tensor, y_true: Tensor, scaler: Any) -> Tuple[Tensor, Tensor]:\n        \"\"\"Rescale y to the original scale.\n\n        Args:\n            y_pred: prediction\n            y_true: ground truth\n            scaler: scaling object\n\n        Returns:\n            y_pred: rescaled prediction\n            y_true: rescaled ground truth\n        \"\"\"\n        # Do inverse transform...\n\n        return y_pred, y_true","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.564597Z","iopub.execute_input":"2024-02-20T09:37:49.564888Z","iopub.status.idle":"2024-02-20T09:37:49.577161Z","shell.execute_reply.started":"2024-02-20T09:37:49.564858Z","shell.execute_reply":"2024-02-20T09:37:49.576339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"trainer\"></a>\n## 6. Define Trainer\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"class _ModelCheckpoint(object):\n    \"\"\"Model checkpooint.\n\n    Args:\n        ckpt_path: path to save model checkpoint\n        ckpt_metric: quantity to monitor during training process\n        ckpt_mode: determine the direction of metric improvement\n        best_ckpt_mid: model identifier of the probably best checkpoint\n            used to do the final evaluation\n    \"\"\"\n\n    def __init__(self, ckpt_path: Path, ckpt_metric: str, ckpt_mode: str, best_ckpt_mid: str) -> None:\n        self.ckpt_path = ckpt_path\n        self.ckpt_metric = ckpt_metric\n        self.ckpt_mode = ckpt_mode\n        self.best_ckpt_mid = best_ckpt_mid\n\n        # Specify checkpoint direction\n        self.ckpt_dir = -1 if ckpt_mode == \"max\" else 1\n\n        # Initialize checkpoint status\n        self.best_val_score = 1e18\n        self.best_epoch = 0\n\n    def step(\n        self, epoch: int, model: nn.Module, val_loss: float, val_result: Dict[str, float], last_epoch: bool = False\n    ) -> None:\n        \"\"\"Update checkpoint status for the current epoch.\n\n        Args:\n            epoch: current epoch\n            model: current model instance\n            val_loss: validation loss\n            val_result: evaluation result on validation set\n            last_epoch: if True, current epoch is the last one\n        \"\"\"\n        val_score = val_loss if self.ckpt_metric is None else val_result[self.ckpt_metric]\n        val_score = val_score * self.ckpt_dir\n        if val_score < self.best_val_score:  # type: ignore\n            logging.info(f\"Validation performance improves at epoch {epoch}!!\")\n            self.best_val_score = val_score\n            self.best_epoch = epoch\n\n            # Save model checkpoint\n            mid = \"loss\" if self.ckpt_metric is None else self.ckpt_metric\n            self._save_ckpt(model, mid)\n\n        if last_epoch:\n            self._save_ckpt(model, \"last\")\n\n    def save_ckpt(self, model: nn.Module, mid: Optional[str] = None) -> None:\n        \"\"\"Save the checkpoint.\n\n        Args:\n            model: current model instance\n            mid: model identifer\n        \"\"\"\n        self._save_ckpt(model, mid)\n\n    def load_best_ckpt(self, model: nn.Module, device: torch.device) -> nn.Module:\n        \"\"\"Load and return the best model checkpoint for final evaluation.\n\n        Args:\n            model: current model instance\n                *Note: Model weights are overrided by the best checkpoint.\n            device: device of the model instance\n\n        Returns:\n            best_model: best model checkpoint\n        \"\"\"\n        model = self._load_ckpt(model, device, self.best_ckpt_mid)\n\n        return model\n\n    def _save_ckpt(self, model: nn.Module, mid: Optional[str] = None) -> None:\n        \"\"\"Save the model checkpoint.\n\n        Args:\n            model: current model instance\n            mid: model identifer\n        \"\"\"\n        model_file = \"model.pth\" if mid is None else f\"model-{mid}.pth\"\n        torch.save(model.state_dict(), os.path.join(self.ckpt_path, model_file))\n\n    def _load_ckpt(self, model: nn.Module, device: torch.device, mid: str = \"last\") -> nn.Module:\n        \"\"\"Load the model checkpoint.\n\n        Args:\n            model: current model instance\n                *Note: Model weights are overrided by the best checkpoint.\n            device: device of the model instance\n            mid: model identifier\n\n        Returns:\n            model: model instance with the loaded weights\n        \"\"\"\n        model_file = f\"model-{mid}.pth\"\n        model.load_state_dict(torch.load(os.path.join(self.ckpt_path, model_file), map_location=device))\n\n        return model","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:37:49.578354Z","iopub.execute_input":"2024-02-20T09:37:49.578616Z","iopub.status.idle":"2024-02-20T09:37:49.592631Z","shell.execute_reply.started":"2024-02-20T09:37:49.578585Z","shell.execute_reply":"2024-02-20T09:37:49.591787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class _BaseTrainer:\n    \"\"\"Base class for all customized trainers.\n\n    Args:\n        logger: message logger\n        trainer_cfg: hyperparameters for training and evaluation processes\n        model: model instance\n        loss_fn: loss criterion\n        optimizer: optimization algorithm\n        lr_skd: learning rate scheduler\n        ckpt_path: path to save model checkpoints\n        es: early stopping tracker\n        evaluator: task-specific evaluator\n        use_wandb: if True, training and evaluation processes are\n            tracked with wandb\n    \"\"\"\n\n    train_loader: DataLoader  # Tmp. workaround\n    eval_loader: DataLoader  # Tmp. workaround\n\n    def __init__(\n        self,\n        logger: _Logger,\n        trainer_cfg: Dict[str, Any],\n        model: nn.Module,\n        loss_fn: _Loss,\n        optimizer: Optimizer,\n        lr_skd: Union[_LRScheduler, lr_scheduler.ReduceLROnPlateau],\n        ckpt_path: Path,\n        evaluator: Evaluator,\n        use_wandb: bool = False,\n    ):\n        self.logger = logger\n        self.trainer_cfg = trainer_cfg\n        self.model = model\n        self.loss_fn = loss_fn\n        self.optimizer = optimizer\n        self.lr_skd = lr_skd\n        self.ckpt_path = ckpt_path\n        self.evaluator = evaluator\n        self.use_wandb = use_wandb\n\n        self.device = CFG.device\n        self.epochs = trainer_cfg[\"epochs\"]\n        self.use_amp = trainer_cfg[\"use_amp\"]\n        self.grad_accum_steps = trainer_cfg[\"grad_accum_steps\"]\n        self.step_per_batch = trainer_cfg[\"step_per_batch\"]\n\n        # Debug options\n        self.one_batch_only = trainer_cfg[\"one_batch_only\"]\n\n        # Model checkpoint\n        self.model_ckpt = _ModelCheckpoint(ckpt_path, **trainer_cfg[\"model_ckpt\"])\n\n        # Early stopping\n        if trainer_cfg[\"es\"][\"patience\"] != 0:\n            self.logger.info(\"Please disable early stop!\")\n#             self.es = EarlyStopping(**trainer_cfg[\"es\"])\n        else:\n            self.es = None\n\n        self._iter = 0\n        self._track_best_model = True  # (Deprecated)\n\n    def train_eval(self, proc_id: int) -> Dict[str, np.ndarray]:\n        \"\"\"Run training and evaluation processes.\n\n        Args:\n            proc_id: identifier of the current process\n        \"\"\"\n        self.logger.info(\"Start training and evaluation processes...\")\n        for epoch in range(self.epochs):\n            self.epoch = epoch  # For interior use\n            train_loss = self._train_epoch()\n            val_loss, val_result, _ = self._eval_epoch()\n\n            # Adjust learning rate\n            if self.lr_skd is not None and not self.step_per_batch:\n                if isinstance(self.lr_skd, lr_scheduler.ReduceLROnPlateau):\n                    self.lr_skd.step(val_loss)\n                else:\n                    self.lr_skd.step()\n\n            # Track and log process result (by epoch)\n            self._log_proc(epoch, train_loss, val_loss, val_result)\n\n            # Record the best checkpoint\n            self.model_ckpt.step(\n                epoch, self.model, val_loss, val_result, last_epoch=False if epoch != self.epochs - 1 else True\n            )\n\n            # Check early stopping is triggered or not\n            if self.es is not None:\n                self.es.step(val_loss)\n                if self.es.stop:\n                    self.logger.info(f\"Early stopping is triggered at epoch {epoch}, training process is halted.\")\n                    break\n        if self.use_wandb:\n            wandb.log({\"best_epoch\": self.model_ckpt.best_epoch + 1})  # `epoch` starts from 0\n\n        # Run final evaluation\n        final_prf_report, y_preds = self._run_final_eval()\n        self._log_best_prf(final_prf_report)\n\n        return y_preds\n\n    @abstractmethod\n    def _train_epoch(self) -> Union[float, Dict[str, float]]:\n        \"\"\"Run training process for one epoch.\n\n        Returns:\n            train_loss_avg: average training loss over batches\n                *Note: If MTL is used, returned object will be dict\n                    containing losses of sub-tasks and the total loss.\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    def _eval_epoch(self, return_output: bool = False) -> Tuple[float, Dict[str, float], Optional[Tensor]]:\n        \"\"\"Run evaluation process for one epoch.\n\n        Args:\n            return_output: whether to return prediction\n\n        Returns:\n            eval_loss_avg: average evaluation loss over batches\n            eval_result: evaluation performance report\n            y_pred: prediction\n        \"\"\"\n        raise NotImplementedError\n\n    def _log_proc(\n        self,\n        epoch: int,\n        train_loss: Union[float, Dict[str, float]],\n        val_loss: Optional[float] = None,\n        val_result: Optional[Dict[str, float]] = None,\n        proc_id: Optional[str] = None,\n    ) -> None:\n        \"\"\"Log message of training process.\n\n        Args:\n            epoch: current epoch number\n            train_loss: training loss\n            val_loss: validation loss\n            val_result: evaluation performance report\n            proc_id: identifier of the current process\n        \"\"\"\n        proc_msg = [f\"Epoch{epoch} [{epoch+1}/{self.epochs}]\"]\n\n        # Construct training loss message\n        if isinstance(train_loss, float):\n            proc_msg.append(f\"Training loss {train_loss:.4f}\")\n        else:\n            for loss_k, loss_v in train_loss.items():\n                loss_name = loss_k.split(\"_\")[0].capitalize()\n                proc_msg.append(f\"{loss_name} loss {round(loss_v, 4)}\")\n\n        # Construct eval prf message\n        if val_loss is not None:\n            proc_msg.append(f\"Validation loss {val_loss:.4f}\")\n        if val_result is not None:\n            for metric, score in val_result.items():\n                proc_msg.append(f\"{metric.upper()} {round(score, 4)}\")\n\n        proc_msg = \" | \".join(proc_msg)\n        self.logger.info(proc_msg)\n\n        if self.use_wandb:\n            # Process loss dict and log\n            log_dict = train_loss if isinstance(train_loss, dict) else {\"train_loss\": train_loss}\n            if val_loss is not None:\n                log_dict[\"val_loss\"] = val_loss\n            if val_result is not None:\n                for metric, score in val_result.items():\n                    log_dict[metric] = score\n\n            if proc_id is not None:\n                log_dict = {f\"{k}_{proc_id}\": v for k, v in log_dict.items()}\n\n            wandb.log(log_dict)\n\n    def _run_final_eval(self) -> Tuple[Dict[str, Dict[str, float]], Dict[str, np.ndarray]]:\n        \"\"\"Run final evaluation process with designated model checkpoint.\n\n        Returns:\n            final_prf_report: performance report of final evaluation\n            y_preds: prediction on different datasets\n        \"\"\"\n        # Load the best model checkpoint\n        self.model = self.model_ckpt.load_best_ckpt(self.model, self.device)\n\n        # Reconstruct dataloaders\n        self._disable_shuffle()\n        val_loader = self.eval_loader\n\n        final_prf_report, y_preds = {}, {}\n        for data_split, dataloader in {\n            # \"train\": self.train_loader,\n            \"val\": val_loader,\n        }.items():\n            self.eval_loader = dataloader\n            _, eval_result, y_pred = self._eval_epoch(return_output=True)\n            final_prf_report[data_split] = eval_result\n            y_preds[data_split] = y_pred.numpy()\n\n        return final_prf_report, y_preds\n\n    def _disable_shuffle(self) -> None:\n        \"\"\"Disable shuffle in train dataloader for final evaluation.\"\"\"\n        self.train_loader = DataLoader(\n            self.train_loader.dataset,\n            batch_size=self.train_loader.batch_size,\n            shuffle=False,  # Reset shuffle to False\n            num_workers=self.train_loader.num_workers,\n            collate_fn=self.train_loader.collate_fn,\n        )\n\n    def _log_best_prf(self, prf_report: Dict[str, Any]) -> None:\n        \"\"\"Log performance evaluated with the best model checkpoint.\n\n        Args:\n            prf_report: performance report\n        \"\"\"\n        self.logger.info(\">>>>> Performance Report - Best Ckpt <<<<<\")\n        self.logger.info(json.dumps(prf_report, indent=4))\n\n        if self.use_wandb:\n            wandb.log(prf_report)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-02-20T09:37:49.596119Z","iopub.execute_input":"2024-02-20T09:37:49.596393Z","iopub.status.idle":"2024-02-20T09:37:49.627696Z","shell.execute_reply.started":"2024-02-20T09:37:49.596358Z","shell.execute_reply":"2024-02-20T09:37:49.627002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MainTrainer(_BaseTrainer):\n    \"\"\"Main trainer.\n\n    Args:\n        logger: message logger\n        trainer_cfg: hyperparameters for training and evaluation processes\n        model: model instance\n        loss_fn: loss criterion\n        optimizer: optimization algorithm\n        lr_scheduler: learning rate scheduler\n        scaler: scaling object\n        train_loader: training data loader\n        eval_loader: validation data loader\n        use_wandb: if True, training and evaluation processes are\n            tracked with wandb\n    \"\"\"\n\n    def __init__(\n        self,\n        logger: _Logger,\n        trainer_cfg: Dict[str, Any],\n        model: nn.Module,\n        loss_fn: _Loss,\n        optimizer: Optimizer,\n        lr_skd: Union[_LRScheduler, lr_scheduler.ReduceLROnPlateau],\n        ckpt_path: Path,\n        evaluator: Evaluator,\n        scaler: Any,\n        train_loader: DataLoader,\n        eval_loader: Optional[DataLoader] = None,\n        use_wandb: bool = False,\n    ):\n        super(MainTrainer, self).__init__(\n            logger,\n            trainer_cfg,\n            model,\n            loss_fn,\n            optimizer,\n            lr_skd,\n            ckpt_path,\n            evaluator,\n            use_wandb,\n        )\n        self.train_loader = train_loader\n        self.eval_loader = eval_loader if eval_loader else train_loader\n        self.scaler = scaler\n\n        self.loss_name = self.loss_fn.__class__.__name__\n\n        # Mixed precision training\n        self.grad_scaler = GradScaler(enabled=self.use_amp)\n\n    def _train_epoch(self) -> float:\n        \"\"\"Run training process for one epoch.\n\n        Returns:\n            train_loss_avg: average training loss over batches\n        \"\"\"\n        train_loss_total = 0\n\n        self.model.train()\n        for i, batch_data in enumerate(tqdm(self.train_loader)):\n            if i % self.grad_accum_steps == 0:\n                self.optimizer.zero_grad(set_to_none=True)\n\n            # Retrieve batched raw data\n            inputs = {}\n            for k, v in batch_data.items():\n                if k != \"y\":\n                    inputs[k] = v.to(self.device)\n                else:\n                    y = v.to(self.device)\n\n            with autocast(enabled=self.use_amp):\n                # Forward pass and derive loss\n                output = self.model(inputs)\n                loss = self.loss_fn(output, y)\n            train_loss_total += loss.item()\n            loss = loss / self.grad_accum_steps\n\n            # Backpropagation\n            self.grad_scaler.scale(loss).backward()\n            if (i + 1) % self.grad_accum_steps == 0:\n                self.grad_scaler.step(self.optimizer)\n                self.grad_scaler.update()\n                if self.step_per_batch:\n                    self.lr_skd.step()\n\n            self._iter += 1\n\n            # Free mem.\n            del inputs, y, output\n            _ = gc.collect()\n\n            if self.one_batch_only:\n                break\n\n        train_loss_avg = train_loss_total / len(self.train_loader)\n\n        return train_loss_avg\n\n    @torch.no_grad()\n    def _eval_epoch(\n        self,\n        return_output: bool = False,\n    ) -> Tuple[float, Dict[str, float], Optional[Tensor]]:\n        \"\"\"Run evaluation process for one epoch.\n\n        Args:\n            return_output: whether to return prediction\n\n        Returns:\n            eval_loss_avg: average evaluation loss over batches\n            eval_result: evaluation performance report\n            y_pred: prediction\n        \"\"\"\n        eval_loss_total = 0\n        y_true, y_pred = [], []\n\n        self.model.eval()\n        for i, batch_data in enumerate(self.eval_loader):\n            # Retrieve batched raw data\n            inputs = {}\n            for k, v in batch_data.items():\n                if k != \"y\":\n                    inputs[k] = v.to(self.device)\n                else:\n                    y = v.to(self.device)\n\n            # Forward pass\n            output = self.model(inputs)\n\n            # Derive loss\n            loss = self.loss_fn(output, y)\n            eval_loss_total += loss.item()\n\n            # Record batched output\n            y_true.append(y.detach().cpu())\n            y_pred.append(output.detach().cpu())\n\n            del inputs, y, output\n            _ = gc.collect()\n\n        eval_loss_avg = eval_loss_total / len(self.eval_loader)\n\n        # Run evaluation with the specified evaluation metrics\n        y_true = torch.cat(y_true, dim=0)\n        y_pred = torch.cat(y_pred, dim=0)\n        eval_result = self.evaluator.evaluate(y_true, y_pred, self.scaler)\n\n        if return_output:\n            return eval_loss_avg, eval_result, y_pred\n        else:\n            return eval_loss_avg, eval_result, None","metadata":{"execution":{"iopub.status.busy":"2024-02-20T09:37:49.629051Z","iopub.execute_input":"2024-02-20T09:37:49.629345Z","iopub.status.idle":"2024-02-20T09:37:49.651714Z","shell.execute_reply.started":"2024-02-20T09:37:49.629313Z","shell.execute_reply":"2024-02-20T09:37:49.650976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"cv\"></a>\n## 7. Run Cross-Validation\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nNow, all componentst have been settled down, let's train our models.","metadata":{}},{"cell_type":"code","source":"if CFG.train_models:\n    oof = np.zeros((len(train), N_CLASSES))\n    prfs = []\n\n    cv = GroupKFold(n_splits=5)\n    for fold, (tr_idx, val_idx) in enumerate(cv.split(train, train[TGT_COL], train[\"patient_id\"])):\n        logger.info(f\"== Train and Eval Process - Fold{fold} ==\")\n\n        # Build dataloaders\n        data_tr, data_val = train.iloc[tr_idx].reset_index(drop=True), train.iloc[val_idx].reset_index(drop=True)\n        train_loader = DataLoader(\n            EEGDataset({\"meta\": data_tr, \"eeg\": all_eegs}, \"train\", **CFG.dataset),\n            shuffle=CFG.trainer[\"dataloader\"][\"shuffle\"],\n            batch_size=CFG.trainer[\"dataloader\"][\"batch_size\"],\n            num_workers=CFG.trainer[\"dataloader\"][\"num_workers\"]\n        )\n        val_loader = DataLoader(\n            EEGDataset({\"meta\": data_val, \"eeg\": all_eegs}, \"valid\", **CFG.dataset),\n            shuffle=False,\n            batch_size=CFG.trainer[\"dataloader\"][\"batch_size\"],\n            num_workers=CFG.trainer[\"dataloader\"][\"num_workers\"]\n        )\n\n        # Build model\n        logger.info(f\"Build model...\")\n        model = DilatedInceptionWaveNet()\n        model.to(CFG.device)\n\n        # Build criterion\n        loss_fn = KLDivWithLogitsLoss()\n\n        # Build solvers\n        optimizer = torch.optim.Adam(model.parameters(), lr=CFG.trainer[\"lr\"])\n        num_training_steps = (\n            math.ceil(\n                len(train_loader.dataset)\n                / (CFG.trainer[\"dataloader\"][\"batch_size\"] * CFG.trainer[\"grad_accum_steps\"])\n            )\n            * CFG.trainer[\"epochs\"]\n        )\n        lr_skd = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)\n\n        # Build evaluator\n        evaluator = Evaluator(metric_names=[\"kldiv\"])\n\n        # Build trainer\n        trainer: _BaseTrainer = None\n        trainer = MainTrainer(\n            logger=logger,\n            trainer_cfg=CFG.trainer,\n            model=model,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            lr_skd=lr_skd,\n            ckpt_path=CFG.exp_dump_path,\n            evaluator=evaluator,\n            scaler=None,\n            train_loader=train_loader,\n            eval_loader=val_loader,\n            use_wandb=False\n        )\n\n        # Run main training and evaluation for one fold\n        y_preds = trainer.train_eval(fold)\n        oof[val_idx, :] = y_preds[\"val\"]\n\n        # Dump output objects\n        for model_path in CFG.exp_dump_path.glob(\"*.pth\"):\n            if \"seed\" in str(model_path) or \"fold\" in str(model_path):\n                continue\n\n            # Rename model file\n            model_file_name_dst = f\"{model_path.stem}_fold{fold}.pth\"\n            model_path_dst = exp.ckpt_path / model_file_name_dst\n            model_path.rename(model_path_dst)\n\n        # Free mem.\n        del (data_tr, data_val, train_loader, val_loader, model, optimizer, lr_skd, evaluator, trainer)\n        _ = gc.collect()\n\n        if CFG.one_fold_only:\n            logger.info(\"Cross-validatoin stops at first fold!!!\")\n            break\n\n    np.save(CFG.exp_dump_pat / \"oof.npy\", oof)\nelse:\n    oof = np.load(\"/kaggle/input/hms-oof-demo/oof_seed0.npy\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"cv_score\"></a>\n## 8. Derive CV Score\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nFinally, let's see how well this model can perform on local cross-validation. It's exciting to see we can go below CV 0.6772 (below 0.7) using raw EEG signals only!","metadata":{}},{"cell_type":"code","source":"sys.path.append('/kaggle/input/kaggle-kl-div')\nfrom kaggle_kl_div import score\n\n# OOf prediction\ny_pred = pd.DataFrame(F.softmax(torch.tensor(oof), dim=1))\ny_pred[\"id\"] = np.arange(len(oof))\n\n# Ground truth\ny_true = pd.DataFrame(train[TGT_VOTE_COLS].values)\ny_true[\"id\"] = np.arange(len(y_true))\n\ncv_score = score(solution=y_true, submission=y_pred, row_id_column_name=\"id\")\nlogger.info(\">>>>> Performance Report <<<<<\")\nlogger.info(f\"-> KL Divergence: {cv_score:.4f}\")","metadata":{"execution":{"iopub.status.busy":"2024-02-19T15:10:57.192172Z","iopub.execute_input":"2024-02-19T15:10:57.192764Z","iopub.status.idle":"2024-02-19T15:10:57.254016Z","shell.execute_reply.started":"2024-02-19T15:10:57.192735Z","shell.execute_reply":"2024-02-19T15:10:57.252938Z"},"trusted":true},"execution_count":null,"outputs":[]}]}