{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7898450,"sourceType":"datasetVersion","datasetId":4638387}],"dockerImageVersionId":30665,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# imports","metadata":{}},{"cell_type":"code","source":"# %pip install --upgrade wandb8","metadata":{"execution":{"iopub.status.busy":"2024-03-21T21:07:10.926539Z","iopub.execute_input":"2024-03-21T21:07:10.927206Z","iopub.status.idle":"2024-03-21T21:07:10.934410Z","shell.execute_reply.started":"2024-03-21T21:07:10.927158Z","shell.execute_reply":"2024-03-21T21:07:10.932860Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport wandb\nimport cv2\nimport pandas as pd\nimport numpy as np\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt \nfrom torch.utils.data import Dataset, DataLoader, random_split\nimport torch as t\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ndevice = 'cuda' if t.cuda.is_available() else 'cpu'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-03-21T21:07:10.936510Z","iopub.execute_input":"2024-03-21T21:07:10.937317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# utils","metadata":{}},{"cell_type":"code","source":"import gc \ndef GC():\n    gc.collect()\n    t.cuda.empty_cache()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## configs and data","metadata":{}},{"cell_type":"code","source":"batch_size = 150 // 5\nprefetch_factor = 2\nnum_workers = 4","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_path = '/kaggle/input/hms-harmful-brain-activity-classification/test_eegs/'\ntrain_path = '/kaggle/input/hms-harmful-brain-activity-classification/train_eegs/'\nBASE_PATH = \"/kaggle/input/hms-harmful-brain-activity-classification\"\nclass_names = ['Seizure', 'LPD', 'GPD', 'LRDA','GRDA', 'Other']\nFEATS_FOR_REAL = ['Fp1', 'F3', 'C3', 'P3', 'F7', 'T3', 'T5', 'O1', 'Fz', 'Cz', 'Pz', 'Fp2', 'F4', 'C4', 'P4', 'F8', 'T4', 'T6', 'O2', 'EKG']\nTARS = {'Seizure':0, 'LPD':1, 'GPD':2, 'LRDA':3, 'GRDA':4, 'Other':5}\nTARGETS = ['seizure_vote', 'lpd_vote', 'gpd_vote', 'lrda_vote', 'grda_vote','other_vote']","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.read_csv(f'{BASE_PATH}/test.csv')\ntest_df['eeg_path'] = f'{BASE_PATH}/test_eegs/'+test_df['eeg_id'].astype(str)+'.parquet'\ntest_df['spec_path'] = f'{BASE_PATH}/test_spectrograms/'+test_df['spectrogram_id'].astype(str)+'.parquet'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(f'{BASE_PATH}/train.csv')\neeg_path = f'{BASE_PATH}/train_eegs/'+train_df['eeg_id'].astype(str)+'.parquet'\nclass_name = train_df.expert_consensus.copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dataset(Dataset):\n    def __init__(self, transform=None):\n        super().__init__()\n        self.dataframe = train_df\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def __getitem__(self, idx):\n        row = self.dataframe.iloc[idx]\n        eeg_id = row['eeg_id']\n        parq_path = f'{train_path}{eeg_id}.parquet'\n        eeg = pd.read_parquet(parq_path)\n        start_time_second = row['eeg_label_offset_seconds']\n        offset_dp = int(start_time_second * 200)\n        duration = 10_000\n    \n        eeg = eeg.iloc[offset_dp:offset_dp+duration]\n        eeg = eeg.ffill(axis=0)\n        eeg = eeg.fillna(0)\n        labels = row[TARGETS].values.astype(np.float64)\n        labels = labels/np.sum(labels)\n        samples = t.tensor(eeg[FEATS_FOR_REAL].values)\n        labels_out = t.tensor(labels,dtype=t.float64)\n        \n        assert not samples.isnan().any()\n        assert not labels_out.isnan().any()\n        \n        return samples, labels_out","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = Dataset()\ntrain_size = int(len(dataset) * 0.9)\ntest_size = len(dataset) - train_size\ntrain_dataset, test_dataset = random_split(dataset, [train_size, test_size])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataloader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers, prefetch_factor=prefetch_factor, shuffle=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, prefetch_factor=prefetch_factor, shuffle=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# model 👯‍♀️","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self, d_in, d_out, kernel_size, drop):\n        super().__init__()\n        self.model = nn.Sequential(\n            nn.Conv1d(d_in, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.MaxPool1d(kernel_size=2, stride=2, padding=0), # reduce sequence size by 2\n        )\n    def forward(self, x):\n        # TODO: add skip for training speed\n        return self.model(x)\n        \nclass Model(nn.Module):\n    def __init__(self, in_channels=20, gru_hidden_size=128, drop=0.2):\n        super().__init__()\n        self.pre_out = in_channels * 4\n        self.gru_hidden_size = gru_hidden_size\n        \n        self.pre_process = nn.Sequential(\n            nn.BatchNorm1d(in_channels, momentum=None),\n            # use conv1d as a denoiser\n            # block 1\n            ConvBlock(in_channels, in_channels * 2, kernel_size=3, drop=drop),\n            nn.BatchNorm1d(in_channels * 2, momentum=None),\n            \n            # block 2\n            ConvBlock(in_channels * 2, in_channels * 4, kernel_size=5, drop=drop),\n            nn.BatchNorm1d(self.pre_out, momentum=None),\n\n            # block 3\n            ConvBlock(in_channels * 4, in_channels * 4, kernel_size=7, drop=drop),\n            nn.BatchNorm1d(self.pre_out, momentum=None),\n        )\n        \n        # TODO: add a learnable first state for GRU or check what is the default\n        self.gru = nn.GRU(self.pre_out, self.gru_hidden_size, num_layers=1, batch_first=True, bidirectional=True)\n\n        self.head = nn.Sequential(\n            nn.Linear(self.gru_hidden_size * 2, self.gru_hidden_size * 4),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.Linear(self.gru_hidden_size * 4, 6)\n        )\n\n    def forward(self, x: ('batch', 'seq', 'channel')):\n        # pre_process: (batch, channel, seq) → (batch / 4, channel * 4, seq)\n        x = x.permute((0, 2, 1))\n        x = self.pre_process(x)\n        x = x.permute((0, 2, 1))\n\n        # GRU: (batch, seq, input_size), [(2 * num_layers, batch, hidden_size)] → (batch, seq, 2 * hidden_size)\n        x, _ = self.gru(x)\n        x = x[:, -1, :]\n\n        # head: (batch, 2 * hidden_size) → (batch, 6)\n        x = self.head(x)\n\n        # out: → (batch, 6)\n        return x\n\ndef scope():\n    m = Model().to(device)\n    x, y = next(train_dataloader.__iter__())\n    r = m(x.to(device))\n    print(f'{r.shape=}')\n    \nscope()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"GC()\nmodel = Model().to(device)\nopt = t.optim.Adam(model.parameters(), lr=1e-3)\nprint(f'model has {sum(p.numel() for p in model.parameters())} params')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# train","metadata":{}},{"cell_type":"code","source":"def train(model, opt, wnb=True, replay_buffer_size=5):\n    model.train()\n    if wnb: wandb.init(project='kaggle-eeg-rc')\n    for epoch in range(1):\n        x_replay_buffer = []\n        y_replay_buffer = []\n        i = 0\n        tq = tqdm(train_dataloader)\n        for x_train, y_train in tq:\n            x_replay_buffer = x_replay_buffer[-replay_buffer_size:]\n            x_replay_buffer.append(x_train.to(device))\n            y_replay_buffer = y_replay_buffer[-replay_buffer_size:]\n            y_replay_buffer.append(y_train.to(device))\n            x_train = t.cat(x_replay_buffer, dim=0)\n            y_train = t.cat(y_replay_buffer, dim=0)\n            \n            i+=1\n            for k in range(1): # the data reading is too slow, so force the GPU to spin\n                logs = model(x_train.to(device)).log_softmax(-1)\n                kl_loss = nn.KLDivLoss(reduction=\"batchmean\")\n                loss = kl_loss(logs, y_train.to(device))\n                opt.zero_grad()\n                loss.backward()\n                opt.step()\n                tq.set_description(f'loss = {loss:.4f}')\n                if wnb: wandb.log({'loss': loss.item()})\n                if k == 0 and i % 50 == 0:\n                    print(f'{epoch} {loss.item()=}')\n        #             if i % 50 == 0:\n        #                 t.save(model.state_dict(),'model-weights.pt')\n        #                 print(\"model-weights.pt\")\n\ntrain(model, opt, wnb=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# kill the dropout\ndef kill_dropout(model):\n    model.pre_process[1].model[2].p = 0.\n    model.pre_process[1].model[5].p = 0.\n    model.pre_process[1].model[8].p = 0.\n    model.pre_process[3].model[2].p = 0.\n    model.pre_process[3].model[5].p = 0.\n    model.pre_process[3].model[8].p = 0.\n    model.pre_process[5].model[2].p = 0.\n    model.pre_process[5].model[5].p = 0.\n    model.pre_process[5].model[8].p = 0.\n    model.head[1].p = 0.","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# t.save(model.state_dict(),'model-weights2.pt')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = Model().to(device)\n# model.load_state_dict(t.load('/kaggle/working/model-weights2.pt', map_location=device))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val, label = next(train_dataloader.__iter__())\nval2, label2 = next(test_dataloader.__iter__())\n\nx_val = val2.to(device)\nx_train = val.to(device)\n\ny_val =label2.to(device)\ny_train = label.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.pre_process[0].track_running_stats","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# val, label = next(train_dataloader.__iter__())\n# val2, label2 = next(test_dataloader.__iter__())\n\n# x_val = val2.to(device)\n# x_train = val.to(device)\n\n# y_val =label2.to(device)\n# y_train = label.to(device)\n\ndef eval(model, x=x_val, y=y_val):\n    assert not x.isnan().any()\n    assert not y.isnan().any()\n    model.train()\n#     model.gru.eval()\n#     model.pre_process[0].train()\n#     model.pre_process[2].train()\n#     model.pre_process[4].train()\n#     model.pre_process[6].train()\n\n#     model.pre_process[0].track_running_stats = False\n#     model.pre_process[2].track_running_stats = False\n#     model.pre_process[4].track_running_stats = False\n#     model.pre_process[6].track_running_stats = False\n    logs = model(x.to(device)).log_softmax(-1)\n    kl_loss = nn.KLDivLoss(reduction=\"batchmean\")\n    loss = kl_loss(logs, y.to(device))\n    model.train()\n    return loss\n\nprint(f'train: {eval(model, x_train, y_train)}')\nprint(f'eval:  {eval(model, x_val, y_val)}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Validation Stuff kthx\n\n@t.no_grad()\ndef submit():\n    model.eval()\n    res = []\n    # TODO: fix, read from the correct place instead\n    for batch, labels in test_dataloader:\n        batch = batch.to(device)\n        prob = model(batch).softmax(-1)\n        res.append(prob.detach().cpu())\n        break\n    res = t.cat(res, dim=0)\n    print(res[0])\n    \n    pred_df = test_df[[\"eeg_id\"]].copy()\n    target_cols = [x.lower()+'_vote' for x in class_names]\n    pred_df[target_cols] = res.tolist()\n    sub_df = pd.read_csv(f'{BASE_PATH}/sample_submission.csv')\n    sub_df = sub_df[[\"eeg_id\"]].copy()\n    sub_df = sub_df.merge(pred_df, on=\"eeg_id\", how=\"left\")\n    sub_df.to_csv(\"submission.csv\", index=False)\n    sub_df.head()\n    \n# submit()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}