{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7392733,"sourceType":"datasetVersion","datasetId":4297749},{"sourceId":7465251,"sourceType":"datasetVersion","datasetId":4317718},{"sourceId":7517324,"sourceType":"datasetVersion","datasetId":4378712}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"},"papermill":{"default_parameters":{},"duration":5345.179582,"end_time":"2024-02-20T15:08:09.804733","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-02-20T13:39:04.625151","version":"2.4.0"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"088fc3f5e4264683bbeaa2d929fab599":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"danger","description":"","description_tooltip":null,"layout":"IPY_MODEL_2c3d333449fb412dbc86ca31835ef4b0","max":1,"min":0,"orientation":"horizontal","style":"IPY_MODEL_3ba5491929cf4cb1bbb8e28514933aa6","value":1}},"11bfd3544cf8411e96cd7ede1451518c":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_5246ec5881c147f09c90846e9911dda6","placeholder":"​","style":"IPY_MODEL_db8fdfb810194e03b66c7c5e5077ccf5","value":" 1/? [00:00&lt;00:00,  1.37it/s]"}},"2c3d333449fb412dbc86ca31835ef4b0":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":"20px"}},"3ba5491929cf4cb1bbb8e28514933aa6":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"5246ec5881c147f09c90846e9911dda6":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"58eb07f979fd49e8a9155ad60c77923c":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"993085d825a542b1b5c44cd42f7f8c89":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_c818523154ec45689615162b4c03bf96","IPY_MODEL_088fc3f5e4264683bbeaa2d929fab599","IPY_MODEL_11bfd3544cf8411e96cd7ede1451518c"],"layout":"IPY_MODEL_fa70a85db4944741a7d7b13cef19bb70"}},"c818523154ec45689615162b4c03bf96":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_58eb07f979fd49e8a9155ad60c77923c","placeholder":"​","style":"IPY_MODEL_ea4dcc5f4fcf4011a7ba7dc587109a9b","value":""}},"db8fdfb810194e03b66c7c5e5077ccf5":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"ea4dcc5f4fcf4011a7ba7dc587109a9b":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"fa70a85db4944741a7d7b13cef19bb70":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}}},"version_major":2,"version_minor":0}}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Libraries","metadata":{}},{"cell_type":"code","source":"# Importing necessary modules\nimport os\nimport gc\nimport sys\nimport math\nimport time\nimport random\nimport datetime as dt\nimport numpy as np\nimport pandas as pd\nimport wandb\n\nfrom glob import glob\nfrom pathlib import Path\nfrom typing import Dict, List, Union\nimport scipy.signal as scisig\nfrom scipy.signal import butter, lfilter, freqz\nfrom matplotlib import pyplot as plt\nfrom tqdm.auto import tqdm\n\n# Importing PyTorch modules\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import Adam, SGD, AdamW\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.optim.lr_scheduler import (\n    ReduceLROnPlateau,\n    OneCycleLR,\n    CosineAnnealingLR,\n    CosineAnnealingWarmRestarts,\n)\nfrom torch.optim.optimizer import Optimizer\nfrom sklearn.model_selection import GroupKFold\n\n# Adding path for Kaggle kernel\nsys.path.append(\"/kaggle/input/kaggle-kl-div\")\nimport kaggle_kl_div\n\n# Ignoring warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n# Setting device to CUDA\ndevice = torch.device(\"cuda\")\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\"\n\n# Printing system information\n!cat /etc/os-release | grep -oP \"PRETTY_NAME=\\\"\\K([^\\\"]*)\"\nprint(f\"BUILD_DATE={os.environ['BUILD_DATE']}, CONTAINER_NAME={os.environ['CONTAINER_NAME']}\")\n\ntry:\n    print(f\"PyTorch Version:{torch.__version__}, CUDA is available:{torch.cuda.is_available()}, Version CUDA:{torch.version.cuda}\")\n    print(f\"Device Capability:{torch.cuda.get_device_capability()}, {torch.cuda.get_arch_list()}\")\n    print(f\"CuDNN Enabled:{torch.backends.cudnn.enabled}, Version:{torch.backends.cudnn.version()}\")\nexcept Exception:\n    pass","metadata":{"execution":{"iopub.status.busy":"2024-04-02T06:59:20.241583Z","iopub.status.idle":"2024-04-02T06:59:28.174340Z","shell.execute_reply":"2024-04-02T06:59:28.173153Z","shell.execute_reply.started":"2024-04-02T06:59:20.241901Z"},"papermill":{"duration":13.898864,"end_time":"2024-02-20T13:39:22.393405","exception":false,"start_time":"2024-02-20T13:39:08.494541","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Directory settings","metadata":{}},{"cell_type":"code","source":"class APP:\n    # Check if running in Jupyter notebook\n    jupyter = \"ipykernel\" in globals()\n    if not jupyter:\n        try:\n            if \"IPython\" in globals().get(\"__doc__\", \"\"):\n                jupyter = True\n        except Exception as inst:\n            print(inst)\n\n    # Check if running in Kaggle kernel\n    kaggle = os.environ.get(\"KAGGLE_KERNEL_RUN_TYPE\", \"\") != \"\"\n    \n    # Check if running locally\n    local = os.environ.get(\"DOCKER_USING\", \"\") == \"LOCAL\"\n    \n    # Get the current date and time\n    date_time_start = dt.datetime.now()\n    dt_start_ymd_hms = date_time_start.strftime(\"%Y.%m.%d_%H-%M-%S\")\n\n    # Get the file run path\n    file_run_path = \"\"\n    if jupyter:\n        try:\n            file_run_path = Path(globals().get(\"__vsc_ipynb_file__\", \"\"))\n        except Exception as inst:\n            print(inst)\n    else:\n        try:\n            file_run_path = Path(__file__)\n        except Exception as inst:\n            print(inst)\n\n    # Get the file run name, app path, run path, and output path\n    file_run_name = file_run_path.stem\n    path_app = file_run_path.parent\n    path_run = Path(os.getcwd())\n    path_out = (\n        Path(\"/kaggle/working\")\n        if kaggle\n        else file_run_path / f\"{file_run_name}_{dt_start_ymd_hms}\"\n    )\n\n# Set the output directory\nOUTPUT_DIR = \"./\"\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)\n\n# Print the environment information\nprint(f\"jupyter:{APP.jupyter}, kaggle:{APP.kaggle}, local:{APP.local}\")\nprint(APP.file_run_path)\nprint(APP.path_out)","metadata":{"execution":{"iopub.execute_input":"2024-04-02T06:59:28.177365Z","iopub.status.busy":"2024-04-02T06:59:28.176617Z","iopub.status.idle":"2024-04-02T06:59:28.188633Z","shell.execute_reply":"2024-04-02T06:59:28.187618Z","shell.execute_reply.started":"2024-04-02T06:59:28.177332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"# This is a configuration class for the project\n# It contains various parameters and settings for the model training process\n\nclass CFG:\n    VERSION = '86'  # Version of the project\n\n    wandb = False  # Flag for using Weights and Biases for experiment tracking\n    debug = False  # Flag for running in debug mode\n    create_eegs = False  # Flag for creating EEGs\n    apex = True  # Flag for using Apex for mixed precision training\n    visualize = False  # Flag for visualizing data\n    save_all_models = True  # Flag for saving all models during training\n\n    # Set the number of workers for data loading and parallel processing\n    if debug:\n        num_workers = 0\n        parallel = False\n    else:\n        num_workers = os.cpu_count()\n        parallel = True\n\n    model_name = \"resnet1d_gru\"  # Name of the model architecture\n    optimizer = \"AdamW\"  # Optimizer for training\n\n    factor = 0.9  # Factor for reducing learning rate\n    eps = 1e-6  # Epsilon value for numerical stability\n    lr = 8e-3  # Learning rate\n    min_lr = 1e-6  # Minimum learning rate\n\n    batch_size = 64  # Batch size for training\n    batch_koef_valid = 2  # Batch size coefficient for validation\n    batch_scheduler = True  # Flag for using batch scheduler\n    weight_decay = 1e-2  # Weight decay for regularization\n    gradient_accumulation_steps = 1  # Number of gradient accumulation steps\n    max_grad_norm = 1e7  # Maximum gradient norm\n\n    fixed_kernel_size = 5  # Fixed kernel size for the model\n    linear_layer_features = 304  # Number of features in the linear layer\n    kernels = [3, 5, 7, 9, 11]  # List of kernel sizes\n\n    seq_length = 50  # Length of the sequence in seconds\n    sampling_rate = 200  # Sampling rate in Hz\n    nsamples = seq_length * sampling_rate  # Total number of samples\n    n_split_samples = 5  # Number of split samples\n    out_samples = nsamples // n_split_samples  # Number of output samples\n    sample_delta = nsamples - out_samples  # Sample delta\n    sample_offset = sample_delta // 2  # Sample offset\n    multi_validation = False  # Flag for multi-validation\n\n    train_by_stages = False  # Flag for training by stages\n    train_by_folds = True  # Flag for training by folds\n\n    # Define the training stages and corresponding parameters\n    n_stages = 2\n    match n_stages:\n        case 1:\n            train_stages = [0]\n            epochs = [100]\n            test_total_eval = 2\n            total_evals_old = [[(2, 3), (6, 29)]]\n            total_evaluators = [\n                [\n                    {'band':(2, 2), 'excl_evals':[]},\n                    {'band':(6, 28), 'excl_evals':[]},\n                ],\n            ]\n        case 2:\n            train_stages = [0, 1]\n            epochs = [50, 100]\n            test_total_eval = 0\n            total_evals_old = [[(1, 2),(4, 5)], (6, 29)]\n            total_evaluators = [\n                [\n                    {'band':(3, 3), 'excl_evals':[]},\n                    {'band':(6, 28), 'excl_evals':[]},\n                ],\n                [\n                    {'band':(1, 2), 'excl_evals':[]},\n                    {'band':(4, 5), 'excl_evals':[]},\n                ],\n            ]\n        case 3:\n            train_stages = [0, 1, 2]\n            epochs = [20, 50, 100]\n            test_total_eval = 0\n            total_evals_old = [(0, 3), (3, 6), (6, 29)]\n            total_evaluators = [\n                [\n                    {'band':(0, 2), 'excl_evals':[]},\n                ],\n                [\n                    {'band':(3, 5), 'excl_evals':[]},\n                ],\n                [\n                    {'band':(6, 28), 'excl_evals':[]},\n                ],\n            ]\n\n    n_fold = 5  # Number of folds for cross-validation\n    train_folds = [0, 1, 2, 3, 4]  # List of train folds\n\n    patience = 11  # Patience for early stopping\n    seed = 2024  # Random seed\n\n    bandpass_filter = {\"low\": 0.5, \"high\": 20, \"order\": 2}  # Bandpass filter parameters\n    rand_filter = {\"probab\": 0.1, \"low\": 10, \"high\": 20, \"band\": 1.0, \"order\": 2}  # Random filter parameters\n    freq_channels = []  # Frequency channels\n    filter_order = 2  # Filter order\n\n    random_divide_signal = 0.05  # Random divide signal\n    random_close_zone = 0.05  # Random close zone\n    random_common_negative_signal = 0.0  # Random common negative signal\n    random_common_reverse_signal = 0.0  # Random common reverse signal\n    random_negative_signal = 0.05  # Random negative signal\n    random_reverse_signal = 0.05  # Random reverse signal\n\n    log_step = 100  # Log step\n    log_show = False  # Flag for showing logs\n\n    scheduler = \"CosineAnnealingWarmRestarts\"  # Scheduler for learning rate\n\n    # Parameters for CosineAnnealingLR scheduler\n    cosanneal_params = {\n        \"T_max\": 6,\n        \"eta_min\": 1e-5,\n        \"last_epoch\": -1,\n    }\n\n    # Parameters for ReduceLROnPlateau scheduler\n    reduce_params = {\n        \"mode\": \"min\",\n        \"factor\": 0.2,\n        \"patience\": 4,\n        \"eps\": 1e-6,\n        \"verbose\": True,\n    }\n\n    # Parameters for CosineAnnealingWarmRestarts scheduler\n    cosanneal_res_params = {\n        \"T_0\": 20,\n        \"eta_min\": 1e-6,\n        \"T_mult\": 1,\n        \"last_epoch\": -1,\n    }\n\n    # Target columns\n    target_cols = [\n        \"seizure_vote\",\n        \"lpd_vote\",\n        \"gpd_vote\",\n        \"lrda_vote\",\n        \"grda_vote\",\n        \"other_vote\",\n    ]  \n\n    pred_cols = [x + \"_pred\" for x in target_cols]  # Predicted columns\n    \n    # Map features\n    map_features = [\n        (\"Fp1\", \"T3\"),\n        (\"T3\", \"O1\"),\n        (\"Fp1\", \"C3\"),\n        (\"C3\", \"O1\"),\n        (\"Fp2\", \"C4\"),\n        (\"C4\", \"O2\"),\n        (\"Fp2\", \"T4\"),\n        (\"T4\", \"O2\"),\n    ]  \n\n    eeg_features = [\"Fp1\", \"T3\", \"C3\", \"O1\", \"Fp2\", \"C4\", \"T4\", \"O2\"]  # EEG features\n    feature_to_index = {x: y for x, y in zip(eeg_features, range(len(eeg_features)))}  # Feature to index mapping\n    simple_features = []  # Simple features\n    n_map_features = len(map_features)  # Number of map features\n    in_channels = n_map_features + n_map_features * len(freq_channels) + len(simple_features)  # Number of input channels\n    target_size = len(target_cols)  # Size of the target\n\n    # Define the file paths\n    path_inp = Path(\"/kaggle/input\")\n    path_src = path_inp / \"hms-harmful-brain-activity-classification/\"\n    file_train = path_src / \"train.csv\"\n    path_train = path_src / \"train_eegs\"\n    file_features_test = path_train / \"100261680.parquet\"\n    file_eeg_specs = path_inp / \"eeg-spectrogram-by-lead-id-unique/eeg_specs.npy\"\n    file_raw_eeg = path_inp / \"brain-eegs/eegs.npy\"\n\n    if APP.kaggle:\n        num_workers = 2\n        parallel = True\n\n# Print the feature to index mapping and the EEG features\nprint(CFG.feature_to_index)\nprint(CFG.eeg_features)","metadata":{"execution":{"iopub.execute_input":"2024-04-02T06:59:28.190168Z","iopub.status.busy":"2024-04-02T06:59:28.189871Z","iopub.status.idle":"2024-04-02T06:59:28.216575Z","shell.execute_reply":"2024-04-02T06:59:28.215735Z","shell.execute_reply.started":"2024-04-02T06:59:28.190145Z"},"papermill":{"duration":0.029582,"end_time":"2024-02-20T13:39:22.435956","exception":false,"start_time":"2024-02-20T13:39:22.406374","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{"papermill":{"duration":0.012027,"end_time":"2024-02-20T13:39:22.461399","exception":false,"start_time":"2024-02-20T13:39:22.449372","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Initialize the logger for logging messages\ndef init_logger(log_file=OUTPUT_DIR + \"train.log\"):\n    from logging import getLogger, INFO, FileHandler, Formatter, StreamHandler\n\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\n\nLOGGER = init_logger()\n\n\n# Function to convert seconds to minutes and seconds format\ndef asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return \"%dm %ds\" % (m, s)\n\n\n# Function to calculate the time remaining\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\n# Function to quantize data using mu-law encoding\ndef quantize_data(data, classes):\n    mu_x = mu_law_encoding(data, classes)\n    return mu_x  # quantized\n\n\n# Function to perform mu-law encoding\ndef mu_law_encoding(data, mu):\n    mu_x = np.sign(data) * np.log(1 + mu * np.abs(data)) / np.log(mu + 1)\n    return mu_x\n\n\n# Function to perform mu-law expansion\ndef mu_law_expansion(data, mu):\n    s = np.sign(data) * (np.exp(np.abs(data) * np.log(mu + 1)) - 1) / mu\n    return s\n\n\n# Function to create a Butterworth bandpass filter\ndef butter_bandpass(lowcut, highcut, fs, order=5):\n    return butter(order, [lowcut, highcut], fs=fs, btype=\"band\")\n\n\n# Function to apply a Butterworth bandpass filter to data\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\n# Function to apply a Butterworth lowpass filter to data\ndef butter_lowpass_filter(data, cutoff_freq=20, sampling_rate=CFG.sampling_rate, order=4):\n    nyquist = 0.5 * sampling_rate\n    normal_cutoff = cutoff_freq / nyquist\n    b, a = butter(order, normal_cutoff, btype=\"low\", analog=False)\n    filtered_data = lfilter(b, a, data, axis=0)\n    return filtered_data\n\n\n# Function to denoise the data using a bandpass filter\ndef denoise_filter(x):\n    y = butter_bandpass_filter(x, CFG.lowcut, CFG.highcut, CFG.sampling_rate, 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    return y","metadata":{"papermill":{"duration":0.045586,"end_time":"2024-02-20T13:39:22.519244","exception":false,"start_time":"2024-02-20T13:39:22.473658","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Parquet to EEG Signals Numpy Processing","metadata":{"papermill":{"duration":0.011709,"end_time":"2024-02-20T13:40:46.512143","exception":false,"start_time":"2024-02-20T13:40:46.500434","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def eeg_from_parquet(\n    parquet_path: str, display: bool = False, seq_length=CFG.seq_length\n) -> np.ndarray:\n    # Read the EEG data from the parquet file\n    eeg = pd.read_parquet(parquet_path, columns=CFG.eeg_features)\n    rows = len(eeg)\n\n    # Calculate the offset for selecting a subset of data\n    offset = (rows - CFG.nsamples) // 2\n\n    # Select a subset of data based on the offset and CFG.nsamples\n    eeg = eeg.iloc[offset : offset + CFG.nsamples]\n\n    if display:\n        # Display the EEG data\n        plt.figure(figsize=(10, 5))\n        offset = 0\n\n    # Create an empty array to store the processed EEG data\n    data = np.zeros((CFG.nsamples, len(CFG.eeg_features)))\n\n    for index, feature in enumerate(CFG.eeg_features):\n        # Get the values of the current feature\n        x = eeg[feature].values.astype(\"float32\")\n\n        # Calculate the mean and percentage of NaN values in the feature\n        mean = np.nanmean(x)\n        nan_percentage = np.isnan(x).mean()\n\n        if nan_percentage < 1:\n            # Replace NaN values with the mean value\n            x = np.nan_to_num(x, nan=mean)\n        else:\n            # If all values are NaN, set all values to 0\n            x[:] = 0\n\n        # Store the processed feature data in the data array\n        data[:, index] = x\n\n        if display:\n            if index != 0:\n                offset += x.max()\n            # Plot the feature data\n            plt.plot(range(CFG.nsamples), x - offset, label=feature)\n            offset -= x.min()\n\n    if display:\n        # Display the plot\n        plt.legend()\n        name = parquet_path.split(\"/\")[-1].split(\".\")[0]\n        plt.yticks([])\n        plt.title(f\"EEG {name}\", size=16)\n        plt.show()\n\n    # Return the processed EEG data\n    return data","metadata":{"papermill":{"duration":0.026463,"end_time":"2024-02-20T13:40:46.550473","exception":false,"start_time":"2024-02-20T13:40:46.52401","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{"papermill":{"duration":0.014802,"end_time":"2024-02-20T13:42:17.914507","exception":false,"start_time":"2024-02-20T13:42:17.899705","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class EEGDataset(Dataset):\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        batch_size: int,\n        eegs: Dict[int, np.ndarray],\n        mode: str = \"train\",\n        downsample: int = None,\n        bandpass_filter: Dict[str, Union[int, float]] = None,\n        rand_filter: Dict[str, Union[int, float]] = None,\n    ):\n        self.df = df\n        self.batch_size = batch_size\n        self.mode = mode\n        self.eegs = eegs\n        self.downsample = downsample\n        self.offset = None\n        self.bandpass_filter = bandpass_filter\n        self.rand_filter = rand_filter\n        \n    def __len__(self):\n        \"\"\"\n        Length of dataset.\n        \"\"\"\n        return len(self.df)\n\n    # Function to get one item from the dataset\n    def __getitem__(self, index):\n        \"\"\"\n        Get one item.\n        \"\"\"\n        X, y_prob = self.__data_generation(index)\n        if self.downsample is not None:\n            X = X[:: self.downsample, :]\n        output = {\n            \"eeg\": torch.tensor(X, dtype=torch.float32),\n            \"labels\": torch.tensor(y_prob, dtype=torch.float32),\n        }\n        return output\n\n    # Function to set the offset\n    def set_offset(self, offset: int):\n        self.offset = offset\n\n    # Function to generate data\n    def __data_generation(self, index):\n        \"\"\"\n        This function generates data for training or testing the model.\n\n        Args:\n            index: The index of the data sample to generate.\n\n        Returns:\n            A tuple containing the generated data (X) and target probabilities (y_prob).\n        \"\"\"\n\n        # Initialize output array with zeros\n        X = np.zeros((CFG.out_samples, CFG.in_channels), dtype=\"float32\")  # Size=(10000, 14)\n        random_divide_signal = False\n\n        # Get the data for the specific index\n        row = self.df.iloc[index]\n        data = self.eegs[row.eeg_id]  # Size=(10000, 8)\n\n        # Handle potential difference between requested and available samples\n        if CFG.nsamples != CFG.out_samples:\n            if self.mode == \"train\":\n                # Introduce random offset in training mode\n                offset = (CFG.sample_delta * random.randint(0, 1000)) // 1000\n            elif not self.offset is None:\n                # Use predefined offset if available\n                offset = self.offset\n            else:\n                # Use default offset\n                offset = CFG.sample_offset\n\n            # Randomly divide signal in training mode with some probability\n            if self.mode == \"train\" and CFG.random_divide_signal > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_divide_signal:\n                random_divide_signal = True\n                multipliers = [(1, 2), (2, 3), (3, 4), (3, 5)]\n                koef_1, koef_2 = multipliers[random.randint(0, 3)]\n                offset = (koef_1 * offset) // koef_2\n                data = data[offset:offset+(CFG.out_samples * koef_2) // koef_1,:]\n            else:\n                # Select data with specified offset and number of samples\n                data = data[offset:offset+CFG.out_samples,:]\n\n        reverse_signal = False\n        negative_signal = False\n        if self.mode == \"train\":\n            if CFG.random_common_reverse_signal > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_common_reverse_signal:\n                reverse_signal = True\n            if CFG.random_common_negative_signal > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_common_negative_signal:\n                negative_signal = True\n\n        # Loop through feature pairs defined in configuration\n        for i, (feat_a, feat_b) in enumerate(CFG.map_features):\n            # Skip this feature pair with some probability in training mode\n            if self.mode == \"train\" and CFG.random_close_zone > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_close_zone:\n                continue\n            \n            # Calculate difference between features\n            diff_feat = (data[:, CFG.feature_to_index[feat_a]] - data[:, CFG.feature_to_index[feat_b]])  # Size=(10000,)\n\n            # Apply random reverse or negation in training mode\n            if self.mode == \"train\":\n                if reverse_signal or CFG.random_reverse_signal > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_reverse_signal:\n                    diff_feat = np.flip(diff_feat)\n                if negative_signal or CFG.random_negative_signal > 0.0 and random.uniform(0.0, 1.0) <= CFG.random_negative_signal:\n                    diff_feat = -diff_feat\n\n            # Apply pre-defined bandpass filter if available\n            if not self.bandpass_filter is None:\n                diff_feat = butter_bandpass_filter(\n                    diff_feat,\n                    self.bandpass_filter[\"low\"],\n                    self.bandpass_filter[\"high\"],\n                    CFG.sampling_rate,\n                    order=self.bandpass_filter[\"order\"],\n                )\n            \n            # Handle random signal division\n            if random_divide_signal:\n                #diff_feat = cp.asnumpy(cpsig.upfirdn([1.0, 1, 1.0], diff_feat, 2, 3))  # linear interp, rate 2/3\n                diff_feat = scisig.upfirdn([1.0, 1, 1.0], diff_feat, koef_1, koef_2)  # linear interp, rate 2/3\n                diff_feat = diff_feat[0:CFG.out_samples]\n\n            if (\n                self.mode == \"train\"\n                and not self.rand_filter is None\n                and random.uniform(0.0, 1.0) <= self.rand_filter[\"probab\"]\n            ):\n                lowcut = random.randint(\n                    self.rand_filter[\"low\"], self.rand_filter[\"high\"]\n                )\n                highcut = lowcut + self.rand_filter[\"band\"]\n                diff_feat = butter_bandpass_filter(\n                    diff_feat,\n                    lowcut,\n                    highcut,\n                    CFG.sampling_rate,\n                    order=self.rand_filter[\"order\"],\n                )\n\n            X[:, i] = diff_feat\n\n        n = CFG.n_map_features\n        if len(CFG.freq_channels) > 0:\n            for i in range(CFG.n_map_features):\n                diff_feat = X[:, i]\n                # Apply bandpass filter for each frequency channel\n                for j, (lowcut, highcut) in enumerate(CFG.freq_channels):\n                    band_feat = butter_bandpass_filter(\n                        diff_feat, lowcut, highcut, CFG.sampling_rate, order=CFG.filter_order,  # 6\n                    )\n                    X[:, n] = band_feat\n                    n += 1\n\n        for spml_feat in CFG.simple_features:\n            feat_val = data[:, CFG.feature_to_index[spml_feat]]\n\n            # Apply bandpass filter if available\n            if not self.bandpass_filter is None:\n                feat_val = butter_bandpass_filter(\n                    feat_val,\n                    self.bandpass_filter[\"low\"],\n                    self.bandpass_filter[\"high\"],\n                    CFG.sampling_rate,\n                    order=self.bandpass_filter[\"order\"],\n                )\n\n            # Apply random bandpass filter in training mode with some probability\n            if (\n                self.mode == \"train\"\n                and not self.rand_filter is None\n                and random.uniform(0.0, 1.0) <= self.rand_filter[\"probab\"]\n            ):\n                lowcut = random.randint(\n                    self.rand_filter[\"low\"], self.rand_filter[\"high\"]\n                )\n                highcut = lowcut + self.rand_filter[\"band\"]\n                feat_val = butter_bandpass_filter(\n                    feat_val,\n                    lowcut,\n                    highcut,\n                    CFG.sampling_rate,\n                    order=self.rand_filter[\"order\"],\n                )\n\n            # Add feature to output array\n            X[:, n] = feat_val\n            n += 1\n            \n        # Clip values in output array\n        X = np.clip(X, -1024, 1024)\n\n        # Handle potential NaN values\n        X = np.nan_to_num(X, nan=0) / 32.0\n\n        # Apply lowpass filter to entire output array\n        X = butter_lowpass_filter(X, order=CFG.filter_order)\n\n        # Initialize target probabilities array\n        y_prob = np.zeros(CFG.target_size, dtype=\"float32\")\n\n        # If not in test mode, get target probabilities from the data row\n        if self.mode != \"test\":\n            y_prob = row[CFG.target_cols].values.astype(np.float32)\n\n        # Return generated data and target probabilities\n        return X, y_prob","metadata":{"papermill":{"duration":0.0354,"end_time":"2024-02-20T13:42:17.964068","exception":false,"start_time":"2024-02-20T13:42:17.928668","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper functions","metadata":{}},{"cell_type":"code","source":"class KLDivLossWithLogits(nn.KLDivLoss):\n    \"\"\"\n    Custom KL divergence loss function with logits as input.\n\n    This class inherits from the nn.KLDivLoss class but applies a log_softmax function\n    to the input (y) before calculating the loss. This is because the KL divergence\n    loss expects log probabilities as input, while the network typically outputs logits.\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) # Apply log_softmax to y\n        loss = super().forward(y, t)\n        return loss\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n\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 seed_torch(seed):\n    \"\"\"Seeds all relevant libraries for deterministic training.\"\"\"\n    \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    torch.backends.cudnn.deterministic = True\n    # torch.backends.cudnn.benchmark = True\n    # pl.seed_everything(seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{"papermill":{"duration":0.031018,"end_time":"2024-02-20T13:42:28.75746","exception":false,"start_time":"2024-02-20T13:42:28.726442","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class ResNet_1D_Block(nn.Module):\n    \"\"\"\n    This class implements a basic residual block for 1D CNNs, which is a commonly\n    used building block for image and time series classification tasks. It consists\n    of two convolutional layers with batch normalization, activation functions (SiLU in this case),\n    and dropout for regularization. A skip connection is added between the input and output\n    of the block, allowing the network to learn identity mappings for better gradient flow.\n    \"\"\"\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size,\n        stride,\n        padding,\n        downsampling,\n        dilation=1,\n        groups=1,\n        dropout=0.0,\n    ):\n        super(ResNet_1D_Block, self).__init__()\n\n        # Initialize the block layers\n        self.bn1 = nn.BatchNorm1d(num_features=in_channels)\n        self.relu_1 = nn.Hardswish()\n        self.relu_2 = nn.Hardswish()\n\n        self.dropout = nn.Dropout(p=dropout, inplace=False)\n        self.conv1 = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            stride=stride,\n            padding=padding,\n            dilation=dilation,\n            groups=groups,\n            bias=False,\n        )\n\n        self.bn2 = nn.BatchNorm1d(num_features=out_channels)\n        self.conv2 = nn.Conv1d(\n            in_channels=out_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            stride=stride,\n            padding=padding,\n            dilation=dilation,\n            groups=groups,\n            bias=False,\n        )\n\n        self.maxpool = nn.MaxPool1d(\n            kernel_size=2,\n            stride=2,\n            padding=0,\n            dilation=dilation,\n        )\n        self.downsampling = downsampling\n\n    def forward(self, x):\n        # Apply the block layers to the input\n        identity = x\n\n        out = self.bn1(x)\n        out = self.relu_1(out)\n        out = self.dropout(out)\n        out = self.conv1(out)\n        out = self.bn2(out)\n        out = self.relu_2(out)\n        out = self.dropout(out)\n        out = self.conv2(out)\n\n        out = self.maxpool(out)\n        identity = self.downsampling(x)\n\n        out += identity\n        return out\n\n\nclass EEGNet(nn.Module):\n    \"\"\"\n    This class implements the EEGNet model, which is a compact and efficient\n    convolutional neural network architecture for EEG classification tasks.\n    The model consists of a series of convolutional layers with batch normalization,\n    activation functions (SiLU in this case), and max pooling operations. The output\n    of the convolutional layers is flattened and passed through a fully connected layer\n    to obtain the final output. The model also includes a GRU layer for processing\n    sequential data.\n    \"\"\"\n    def __init__(\n        self,\n        kernels,\n        in_channels,\n        fixed_kernel_size,\n        num_classes,\n        linear_layer_features,\n        dilation=1,\n        groups=1,\n    ):\n        super(EEGNet, self).__init__()\n        self.kernels = kernels\n        self.planes = 24\n        self.parallel_conv = nn.ModuleList()\n        self.in_channels = in_channels\n\n        for i, kernel_size in enumerate(list(self.kernels)):\n            sep_conv = nn.Conv1d(\n                in_channels=in_channels,\n                out_channels=self.planes,\n                kernel_size=(kernel_size),\n                stride=1,\n                padding=0,\n                dilation=dilation,\n                groups=groups,\n                bias=False,\n            )\n            self.parallel_conv.append(sep_conv)\n\n        # Initialize the model layers\n        self.bn1 = nn.BatchNorm1d(num_features=self.planes)\n        self.relu_1 = nn.SiLU()\n        self.relu_2 = nn.SiLU()\n\n        self.conv1 = nn.Conv1d(\n            in_channels=self.planes,\n            out_channels=self.planes,\n            kernel_size=fixed_kernel_size,\n            stride=2,\n            padding=2,\n            dilation=dilation,\n            groups=groups,\n            bias=False,\n        )\n\n        self.block = self._make_resnet_layer(\n            kernel_size=fixed_kernel_size,\n            stride=1,\n            dilation=dilation,\n            groups=groups,\n            padding=fixed_kernel_size // 2,\n        )\n        self.bn2 = nn.BatchNorm1d(num_features=self.planes)\n        self.avgpool = nn.AvgPool1d(kernel_size=6, stride=6, padding=2)\n\n        # Initialize the GRU layer\n        self.rnn = nn.GRU(\n            input_size=self.in_channels,\n            hidden_size=128,\n            num_layers=1,\n            bidirectional=True,\n            # dropout=0.2,\n        )\n\n        # Initialize the linear layer\n        self.fc = nn.Linear(in_features=linear_layer_features, out_features=num_classes)\n\n\n    # Function to create a ResNet layer\n    def _make_resnet_layer(\n        self,\n        kernel_size,\n        stride,\n        dilation=1,\n        groups=1,\n        blocks=9,\n        padding=0,\n        dropout=0.0,\n    ):\n        layers = []\n        downsample = None\n        base_width = self.planes\n\n        for i in range(blocks):\n            downsampling = nn.Sequential(\n                nn.MaxPool1d(kernel_size=2, stride=2, padding=0)\n            )\n            layers.append(\n                ResNet_1D_Block(\n                    in_channels=self.planes,\n                    out_channels=self.planes,\n                    kernel_size=kernel_size,\n                    stride=stride,\n                    padding=padding,\n                    downsampling=downsampling,\n                    dilation=dilation,\n                    groups=groups,\n                    dropout=dropout,\n                )\n            )\n        return nn.Sequential(*layers)\n\n    # Function to extract features from the input\n    def extract_features(self, x):\n        x = x.permute(0, 2, 1)\n\n        out_sep = []\n        for i in range(len(self.kernels)):\n            sep = self.parallel_conv[i](x)\n            out_sep.append(sep)\n\n        out = torch.cat(out_sep, dim=2)\n        out = self.bn1(out)\n        out = self.relu_1(out)\n        out = self.conv1(out)\n\n        out = self.block(out)\n        out = self.bn2(out)\n        out = self.relu_2(out)\n        out = self.avgpool(out)\n\n        out = out.reshape(out.shape[0], -1)\n        rnn_out, _ = self.rnn(x.permute(0, 2, 1))\n        new_rnn_h = rnn_out[:, -1, :]\n\n        new_out = torch.cat([out, new_rnn_h], dim=1)\n        return new_out\n\n    # Function to forward pass through the model\n    def forward(self, x):\n        new_out = self.extract_features(x)\n        result = self.fc(new_out)\n        return result","metadata":{"papermill":{"duration":0.05674,"end_time":"2024-02-20T13:42:28.844768","exception":false,"start_time":"2024-02-20T13:42:28.788028","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Adan Optimizer","metadata":{}},{"cell_type":"code","source":"class 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__(\n        self,\n        params,\n        lr=1e-3,\n        betas=(0.98, 0.92, 0.99),\n        eps=1e-8,\n        weight_decay=0.2,\n        max_grad_norm=0.0,\n        no_prox=False,\n    ):\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(\n            lr=lr,\n            betas=betas,\n            eps=eps,\n            weight_decay=weight_decay,\n            max_grad_norm=max_grad_norm,\n            no_prox=no_prox,\n        )\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(\n                max_grad_norm / (global_grad_norm + group[\"eps\"]), max=1.0\n            )\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            bias_correction2 = 1.0 - beta2 ** group[\"step\"]\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 = (\n                    state[\"exp_avg\"],\n                    state[\"exp_avg_sq\"],\n                    state[\"exp_avg_diff\"],\n                )\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_(\n                    group[\"eps\"]\n                )\n                update = (\n                    (\n                        exp_avg / bias_correction1\n                        + beta2 * exp_avg_diff / bias_correction2\n                    )\n                ).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":{"papermill":{"duration":0.058157,"end_time":"2024-02-20T13:42:29.60282","exception":false,"start_time":"2024-02-20T13:42:29.544663","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Func","metadata":{}},{"cell_type":"code","source":"def train_fn(\n    stage, fold, train_loader, model, criterion, optimizer, epoch, scheduler, device\n):\n    \"\"\"\n    This function performs one epoch of training on the provided dataset.\n    It iterates through the training data loader, calculates the loss for each batch,\n    performs backpropagation, updates the model weights using the optimizer,\n    and tracks various metrics like average loss and learning rate.\n    \"\"\"\n    model.train()\n\n    scaler = torch.cuda.amp.GradScaler(enabled=CFG.apex)\n    losses = AverageMeter()\n    start = end = time.time()\n    global_step = 0\n\n    for step, batch in enumerate(train_loader):\n        eegs = batch[\"eeg\"].to(device)\n        labels = batch[\"labels\"].to(device)\n        batch_size = labels.size(0)\n\n        with torch.cuda.amp.autocast(enabled=CFG.apex):\n            y_preds = model(eegs)\n            # Calculate loss using log_softmax and the provided criterion\n            loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n\n        if CFG.gradient_accumulation_steps > 1:\n            # Gradient accumulation (accumulate gradients for multiple batches before update)\n            loss = loss / CFG.gradient_accumulation_steps\n\n        losses.update(loss.item(), batch_size)\n\n        # Perform backward pass and gradient scaling (if mixed precision)\n        scaler.scale(loss).backward()\n\n        # Clip gradients to prevent exploding gradients\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.max_grad_norm)\n\n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            # Perform optimizer step and update scaler after gradient accumulation steps\n            scaler.step(optimizer)\n            scaler.update()\n            global_step += 1\n            if CFG.batch_scheduler: # Update learning rate scheduler if configured\n                scheduler.step()\n        end = time.time() # Update end time for logging\n\n        if CFG.log_show and (\n            step % CFG.log_step == 0 or step == (len(train_loader) - 1)\n        ):\n            # remain=timeSince(start, float(step + 1) / len(train_loader))\n            LOGGER.info(\n                f\"Epoch {epoch+1} [{step}/{len(train_loader)}] Loss: {losses.val:.4f} Loss Avg:{losses.avg:.4f}\"\n            )\n            # \"Elapsed {remain:s} Grad: {grad_norm:.4f}  LR: {cheduler.get_lr()[0]:.8f}\"\n\n        if CFG.wandb:\n            wandb.log(\n                {\n                    f\"[fold{fold}] loss\": losses.val,\n                    f\"[fold{fold}] lr\": scheduler.get_lr()[0],\n                }\n            )\n    return losses.avg","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Valid Func","metadata":{}},{"cell_type":"code","source":"def valid_fn(stage, epoch, valid_loader, model, criterion, device):\n    \"\"\"\n    This function iterates through the validation data loader,\n    evaluates the model's performance on the validation set,\n    calculates the average validation loss, and returns the loss and predictions.\n    \"\"\"\n    losses = AverageMeter()  # Object to track and store average loss\n    model.eval()  # Set the model to evaluation mode (disable dropout etc.)\n\n    preds = []  # List to store model predictions\n    targets = []  # List to store ground truth labels\n\n    start = time.time()  # Record validation start time\n    end = start  # Initialize end time (updated later)\n\n    for step, batch in enumerate(valid_loader):\n        eegs = batch[\"eeg\"].to(device)\n        labels = batch[\"labels\"].to(device)\n        batch_size = labels.size(0)\n\n        with torch.no_grad(): \n            # Disable gradient calculation for validation (saves memory)\n            y_preds = model(eegs)\n            # Calculate loss using log_softmax and the provided criterion\n            loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n\n        if CFG.gradient_accumulation_steps > 1:\n            # Gradient accumulation (if used during training, adjust loss here)\n            loss = loss / CFG.gradient_accumulation_steps\n\n        losses.update(loss.item(), batch_size) # Update average loss meter\n\n        # Collect predictions and targets for later evaluation metrics\n        preds.append(nn.Softmax(dim=1)(y_preds).to(\"cpu\").numpy())\n        targets.append(labels.to(\"cpu\").numpy())\n\n        end = time.time()  # Update end time for logging\n\n        if CFG.log_show and (\n            step % CFG.log_step == 0 or step == (len(valid_loader) - 1)\n        ):\n            LOGGER.info(\n                f\"Epoch {epoch+1} VALIDATION: [{step}/{len(valid_loader)}] Val Loss: {losses.val:.4f} Val Loss Avg: {losses.avg:.4f}\"\n            )\n    \n    # Concatenate predictions and targets from all batches\n    predictions = np.concatenate(preds)\n    targets = np.concatenate(targets)\n\n    return losses.avg, predictions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Build Optimizer","metadata":{"papermill":{"duration":0.031254,"end_time":"2024-02-20T13:42:29.808216","exception":false,"start_time":"2024-02-20T13:42:29.776962","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def build_optimizer(cfg, model, device, epochs, num_batches_per_epoch):\n    \"\"\"\n    This function reads the optimizer configuration from the provided dictionary\n    and creates the corresponding optimizer object. It supports various optimizers\n    like SGD, Adam, AdamW, Ranger21, Lion, Adan, and SAM (Sharpness-Aware Minimization).\n    The optimizer is configured with the learning rate, weight decay, and other\n    parameters specified in the configuration.\n    \"\"\"\n    lr = cfg.lr  # Extract learning rate from the configuration\n    if cfg.optimizer == \"SAM\":\n        # Build SAM optimizer (Sharpness-Aware Minimization)\n        base_optimizer = (\n            torch.optim.SGD\n        )\n        optimizer_model = SAM(\n            model.parameters(),\n            base_optimizer,\n            lr=lr,\n            momentum=0.9,\n            weight_decay=cfg.weight_decay,\n            adaptive=True,\n        )\n    elif cfg.optimizer == \"Ranger21\":\n        # Build Ranger21 optimizer\n        optimizer_model = Ranger21(\n            model.parameters(),\n            lr=lr,\n            weight_decay=cfg.weight_decay,\n            num_epochs=epochs,\n            num_batches_per_epoch=num_batches_per_epoch,\n        )\n    elif cfg.optimizer == \"SGD\":\n        # Build SGD optimizer\n        optimizer_model = torch.optim.SGD(\n            model.parameters(), lr=lr, weight_decay=cfg.weight_decay, momentum=0.9\n        )\n    elif cfg.optimizer == \"Adam\":\n        # Build Adam optimizer\n        optimizer_model = Adam(model.parameters(), lr=lr, weight_decay=CFG.weight_decay)\n    elif cfg.optimizer == \"AdamW\":\n        # Build AdamW optimizer\n        optimizer_model = AdamW(\n            model.parameters(), lr=lr, weight_decay=CFG.weight_decay\n        )\n    elif cfg.optimizer == \"Lion\":\n        # Build Lion optimizer\n        optimizer_model = Lion(model.parameters(), lr=lr, weight_decay=cfg.weight_decay)\n    elif cfg.optimizer == \"Adan\":\n        # Build Adan optimizer\n        optimizer_model = Adan(model.parameters(), lr=lr, weight_decay=cfg.weight_decay)\n    else:\n        raise NotImplementedError(f\"Optimizer {cfg.optimizer} not supported!\")\n\n    return optimizer_model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Scheduler","metadata":{}},{"cell_type":"code","source":"def get_scheduler(optimizer, epochs, steps_per_epoch):\n    \"\"\"\n    This function creates a learning rate scheduler based on the provided configuration.\n    It supports various schedulers like ReduceLROnPlateau, CosineAnnealingLR,\n    CosineAnnealingWarmRestarts, and OneCycleLR. The scheduler is configured with\n    the optimizer, learning rate, and other parameters specified in the configuration.\n    \"\"\"\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        scheduler = OneCycleLR(\n            optimizer=optimizer,\n            epochs=epochs,\n            pct_start=0.0,\n            steps_per_epoch=steps_per_epoch,\n            max_lr=CFG.lr,\n            div_factor=25,\n            final_div_factor=4.0e-01,\n        )\n    return scheduler","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Loop","metadata":{}},{"cell_type":"code","source":"def train_loop(stage, epochs, folds, fold, directory, prev_dir, eggs):\n    \"\"\"\n    This function performs the following steps for a single fold in a k-fold\n    cross-validation setting:\n\n    1. Splits the data into training and validation sets based on the fold.\n    2. Creates training and validation data loaders.\n    3. Initializes the model (potentially loading weights from previous stages).\n    4. Defines the optimizer, scheduler, and loss criterion.\n    5. Trains the model for the specified number of epochs.\n    6. Validates the model on the validation set after each epoch.\n    7. Tracks the best validation loss and saves the model with the best performance.\n    8. Returns the updated validation folds data and the best validation loss.\n    \"\"\"\n    train_folds = folds[folds[\"fold\"] != fold].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    # Creating training dataset\n    train_dataset = EEGDataset(\n        train_folds,\n        batch_size=CFG.batch_size,\n        mode=\"train\",\n        eegs=eggs,\n        bandpass_filter=CFG.bandpass_filter,\n        rand_filter=CFG.rand_filter,\n    )\n    \n    # Creating validation dataset\n    valid_dataset = EEGDataset(\n        valid_folds,\n        batch_size=CFG.batch_size,\n        mode=\"valid\",\n        eegs=eggs,\n        bandpass_filter=CFG.bandpass_filter,\n    )\n\n    # Creating training and validation data loaders\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=CFG.batch_size,\n        shuffle=True,\n        num_workers=CFG.num_workers,\n        pin_memory=True,\n        drop_last=True,\n    )\n\n    valid_loader = DataLoader(\n        valid_dataset,\n        batch_size=CFG.batch_size * CFG.batch_koef_valid,\n        shuffle=False,\n        num_workers=CFG.num_workers,\n        pin_memory=True,\n        drop_last=False,\n    )\n\n    LOGGER.info(\n        f\"========== stage: {stage} fold: {fold} training {len(train_loader)} / {len(valid_loader)} ==========\"\n    )\n\n    # Initialize the model\n    model = EEGNet(\n        kernels=CFG.kernels,\n        in_channels=CFG.in_channels,\n        fixed_kernel_size=CFG.fixed_kernel_size,\n        num_classes=CFG.target_size,\n        linear_layer_features=CFG.linear_layer_features,\n    )\n\n    # Load model weights from previous stage if available\n    if stage > 1:\n        model_weight = f\"{prev_dir}{CFG.model_name}_ver-{CFG.VERSION}_stage-{stage-1}_fold-{fold}_best.pth\"\n        checkpoint = torch.load(model_weight, map_location=device)\n        model.load_state_dict(checkpoint[\"model\"])\n\n    model.to(device)\n\n    # CPMP: wrap the model to use all GPUs\n    if CFG.parallel:\n        model = nn.DataParallel(model)\n\n    # Initialize the optimizer, scheduler, and loss criterion\n    optimizer = build_optimizer(\n        CFG, model, device, epochs=epochs, num_batches_per_epoch=len(train_loader)\n    )\n    scheduler = get_scheduler(\n        optimizer, epochs=epochs, steps_per_epoch=len(train_loader)\n    )\n    criterion = nn.KLDivLoss(reduction=\"batchmean\")\n\n    # Train the model for the specified number of epochs\n    best_score = np.inf\n    for epoch in range(epochs):\n        start_time = time.time()\n\n        # train\n        avg_loss = train_fn(\n            stage,\n            fold,\n            train_loader,\n            model,\n            criterion,\n            optimizer,\n            epoch,\n            scheduler,\n            device,\n        )\n\n        # eval\n        valid_dataset.set_offset(CFG.sample_offset)\n        avg_val_loss, predictions = valid_fn(\n            stage,\n            epoch,\n            valid_loader,\n            model,\n            criterion,\n            device,\n        )\n        \n        avg_loss_line = ''\n\n        # multi-validation\n        if CFG.multi_validation:\n            multi_avg_val_loss = np.zeros(CFG.n_split_samples)\n            start = (2 * CFG.sample_delta) // CFG.n_split_samples\n            finish = (3 * CFG.sample_delta) // CFG.n_split_samples\n            delta = (finish - start) // 5\n            for i in range(CFG.n_split_samples):\n                valid_dataset.set_offset(start)\n                multi_avg_val_loss[i], _ = valid_fn(\n                    stage,\n                    epoch,\n                    valid_loader,\n                    model,\n                    criterion,\n                    device,\n                )\n                avg_loss_line += f\" {multi_avg_val_loss[i]:.4f}\"\n                start += delta\n            avg_loss_line += f\" mean={np.mean(multi_avg_val_loss):.4f}\"\n\n        elapsed = time.time() - start_time\n\n        LOGGER.info(\n            f\"Epoch {epoch+1} Avg Train Loss: {avg_loss:.4f} Avg Valid Loss: {avg_val_loss:.4f} / {avg_loss_line}\"\n        )\n        #   time: {elapsed:.0f}s\n        if CFG.wandb:\n            wandb.log(\n                {\n                    f\"[fold{fold}] stage\": stage,\n                    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            )\n\n        # Save the model with the best validation loss\n        if CFG.save_all_models:\n            torch.save(\n                {\"model\": model.module.state_dict(), \"predictions\": predictions},\n                f\"{directory}{CFG.model_name}_ver-{CFG.VERSION}_stage-{stage}_fold-{fold}_epoch-{epoch}_val-{avg_val_loss:.4f}_train-{avg_loss:.4f}.pth\",\n            )\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}\")\n            # CPMP: save the original model. It is stored as the module attribute of the DP model.\n            torch.save(\n                {\"model\": model.module.state_dict(), \"predictions\": predictions},\n                f\"{directory}{CFG.model_name}_ver-{CFG.VERSION}_stage-{stage}_fold-{fold}_best.pth\",\n            )\n\n    # Load best model to use it for inference\n    predictions = torch.load(\n        f\"{directory}{CFG.model_name}_ver-{CFG.VERSION}_stage-{stage}_fold-{fold}_best.pth\",\n        map_location=torch.device(\"cpu\"),\n    )[\"predictions\"]\n\n    # valid_folds[[f\"pred_{c}\" for c in CFG.target_cols]] = predictions\n    valid_folds[CFG.pred_cols] = predictions\n    valid_folds[CFG.target_cols] = valid_labels\n\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # Return updated validation folds data and best validation loss\n    return valid_folds, best_score","metadata":{"papermill":{"duration":0.05868,"end_time":"2024-02-20T13:42:29.897814","exception":false,"start_time":"2024-02-20T13:42:29.839134","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load train data","metadata":{}},{"cell_type":"code","source":"# Read the training data\ntrain = pd.read_csv(CFG.file_train)\nTARGETS = train.columns[-6:]\nprint(\"Train shape:\", train.shape)\nprint(\"Targets\", list(TARGETS))\n\ntrain[\"total_evaluators\"] = train[CFG.target_cols].sum(axis=1)\n\ntrain_uniq = train.drop_duplicates(subset=[\"eeg_id\"] + list(TARGETS))\n\n# Display some basic information about the training data\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.\")\nprint(f\"There are {train_uniq.shape[0]} unique eeg_id + votes in the training data.\")\n\n# Display the distribution of the number of unique votes per EEG ID\nif CFG.visualize:\n    train_uniq.eeg_id.value_counts().value_counts().plot(\n        kind=\"bar\",\n        title=f\"Distribution of Count of EEG w Unique Vote: \"\n        f\"{train_uniq.shape[0]} examples\",\n    )\n\ndel train_uniq\n_ = gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualize the distribution based on the number of total evaluators\nif CFG.visualize:\n    plt.figure(figsize=(10, 6))\n    plt.hist(train[\"total_evaluators\"], bins=10, color=\"blue\", edgecolor=\"black\")\n    plt.title(\"Histogram of Total Evaluators\")\n    plt.xlabel(\"Total Evaluators\")\n    plt.ylabel(\"Frequency\")\n    plt.grid(True)\n    plt.show()\n\n# Read the test data\ntst_eeg_df = pd.read_parquet(CFG.file_features_test)\ntst_eeg_features = tst_eeg_df.columns\nprint(f\"There are {len(tst_eeg_features)} raw eeg features\")\nprint(list(tst_eeg_features))\ndel tst_eeg_df\n_ = gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Split Data","metadata":{}},{"cell_type":"code","source":"# %%time\n# Load EEG specifications\nall_eeg_specs = np.load(CFG.file_eeg_specs, allow_pickle=True).item()\n\n# Filter training data based on EEG specifications\ntrain = train[train[\"label_id\"].isin(all_eeg_specs.keys())].copy()\nprint(train.shape[0])\n\n# Prepare target data with regularization\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\n# Add a new column \"target\" containing the expert consensus\ntrain[\"target\"] = train[\"expert_consensus\"]\n\ntrain[train['total_evaluators'] == CFG.test_total_eval].groupby(['expert_consensus','total_evaluators']).count()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.test_total_eval > 0:\n    # Assign a unique key identifier to each training sample\n    train['key_id'] = range(train.shape[0])\n\n    # List to store data subsets for different total_evaluators ranges (old)\n    train_pop_olds = []\n\n    # Loop through different total_evaluators ranges defined in CFG.total_evals_old\n    for total_eval in CFG.total_evals_old:\n        if type(total_eval) is list:\n            pop_idx = (train[\"total_evaluators\"] >= total_eval[0][0]) & (\n                train[\"total_evaluators\"] < total_eval[0][1]\n            ) | (train[\"total_evaluators\"] >= total_eval[1][0]) & (\n                train[\"total_evaluators\"] < total_eval[1][1]\n            )\n        else:\n            pop_idx = (train[\"total_evaluators\"] >= total_eval[0]) & (\n                train[\"total_evaluators\"] < total_eval[1]\n            )\n\n        # Filter training data based on the indexing and create a copy\n        train_pop = train[pop_idx].copy().reset_index()\n\n        # Initialize stratified group k-fold cross-validation (considering patient_id)\n        sgkf = GroupKFold(n_splits=CFG.n_fold)\n        train_pop[\"fold\"] = -1\n        for fold_id, (_, val_idx) in enumerate(\n            sgkf.split(train_pop, y=train_pop[\"target\"], groups=train_pop[\"patient_id\"])\n        ):\n            train_pop.loc[val_idx, \"fold\"] = fold_id\n\n        # Append the filtered and fold-assigned data subset to the list\n        train_pop_olds.append(train_pop)\n        print(train_pop.shape[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# List to store data subsets for different total_evaluators configurations\ntrain_pops = []\nfor eval_list in CFG.total_evaluators:\n    # List to store filtered subsets based on individual configurations within the list\n    result = []\n    # Start with the entire training data for subsetting\n    train_pop = train.copy()\n    \n    # Loop through each evaluation dictionary within the current configuration list\n    for eval_dict in eval_list:\n        # Extract the band range from the dictionary\n        band = eval_dict['band']\n\n        # Create boolean indexing to filter based on total_evaluators range\n        pop_idx = (train_pop[\"total_evaluators\"] >= band[0]) \n        pop_idx &= (train_pop[\"total_evaluators\"] <= band[1])\n\n        # Loop through evaluators to exclude (if any)\n        for exclude in eval_dict['excl_evals']:\n            pop_idx &= ~(train_pop['expert_consensus'] == exclude)\n            pass\n\n        # Filter the training data based on the current configuration and append to the result list\n        result.append(train_pop[pop_idx])\n\n    # Concatenate the filtered subsets from the current configuration list\n    train_pop = pd.concat(result).copy().reset_index()\n\n    # Initialize stratified group k-fold cross-validation (considering patient_id)\n    sgkf = GroupKFold(n_splits=CFG.n_fold)\n\n    # Add a new column \"fold\" to assign folds for cross-validation within the subset\n    train_pop[\"fold\"] = -1\n\n    for fold_id, (_, val_idx) in enumerate(\n        sgkf.split(train_pop, y=train_pop[\"target\"], groups=train_pop[\"patient_id\"])\n    ):\n        train_pop.loc[val_idx, \"fold\"] = fold_id\n\n    # Append the filtered and fold-assigned data subset to the list\n    train_pops.append(train_pop)\n\n    # Print the number of samples in the current subset\n    print(train_pop.shape[0])\n\ntrain_0 = train_pops[0]\ntrain_0[train_0['total_evaluators'] == CFG.test_total_eval].groupby(['expert_consensus','total_evaluators']).count()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.test_total_eval > 0:\n    df_old = train_pop_olds[0].copy(deep=True).set_index(['key_id'], drop=True).drop(columns=['fold'])\n    df_new = train_pops[0].copy(deep=True).set_index(['key_id'], drop=True).drop(columns=['fold'])\n\n    #outer merge the two DataFrames, adding an indicator column called 'Exist'\n    diff_df = pd.merge(df_old, df_new, how='outer', indicator='Exist')\n\n    #find which rows don't exist in both DataFrames\n    diff_df = diff_df.loc[diff_df['Exist'] != 'both']\n    display(diff_df)\n\n    del df_old, df_new, diff_df, train_pop_olds\n    _ = gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.visualize:\n    print(\"Pop 1: train unique eeg_id + votes shape:\", train_pops[0].shape)\n    plt.figure(figsize=(10, 6))\n    plt.hist(train[\"total_evaluators\"], bins=10, color=\"blue\", edgecolor=\"black\")\n    plt.title(\"Histogram of Total Evaluators\")\n    plt.xlabel(\"Total Evaluators\")\n    plt.ylabel(\"Frequency\")\n    plt.grid(True)\n    plt.show()\n\ndel all_eeg_specs\n_ = gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Deduplicate Train EEG Id","metadata":{}},{"cell_type":"code","source":"# %%time\nif CFG.create_eegs:\n    all_eegs = {}  # Dictionary to store loaded EEG data\n    visualize = 1 if CFG.visualize else 0  # Set visualization flag based on configuration\n    eeg_ids = train.eeg_id.unique()  # Get unique EEG IDs from the training data\n\n    for i, eeg_id in tqdm(enumerate(eeg_ids)):  # Loop through unique EEG IDs with progress bar\n        eeg_path = CFG.path_train / f\"{eeg_id}.parquet\"  # Construct path to the EEG file\n\n        # Load EEG data from parquet file\n        data = eeg_from_parquet(eeg_path, display=i < visualize)\n        all_eegs[eeg_id] = data  # Store loaded data in the dictionary\n\n        if i == visualize:  # After processing a certain number for visualization (if set)\n            if CFG.create_eegs:  # If creating EEG data\n                print(f\"Processing {train['eeg_id'].nunique()} eeg parquets... \", end=\"\")\n            else:  # If loading pre-created EEG data\n                print(f\"Reading {len(eeg_ids)} eeg NumPys from disk.\")\n                break  # Exit the loop after processing the desired number for visualization\n\n    # Save the dictionary containing all EEG data to a NumPy file\n    np.save(\"./eegs\", all_eegs)\nelse:  # If not creating EEG data, load from a pre-existing file\n    all_eegs = np.load(CFG.file_raw_eeg, allow_pickle=True).item()  # Load EEG data from file\n\nif CFG.visualize:  # Check if visualization flag is set\n    frequencies = [1, 2, 4, 8, 16][::-1]  # Define list of frequencies for filtering (Hz)\n    x = [all_eegs[eeg_ids[0]][:, 0]]  # Select the first EEG feature from the first EEG\n\n    # Apply butter low-pass filter to the selected feature with different frequencies\n    for frequency in frequencies:\n        x.append(butter_lowpass_filter(x[0], cutoff_freq=frequency))\n\n    # Create a plot to visualize the filter effects\n    plt.figure(figsize=(12, 8))\n    plt.plot(range(CFG.nsamples), x[0], label=\"without filter\")\n    for k in range(1, len(x)):\n        # Shift each filtered signal vertically to avoid overlapping lines\n        plt.plot(\n            range(CFG.nsamples),\n            x[k] - k * (x[0].max() - x[0].min()),\n            label=f\"with filter {frequencies[k - 1]}Hz\",\n        )\n\n    plt.legend()\n    plt.yticks([])\n    plt.title(\"Butter Low-Pass Filter Examples\", size=18)\n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nif CFG.visualize:\n    # Create an EEG dataset object from the first element of train_pops\n    train_dataset = EEGDataset(\n        train_pops[0], batch_size=CFG.batch_size, eegs=all_eegs, mode=\"train\"\n    )\n\n    # Create a data loader for efficient batch processing\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=CFG.batch_size,\n        shuffle=False,  # Don't shuffle data for visualization\n        num_workers=CFG.num_workers,\n        pin_memory=True,  # Pin data to GPU memory for faster access (if applicable)\n        drop_last=True,  # Drop the last incomplete batch if necessary\n    )\n\n    output = train_dataset[0]\n    X, y = output[\"eeg\"], output[\"labels\"]\n    print(f\"X shape: {X.shape}, y shape: {y.shape}\")\n\n    # Create a random tensor as input for model testing (replace with real input later)\n    iot = torch.randn(2, CFG.nsamples, CFG.in_channels)  # .cuda()\n\n    # Define the EEGNet model architecture\n    model = EEGNet(\n        kernels=CFG.kernels,\n        in_channels=CFG.in_channels,\n        fixed_kernel_size=CFG.fixed_kernel_size,\n        num_classes=CFG.target_size,\n        linear_layer_features=CFG.linear_layer_features,\n    )\n\n    # Test the model output with the random tensor\n    output = model(iot)\n    print(f\"Model output shape: {output.shape}\")\n\n    # Loop through a batch of data from the data loader\n    for batch in train_loader:\n        # Extract the EEG data (X) and labels (y) from the batch\n        X = batch.pop(\"eeg\")\n        y = batch.pop(\"labels\")\n\n        # Visualize the first 4 samples in the batch\n        for item in range(4):\n            plt.figure(figsize=(20, 4))\n            offset = 0\n            for col in range(X.shape[-1]):\n                # Avoid overlapping lines by adjusting offset based on min/max values\n                if col != 0:\n                    offset -= X[item, :, col].min()\n                plt.plot(\n                    range(CFG.nsamples),\n                    X[item, :, col] + offset,\n                    label=f\"feature {col+1}\",\n                )\n                offset += X[item, :, col].max()\n\n            # Format and display the target labels\n            tt = f\"{y[col][0]:0.1f}\"\n            for t in y[col][1:]:\n                tt += f\", {t:0.1f}\"\n            plt.title(f\"EEG_Id = {eeg_ids[item]}\\nTarget = {tt}\", size=14)\n            plt.legend()\n            plt.show()\n\n        # Break out of the loop after visualizing 4 samples\n        break\n\n    # Delete temporary tensors and model to free memory\n    del iot, model\n    gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Stages","metadata":{}},{"cell_type":"code","source":"def get_score(preds, targets):\n    \"\"\"\n    Calculates the score (metric) between predicted and target values.\n    \"\"\"\n\n    # Create a copy of the predictions\n    oof = pd.DataFrame(preds.copy())\n\n    # Add an \"id\" column with sequential indices for each data point in predictions\n    oof[\"id\"] = np.arange(len(oof))\n\n    # Create a copy of the target\n    true = pd.DataFrame(targets.copy())\n\n    # Add an \"id\" column with sequential indices for each data point in targets\n    true[\"id\"] = np.arange(len(true))\n\n    # Calculate the score (metric) using the `kaggle_kl_div` function\n    cv = kaggle_kl_div.score(solution=true, submission=oof, row_id_column_name=\"id\")\n\n    # Return the calculated score\n    return cv\n\n\ndef get_result(result_df):\n    \"\"\"\n    Processes and evaluates the final results DataFrame.\n    \"\"\"\n\n    # Select columns from the result DataFrame containing EEG IDs and target columns as defined in CFG\n    gt = result_df[[\"eeg_id\"] + CFG.target_cols]\n\n    # Sort the data by \"eeg_id\" to ensure consistent order\n    gt.sort_values(by=\"eeg_id\", inplace=True)\n\n    # Reset the index after sorting\n    gt.reset_index(inplace=True, drop=True)\n\n    # Select columns from the result DataFrame containing EEG IDs and predicted target columns\n    preds = result_df[[\"eeg_id\"] + CFG.pred_cols]\n\n    # Rename the predicted target columns to match the target columns format\n    preds.columns = [\"eeg_id\"] + CFG.target_cols\n\n    # Sort the data by \"eeg_id\" to ensure consistent order (same as ground truth)\n    preds.sort_values(by=\"eeg_id\", inplace=True)\n\n    # Reset the index after sorting\n    preds.reset_index(inplace=True, drop=True)\n\n    # Calculate the score using the `get_score` function\n    score_loss = get_score(gt[CFG.target_cols], preds[CFG.target_cols])\n\n    # Log the score information\n    LOGGER.info(f\"Score with best loss weights: {score_loss}\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == \"__main__\" and CFG.train_by_stages:\n    # Check if script is executed directly and training by stages is enabled\n\n    # Set random seed for reproducibility\n    seed_torch(seed=CFG.seed)\n\n    # Variable to store the directory path from the previous stage\n    prev_dir = \"\"\n\n    # Loop through all defined stages\n    for stage in range(len(CFG.total_evaluators)):\n        # Construct directory path for the current stage outputs\n        pop_dir = f\"{OUTPUT_DIR}pop_{stage+1}_weight_oof/\"\n        # Create the directory if it doesn't exist\n        if not os.path.exists(pop_dir):\n            os.makedirs(pop_dir)\n\n        # Skip stages not included in the training stages list (potentially for fine-tuning)\n        if stage not in CFG.train_stages:\n            prev_dir = pop_dir\n            continue\n\n        # Initialize empty DataFrame to store out-of-fold (OOF) predictions\n        oof_df = pd.DataFrame()\n        # Initialize empty list to store scores for each fold\n        scores = []\n\n        # Loop through each fold in the defined training folds\n        for fold in CFG.train_folds:\n            # Call the `train_loop` function to perform training for the current stage, fold, etc.\n            train_oof_df, score = train_loop(\n                stage=stage + 1,\n                epochs=CFG.epochs[stage],\n                fold=fold,\n                folds=train_pops[stage],\n                directory=pop_dir,\n                prev_dir=prev_dir,\n                eggs=all_eegs,\n            )\n\n            # Concatenate OOF predictions from the current fold to the main OOF DataFrame\n            oof_df = pd.concat([oof_df, train_oof_df])\n            # Append the score for the current fold\n            scores.append(score)\n\n            # Log information about the fold result\n            LOGGER.info(f\"========== stage: {stage+1} fold: {fold} result ==========\")\n            LOGGER.info(f\"Score with best loss weights stage{stage+1}: {score:.4f}\")\n\n        # Log overall CV (Cross-Validation) score (average of fold scores)\n        LOGGER.info(f\"==================== CV ====================\")\n        LOGGER.info(f\"Score with best loss weights: {np.mean(scores):.4f}\")\n\n        # Reset index of the OOF DataFrame\n        oof_df.reset_index(drop=True, inplace=True)\n\n        # Save the OOF DataFrame to a CSV file with specific naming convention\n        oof_df.to_csv(\n            f\"{pop_dir}{CFG.model_name}_oof_df_ver-{CFG.VERSION}_stage-{stage+1}.csv\",\n            index=False,\n        )\n\n        # Update the previous directory path for referencing in subsequent stages\n        prev_dir = pop_dir\n\n    # If using Weights & Biases logging, finish the session\n    if CFG.wandb:\n        wandb.finish()\n","metadata":{"papermill":{"duration":2650.233098,"end_time":"2024-02-20T14:26:40.161837","exception":false,"start_time":"2024-02-20T13:42:29.928739","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == \"__main__\" and CFG.train_by_folds:\n    # Check if script is executed directly and training by folds is enabled\n\n    # Set random seed for reproducibility\n    seed_torch(seed=CFG.seed)\n\n    # Initialize dictionaries to store scores and OOF DataFrames for each stage\n    stages_scores = {i: [] for i in CFG.train_stages}\n    stages_oof_df = {i: pd.DataFrame() for i in CFG.train_stages}\n\n    # Loop through each fold in the defined training folds\n    for fold in CFG.train_folds:\n\n        # Variable to store the directory path from the previous stage\n        prev_dir = \"\"\n\n        # Loop through all defined stages\n        for stage in range(len(CFG.total_evaluators)):\n            # Construct directory path for the current stage outputs\n            pop_dir = f\"{OUTPUT_DIR}pop_{stage+1}_weight_oof/\"\n            # Create the directory if it doesn't exist\n            if not os.path.exists(pop_dir):\n                os.makedirs(pop_dir)\n\n            # Skip stages not included in the training stages list\n            if stage not in CFG.train_stages:\n                prev_dir = pop_dir\n                continue\n\n            # Call the `train_loop` function to perform training for the current stage, fold, etc.\n            train_oof_df, score = train_loop(\n                stage=stage + 1,\n                epochs=CFG.epochs[stage],\n                fold=fold,\n                folds=train_pops[stage],\n                directory=pop_dir,\n                prev_dir=prev_dir,\n                eggs=all_eegs,\n            )\n\n            # Concatenate OOF predictions from the current fold to the stage's OOF DataFrame\n            stages_oof_df[stage] = pd.concat([stages_oof_df[stage], train_oof_df])\n            # Append the score for the current fold to the stage's score list\n            stages_scores[stage].append(score)\n\n            # Update the previous directory path for referencing in subsequent stages\n            prev_dir = pop_dir\n\n            # Log information about the fold result\n            LOGGER.info(f\"========== fold: {fold} stage: {stage+1} result ==========\")\n            LOGGER.info(f\"Score with best loss weights stage{stage+1}: {score:.4f}\")\n\n    # Log overall CV (Cross-Validation) score (average score) for each stage\n    for stage, scores in stages_scores.items():\n        LOGGER.info(f\"============ CV score with best loss weights ============\")\n        LOGGER.info(f\"Stage {stage}: {np.mean(scores):.4f}\")\n\n    # Save OOF DataFrames for each stage to separate CSV files\n    for stage, oof_df in stages_oof_df.items():\n        pop_dir = f\"{OUTPUT_DIR}pop_{stage+1}_weight_oof/\"\n        # Reset index of the OOF DataFrame\n        oof_df.reset_index(drop=True, inplace=True)\n        # Save the OOF DataFrame to a CSV file with specific naming convention\n        oof_df.to_csv(\n            f\"{pop_dir}{CFG.model_name}_oof_df_ver-{CFG.VERSION}_stage-{stage+1}.csv\",\n            index=False,\n        )\n\n    # If using Weights & Biases logging, finish the session\n    if CFG.wandb:\n        wandb.finish()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"# === Pre-process OOF ===\n# Prepare ground truth data\ngt = oof_df[[\"eeg_id\"] + CFG.target_cols]  # Select relevant columns from OOF DataFrame\ngt.sort_values(by=\"eeg_id\", inplace=True)  # Sort by \"eeg_id\" for consistent order\ngt.reset_index(inplace=True, drop=True)  # Reset index after sorting\n\n# Prepare predicted data\npreds = oof_df[[\"eeg_id\"] + CFG.pred_cols]  # Select predicted target columns\npreds.columns = [\"eeg_id\"] + CFG.target_cols  # Rename columns to match target format\npreds.sort_values(by=\"eeg_id\", inplace=True)  # Sort by \"eeg_id\" for consistency\npreds.reset_index(inplace=True, drop=True)  # Reset index after sorting\n\n# Extract target values (ground truth and predicted)\ny_trues = gt[CFG.target_cols]  # Select target columns from ground truth DataFrame\ny_preds = preds[CFG.target_cols]  # Select target columns from predicted DataFrame\n\n# Create DataFrames for calculating score (metric)\noof = pd.DataFrame(y_preds.copy())  # Copy predicted values to avoid modifying original DataFrame\noof[\"id\"] = np.arange(len(oof))  # Add \"id\" column with sequential indices\n\ntrue = pd.DataFrame(y_trues.copy())  # Copy ground truth values to avoid modifying original DataFrame\ntrue[\"id\"] = np.arange(len(true))  # Add \"id\" column with sequential indices\n\n# Calculate CV (Cross-Validation) score using kaggle_kl_div library\ncv = kaggle_kl_div.score(solution=true, submission=oof, row_id_column_name=\"id\")\n\n# Print the CV score with informative message\nprint(f\"CV Score with resnet1D_gru Raw EEG = {cv:.4f}\")\n","metadata":{"papermill":{"duration":0.167828,"end_time":"2024-02-20T15:08:05.860071","exception":false,"start_time":"2024-02-20T15:08:05.692243","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}