{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7447509,"sourceType":"datasetVersion","datasetId":4334995},{"sourceId":7392775,"sourceType":"datasetVersion","datasetId":4297782}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Code referenced from https://www.kaggle.com/code/alejopaullier/hms-efficientnetb0-pytorch-train. Adapted for 8*256*256 resolution and transformer models instead of EfficientNet","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nimport gc\nimport matplotlib.pyplot as plt\nimport math\nimport multiprocessing\nimport numpy as np\nimport os\nimport pandas as pd\nimport random\nimport time\nimport timm\nimport torch\nimport torch.nn as nn\n\n\nfrom albumentations.pytorch import ToTensorV2\nfrom glob import glob\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm import tqdm\nfrom typing import Dict, List\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\"\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint('Using', torch.cuda.device_count(), 'GPU(s)')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration","metadata":{}},{"cell_type":"code","source":"class config:\n    AMP = True\n    BATCH_SIZE_TRAIN = 32\n    BATCH_SIZE_VALID = 32\n    EPOCHS = 4\n    FOLDS = 3\n    FREEZE = False\n    GRADIENT_ACCUMULATION_STEPS = 1\n    MAX_GRAD_NORM = 1e7\n    MODEL = \"swinv2_base_window8_256\" # Tried swinv2_base_window8_256, vit_relpos_base_patch32_plus_rpn_256, lambdaresnet26_rpt, vit_medium_patch16_reg4_gap_256\n    NUM_WORKERS = 0 \n    PRINT_FREQ = 20\n    IN_CHANS = 4\n    SEED = 20\n    TRAIN_FULL_DATA = False\n    VISUALIZE = True\n    WEIGHT_DECAY = 0.01\n    NUM_CLASSES = 6\n    \n    \nclass paths:\n    OUTPUT_DIR = \"/kaggle/working/\"  \n    PRE_LOADED_EEGS = '/kaggle/input/brain-eeg-spectrograms/eeg_specs.npy'\n    PRE_LOADED_SPECTOGRAMS = '/kaggle/input/brain-spectrograms/specs.npy'\n    TRAIN_CSV = \"/kaggle/input/hms-harmful-brain-activity-classification/train.csv\"\n    TRAIN_EEGS = \"/kaggle/input/brain-eeg-spectrograms/EEG_Spectrograms/\"\n    TRAIN_SPECTOGRAMS = \"/kaggle/input/hms-harmful-brain-activity-classification/train_spectrograms/\"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utility Functions","metadata":{}},{"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\ndef asMinutes(s: float):\n    \"Convert to minutes.\"\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since: float, percent: float):\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 get_logger(filename=paths.OUTPUT_DIR):\n    from logging import getLogger, INFO, StreamHandler, FileHandler, Formatter\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=f\"{filename}.log\")\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\n\ndef plot_spectrogram(spectrogram_path: str):\n    \"\"\"\n    Source: https://www.kaggle.com/code/mvvppp/hms-eda-and-domain-journey\n    Visualize spectogram recordings from a parquet file.\n    :param spectrogram_path: path to the spectogram parquet.\n    \"\"\"\n    sample_spect = pd.read_parquet(spectrogram_path)\n    \n    split_spect = {\n        \"LL\": sample_spect.filter(regex='^LL', axis=1),\n        \"RL\": sample_spect.filter(regex='^RL', axis=1),\n        \"RP\": sample_spect.filter(regex='^RP', axis=1),\n        \"LP\": sample_spect.filter(regex='^LP', axis=1),\n    }\n    \n    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(15, 12))\n    axes = axes.flatten()\n    label_interval = 5\n    for i, split_name in enumerate(split_spect.keys()):\n        ax = axes[i]\n        img = ax.imshow(np.log(split_spect[split_name]).T, cmap='viridis', aspect='auto', origin='lower')\n        cbar = fig.colorbar(img, ax=ax)\n        cbar.set_label('Log(Value)')\n        ax.set_title(split_name)\n        ax.set_ylabel(\"Frequency (Hz)\")\n        ax.set_xlabel(\"Time\")\n\n        ax.set_yticks(np.arange(len(split_spect[split_name].columns)))\n        ax.set_yticklabels([column_name[3:] for column_name in split_spect[split_name].columns])\n        frequencies = [column_name[3:] for column_name in split_spect[split_name].columns]\n        ax.set_yticks(np.arange(0, len(split_spect[split_name].columns), label_interval))\n        ax.set_yticklabels(frequencies[::label_interval])\n    plt.tight_layout()\n    plt.show()\n    \n    \ndef seed_everything(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed) \n\n    \ndef sep():\n    print(\"-\"*100)\n    \n\ntarget_preds = [x + \"_pred\" for x in ['seizure_vote', 'lpd_vote', 'gpd_vote', 'lrda_vote', 'grda_vote', 'other_vote']]\nlabel_to_num = {'Seizure': 0, 'LPD': 1, 'GPD': 2, 'LRDA': 3, 'GRDA': 4, 'Other':5}\nnum_to_label = {v: k for k, v in label_to_num.items()}\nLOGGER = get_logger()\nseed_everything(config.SEED)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Dataset","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(paths.TRAIN_CSV)\nlabel_cols = df.columns[-6:]\nprint(f\"Train cataframe shape is: {df.shape}\")\nprint(f\"Labels: {list(label_cols)}\")\ndf.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Preprocessing","metadata":{}},{"cell_type":"code","source":"train_df = df.groupby('eeg_id')[['spectrogram_id','spectrogram_label_offset_seconds']].agg({\n    'spectrogram_id':'first',\n    'spectrogram_label_offset_seconds':'min'\n})\ntrain_df.columns = ['spectogram_id','min']\n\naux = df.groupby('eeg_id')[['spectrogram_id','spectrogram_label_offset_seconds']].agg({\n    'spectrogram_label_offset_seconds':'max'\n})\ntrain_df['max'] = aux\n\naux = df.groupby('eeg_id')[['patient_id']].agg('first')\ntrain_df['patient_id'] = aux\n\naux = df.groupby('eeg_id')[label_cols].agg('sum')\nfor label in label_cols:\n    train_df[label] = aux[label].values\n    \ny_data = train_df[label_cols].values\ny_data = y_data / y_data.sum(axis=1,keepdims=True)\ntrain_df[label_cols] = y_data\n\naux = df.groupby('eeg_id')[['expert_consensus']].agg('first')\ntrain_df['target'] = aux\n\ntrain_df = train_df.reset_index()\nprint('Train non-overlapp eeg_id shape:', train_df.shape )\ntrain_df.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Read Train Spectrograms","metadata":{}},{"cell_type":"code","source":"%%time\nREAD_SPEC_FILES = False\n\npaths_spectograms = glob(paths.TRAIN_SPECTOGRAMS + \"*.parquet\")\nprint(f'There are {len(paths_spectograms)} spectrogram parquets')\n\nif READ_SPEC_FILES:    \n    all_spectrograms = {}\n    for file_path in tqdm(paths_spectograms):\n        aux = pd.read_parquet(file_path)\n        name = int(file_path.split(\"/\")[-1].split('.')[0])\n        all_spectrograms[name] = aux.iloc[:,1:].values\n        del aux\nelse:\n    all_spectrograms = np.load(paths.PRE_LOADED_SPECTOGRAMS, allow_pickle=True).item()\n    \nif config.VISUALIZE:\n    idx = np.random.randint(0,len(paths_spectograms))\n    spectrogram_path = paths_spectograms[idx]\n    plot_spectrogram(spectrogram_path)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Read EEG Spectrograms","metadata":{}},{"cell_type":"code","source":"%%time\nREAD_EEG_SPEC_FILES = False\n\npaths_eegs = glob(paths.TRAIN_EEGS + \"*.npy\")\nprint(f'There are {len(paths_eegs)} EEG spectograms')\n\nif READ_EEG_SPEC_FILES:\n    all_eegs = {}\n    for file_path in tqdm(paths_eegs):\n        eeg_id = file_path.split(\"/\")[-1].split(\".\")[0]\n        eeg_spectogram = np.load(file_path)\n        all_eegs[eeg_id] = eeg_spectogram\nelse:\n    all_eegs = np.load(paths.PRE_LOADED_EEGS, allow_pickle=True).item()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Validation","metadata":{}},{"cell_type":"code","source":"from sklearn.model_selection import KFold, GroupKFold\n\n\ngkf = GroupKFold(n_splits=config.FOLDS)\nfor fold, (train_index, valid_index) in enumerate(gkf.split(train_df, train_df.target, train_df.patient_id)):\n    train_df.loc[valid_index, \"fold\"] = int(fold)\n    \ndisplay(train_df.groupby('fold').size()), sep()\ndisplay(train_df.head())","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(\n        self, df: pd.DataFrame, config,\n        augment: bool = False, mode: str = 'train',\n        specs: Dict[int, np.ndarray] = all_spectrograms,\n        eeg_specs: Dict[int, np.ndarray] = all_eegs\n    ): \n        self.df = df\n        self.config = config\n        self.batch_size = self.config.BATCH_SIZE_TRAIN\n        self.augment = augment\n        self.mode = mode\n        self.spectograms = all_spectrograms\n        self.eeg_spectograms = eeg_specs\n        \n    def __len__(self):\n        \"\"\"\n        Denotes the number of batches per epoch.\n        \"\"\"\n        return len(self.df)\n        \n    def __getitem__(self, index):\n        \"\"\"\n        Generate one batch of data.\n        \"\"\"\n        X, y = self.__data_generation(index)\n        if self.augment:\n            X = self.__transform(X) \n        return torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)\n                        \n    def __data_generation(self, index):\n        \"\"\"\n        Generates data containing batch_size samples.\n        \"\"\"\n        X = np.zeros((128, 256, 8), dtype='float32')\n        y = np.zeros(6, dtype='float32')\n        img = np.ones((128,256), dtype='float32')\n        row = self.df.iloc[index]\n        if self.mode=='test': \n            r = 0\n        else: \n            r = int((row['min'] + row['max']) // 4)\n            \n        for region in range(4):\n            img = self.spectograms[row.spectogram_id][r:r+300, region*100:(region+1)*100].T\n            \n            # Log transform spectogram\n            img = np.clip(img, np.exp(-4), np.exp(8))\n            img = np.log(img)\n\n            # Standarize per image\n            ep = 1e-6\n            mu = np.nanmean(img.flatten())\n            std = np.nanstd(img.flatten())\n            img = (img-mu)/(std+ep)\n            img = np.nan_to_num(img, nan=0.0)\n            X[14:-14, :, region] = img[:, 22:-22] / 2.0\n            img = self.eeg_spectograms[row.eeg_id]\n            X[:, :, 4:] = img\n                \n            if self.mode != 'test':\n                y = row[label_cols].values.astype(np.float32)\n            \n        return X, y\n    \n    def __transform(self, img):\n        transforms = A.Compose([\n            A.HorizontalFlip(p=0.5),\n        ])\n        return transforms(image=img)['image']","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DataLoader","metadata":{}},{"cell_type":"code","source":"train_dataset = CustomDataset(train_df, config, mode=\"train\")\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=config.BATCH_SIZE_TRAIN,\n    shuffle=False,\n    num_workers=config.NUM_WORKERS, pin_memory=True, drop_last=True\n)\nX, y = train_dataset[0]\nprint(f\"X shape: {X.shape}\")\nprint(f\"y shape: {y.shape}\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualize DataLoader","metadata":{}},{"cell_type":"code","source":"if config.VISUALIZE:\n    ROWS = 2\n    COLS = 3\n    for (X, y) in train_loader:\n        plt.figure(figsize=(20,8))\n        for row in range(ROWS):\n            for col in range(COLS):\n                plt.subplot(ROWS, COLS, row*COLS + col+1)\n                t = y[row*COLS + col]\n                img = X[row*COLS + col, :, :, 0]\n                mn = img.flatten().min()\n                mx = img.flatten().max()\n                img = (img-mn)/(mx-mn)\n                plt.imshow(img)\n                tars = f'[{t[0]:0.2f}'\n                for s in t[1:]:\n                    tars += f', {s:0.2f}'\n                eeg = train_df.eeg_id.values[row*config.BATCH_SIZE_TRAIN + row*COLS + col]\n                plt.title(f'EEG = {eeg}\\nTarget = {tars}',size=12)\n                plt.yticks([])\n                plt.ylabel('Frequencies (Hz)',size=14)\n                plt.xlabel('Time (sec)',size=16)\n        plt.show()\n        break","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class CustomModel(nn.Module):\n    def __init__(self, config, num_classes: int = 6, pretrained: bool = True):\n        super(CustomModel, self).__init__()\n        self.USE_KAGGLE_SPECTROGRAMS = True\n        self.USE_EEG_SPECTROGRAMS = True\n        try:\n            self.model = timm.create_model(\n                config.MODEL,\n                pretrained=pretrained,\n                in_chans = config.IN_CHANS,\n                num_classes = 6,\n                drop_rate = 0.1,\n                drop_path_rate = 0.2,\n            )\n        except:\n                self.model = timm.create_model(\n                config.MODEL,\n                pretrained=False,\n                in_chans = config.IN_CHANS,\n                num_classes = 6,\n                drop_rate = 0.1,\n                drop_path_rate = 0.2,\n            )\n            \n        if config.FREEZE:\n            for i,(name, param) in enumerate(list(self.model.named_parameters())\\\n                                             [0:config.NUM_FROZEN_LAYERS]):\n                param.requires_grad = False\n\n        self.model = self.model\n\n    def __reshape_input(self, x):\n        \"\"\"\n        Reshapes input (128, 256, 8) -> (512, 512, 3) monotone image.\n        \"\"\" \n        # === Get spectograms ===\n        spectograms = [x[:, :, :, i:i+1] for i in range(4)]\n        spectograms = torch.cat(spectograms, dim=-1)\n        \n        # === Get EEG spectograms ===\n        eegs = [x[:, :, :, i:i+1] for i in range(4,8)]\n        eegs = torch.cat(eegs, dim=-1)\n        \n        # === Reshape (512,512,3) ===\n        if self.USE_KAGGLE_SPECTROGRAMS & self.USE_EEG_SPECTROGRAMS:\n            x = torch.cat([spectograms, eegs], dim=1)\n        elif self.USE_EEG_SPECTROGRAMS:\n            x = eegs\n        else:\n            x = spectograms\n        x = x.permute(0, 3, 1, 2)\n\n        \n        return x.to('cuda')\n    \n    def forward(self, x):\n        x = self.__reshape_input(x)\n        x = self.model(x)\n        return x","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Scheduler","metadata":{}},{"cell_type":"code","source":"from torch.optim.lr_scheduler import OneCycleLR\n\nEPOCHS = config.EPOCHS\nBATCHES = len(train_loader)\nsteps = []\nlrs = []\noptim_lrs = []\nmodel = CustomModel(config)\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\nscheduler = OneCycleLR(\n    optimizer,\n    max_lr=1e-3,\n    epochs=config.EPOCHS,\n    steps_per_epoch=len(train_loader),\n    pct_start=0.05,\n    anneal_strategy=\"cos\",\n    final_div_factor=100,\n)\nfor epoch in range(EPOCHS):\n    for batch in range(BATCHES):\n        scheduler.step()\n        lrs.append(scheduler.get_last_lr()[0])\n        steps.append(epoch * BATCHES + batch)\n\nmax_lr = max(lrs)\nmin_lr = min(lrs)\nprint(f\"Maximum LR: {max_lr} | Minimum LR: {min_lr}\")\nplt.figure()\nplt.plot(steps, lrs, label='OneCycle')\nplt.ticklabel_format(axis='y', style='sci', scilimits=(0,0))\nplt.xlabel(\"Step\")\nplt.ylabel(\"Learning Rate\")\nplt.show()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Function","metadata":{}},{"cell_type":"code","source":"import torch.nn.functional as F\n\n# === Reduction = \"mean\" ===\ncriterion = nn.KLDivLoss(reduction=\"mean\")\ny_pred = F.log_softmax(torch.randn(6, 2, requires_grad=True), dim=1)\ny_true = F.softmax(torch.rand(6, 2), dim=1)\nprint(f\"Predictions: {y_pred}\")\nprint(f\"Targets: {y_true}\")\noutput = criterion(y_pred, y_true)\nprint(f\"Output: {output}\")\n\nprint(\"\\n\", \"=\"*100, \"\\n\")\n\n# === Reduction = \"batchmean\" ===\ncriterion = nn.KLDivLoss(reduction=\"batchmean\")\ny_pred = F.log_softmax(torch.randn(2, 6, requires_grad=True), dim=1)\ny_true = F.softmax(torch.rand(2, 6), dim=1)\nprint(f\"Predictions: {y_pred}\")\nprint(f\"Targets: {y_true}\")\noutput = criterion(y_pred, y_true)\nprint(f\"Output: {output}\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train and Validation Functions","metadata":{}},{"cell_type":"code","source":"def train_epoch(train_loader, model, criterion, optimizer, epoch, scheduler, device):\n    \"\"\"One epoch training pass.\"\"\"\n    model.train() \n    criterion = nn.KLDivLoss(reduction=\"batchmean\")\n    scaler = torch.cuda.amp.GradScaler(enabled=config.AMP)\n    losses = AverageMeter()\n    start = end = time.time()\n    global_step = 0\n    \n    # ========== ITERATE OVER TRAIN BATCHES ============\n    with tqdm(train_loader, unit=\"train_batch\", desc='Train') as tqdm_train_loader:\n        for step, (X, y) in enumerate(tqdm_train_loader):\n            X = X.to(device)\n            y = y.to(device)\n            batch_size = y.size(0)\n            with torch.cuda.amp.autocast(enabled=config.AMP):\n                y_preds = model(X) \n                loss = criterion(F.log_softmax(y_preds, dim=1), y)\n            if config.GRADIENT_ACCUMULATION_STEPS > 1:\n                loss = loss / config.GRADIENT_ACCUMULATION_STEPS\n            losses.update(loss.item(), batch_size)\n            scaler.scale(loss).backward()\n            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.MAX_GRAD_NORM)\n\n            if (step + 1) % config.GRADIENT_ACCUMULATION_STEPS == 0:\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad()\n                global_step += 1\n                scheduler.step()\n            end = time.time()\n\n            # ========== LOG INFO ==========\n            if step % config.PRINT_FREQ == 0 or step == (len(train_loader)-1):\n                print('Epoch: [{0}][{1}/{2}] '\n                      'Elapsed {remain:s} '\n                      'Loss: {loss.avg:.4f} '\n                      'Grad: {grad_norm:.4f}  '\n                      'LR: {lr:.8f}  '\n                      .format(epoch+1, step, len(train_loader), \n                              remain=timeSince(start, float(step+1)/len(train_loader)),\n                              loss=losses,\n                              grad_norm=grad_norm,\n                              lr=scheduler.get_last_lr()[0]))\n\n    return losses.avg\n\n\ndef valid_epoch(valid_loader, model, criterion, device):\n    model.eval()\n    softmax = nn.Softmax(dim=1)\n    losses = AverageMeter()\n    prediction_dict = {}\n    preds = []\n    start = end = time.time()\n    with tqdm(valid_loader, unit=\"valid_batch\", desc='Validation') as tqdm_valid_loader:\n        for step, (X, y) in enumerate(tqdm_valid_loader):\n            X = X.to(device)\n            y = y.to(device)\n            batch_size = y.size(0)\n            with torch.no_grad():\n                y_preds = model(X)\n                loss = criterion(F.log_softmax(y_preds, dim=1), y)\n            if config.GRADIENT_ACCUMULATION_STEPS > 1:\n                loss = loss / config.GRADIENT_ACCUMULATION_STEPS\n            losses.update(loss.item(), batch_size)\n            y_preds = softmax(y_preds)\n            preds.append(y_preds.to('cpu').numpy())\n            end = time.time()\n\n            # ========== LOG INFO ==========\n            if step % config.PRINT_FREQ == 0 or step == (len(valid_loader)-1):\n                print('EVAL: [{0}/{1}] '\n                      'Elapsed {remain:s} '\n                      'Loss: {loss.avg:.4f} '\n                      .format(step, len(valid_loader),\n                              remain=timeSince(start, float(step+1)/len(valid_loader)),\n                              loss=losses))\n                \n    prediction_dict[\"predictions\"] = np.concatenate(preds)\n    return losses.avg, prediction_dict","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Loop","metadata":{}},{"cell_type":"code","source":"def train_loop(df, fold):\n    \n    LOGGER.info(f\"========== Fold: {fold} training ==========\")\n\n    # ======== SPLIT ==========\n    train_folds = df[df['fold'] != fold].reset_index(drop=True)\n    valid_folds = df[df['fold'] == fold].reset_index(drop=True)\n    \n    # ======== DATASETS ==========\n    train_dataset = CustomDataset(train_folds, config, mode=\"train\", augment=True)\n    valid_dataset = CustomDataset(valid_folds, config, mode=\"train\", augment=False)\n    \n    # ======== DATALOADERS ==========\n    train_loader = DataLoader(train_dataset,\n                              batch_size=config.BATCH_SIZE_TRAIN,\n                              shuffle=False,\n                              num_workers=config.NUM_WORKERS, pin_memory=True, drop_last=True)\n    valid_loader = DataLoader(valid_dataset,\n                              batch_size=config.BATCH_SIZE_VALID,\n                              shuffle=False,\n                              num_workers=config.NUM_WORKERS, pin_memory=True, drop_last=False)\n    \n    # ======== MODEL ==========\n    model = CustomModel(config)\n    model.to(device)\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=0.1, weight_decay=config.WEIGHT_DECAY)\n    scheduler = OneCycleLR(\n        optimizer,\n        max_lr=1e-3,\n        epochs=config.EPOCHS,\n        steps_per_epoch=len(train_loader),\n        pct_start=0.1,\n        anneal_strategy=\"cos\",\n        final_div_factor=100,\n    )\n\n    # ======= LOSS ==========\n    criterion = nn.KLDivLoss(reduction=\"batchmean\")\n    \n    best_loss = np.inf\n    # ====== ITERATE EPOCHS ========\n    for epoch in range(config.EPOCHS):\n        start_time = time.time()\n\n        # ======= TRAIN ==========\n        avg_train_loss = train_epoch(train_loader, model, criterion, optimizer, epoch, scheduler, device)\n\n        # ======= EVALUATION ==========\n        avg_val_loss, prediction_dict = valid_epoch(valid_loader, model, criterion, device)\n        predictions = prediction_dict[\"predictions\"]\n        \n        # ======= SCORING ==========\n        elapsed = time.time() - start_time\n\n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_train_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n        \n        if avg_val_loss < best_loss:\n            best_loss = avg_val_loss\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Loss: {best_loss:.4f} Model')\n            torch.save({'model': model.state_dict(),\n                        'predictions': predictions},\n                        paths.OUTPUT_DIR + f\"/{config.MODEL.replace('/', '_')}_fold_{fold}_best.pth\")\n\n    predictions = torch.load(paths.OUTPUT_DIR + f\"/{config.MODEL.replace('/', '_')}_fold_{fold}_best.pth\", \n                             map_location=torch.device('cpu'))['predictions']\n    valid_folds[target_preds] = predictions\n\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return valid_folds","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Full Data","metadata":{}},{"cell_type":"code","source":"def train_loop_full_data(df):\n    train_dataset = CustomDataset(df, config, mode=\"train\", augment=True)\n    train_loader = DataLoader(train_dataset,\n                              batch_size=config.BATCH_SIZE_TRAIN,\n                              shuffle=False,\n                              num_workers=config.NUM_WORKERS, pin_memory=True, drop_last=True)\n    model = CustomModel(config)\n    model.to(device)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=0.1, weight_decay=config.WEIGHT_DECAY)\n    scheduler = OneCycleLR(\n        optimizer,\n        max_lr=1e-3,\n        epochs=config.EPOCHS,\n        steps_per_epoch=len(train_loader),\n        pct_start=0.1,\n        anneal_strategy=\"cos\",\n        final_div_factor=100,\n    )\n    criterion = nn.KLDivLoss(reduction=\"batchmean\")\n    best_loss = np.inf\n    for epoch in range(config.EPOCHS):\n        start_time = time.time()\n        avg_train_loss = train_epoch(train_loader, model, criterion, optimizer, epoch, scheduler, device)\n        elapsed = time.time() - start_time\n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_train_loss:.4f}  time: {elapsed:.0f}s')\n        torch.save(\n            {'model': model.state_dict()},\n            paths.OUTPUT_DIR + f\"/{config.MODEL.replace('/', '_')}_epoch_{epoch}.pth\")\n    torch.cuda.empty_cache()\n    gc.collect()\n    return _","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def get_result(oof_df):\n    kl_loss = nn.KLDivLoss(reduction=\"batchmean\")\n    labels = torch.tensor(oof_df[label_cols].values)\n    preds = torch.tensor(oof_df[target_preds].values)\n    preds = F.log_softmax(preds, dim=1)\n    result = kl_loss(preds, labels)\n    return result\n\nif not config.TRAIN_FULL_DATA:\n    oof_df = pd.DataFrame()\n    for fold in range(config.FOLDS):\n        if fold in [0, 1, 2, 3, 4]:\n            _oof_df = train_loop(train_df, fold)\n            oof_df = pd.concat([oof_df, _oof_df])\n            LOGGER.info(f\"========== Fold {fold} result: {get_result(_oof_df)} ==========\")\n            print(f\"========== Fold {fold} result: {get_result(_oof_df)} ==========\")\n    oof_df = oof_df.reset_index(drop=True)\n    LOGGER.info(f\"========== CV: {get_result(oof_df)} ==========\")\n    oof_df.to_csv(paths.OUTPUT_DIR + '/oof_df.csv', index=False)\nelse:\n    train_loop_full_data(train_df)","metadata":{},"execution_count":null,"outputs":[]}]}