{"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":7392775,"sourceType":"datasetVersion","datasetId":4297782},{"sourceId":7517324,"sourceType":"datasetVersion","datasetId":4378712},{"sourceId":7652061,"sourceType":"datasetVersion","datasetId":4216847}],"dockerImageVersionId":30646,"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\nSo I looked in the internet to find other CNN architecture and I stumbled across this [github](https://github.com/keyu-tian/SparK/tree/main). Their work consists of pretraining any CNN (e.g., ResNet) in a BERT-style self-supervised manner which was curious for me, so I wanted to experiment on these weights and see if I can go below 0.4 LB.\n\nI share with you a detailed table of the pretrained weights and [dataset](https://www.kaggle.com/datasets/medali1992/resnet50d-spark) I created containing these weights.\n<details>\n\n<summary> <b> Pretrained weights (with SparK's UNet-style decoder; can be used to reconstruct images) </b> </summary>\n\n<br>\n\n|   arch.    | reso. | acc@1 | #params | flops  | weights (self-supervised, with SparK's decoder)                                                                                           |\n|:----------:|:-----:|:-----:|:-------:|:------:|:------------------------------------------------------------------------------------------------------------------------------------------|\n|  ResNet50  |  224  | 80.6  |   26M   |  4.1G  | [res50_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/1STt3w3e5q9eCPZa8VzcJj1zG6p3jLeSF/view?usp=share_link)   |\n| ResNet101  |  224  | 82.2  |   45M   |  7.9G  | [res101_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/1GjN48LKtlop2YQre6---7ViCWO-3C0yr/view?usp=share_link)  |\n| ResNet152  |  224  | 82.7  |   60M   | 11.6G  | [res152_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/1U3Cd94j4ZHfYR2dUjWmsEWfjP6Opx4oo/view?usp=share_link)  |\n| ResNet200  |  224  | 83.1  |   65M   | 15.1G  | [res200_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/13AFSqvIr0v-2hmb4DzVza45t_lhf2CnD/view?usp=share_link)  |\n| ConvNeXt-S |  224  | 84.1  |   50M   |  8.7G  | [cnxS224_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/1bKvrE4sNq1PfzhWlQJXEPrl2kHqHRZM-/view?usp=share_link) |\n| ConvNeXt-L |  384  | 86.0  |  198M   | 101.0G | [cnxL384_withdecoder_1kpretrained_spark_style.pth](https://drive.google.com/file/d/1ZI9Jgtb3fKWE_vDFEly29w-1FWZSNwa0/view?usp=share_link) |\n\n</details>\n\n<br>\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\n* I train a resnet50 models using SparK weights.\n\n### Hyperparams\n\n```\n  scheduler='OneCycleLR' \n  print_freq=50\n  num_workers = 1\n  model_name = 'resnet50'\n  optimizer='Adam'\n  stage1_epochs = 10\n  stage1_epochs = 20\n  eps = 1e-6\n  lr = 1e-3\n  batch_size = 64\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\n### Hyperparams\n\n```\n  scheduler='OneCycleLR' \n  print_freq=50\n  num_workers = 1\n  model_name = 'resnet50'\n  optimizer='Adam'\n  stage1_epochs = 8\n  stage1_epochs = 8\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":"# ====================================================\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-02-29T03:32:38.241847Z","iopub.execute_input":"2024-02-29T03:32:38.242444Z","iopub.status.idle":"2024-02-29T03:32:38.254734Z","shell.execute_reply.started":"2024-02-29T03:32:38.242418Z","shell.execute_reply":"2024-02-29T03:32:38.253925Z"},"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')\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nfrom matplotlib import pyplot as plt\nimport joblib\nos.environ['CUDA_VISIBLE_DEVICES'] = \"0,1\"\n\nVERSION=2","metadata":{"execution":{"iopub.status.busy":"2024-02-29T03:32:38.256196Z","iopub.execute_input":"2024-02-29T03:32:38.256676Z","iopub.status.idle":"2024-02-29T03:33:00.052364Z","shell.execute_reply.started":"2024-02-29T03:32:38.256646Z","shell.execute_reply":"2024-02-29T03:33:00.051510Z"},"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=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 = 'resnet50'\n    optimizer='Adan'\n    epochs = 8\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-02-29T03:33:00.053980Z","iopub.execute_input":"2024-02-29T03:33:00.054320Z","iopub.status.idle":"2024-02-29T03:33:00.065158Z","shell.execute_reply.started":"2024-02-29T03:33:00.054294Z","shell.execute_reply":"2024-02-29T03:33:00.064291Z"},"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-02-29T03:33:00.070492Z","iopub.execute_input":"2024-02-29T03:33:00.070759Z","iopub.status.idle":"2024-02-29T03:33:00.093545Z","shell.execute_reply.started":"2024-02-29T03:33:00.070735Z","shell.execute_reply":"2024-02-29T03:33:00.092659Z"},"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-02-29T03:33:00.094854Z","iopub.execute_input":"2024-02-29T03:33:00.095203Z","iopub.status.idle":"2024-02-29T03:33:00.429534Z","shell.execute_reply.started":"2024-02-29T03:33:00.095173Z","shell.execute_reply":"2024-02-29T03:33:00.428608Z"},"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-02-29T03:33:00.431071Z","iopub.execute_input":"2024-02-29T03:33:00.431390Z","iopub.status.idle":"2024-02-29T03:33:00.715649Z","shell.execute_reply.started":"2024-02-29T03:33:00.431364Z","shell.execute_reply":"2024-02-29T03:33:00.714687Z"},"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-02-29T03:33:00.717105Z","iopub.execute_input":"2024-02-29T03:33:00.717805Z","iopub.status.idle":"2024-02-29T03:34:10.886531Z","shell.execute_reply.started":"2024-02-29T03:33:00.717758Z","shell.execute_reply":"2024-02-29T03:34:10.885571Z"},"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-02-29T03:34:10.887587Z","iopub.execute_input":"2024-02-29T03:34:10.887871Z","iopub.status.idle":"2024-02-29T03:35:52.368336Z","shell.execute_reply.started":"2024-02-29T03:34:10.887845Z","shell.execute_reply":"2024-02-29T03:35:52.367288Z"},"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 +  0.166666667 # 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-02-29T03:35:52.369541Z","iopub.execute_input":"2024-02-29T03:35:52.369834Z","iopub.status.idle":"2024-02-29T03:35:52.846142Z","shell.execute_reply.started":"2024-02-29T03:35:52.369810Z","shell.execute_reply":"2024-02-29T03:35:52.845202Z"},"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-02-29T03:35:52.850333Z","iopub.execute_input":"2024-02-29T03:35:52.850622Z","iopub.status.idle":"2024-02-29T03:35:53.081171Z","shell.execute_reply.started":"2024-02-29T03:35:52.850598Z","shell.execute_reply":"2024-02-29T03:35:53.080179Z"},"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        \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        \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        \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           \n        transforms = A.Compose([\n            A.VerticalFlip(p=0.3),\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-02-29T03:35:53.082378Z","iopub.execute_input":"2024-02-29T03:35:53.082634Z","iopub.status.idle":"2024-02-29T03:35:53.097828Z","shell.execute_reply.started":"2024-02-29T03:35:53.082612Z","shell.execute_reply":"2024-02-29T03:35:53.096882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataloader","metadata":{}},{"cell_type":"code","source":"dataset = CustomDataset(train, 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-02-29T03:35:53.099103Z","iopub.execute_input":"2024-02-29T03:35:53.099386Z","iopub.status.idle":"2024-02-29T03:35:53.377305Z","shell.execute_reply.started":"2024-02-29T03:35:53.099363Z","shell.execute_reply":"2024-02-29T03:35:53.376182Z"},"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-02-29T03:35:53.378476Z","iopub.execute_input":"2024-02-29T03:35:53.378761Z","iopub.status.idle":"2024-02-29T03:35:55.057042Z","shell.execute_reply.started":"2024-02-29T03:35:53.378736Z","shell.execute_reply":"2024-02-29T03:35:55.055635Z"},"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        # Optionally load state from checkpoint\n        if config.SparK:\n            state = torch.load('/kaggle/input/resnet50d-spark/resnet50_1kpretrained_timm_style.pth', 'cpu')\n            self.model.load_state_dict(state.get('module', state), strict=False)\n        if config.FREEZE:\n            for i,(name, param) in enumerate(list(self.model.named_parameters())\\\n                                             [0:config.NUM_FROZEN_LAYERS]):\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    \n    def forward(self, x):\n        x = self.__reshape_input(x)\n        x = self.features(x)\n        x = self.custom_layers(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-02-29T03:35:55.059096Z","iopub.execute_input":"2024-02-29T03:35:55.059405Z","iopub.status.idle":"2024-02-29T03:35:55.072466Z","shell.execute_reply.started":"2024-02-29T03:35:55.059380Z","shell.execute_reply":"2024-02-29T03:35:55.071539Z"},"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()","metadata":{"execution":{"iopub.status.busy":"2024-02-29T03:35:55.073486Z","iopub.execute_input":"2024-02-29T03:35:55.073776Z","iopub.status.idle":"2024-02-29T03:35:59.583984Z","shell.execute_reply.started":"2024-02-29T03:35:55.073750Z","shell.execute_reply":"2024-02-29T03:35:59.583116Z"},"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-02-29T03:35:59.585323Z","iopub.execute_input":"2024-02-29T03:35:59.585667Z","iopub.status.idle":"2024-02-29T03:35:59.612294Z","shell.execute_reply.started":"2024-02-29T03:35:59.585635Z","shell.execute_reply":"2024-02-29T03:35:59.611489Z"},"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        with torch.cuda.amp.autocast(enabled=CFG.apex):\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        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    targets = []\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        targets.append(labels.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    targets = np.concatenate(targets)\n    return losses.avg, predictions","metadata":{"execution":{"iopub.status.busy":"2024-02-29T03:35:59.613339Z","iopub.execute_input":"2024-02-29T03:35:59.613605Z","iopub.status.idle":"2024-02-29T03:35:59.635833Z","shell.execute_reply.started":"2024-02-29T03:35:59.613582Z","shell.execute_reply":"2024-02-29T03:35:59.635101Z"},"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, mode=\"train\")\n    valid_dataset = CustomDataset(valid_folds, 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    model.to(device)\n    # CPMP: wrap the model to use all GPUs\n    model = nn.DataParallel(model)\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                \n                torch.save({'model': model.module.state_dict(),\n                            'predictions': predictions},\n                             directory+f\"{CFG.model_name}_fold{fold}_best_version{VERSION}_stage1.pth\")\n            else:\n                \n                torch.save({'model': model.module.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-02-29T03:35:59.637095Z","iopub.execute_input":"2024-02-29T03:35:59.637409Z","iopub.status.idle":"2024-02-29T03:35:59.663451Z","shell.execute_reply.started":"2024-02-29T03:35:59.637386Z","shell.execute_reply":"2024-02-29T03:35:59.662617Z"},"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-02-29T03:35:59.664592Z","iopub.execute_input":"2024-02-29T03:35:59.664860Z","iopub.status.idle":"2024-02-29T03:40:57.274330Z","shell.execute_reply.started":"2024-02-29T03:35:59.664838Z","shell.execute_reply":"2024-02-29T03:40:57.272871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CFG.stage1_pop1 = False\nCFG.stage2_pop2 = True\nCFG.epochs = 8\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-02-29T03:40:57.275335Z","iopub.status.idle":"2024-02-29T03:40:57.275660Z","shell.execute_reply.started":"2024-02-29T03:40:57.275503Z","shell.execute_reply":"2024-02-29T03:40:57.275518Z"},"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-02-29T03:40:57.276933Z","iopub.status.idle":"2024-02-29T03:40:57.277309Z","shell.execute_reply.started":"2024-02-29T03:40:57.277124Z","shell.execute_reply":"2024-02-29T03:40:57.277153Z"},"trusted":true},"execution_count":null,"outputs":[]}]}