{"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":"gpu","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7392733,"sourceType":"datasetVersion","datasetId":4297749},{"sourceId":7392775,"sourceType":"datasetVersion","datasetId":4297782},{"sourceId":7517324,"sourceType":"datasetVersion","datasetId":4378712},{"sourceId":7652061,"sourceType":"datasetVersion","datasetId":4216847}],"dockerImageVersionId":30648,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Directory settings","metadata":{}},{"cell_type":"markdown","source":"# About this notebook\n\nThe goal of this notebook is to improve the results of the [notebook](https://www.kaggle.com/code/cdeotte/efficientnetb0-starter-lb-0-43#Train-DataLoader) shared by @cdeotte and [@alejopaullier](https://www.kaggle.com/code/alejopaullier/hms-efficientnetb0-pytorch-train), please check them out great notebooks.\n\n\n**Important note**:\n\nI shared in the last days a [discussion](https://www.kaggle.com/competitions/hms-harmful-brain-activity-classification/discussion/478474) which was an Improvement of this [notebook](https://www.kaggle.com/code/nischaydnk/hms-submission-1d-eegnet-pipeline-lightning) created by @nischaydnk. The tricks I used were notably changing the optimizer (I used Adan) and I used the two stage training as stated by @seanbearden [here](https://www.kaggle.com/code/seanbearden/effnetb0-2-pop-model-train-twice-lb-0-39) and not using downsampling which worked pretty well.\nBut the two stage training an issue which is data leakage. The idea was to seperate the data with few votes because the kl will hardly penalize the model if it mislabel them. So in my previous experiments I use two groupkfold CV on the two datasets and the samples with few votes are present in both datasets.\n\nSo in this notebook, I use one CV scheme and in each stage I filter the data then validate on the data that contains both population to prevent data leakage. Let me know in the comment if this approach is correct more info can be found [here](https://www.kaggle.com/competitions/hms-harmful-brain-activity-classification/discussion/477461).\n\nAs stated [here](https://www.kaggle.com/competitions/hms-harmful-brain-activity-classification/discussion/477498), adding 0.166666667 to the targets will reduce the CV/LB gap.\n\n**Consider upvoting this notebook if you find it useful**\n\n## Version 1\n* I train a tf_efficientnet_b0_ns model.\n\n### Hyperparams\n\n```\nscheduler='OneCycleLR' \n # OneCycleLR params\n  cosanneal_res_params={\n      'T_0':20,\n      'eta_min':1e-6,\n      'T_mult':1,\n      'last_epoch':-1}\n  print_freq=50\n  num_workers = 1\n  model_name = 'tf_efficientnet_b0_ns'\n  optimizer='Adam'\n  stage1_epochs = 5\n  stage1_epochs = 7\n  eps = 1e-6\n  lr = 1e-3\n  batch_size = 128\n  weight_decay = 1e-2\n  seed = 2024\n```\n\n## Version2\n\n* I changed the CV sheme, first stage train on all data second stage train on data with `total_evaluators >= 10`\n* Added Time masking augmentation from [here](https://www.kaggle.com/code/iglovikov/xymasking-aug).\n\n### Hyperparams\n\n```\nscheduler='OneCycleLR' \n # OneCycleLR params\n  cosanneal_res_params={\n      'T_0':20,\n      'eta_min':1e-6,\n      'T_mult':1,\n      'last_epoch':-1}\n  print_freq=50\n  num_workers = 1\n  model_name = 'tf_efficientnet_b0_ns'\n  optimizer='Adam'\n  stage1_epochs = 5\n  stage1_epochs = 5\n  eps = 1e-6\n  lr = 1e-3\n  batch_size = 64\n  weight_decay = 1e-2\n  seed = 2024\n```\n\n## Version3\n\n* removed the `0.166667` hack\n* added mixup and hflip augmentations\n\n```\nscheduler='OneCycleLR' \n # OneCycleLR params\n  cosanneal_res_params={\n      'T_0':20,\n      'eta_min':1e-6,\n      'T_mult':1,\n      'last_epoch':-1}\n  print_freq=50\n  num_workers = 1\n  model_name = 'tf_efficientnet_b0_ns'\n  optimizer='Adam'\n  stage1_epochs = 5\n  stage1_epochs = 5\n  eps = 1e-6\n  lr = 1e-3\n  batch_size = 64\n  weight_decay = 1e-2\n  seed = 2024\n```","metadata":{}},{"cell_type":"code","source":"!pip install -U albumentations","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:20.392939Z","iopub.execute_input":"2024-03-04T18:14:20.393611Z","iopub.status.idle":"2024-03-04T18:14:34.117788Z","shell.execute_reply.started":"2024-03-04T18:14:20.393577Z","shell.execute_reply":"2024-03-04T18:14:34.116853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# directory settings\n# ====================================================\n\nimport os\n\nOUTPUT_DIR = './'\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)\n    \nPOP_2_DIR = OUTPUT_DIR + 'pop_2_weight_oof/'\nif not os.path.exists(POP_2_DIR):\n    os.makedirs(POP_2_DIR)\n    \nPOP_1_DIR = OUTPUT_DIR + 'pop_1_weight_oof/'\nif not os.path.exists(POP_1_DIR):\n    os.makedirs(POP_1_DIR)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-03-04T18:14:34.119646Z","iopub.execute_input":"2024-03-04T18:14:34.119949Z","iopub.status.idle":"2024-03-04T18:14:34.126556Z","shell.execute_reply.started":"2024-03-04T18:14:34.119922Z","shell.execute_reply":"2024-03-04T18:14:34.125624Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Libraries","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nfrom glob import glob\nimport sys\nimport math\nimport time\nimport random\nimport shutil\nfrom pathlib import Path\nfrom typing import Dict, List\nfrom scipy.stats import entropy\nfrom scipy.signal import butter, lfilter, freqz\nfrom contextlib import contextmanager\nfrom collections import defaultdict, Counter\nsys.path.append('/kaggle/input/kaggle-kl-div')\nfrom kaggle_kl_div import score\nimport numpy as np\nimport pandas as pd\nfrom sklearn import preprocessing\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.metrics import accuracy_score, log_loss\nfrom tqdm.auto import tqdm\nfrom functools import partial\nimport cv2\nfrom PIL import Image\nimport torch\nimport torch.nn as nn\nimport pytorch_lightning as pl\nimport torch.nn.functional as F\nfrom torch.optim import Adam, SGD, AdamW\nimport torchvision.models as models\nfrom torch.nn.parameter import Parameter\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, OneCycleLR, CosineAnnealingLR, CosineAnnealingWarmRestarts\nfrom sklearn.preprocessing import LabelEncoder\nfrom torchvision.transforms import v2\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.model_selection import train_test_split\nimport albumentations as A\nfrom albumentations import (Compose, Normalize, Resize, RandomResizedCrop, HorizontalFlip, VerticalFlip, ShiftScaleRotate, Transpose)\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\nimport timm\nimport warnings \nwarnings.filterwarnings('ignore')\nos.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1\"\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nfrom matplotlib import pyplot as plt\nimport joblib\nVERSION=3","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:34.127723Z","iopub.execute_input":"2024-03-04T18:14:34.127966Z","iopub.status.idle":"2024-03-04T18:14:44.579181Z","shell.execute_reply.started":"2024-03-04T18:14:34.127945Z","shell.execute_reply":"2024-03-04T18:14:44.578351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# CFG\n# ====================================================\n\nclass CFG:\n    wandb = False\n    debug = False\n    train=True\n    apex=True\n    stage1_pop1=True\n    stage2_pop2=False\n    VISUALIZE=True\n    FREEZE=False\n    SparK=False\n    t4_gpu=False\n    USE_MIXUP = True\n    scheduler='OneCycleLR' # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts','OneCycleLR']\n    # CosineAnnealingLR params\n    cosanneal_params={\n        'T_max':6,\n        'eta_min':1e-5,\n        'last_epoch':-1\n    }\n    #ReduceLROnPlateau params\n    reduce_params={\n        'mode':'min',\n        'factor':0.2,\n        'patience':4,\n        'eps':1e-6,\n        'verbose':True\n    }\n    # CosineAnnealingWarmRestarts params\n    cosanneal_res_params={\n        'T_0':20,\n        'eta_min':1e-6,\n        'T_mult':1,\n        'last_epoch':-1\n    }\n    print_freq=50\n    num_workers = 1\n    model_name = 'tf_efficientnet_b0_ns'\n    optimizer='Adan'\n    epochs = 5\n    factor = 0.9\n    patience = 2\n    eps = 1e-6\n    lr = 1e-3\n    min_lr = 1e-6\n    batch_size = 64\n    weight_decay = 1e-2\n    batch_scheduler=True\n    gradient_accumulation_steps = 1\n    max_grad_norm = 1e7\n    seed = 2024\n    target_cols = ['seizure_vote', 'lpd_vote', 'gpd_vote', 'lrda_vote', 'grda_vote', 'other_vote']\n    target_size = 6\n    pred_cols = ['pred_seizure_vote', 'pred_lpd_vote', 'pred_gpd_vote', 'pred_lrda_vote', 'pred_grda_vote', 'pred_other_vote']\n    n_fold = 5\n    trn_fold = [0, 1, 2, 3, 4]\n    PATH = '/kaggle/input/hms-harmful-brain-activity-classification/'\n    data_root = \"/kaggle/input/hms-harmful-brain-activity-classification/train_eegs/\"\n    raw_eeg_path = \"/kaggle/input/brain-eegs/eegs.npy\"","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:44.581695Z","iopub.execute_input":"2024-03-04T18:14:44.582077Z","iopub.status.idle":"2024-03-04T18:14:44.592411Z","shell.execute_reply.started":"2024-03-04T18:14:44.582028Z","shell.execute_reply":"2024-03-04T18:14:44.591017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"def init_logger(log_file=OUTPUT_DIR+'train.log'):\n    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = init_logger()\n\ndef get_score(preds, targets):\n    oof = pd.DataFrame(preds.copy())\n    oof['id'] = np.arange(len(oof))\n\n    true = pd.DataFrame(targets.copy())\n    true['id'] = np.arange(len(true))\n\n    cv = score(solution=true, submission=oof, row_id_column_name='id')\n    return cv\n\ndef butter_bandpass(lowcut, highcut, fs, order=5):\n    return butter(order, [lowcut, highcut], fs=fs, btype='band')\n\ndef butter_bandpass_filter(data, lowcut, highcut, fs, order=5):\n    b, a = butter_bandpass(lowcut, highcut, fs, order=order)\n    y = lfilter(b, a, data)\n    return y\n\n\ndef denoise_filter(x):\n    # Sample rate and desired cutoff frequencies (in Hz).\n    fs = 200.0\n    lowcut = 1.0\n    highcut = 25.0\n    \n    # Filter a noisy signal.\n    T = 50\n    nsamples = T * fs\n    t = np.arange(0, nsamples) / fs\n    y = butter_bandpass_filter(x, lowcut, highcut, fs, order=6)\n    y = (y + np.roll(y,-1)+ np.roll(y,-2)+ np.roll(y,-3))/4\n    y = y[0:-1:4]\n    \n    return y\n\nclass KLDivLossWithLogits(nn.KLDivLoss):\n\n    def __init__(self):\n        super().__init__(reduction=\"batchmean\")\n\n    def forward(self, y, t):\n        y = nn.functional.log_softmax(y,  dim=1)\n        loss = super().forward(y, t)\n\n        return loss\n\n\n\ndef seed_torch(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    \nseed_torch(seed=CFG.seed)","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:44.593583Z","iopub.execute_input":"2024-03-04T18:14:44.593934Z","iopub.status.idle":"2024-03-04T18:14:44.617819Z","shell.execute_reply.started":"2024-03-04T18:14:44.593870Z","shell.execute_reply":"2024-03-04T18:14:44.617054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load train data","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/hms-harmful-brain-activity-classification/train.csv')\nTARGETS = train.columns[-6:]\nprint('Train shape:', train.shape )\nprint('Targets', list(TARGETS))\n\ntrain['total_evaluators'] = train[['seizure_vote', 'lpd_vote', 'gpd_vote', 'lrda_vote', 'grda_vote', 'other_vote']].sum(axis=1)\n\nprint(f'There are {train.patient_id.nunique()} patients in the training data.')\nprint(f'There are {train.eeg_id.nunique()} EEG IDs in the training data.')","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:44.618921Z","iopub.execute_input":"2024-03-04T18:14:44.619183Z","iopub.status.idle":"2024-03-04T18:14:44.935675Z","shell.execute_reply.started":"2024-03-04T18:14:44.619161Z","shell.execute_reply":"2024-03-04T18:14:44.934661Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(10, 6))\nplt.hist(train['total_evaluators'], bins=10, color='blue', edgecolor='black')\nplt.title('Histogram of Total Evaluators')\nplt.xlabel('Total Evaluators')\nplt.ylabel('Frequency')\nplt.grid(True)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:44.936800Z","iopub.execute_input":"2024-03-04T18:14:44.937120Z","iopub.status.idle":"2024-03-04T18:14:45.193509Z","shell.execute_reply.started":"2024-03-04T18:14:44.937093Z","shell.execute_reply":"2024-03-04T18:14:45.192609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nspectrograms = np.load('/kaggle/input/brain-spectrograms/specs.npy', allow_pickle=True).item()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:14:45.195015Z","iopub.execute_input":"2024-03-04T18:14:45.195397Z","iopub.status.idle":"2024-03-04T18:15:56.455386Z","shell.execute_reply.started":"2024-03-04T18:14:45.195362Z","shell.execute_reply":"2024-03-04T18:15:56.454340Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nall_eegs = np.load('/kaggle/input/eeg-spectrogram-by-lead-id-unique/eeg_specs.npy',allow_pickle=True).item()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:15:56.457877Z","iopub.execute_input":"2024-03-04T18:15:56.458236Z","iopub.status.idle":"2024-03-04T18:17:39.292457Z","shell.execute_reply.started":"2024-03-04T18:15:56.458210Z","shell.execute_reply":"2024-03-04T18:17:39.291512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Deduplicate Train EEG Id","metadata":{}},{"cell_type":"code","source":"train = train[train['label_id'].isin(all_eegs.keys())].copy()\n\ny_data = train[TARGETS].values # Regularization value\ny_data = y_data / y_data.sum(axis=1,keepdims=True)\ntrain[TARGETS] = y_data\n\ntrain['target'] = train['expert_consensus']\n\ntrain = train.reset_index(drop=True)\n\nplt.figure(figsize=(10, 6))\nplt.hist(train['total_evaluators'], bins=10, color='blue', edgecolor='black')\nplt.title('Histogram of Total Evaluators')\nplt.xlabel('Total Evaluators')\nplt.ylabel('Frequency')\nplt.grid(True)\nplt.show()\n\ndel y_data\n_ = gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:39.295793Z","iopub.execute_input":"2024-03-04T18:17:39.296078Z","iopub.status.idle":"2024-03-04T18:17:39.775015Z","shell.execute_reply.started":"2024-03-04T18:17:39.296043Z","shell.execute_reply":"2024-03-04T18:17:39.774122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CV Scheme","metadata":{}},{"cell_type":"code","source":"gkf = GroupKFold(n_splits=CFG.n_fold)\n\ntrain[\"fold\"] = -1\n\nfor fold_id, (_, val_idx) in enumerate(\n    gkf.split(train, y=train[\"target\"], groups=train[\"patient_id\"])\n):\n    train.loc[val_idx, \"fold\"] = fold_id\n    \ndel gkf\n_ = gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:39.776565Z","iopub.execute_input":"2024-03-04T18:17:39.776935Z","iopub.status.idle":"2024-03-04T18:17:40.015196Z","shell.execute_reply.started":"2024-03-04T18:17:39.776902Z","shell.execute_reply":"2024-03-04T18:17:40.014436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(\n        self, df: pd.DataFrame,\n        augment: bool = False, mode: str = 'train',\n        specs: Dict[int, np.ndarray] = spectrograms,\n        eeg_specs: Dict[int, np.ndarray] = all_eegs\n    ): \n        self.df = df\n        self.augment = augment\n        self.mode = mode\n        self.spectograms = spectrograms\n        self.eeg_spectograms = eeg_specs\n        \n    def __len__(self):\n        \"\"\"\n        Denotes the number of batches per epoch.\n        \"\"\"\n        return len(self.df)\n        \n    def __getitem__(self, index):\n        \"\"\"\n        Generate one batch of data.\n        \"\"\"\n        X, y = self.__data_generation(index)\n        if self.augment:\n            X = self.__transform(X) \n        return {\"spectrogram\":torch.tensor(X, dtype=torch.float32), \"labels\":torch.tensor(y, dtype=torch.float32)}\n                        \n    def __data_generation(self, index):\n        \"\"\"\n        Generates data containing batch_size samples.\n        \"\"\"\n        X = np.zeros((128, 256, 8), dtype='float32')\n        y = np.zeros(6, dtype='float32')\n        img = np.ones((128,256), dtype='float32')\n        row = self.df.iloc[index]\n        if self.mode=='test': \n            r = 0\n        else: \n            r = int(row['spectrogram_label_offset_seconds'] // 2)\n            \n        for region in range(4):\n            img = self.spectograms[row.spectrogram_id][r:r+300, region*100:(region+1)*100].T\n            \n            # Log transform spectogram\n            img = np.clip(img, np.exp(-4), np.exp(8))\n            img = np.log(img)\n\n            # Standarize per image\n            ep = 1e-6\n            mu = np.nanmean(img.flatten())\n            std = np.nanstd(img.flatten())\n            img = (img-mu)/(std+ep)\n            img = np.nan_to_num(img, nan=0.0)\n            X[14:-14, :, region] = img[:, 22:-22] / 2.0\n            img = self.eeg_spectograms[row.label_id]\n            X[:, :, 4:] = img\n                \n            if self.mode != 'test':\n                y = row[TARGETS].values.astype(np.float32)\n            \n        return X, y\n    \n    def __transform(self, img):\n        params1 = {\n                    \"num_masks_x\": 1,    \n                    \"mask_x_length\": (0, 20), # This line changed from fixed  to a range\n                    \"fill_value\": (0, 1, 2, 3, 4, 5, 6, 7),\n                    }\n        params2 = {    \n                    \"num_masks_y\": 1,    \n                    \"mask_y_length\": (0, 20),\n                    \"fill_value\": (0, 1, 2, 3, 4, 5, 6, 7),    \n                    }\n        params3 = {    \n                    \"num_masks_x\": (2, 4),\n                    \"num_masks_y\": 5,    \n                    \"mask_y_length\": 8,\n                    \"mask_x_length\": (10, 20),\n                    \"fill_value\": (0, 1, 2, 3, 4, 5, 6, 7),  \n                    }\n        \n        transforms = A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.XYMasking(**params1, p=0.3),\n            A.XYMasking(**params2, p=0.3),\n            A.XYMasking(**params3, p=0.3),\n        ])\n        return transforms(image=img)['image']","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:40.016387Z","iopub.execute_input":"2024-03-04T18:17:40.016716Z","iopub.status.idle":"2024-03-04T18:17:40.035763Z","shell.execute_reply.started":"2024-03-04T18:17:40.016687Z","shell.execute_reply":"2024-03-04T18:17:40.034706Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataloader","metadata":{}},{"cell_type":"code","source":"dataset = CustomDataset(train, augment=True, mode=\"train\")\ndataloader = DataLoader(dataset, batch_size=32, shuffle=False)\n\nbatch = dataset[0]\nX, y = batch[\"spectrogram\"], batch[\"labels\"]\nprint(f\"X shape: {X.shape}\")\nprint(f\"y shape: {y.shape}\")\n\ndel dataset, X, y\n_ = gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:40.037348Z","iopub.execute_input":"2024-03-04T18:17:40.037645Z","iopub.status.idle":"2024-03-04T18:17:40.291915Z","shell.execute_reply.started":"2024-03-04T18:17:40.037622Z","shell.execute_reply":"2024-03-04T18:17:40.290964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.VISUALIZE:\n    ROWS = 2\n    COLS = 3\n    for batch in dataloader:\n        X, y = batch[\"spectrogram\"], batch[\"labels\"]\n        plt.figure(figsize=(20,8))\n        for row in range(ROWS):\n            for col in range(COLS):\n                plt.subplot(ROWS, COLS, row*COLS + col+1)\n                t = y[row*COLS + col]\n                img = X[row*COLS + col, :, :, 0]\n                mn = img.flatten().min()\n                mx = img.flatten().max()\n                img = (img-mn)/(mx-mn)\n                plt.imshow(img)\n                tars = f'[{t[0]:0.2f}'\n                for s in t[1:]:\n                    tars += f', {s:0.2f}'\n                eeg = train.eeg_id.values[row*CFG.batch_size + row*COLS + col]\n                plt.title(f'EEG = {eeg}\\nTarget = {tars}',size=12)\n                plt.yticks([])\n                plt.ylabel('Frequencies (Hz)',size=14)\n                plt.xlabel('Time (sec)',size=16)\n        plt.show()\n        break\n        \ndel dataloader\n_ = gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:40.293005Z","iopub.execute_input":"2024-03-04T18:17:40.293291Z","iopub.status.idle":"2024-03-04T18:17:41.896966Z","shell.execute_reply.started":"2024-03-04T18:17:40.293267Z","shell.execute_reply":"2024-03-04T18:17:41.894525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class CustomModel(nn.Module):\n    def __init__(self, config, num_classes: int = 6, pretrained: bool = True):\n        super(CustomModel, self).__init__()\n        self.USE_KAGGLE_SPECTROGRAMS = True\n        self.USE_EEG_SPECTROGRAMS = True\n        self.model = timm.create_model(\n            config.model_name,\n            pretrained=pretrained,\n        )\n        if CFG.FREEZE:\n            total_params = sum(p.numel() for p in self.model.parameters())\n            percentage_to_freeze = 0.1\n            for i,(name, param) in enumerate(list(self.model.named_parameters())\\\n                                             [0:int(total_params * percentage_to_freeze)]):\n                param.requires_grad = False\n\n        self.features = nn.Sequential(*list(self.model.children())[:-2])\n        self.custom_layers = nn.Sequential(\n            nn.AdaptiveAvgPool2d(1),\n            nn.Flatten(),\n            nn.Linear(self.model.num_features, num_classes)\n        )\n\n    def __reshape_input(self, x):\n        \"\"\"\n        Reshapes input (128, 256, 8) -> (512, 512, 3) monotone image.\n        \"\"\" \n        # === Get spectograms ===\n        spectograms = [x[:, :, :, i:i+1] for i in range(4)]\n        spectograms = torch.cat(spectograms, dim=1)\n        \n        # === Get EEG spectograms ===\n        eegs = [x[:, :, :, i:i+1] for i in range(4,8)]\n        eegs = torch.cat(eegs, dim=1)\n        \n        # === Reshape (512,512,3) ===\n        if self.USE_KAGGLE_SPECTROGRAMS & self.USE_EEG_SPECTROGRAMS:\n            x = torch.cat([spectograms, eegs], dim=2)\n        elif self.USE_EEG_SPECTROGRAMS:\n            x = eegs\n        else:\n            x = spectograms\n            \n        x = torch.cat([x,x,x], dim=3)\n        x = x.permute(0, 3, 1, 2)\n        return x\n    def extract_features(self, x):\n        x = self.__reshape_input(x)\n        x = self.features(x)\n        return x\n    \n    def forward(self, x):\n        x = self.extract_features(x)\n        x = self.custom_layers(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:41.898577Z","iopub.execute_input":"2024-03-04T18:17:41.898873Z","iopub.status.idle":"2024-03-04T18:17:41.912024Z","shell.execute_reply.started":"2024-03-04T18:17:41.898845Z","shell.execute_reply":"2024-03-04T18:17:41.910965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"iot = torch.randn(2, 128, 256, 8)\nmodel = CustomModel(CFG)\noutput = model(iot)\nprint(output.shape)\n\ndel iot, model, output\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:41.913124Z","iopub.execute_input":"2024-03-04T18:17:41.913412Z","iopub.status.idle":"2024-03-04T18:17:44.834077Z","shell.execute_reply.started":"2024-03-04T18:17:41.913382Z","shell.execute_reply":"2024-03-04T18:17:44.833123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Mixup Augmentation","metadata":{}},{"cell_type":"code","source":"def mixup_data(\n        X: torch.Tensor, y: torch.Tensor, config, device,\n        alpha: float = 1.0\n    ):\n    \"\"\"\n    Performs MixUp augmentation.\n    :param config: configuration class with param to use mixup or not.\n    :param X: batch with images.\n    :param y: ground truth.\n    :param alpha: parameter to use in the beta distribution.\n    :param device: indicates if using CPU or GPU.\n    :return mixed_X: mixed X.\n    :return y_a: class of the original image.\n    :return y_b: class of the image used for mixing.\n    :return lambda: lamdba parameter that indicates the percentage of the mix.\n    \"\"\"\n    if not config.USE_MIXUP:\n        return X, y, None, None\n    \n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha) # draw random number from beta distribution\n    else:\n        lam = 1\n\n    batch_size = X.size()[0]\n    index = torch.randperm(batch_size).to(device) # torch tensor with shuffled numbers between 1:batch_size\n    mixed_X = lam * X + (1 - lam) * X[index, :] # perform mixup to the whole batch\n    y_a, y_b = y, y[index]\n    return mixed_X, y_a, y_b, lam\n\n\ndef get_criterion(config, criterion):\n    \"\"\"\n    This function computes the criterion/loss depending whether MixUp augmentation was applied or not.\n    If MixUp was applied it returns a weighted average of the loss averaging by the lambda parameter.\n    Otherwise, it returns the regular loss as MixUp was not applied.\n    :param config: configuration class with param to use mixup or not.\n    :param criterion: loss function to use.\n    \"\"\"\n\n    def mixup_criterion(pred, y_a, y_b, lam):\n        return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)\n\n    def single_criterion(pred, y_a, y_b, lam):\n        return criterion(pred, y_a)\n    \n    if config.USE_MIXUP:\n        return mixup_criterion\n    else:\n        return single_criterion\n","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:44.835142Z","iopub.execute_input":"2024-03-04T18:17:44.835419Z","iopub.status.idle":"2024-03-04T18:17:44.846164Z","shell.execute_reply.started":"2024-03-04T18:17:44.835396Z","shell.execute_reply":"2024-03-04T18:17:44.845177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Adan Optimizer","metadata":{}},{"cell_type":"code","source":"import math\nimport torch\nfrom torch.optim.optimizer import Optimizer\n\n\nclass Adan(Optimizer):\n    \"\"\"\n    Implements a pytorch variant of Adan\n    Adan was proposed in\n    Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models[J]. arXiv preprint arXiv:2208.06677, 2022.\n    https://arxiv.org/abs/2208.06677\n    Arguments:\n        params (iterable): iterable of parameters to optimize or dicts defining parameter groups.\n        lr (float, optional): learning rate. (default: 1e-3)\n        betas (Tuple[float, float, flot], optional): coefficients used for computing \n            running averages of gradient and its norm. (default: (0.98, 0.92, 0.99))\n        eps (float, optional): term added to the denominator to improve \n            numerical stability. (default: 1e-8)\n        weight_decay (float, optional): decoupled weight decay (L2 penalty) (default: 0)\n        max_grad_norm (float, optional): value used to clip \n            global grad norm (default: 0.0 no clip)\n        no_prox (bool): how to perform the decoupled weight decay (default: False)\n    \"\"\"\n\n    def __init__(self, params, lr=1e-3, betas=(0.98, 0.92, 0.99), eps=1e-8,\n                 weight_decay=0.2, max_grad_norm=0.0, no_prox=False):\n        if not 0.0 <= max_grad_norm:\n            raise ValueError(\"Invalid Max grad norm: {}\".format(max_grad_norm))\n        if not 0.0 <= lr:\n            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n        if not 0.0 <= eps:\n            raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n        if not 0.0 <= betas[0] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n        if not 0.0 <= betas[1] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n        if not 0.0 <= betas[2] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 2: {}\".format(betas[2]))\n        defaults = dict(lr=lr, betas=betas, eps=eps,\n                        weight_decay=weight_decay,\n                        max_grad_norm=max_grad_norm, no_prox=no_prox)\n        super(Adan, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(Adan, self).__setstate__(state)\n        for group in self.param_groups:\n            group.setdefault('no_prox', False)\n\n    @torch.no_grad()\n    def restart_opt(self):\n        for group in self.param_groups:\n            group['step'] = 0\n            for p in group['params']:\n                if p.requires_grad:\n                    state = self.state[p]\n                    # State initialization\n\n                    # Exponential moving average of gradient values\n                    state['exp_avg'] = torch.zeros_like(p)\n                    # Exponential moving average of squared gradient values\n                    state['exp_avg_sq'] = torch.zeros_like(p)\n                    # Exponential moving average of gradient difference\n                    state['exp_avg_diff'] = torch.zeros_like(p)\n\n    @torch.no_grad()\n    def step(self):\n        \"\"\"\n            Performs a single optimization step.\n        \"\"\"\n        if self.defaults['max_grad_norm'] > 0:\n            device = self.param_groups[0]['params'][0].device\n            global_grad_norm = torch.zeros(1, device=device)\n\n            max_grad_norm = torch.tensor(self.defaults['max_grad_norm'], device=device)\n            for group in self.param_groups:\n\n                for p in group['params']:\n                    if p.grad is not None:\n                        grad = p.grad\n                        global_grad_norm.add_(grad.pow(2).sum())\n\n            global_grad_norm = torch.sqrt(global_grad_norm)\n\n            clip_global_grad_norm = torch.clamp(max_grad_norm / (global_grad_norm + group['eps']), max=1.0)\n        else:\n            clip_global_grad_norm = 1.0\n\n        for group in self.param_groups:\n            beta1, beta2, beta3 = group['betas']\n            # assume same step across group now to simplify things\n            # per parameter step can be easily support by making it tensor, or pass list into kernel\n            if 'step' in group:\n                group['step'] += 1\n            else:\n                group['step'] = 1\n\n            bias_correction1 = 1.0 - beta1 ** group['step']\n\n            bias_correction2 = 1.0 - beta2 ** group['step']\n\n            bias_correction3 = 1.0 - beta3 ** group['step']\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n\n                state = self.state[p]\n                if len(state) == 0:\n                    state['exp_avg'] = torch.zeros_like(p)\n                    state['exp_avg_sq'] = torch.zeros_like(p)\n                    state['exp_avg_diff'] = torch.zeros_like(p)\n\n                grad = p.grad.mul_(clip_global_grad_norm)\n                if 'pre_grad' not in state or group['step'] == 1:\n                    state['pre_grad'] = grad\n\n                copy_grad = grad.clone()\n\n                exp_avg, exp_avg_sq, exp_avg_diff = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_diff']\n                diff = grad - state['pre_grad']\n\n                update = grad + beta2 * diff\n                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)  # m_t\n                exp_avg_diff.mul_(beta2).add_(diff, alpha=1 - beta2)  # diff_t\n                exp_avg_sq.mul_(beta3).addcmul_(update, update, value=1 - beta3)  # n_t\n\n                denom = ((exp_avg_sq).sqrt() / math.sqrt(bias_correction3)).add_(group['eps'])\n                update = ((exp_avg / bias_correction1 + beta2 * exp_avg_diff / bias_correction2)).div_(denom)\n\n                if group['no_prox']:\n                    p.data.mul_(1 - group['lr'] * group['weight_decay'])\n                    p.add_(update, alpha=-group['lr'])\n                else:\n                    p.add_(update, alpha=-group['lr'])\n                    p.data.div_(1 + group['lr'] * group['weight_decay'])\n\n                state['pre_grad'] = copy_grad","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:44.847916Z","iopub.execute_input":"2024-03-04T18:17:44.848317Z","iopub.status.idle":"2024-03-04T18:17:44.875893Z","shell.execute_reply.started":"2024-03-04T18:17:44.848283Z","shell.execute_reply":"2024-03-04T18:17:44.875102Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper functions","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Helper functions\n# ====================================================\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\ndef asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since, percent):\n    now = time.time()\n    s = now - since\n    es = s / (percent)\n    rs = es - s\n    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))\n\n\ndef train_fn(fold, train_loader, model, criterion, optimizer, epoch, scheduler, device):\n    model.train()\n    scaler = torch.cuda.amp.GradScaler(enabled=CFG.apex)\n    losses = AverageMeter()\n    start = end = time.time()\n    global_step = 0\n    for step, batch in enumerate(train_loader):\n        spectrogram = batch['spectrogram'].to(device)\n        labels = batch['labels'].to(device)\n        batch_size = labels.size(0)\n        mixed_X, y_a, y_b, lam = mixup_data(spectrogram, labels, CFG, device)\n        new_criterion = get_criterion(CFG, criterion)\n        with torch.cuda.amp.autocast(enabled=CFG.apex):\n            y_preds= model(mixed_X)\n            #loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n            loss = new_criterion(F.log_softmax(y_preds, dim=1), y_a, y_b, lam)\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        losses.update(loss.item(), batch_size)\n        scaler.scale(loss).backward()\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.max_grad_norm)\n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            global_step += 1\n            if CFG.batch_scheduler:\n                scheduler.step()\n        end = time.time()\n        if step % CFG.print_freq == 0 or step == (len(train_loader)-1):\n            print('Epoch: [{0}][{1}/{2}] '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  'Grad: {grad_norm:.4f}  '\n                  'LR: {lr:.8f}  '\n                  .format(epoch+1, step, len(train_loader), \n                          remain=timeSince(start, float(step+1)/len(train_loader)),\n                          loss=losses,\n                          grad_norm=grad_norm,\n                          lr=scheduler.get_lr()[0]))\n        if CFG.wandb:\n            wandb.log({f\"[fold{fold}] loss\": losses.val,\n                       f\"[fold{fold}] lr\": scheduler.get_lr()[0]})\n    return losses.avg\n\n\ndef valid_fn(valid_loader, model, criterion, device):\n    losses = AverageMeter()\n    model.eval()\n    preds = []\n    start = end = time.time()\n    for step, batch in enumerate(valid_loader):\n        spectrogram = batch['spectrogram'].to(device)\n        labels = batch['labels'].to(device)\n        batch_size = labels.size(0)\n        with torch.no_grad():\n            y_preds = model(spectrogram)\n            loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        losses.update(loss.item(), batch_size)\n        preds.append(nn.Softmax(dim=1)(y_preds).to('cpu').numpy())\n        end = time.time()\n        if step % CFG.print_freq == 0 or step == (len(valid_loader)-1):\n            print('EVAL: [{0}/{1}] '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  .format(step, len(valid_loader),\n                          loss=losses,\n                          remain=timeSince(start, float(step+1)/len(valid_loader))))\n    predictions = np.concatenate(preds)\n    return losses.avg, predictions","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:44.877391Z","iopub.execute_input":"2024-03-04T18:17:44.877779Z","iopub.status.idle":"2024-03-04T18:17:44.900187Z","shell.execute_reply.started":"2024-03-04T18:17:44.877748Z","shell.execute_reply":"2024-03-04T18:17:44.899253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Loop","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# train loop\n# ====================================================\ndef train_loop(folds, fold, directory):\n    \n    LOGGER.info(f\"========== fold: {fold} training ==========\")\n\n    # ====================================================\n    # loader\n    # ====================================================\n    if CFG.stage1_pop1:\n        train_folds = folds[folds['fold'] != fold].reset_index(drop=True)\n    else:\n        train_folds = folds[(folds['fold'] != fold) & (folds['total_evaluators'] >= 10)].reset_index(drop=True)\n    valid_folds = folds[folds['fold'] == fold].reset_index(drop=True)\n    valid_labels = valid_folds[ CFG.target_cols].values\n    \n    train_dataset = CustomDataset(train_folds, augment=True, mode=\"train\")\n    valid_dataset = CustomDataset(valid_folds, augment=False, mode=\"train\")\n\n    train_loader = DataLoader(train_dataset,\n                              batch_size=CFG.batch_size,\n                              shuffle=True,\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=True)\n    valid_loader = DataLoader(valid_dataset,\n                              batch_size=CFG.batch_size * 2,\n                              shuffle=False,\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\n\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n    model = CustomModel(CFG)\n    if CFG.stage2_pop2:\n        model_weight = POP_1_DIR + f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage1.pth\"\n        checkpoint = torch.load(model_weight, map_location=device)\n        model.load_state_dict(checkpoint[\"model\"])\n    # CPMP: wrap the model to use all GPUs\n    model.to(device)\n    if CFG.t4_gpu:\n        model = nn.DataParallel(model)\n    \n    \n    \n    def build_optimizer(cfg, model, device):\n        lr = cfg.lr\n        # lr = default_configs[\"lr\"]\n        if cfg.optimizer == \"SAM\":\n            base_optimizer = torch.optim.SGD  # define an optimizer for the \"sharpness-aware\" update\n            optimizer_model = SAM(model.parameters(), base_optimizer, lr=lr, momentum=0.9, weight_decay=cfg.weight_decay, adaptive=True)\n        elif cfg.optimizer == \"Ranger21\":\n            optimizer_model = Ranger21(model.parameters(), lr=lr, weight_decay=cfg.weight_decay, \n            num_epochs=cfg.epochs, num_batches_per_epoch=len(train_loader))\n        elif cfg.optimizer == \"SGD\":\n            optimizer_model = torch.optim.SGD(model.parameters(), lr=lr, weight_decay=cfg.weight_decay, momentum=0.9)\n        elif cfg.optimizer == \"Adam\":\n            optimizer_model = Adam(model.parameters(), lr=lr, weight_decay=CFG.weight_decay)\n        elif cfg.optimizer == \"Lion\":\n            optimizer_model = Lion(model.parameters(), lr=lr, weight_decay=cfg.weight_decay)\n        elif cfg.optimizer == \"Adan\":\n            optimizer_model = Adan(model.parameters(), lr=lr, weight_decay=cfg.weight_decay)\n    \n        return optimizer_model\n    \n    optimizer = build_optimizer(CFG, model, device)\n    \n    # ====================================================\n    # scheduler\n    # ====================================================\n    # ====================================================\n\n    def get_scheduler(optimizer):\n        if CFG.scheduler=='ReduceLROnPlateau':\n            scheduler = ReduceLROnPlateau(optimizer, **CFG.reduce_params)\n        elif CFG.scheduler=='CosineAnnealingLR':\n            scheduler = CosineAnnealingLR(optimizer, **CFG.cosanneal_params)\n        elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n            scheduler = CosineAnnealingWarmRestarts(optimizer, **CFG.cosanneal_res_params)\n        elif CFG.scheduler=='OneCycleLR':\n            steps_per_epoch=len(train_loader),\n            scheduler = OneCycleLR(optimizer=optimizer, epochs=CFG.epochs, anneal_strategy=\"cos\", pct_start=0.05, steps_per_epoch=len(train_loader),\n        max_lr=CFG.lr, final_div_factor=100)\n        return scheduler\n    \n    scheduler = get_scheduler(optimizer)\n\n    # ====================================================\n    # loop\n    # ====================================================\n    criterion = nn.KLDivLoss(reduction=\"batchmean\")\n\n    \n    best_score = np.inf\n\n    for epoch in range(CFG.epochs):\n\n        start_time = time.time()\n\n        # train\n        avg_loss = train_fn(fold, train_loader, model, criterion, optimizer, epoch, scheduler, device)\n\n        # eval\n        avg_val_loss, predictions = valid_fn(valid_loader, model, criterion, device)\n\n        elapsed = time.time() - start_time\n\n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n        if CFG.wandb:\n            wandb.log({f\"[fold{fold}] epoch\": epoch+1, \n                       f\"[fold{fold}] avg_train_loss\": avg_loss, \n                       f\"[fold{fold}] avg_val_loss\": avg_val_loss,\n                       f\"[fold{fold}] score\": score})\n        \n        if best_score > avg_val_loss:\n            best_score = avg_val_loss\n            LOGGER.info(f'Epoch {epoch+1} - Save Best valid loss: {avg_val_loss:.4f} Model')\n            # CPMP: save the original model. It is stored as the module attribute of the DP model.\n            if CFG.stage1_pop1:\n                state_dict = model.module.state_dict() if CFG.t4_gpu else model.state_dict()\n                torch.save({'model': state_dict ,\n                            'predictions': predictions},\n                             directory+f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage1.pth\")\n            else:\n                state_dict = model.module.state_dict() if CFG.t4_gpu else model.state_dict()\n                torch.save({'model': state_dict,\n                            'predictions': predictions},\n                             directory+f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage2.pth\")\n                \n    if CFG.stage1_pop1:\n        predictions = torch.load(directory+f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage1.pth\", \n                             map_location=torch.device('cpu'))['predictions']\n    else:\n        predictions = torch.load(directory+f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage2.pth\", \n                             map_location=torch.device('cpu'))['predictions']\n    valid_folds[[f\"pred_{c}\" for c in CFG.target_cols]] = predictions\n    valid_folds[CFG.target_cols] = valid_labels \n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return valid_folds, best_score","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:44.901535Z","iopub.execute_input":"2024-03-04T18:17:44.901911Z","iopub.status.idle":"2024-03-04T18:17:44.928741Z","shell.execute_reply.started":"2024-03-04T18:17:44.901881Z","shell.execute_reply":"2024-03-04T18:17:44.927939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == '__main__':\n    \n    if CFG.train:\n        oof_df = pd.DataFrame()\n        scores = []\n        for fold in range(CFG.n_fold):\n            if fold in CFG.trn_fold:\n                _oof_df, score = train_loop(train, fold, POP_1_DIR)\n                oof_df = pd.concat([oof_df, _oof_df])\n                scores.append(score)\n                LOGGER.info(f\"========== fold: {fold} result ==========\")\n                LOGGER.info(f'Score with best loss weights stage1: {score}')\n        oof_df = oof_df.reset_index(drop=True)\n        LOGGER.info(f\"========== CV ==========\")\n        LOGGER.info(f'Score with best loss weights stage1: {np.mean(scores)}')\n        oof_df.to_csv(POP_1_DIR+f'{CFG.model_name}_oof_df_version{VERSION}_stage1.csv', index=False)\n        \n    if CFG.wandb:\n        wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:44.929974Z","iopub.execute_input":"2024-03-04T18:17:44.930357Z","iopub.status.idle":"2024-03-04T18:17:50.925833Z","shell.execute_reply.started":"2024-03-04T18:17:44.930328Z","shell.execute_reply":"2024-03-04T18:17:50.924123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CFG.stage1_pop1 = False\nCFG.stage2_pop2 = True\nCFG.epochs = 5\n\nif __name__ == '__main__':\n    \n    if CFG.train:\n        oof_df = pd.DataFrame()\n        scores = []\n        for fold in range(CFG.n_fold):\n            if fold in CFG.trn_fold:\n                _oof_df, score = train_loop(train, fold, POP_2_DIR)\n                oof_df = pd.concat([oof_df, _oof_df])\n                scores.append(score)\n                LOGGER.info(f\"========== fold: {fold} result ==========\")\n                LOGGER.info(f'Score with best loss weights stage2: {score}')\n        oof_df = oof_df.reset_index(drop=True)\n        LOGGER.info(f\"========== CV ==========\")\n        LOGGER.info(f'Score with best loss weights stage2: {np.mean(scores)}')\n        oof_df.to_csv(POP_2_DIR+f'{CFG.model_name}_oof_df_version{VERSION}_stage2.csv', index=False)\n        \n    if CFG.wandb:\n        wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:50.927217Z","iopub.status.idle":"2024-03-04T18:17:50.927678Z","shell.execute_reply.started":"2024-03-04T18:17:50.927447Z","shell.execute_reply":"2024-03-04T18:17:50.927467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/kaggle-kl-div')\nfrom kaggle_kl_div import score\n\n# === Pre-process OOF ===\nlabel_cols = CFG.target_cols\ngt = oof_df[[\"eeg_id\"] + CFG.target_cols]\ngt.sort_values(by=\"eeg_id\", inplace=True)\ngt.reset_index(inplace=True, drop=True)\n\npreds = oof_df[[\"eeg_id\"] + CFG.pred_cols]\npreds.columns = [\"eeg_id\"] + CFG.target_cols\npreds.sort_values(by=\"eeg_id\", inplace=True)\npreds.reset_index(inplace=True, drop=True)\n\ny_trues = gt[CFG.target_cols]\ny_preds = preds[CFG.target_cols]\n\noof = pd.DataFrame(y_preds.copy())\noof['id'] = np.arange(len(oof))\n\ntrue = pd.DataFrame(y_trues.copy())\ntrue['id'] = np.arange(len(true))\n\ncv = score(solution=true, submission=oof, row_id_column_name='id')\nprint('CV Stage2 Score with SparK resnet50 Spectrogram =',cv)","metadata":{"execution":{"iopub.status.busy":"2024-03-04T18:17:50.929237Z","iopub.status.idle":"2024-03-04T18:17:50.929679Z","shell.execute_reply.started":"2024-03-04T18:17:50.929455Z","shell.execute_reply":"2024-03-04T18:17:50.929473Z"},"trusted":true},"execution_count":null,"outputs":[]}]}