{"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":"none","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}],"dockerImageVersionId":30674,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Based on [HMS Resnet1d GRU Train - 1 / 5 Dataset](https://www.kaggle.com/code/konstantinboyko/hms-resnet1d-gru-train-1-5-dataset) by [Konstantin Boyko](https://www.kaggle.com/konstantinboyko)\n\n**I'm going to make X as a PSD np.array instead of raw EEGs**\n - In Utils added _signal_psd_multitaper but i'm not sure that i correct use it at the end of the Dataset class because for computing needs a lot of time. May be i should make separate dataset with PSD in file and do some preload to get rid of repeating to compute this function. \n - Also i should change dimension from (2000,8) to (1000,8) to get rid of duplicate PSD array, but i don't sure how to correct calculate output of the layers. ","metadata":{}},{"cell_type":"markdown","source":"# **Library**","metadata":{}},{"cell_type":"code","source":"import 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\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#import cupy as cp\n#import cupyx.scipy.signal as cpsig\n\nsys.path.append(\"/kaggle/input/kaggle-kl-div\")\nimport kaggle_kl_div\n\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")\n\ndevice = torch.device(\"cuda\")\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\"\n\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(\n        f\"PyTorch Version:{torch.__version__}, CUDA is available:{torch.cuda.is_available()}, Version CUDA:{torch.version.cuda}\"\n    )\n    print(\n        f\"Device Capability:{torch.cuda.get_device_capability()}, {torch.cuda.get_arch_list()}\"\n    )\n    print(\n        f\"CuDNN Enabled:{torch.backends.cudnn.enabled}, Version:{torch.backends.cudnn.version()}\"\n    )\nexcept Exception:\n    pass","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:23:01.061023Z","iopub.execute_input":"2024-03-28T23:23:01.061406Z","iopub.status.idle":"2024-03-28T23:23:09.794301Z","shell.execute_reply.started":"2024-03-28T23:23:01.061377Z","shell.execute_reply":"2024-03-28T23:23:09.792705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Directory settings**","metadata":{}},{"cell_type":"code","source":"class APP:\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    kaggle = os.environ.get(\"KAGGLE_KERNEL_RUN_TYPE\", \"\") != \"\"\n    local = os.environ.get(\"DOCKER_USING\", \"\") == \"LOCAL\"\n    date_time_start = dt.datetime.now()\n    dt_start_ymd_hms = date_time_start.strftime(\"%Y.%m.%d_%H-%M-%S\")\n\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\n    else:\n        try:\n            file_run_path = Path(__file__)\n        except Exception as inst:\n            print(inst)\n\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\nOUTPUT_DIR = \"./\"\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)\n\nprint(f\"jupyter:{APP.jupyter}, kaggle:{APP.kaggle}, local:{APP.local}\")\nprint(APP.file_run_path)\nprint(APP.path_out)","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:23:09.797217Z","iopub.execute_input":"2024-03-28T23:23:09.798909Z","iopub.status.idle":"2024-03-28T23:23:09.811330Z","shell.execute_reply.started":"2024-03-28T23:23:09.798872Z","shell.execute_reply":"2024-03-28T23:23:09.810113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Config**","metadata":{}},{"cell_type":"code","source":"class CFG:\n    VERSION = '82'\n\n    wandb = False\n    debug = False\n    apex = True\n    visualize = False\n    save_all_models = True\n\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\"\n    # optimizer = \"Adan\"\n    optimizer = \"AdamW\"\n\n    factor = 0.9\n    eps = 1e-6\n    lr = 8e-3\n    min_lr = 1e-6\n\n    batch_size = 64\n    batch_koef_valid = 2\n    batch_scheduler = True\n    weight_decay = 1e-2\n    gradient_accumulation_steps = 1\n    max_grad_norm = 1e7\n\n    fixed_kernel_size = 5\n    # linear_layer_features = 424\n    # kernels = [3, 5, 7, 9]\n    #linear_layer_features = 448  # Full Signal = 10_000\n    #linear_layer_features = 352  # Half Signal = 5_000\n    linear_layer_features = 304   # 1/4, 1/5, 1/6  Signal = 2_000\n    #linear_layer_features = 280  # 1/10  Signal = 1_000\n    kernels = [3, 5, 7, 9, 11]\n    # kernels = [5, 7, 9, 11, 13]\n\n    seq_length = 50  # Second's\n    sampling_rate = 200  # Hz\n    nsamples = seq_length * sampling_rate  # Число семплов 10_000\n    n_split_samples = 1\n    out_samples = nsamples // n_split_samples  # 2_000\n    sample_delta = nsamples - out_samples  # 8000\n    sample_offset = sample_delta // 2\n    multi_validation = False\n\n    train_by_stages = False\n    train_by_folds = True\n\n    # 'GPD', 'GRDA', 'LPD', 'LRDA', 'Other', 'Seizure'\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)]]  # Deprecated\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 = 4\n            total_evals_old = [[(1, 4),(4, 5), (5, 6)], (6, 29)]  # Deprecated\n            total_evaluators = [ \n                [   \n                    {'band':(1, 3), 'excl_evals':[]}, \n                    {'band':(4, 4), 'excl_evals':['GPD']}, \n                    {'band':(5, 5), 'excl_evals':[]},\n                    {'band':(6, 28), 'excl_evals':[]},\n                ], \n                [   \n                    {'band':(6, 28), '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)]  # Deprecated\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\n    train_folds = [0, 1, 2, 3, 4]\n    # train_folds = [0]\n\n    patience = 11\n    seed = 2024\n\n    bandpass_filter = {\"low\": 0.5, \"high\": 20, \"order\": 2}\n    rand_filter = {\"probab\": 0.1, \"low\": 10, \"high\": 20, \"band\": 1.0, \"order\": 2}\n    freq_channels = []  # [(8.0, 12.0)]; [(0.5, 4.5)]\n    filter_order = 2\n\n    random_divide_signal = 0.05\n    random_close_zone = 0.05\n    random_common_negative_signal = 0.0\n    random_common_reverse_signal = 0.0\n    random_negative_signal = 0.05\n    random_reverse_signal = 0.05\n\n    log_step = 100  # Шаг отображения тренировки\n    log_show = False\n\n    scheduler = \"CosineAnnealingWarmRestarts\"  # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts','OneCycleLR']\n\n    # CosineAnnealingLR params\n    cosanneal_params = {\n        \"T_max\": 6,\n        \"eta_min\": 1e-5,\n        \"last_epoch\": -1,\n    }\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\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\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]\n\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        #('Fz', 'Cz'), ('Cz', 'Pz'),\n    ]\n\n    eeg_features = [\"Fp1\", \"T3\", \"C3\", \"O1\", \"Fp2\", \"C4\", \"T4\", \"O2\"]  # 'Fz', 'Cz', 'Pz'\n        # 'F3', 'P3', 'F7', 'T5', 'Fz', 'Cz', 'Pz', 'F4', 'P4', 'F8', 'T6', 'EKG']\n    feature_to_index = {x: y for x, y in zip(eeg_features, range(len(eeg_features)))}\n    simple_features = []  # 'Fz', 'Cz', 'Pz', 'EKG'\n\n    # eeg_features = [row for row in feature_to_index]\n    # eeg_feat_size = len(eeg_features)\n    \n    n_map_features = len(map_features)\n    in_channels = n_map_features + n_map_features * len(freq_channels) + len(simple_features)\n    target_size = len(target_cols)\n\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    #file_raw_eeg = path_inp / \"brain-eegs-plus/eegs.npy\"\n    #file_raw_eeg = path_inp / \"brain-eegs-full/eegs.npy\"\n\n    if APP.kaggle:\n        num_workers = 2\n        parallel = True\n        # GPU_DEVICES = \"auto\"\n\n\n# print(CFG.eeg_feat_size, CFG.in_channels)\nprint(CFG.feature_to_index)\nprint(CFG.eeg_features)","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:23:48.835473Z","iopub.execute_input":"2024-03-28T23:23:48.835867Z","iopub.status.idle":"2024-03-28T23:23:48.866596Z","shell.execute_reply.started":"2024-03-28T23:23:48.835838Z","shell.execute_reply":"2024-03-28T23:23:48.865484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(CFG.out_samples, CFG.in_channels)","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:00.665073Z","iopub.execute_input":"2024-03-28T23:24:00.665786Z","iopub.status.idle":"2024-03-28T23:24:00.672186Z","shell.execute_reply.started":"2024-03-28T23:24:00.665743Z","shell.execute_reply":"2024-03-28T23:24:00.670819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Utils**","metadata":{}},{"cell_type":"code","source":"def _signal_psd_multitaper(\n    signal, sampling_rate=200, min_frequency=0.02, max_frequency=20\n):\n    try:\n        import mne\n\n        power, frequency = mne.time_frequency.psd_array_multitaper(\n                signal,\n                sfreq=sampling_rate,\n                fmin=min_frequency,\n                fmax=max_frequency,\n                adaptive=True,\n                normalization=\"full\",\n                verbose=False,\n        )\n    except ImportError:\n        raise ImportError(\n            \"NeuroKit warning: signal_psd(): the 'mne'\",\n            \"module is required for the 'mne' method to run.\",\n            \"Please install it first (`pip install mne`).\",\n        )\n    return frequency, power\n\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\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 quantize_data(data, classes):\n    mu_x = mu_law_encoding(data, classes)\n    return mu_x  # quantized\n\n\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\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\ndef butter_bandpass(lowcut, highcut, fs, order=5):\n    return butter(order, [lowcut, highcut], fs=fs, btype=\"band\")\n\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 butter_lowpass_filter(\n    data, cutoff_freq=20, sampling_rate=CFG.sampling_rate, order=4\n):\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\ndef denoise_filter(x):\n    # Частота дискретизации и желаемые частоты среза (в Гц).\n    # Отфильтруйте шумный сигнал\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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:04.610924Z","iopub.execute_input":"2024-03-28T23:24:04.611678Z","iopub.status.idle":"2024-03-28T23:24:04.633498Z","shell.execute_reply.started":"2024-03-28T23:24:04.611635Z","shell.execute_reply":"2024-03-28T23:24:04.632342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Dataset**","metadata":{}},{"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        # Обозначает количество пакетов за эпоху\n        return len(self.df)\n\n    def __getitem__(self, index):\n        \"\"\"\n        Get one item.\n        \"\"\"\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    def set_offset(self, offset: int):\n        self.offset = offset\n\n    def __data_generation(self, index):\n        # Генерирует данные, содержащие образцы размера партии\n        X = np.zeros(\n            (CFG.nsamples, CFG.in_channels), dtype=\"float32\"\n        )  # Size=(10000, 14)\n\n        random_divide_signal = False\n        row = self.df.iloc[index]  # Строка Pandas\n        data = self.eegs[row.eeg_id]  # Size=(10000, 8)\n        if CFG.nsamples != CFG.out_samples:\n            if self.mode == \"train\":\n                offset = (CFG.sample_delta * random.randint(0, 1000)) // 1000\n            elif not self.offset is None:\n                offset = self.offset\n            else:\n                offset = CFG.sample_offset\n\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                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        for i, (feat_a, feat_b) in enumerate(CFG.map_features):\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            diff_feat = (\n                data[:, CFG.feature_to_index[feat_a]]\n                - data[:, CFG.feature_to_index[feat_b]]\n            )  # Size=(10000,)\n\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            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            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                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            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            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            X[:, n] = feat_val\n            n += 1\n            \n        # Обрезать края превышающие значения [-1024, 1024]\n        X = np.clip(X, -1024, 1024)\n\n        # Замените NaN нулем и разделить все на 32\n        X = np.nan_to_num(X, nan=0) / 32.0\n\n        # обрезать полосовым фильтром верхнюю границу в 20 Hz.\n        X = butter_lowpass_filter(X, order=CFG.filter_order)  # 4\n\n        y_prob = np.zeros(CFG.target_size, dtype=\"float32\")  # Size=(6,)\n        if self.mode != \"test\":\n            y_prob = row[CFG.target_cols].values.astype(np.float32)\n        \n        for i in range(CFG.n_map_features):\n            signal = X[:,i]\n            f, p = _signal_psd_multitaper(signal)\n            X[:1000,i] = p\n            X[1000:2000,i] = p\n        X = X[:2000]\n            \n        \n        return X, y_prob","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:12.790025Z","iopub.execute_input":"2024-03-28T23:24:12.790435Z","iopub.status.idle":"2024-03-28T23:24:12.828529Z","shell.execute_reply.started":"2024-03-28T23:24:12.790407Z","shell.execute_reply":"2024-03-28T23:24:12.827738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Helper function**","metadata":{}},{"cell_type":"code","source":"class KLDivLossWithLogits(nn.KLDivLoss):\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        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    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  # Это опция требует много паямяти GPU\n    # pl.seed_everything(seed)\n","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:14.677824Z","iopub.execute_input":"2024-03-28T23:24:14.678245Z","iopub.status.idle":"2024-03-28T23:24:14.689752Z","shell.execute_reply.started":"2024-03-28T23:24:14.678213Z","shell.execute_reply":"2024-03-28T23:24:14.688567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Model**","metadata":{}},{"cell_type":"code","source":"class ResNet_1D_Block(nn.Module):\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        self.bn1 = nn.BatchNorm1d(num_features=in_channels)\n        # self.relu = nn.ReLU(inplace=False)\n        # self.relu_1 = nn.PReLU()\n        # self.relu_2 = nn.PReLU()\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        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    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        self.bn1 = nn.BatchNorm1d(num_features=self.planes)\n        # self.relu = nn.ReLU(inplace=False)\n        # self.relu_1 = nn.ReLU()\n        # self.relu_2 = nn.ReLU()\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        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        self.fc = nn.Linear(in_features=linear_layer_features, out_features=num_classes)\n\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    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    def forward(self, x):\n        new_out = self.extract_features(x)\n        result = self.fc(new_out)\n        return result","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-03-28T23:24:16.245262Z","iopub.execute_input":"2024-03-28T23:24:16.245652Z","iopub.status.idle":"2024-03-28T23:24:16.278677Z","shell.execute_reply.started":"2024-03-28T23:24:16.245624Z","shell.execute_reply":"2024-03-28T23:24:16.277723Z"},"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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:17.698476Z","iopub.execute_input":"2024-03-28T23:24:17.699250Z","iopub.status.idle":"2024-03-28T23:24:17.728849Z","shell.execute_reply.started":"2024-03-28T23:24:17.699212Z","shell.execute_reply":"2024-03-28T23:24:17.727743Z"},"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    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            loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n\n        losses.update(loss.item(), batch_size)\n\n        scaler.scale(loss).backward()\n\n        grad_norm = torch.nn.utils.clip_grad_norm_(\n            model.parameters(), CFG.max_grad_norm\n        )\n\n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            global_step += 1\n            if CFG.batch_scheduler:\n                scheduler.step()\n        end = time.time()\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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:19.159707Z","iopub.execute_input":"2024-03-28T23:24:19.160131Z","iopub.status.idle":"2024-03-28T23:24:19.173286Z","shell.execute_reply.started":"2024-03-28T23:24:19.160100Z","shell.execute_reply":"2024-03-28T23:24:19.171914Z"},"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    losses = AverageMeter()\n    model.eval()\n    preds = []\n    targets = []\n    start = end = time.time()\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            y_preds = model(eegs)\n            loss = criterion(F.log_softmax(y_preds, dim=1), labels)\n\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n\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\n        if CFG.log_show and (\n            step % CFG.log_step == 0 or step == (len(valid_loader) - 1)\n        ):\n            # remain=timeSince(start, float(step + 1) / len(valid_loader))\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            # Elapsed {remain:s}\n\n    predictions = np.concatenate(preds)\n    targets = np.concatenate(targets)\n\n    return losses.avg, predictions","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:20.508414Z","iopub.execute_input":"2024-03-28T23:24:20.508811Z","iopub.status.idle":"2024-03-28T23:24:20.520289Z","shell.execute_reply.started":"2024-03-28T23:24:20.508779Z","shell.execute_reply":"2024-03-28T23:24:20.519119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Build Optimizer**","metadata":{}},{"cell_type":"code","source":"def build_optimizer(cfg, model, device, epochs, num_batches_per_epoch):\n    lr = cfg.lr\n    # lr = default_configs[\"lr\"]\n    if cfg.optimizer == \"SAM\":\n        base_optimizer = (\n            torch.optim.SGD\n        )  # define an optimizer for the \"sharpness-aware\" update\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        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        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        optimizer_model = Adam(model.parameters(), lr=lr, weight_decay=CFG.weight_decay)\n    elif cfg.optimizer == \"AdamW\":\n        optimizer_model = AdamW(\n            model.parameters(), lr=lr, weight_decay=CFG.weight_decay\n        )\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","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:21.917046Z","iopub.execute_input":"2024-03-28T23:24:21.918173Z","iopub.status.idle":"2024-03-28T23:24:21.930105Z","shell.execute_reply.started":"2024-03-28T23:24:21.918131Z","shell.execute_reply":"2024-03-28T23:24:21.928356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Sheduler**","metadata":{}},{"cell_type":"code","source":"def get_scheduler(optimizer, epochs, steps_per_epoch):\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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:23.568308Z","iopub.execute_input":"2024-03-28T23:24:23.568740Z","iopub.status.idle":"2024-03-28T23:24:23.576409Z","shell.execute_reply.started":"2024-03-28T23:24:23.568705Z","shell.execute_reply":"2024-03-28T23:24:23.575308Z"},"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    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    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    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        #rand_filter=CFG.rand_filter,\n    )\n\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    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    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    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    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        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        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    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 valid_folds, best_score","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:27.497167Z","iopub.execute_input":"2024-03-28T23:24:27.497556Z","iopub.status.idle":"2024-03-28T23:24:27.523549Z","shell.execute_reply.started":"2024-03-28T23:24:27.497528Z","shell.execute_reply":"2024-03-28T23:24:27.522293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Load train data**","metadata":{}},{"cell_type":"code","source":"train = 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\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\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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:29.902682Z","iopub.execute_input":"2024-03-28T23:24:29.903086Z","iopub.status.idle":"2024-03-28T23:24:30.416198Z","shell.execute_reply.started":"2024-03-28T23:24:29.903042Z","shell.execute_reply":"2024-03-28T23:24:30.414885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.visualize:\n    plt.figure(figsize=(10, 6))\n    plt.hist(train[\"total_evaluators\"], bins=15, 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\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":{"execution":{"iopub.status.busy":"2024-03-28T23:24:33.038427Z","iopub.execute_input":"2024-03-28T23:24:33.038791Z","iopub.status.idle":"2024-03-28T23:24:33.432614Z","shell.execute_reply.started":"2024-03-28T23:24:33.038757Z","shell.execute_reply":"2024-03-28T23:24:33.431469Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Split data**","metadata":{}},{"cell_type":"code","source":"# %%time\nall_eeg_specs = np.load(CFG.file_eeg_specs, allow_pickle=True).item()\n\ntrain = train[train[\"label_id\"].isin(all_eeg_specs.keys())].copy()\nprint(train.shape)\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['total_evaluators'] == CFG.test_total_eval].groupby(['expert_consensus','total_evaluators']).count()","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:24:35.872354Z","iopub.execute_input":"2024-03-28T23:24:35.873029Z","iopub.status.idle":"2024-03-28T23:26:13.574314Z","shell.execute_reply.started":"2024-03-28T23:24:35.872993Z","shell.execute_reply":"2024-03-28T23:26:13.573196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_pops = []\nfor eval_list in CFG.total_evaluators:\n    result=[]\n    train_pop = train  \n    for eval_dict in eval_list:\n        band = eval_dict['band']\n        pop_idx = (train_pop[\"total_evaluators\"] >= band[0]) \n        pop_idx &= (train_pop[\"total_evaluators\"] <= band[1])\n        for exclude in eval_dict['excl_evals']:\n            pop_idx &= ~(train_pop['expert_consensus'] == exclude)\n            pass\n        result.append(train_pop[pop_idx])\n    train_pop = pd.concat(result).copy().reset_index()\n\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    train_pops.append(train_pop)\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()\n","metadata":{"execution":{"iopub.status.busy":"2024-03-28T23:27:29.660252Z","iopub.execute_input":"2024-03-28T23:27:29.660686Z","iopub.status.idle":"2024-03-28T23:27:29.760213Z","shell.execute_reply.started":"2024-03-28T23:27:29.660654Z","shell.execute_reply":"2024-03-28T23:27:29.759000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Dedupicate train EEG ID**","metadata":{}},{"cell_type":"code","source":"# %%time\nall_eegs = np.load(CFG.file_raw_eeg, allow_pickle=True).item()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Train Stage**","metadata":{}},{"cell_type":"code","source":"def get_score(preds, targets):\n    oof = pd.DataFrame(preds.copy())\n    oof[\"id\"] = np.arange(len(oof))\n    true = pd.DataFrame(targets.copy())\n    true[\"id\"] = np.arange(len(true))\n    cv = kaggle_kl_div.score(solution=true, submission=oof, row_id_column_name=\"id\")\n    return cv\n\n\ndef get_result(result_df):\n    gt = result_df[[\"eeg_id\"] + CFG.target_cols]\n    gt.sort_values(by=\"eeg_id\", inplace=True)\n    gt.reset_index(inplace=True, drop=True)\n    preds = result_df[[\"eeg_id\"] + CFG.pred_cols]\n    preds.columns = [\"eeg_id\"] + CFG.target_cols\n    preds.sort_values(by=\"eeg_id\", inplace=True)\n    preds.reset_index(inplace=True, drop=True)\n    score_loss = get_score(gt[CFG.target_cols], preds[CFG.target_cols])\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_folds:\n    seed_torch(seed=CFG.seed)\n\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    for fold in CFG.train_folds:\n\n        prev_dir = \"\"\n        for stage in range(len(CFG.total_evaluators)):\n\n            pop_dir = f\"{OUTPUT_DIR}pop_{stage+1}_weight_oof/\"\n            if not os.path.exists(pop_dir):\n                os.makedirs(pop_dir)\n\n            if stage not in CFG.train_stages:\n                prev_dir = pop_dir\n                continue\n\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            stages_oof_df[stage] = pd.concat([stages_oof_df[stage], train_oof_df])\n            stages_scores[stage].append(score)\n\n            prev_dir = pop_dir\n\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    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    for stage, oof_df in stages_oof_df.items():\n        pop_dir = f\"{OUTPUT_DIR}pop_{stage+1}_weight_oof/\"\n        oof_df.reset_index(drop=True, inplace=True)\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 CFG.wandb:\n        wandb.finish()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Submission**","metadata":{}},{"cell_type":"code","source":"# === Pre-process OOF ===\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 = kaggle_kl_div.score(solution=true, submission=oof, row_id_column_name=\"id\")\nprint(f\"CV Score with resnet1D_gru Raw EEG = {cv:.4f}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}