{"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"![](https://storage.googleapis.com/kaggle-competitions/kaggle/23249/logos/header.png)","metadata":{}},{"cell_type":"markdown","source":"# Motivation","metadata":{}},{"cell_type":"markdown","source":"Motivation: I wanted to implement a notebook where it would be easy enough to run as just editing the config (first cell) and changing models (if you want to) in order to save experiment time. Another major point for me was being able to run it on both the GPU and TPU.\n\nThe following ideas are implemented in this Notebook:\n* Optimizers : Adam, AdamW, AdamP, Ranger\n* Schedulers : ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR' ,'GradualWarmupSchedulerV2']\n* Models    : ['deit_base_patch16_384','vit_large_patch16_384','tf_efficientnet_b4_ns','resnext50_32x4d'] it also works for different effnets and resnexts\n* Loss Fns   : ['BCEwithLogits']\n* Augmentations: HeavyAug+Cutout(Train),LightAug (only rotations), HeavyAug, AutoAugment(imagenet)\n\nAny feedback is appreciated!\nBe sure to also checkout:\n* Inference: tbd","metadata":{}},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Configurations\n# ====================================================\nimport os\nclass CFG:\n    DEBUG = False\n    \n    #Model Params\n    device = 'GPU' #['CPU','GPU','TPU']\n    N_FOLDS = 5\n    MODEL_NAME = 'tf_efficientnet_b4_ns' # Recommended : ['tf_efficientnet_b4_ns','resnext50_32x4d']\n    pretrained = True   \n    EPOCHS = 10 if not DEBUG else 3 # more is definitely plausible\n    TRAIN_FOLDS = [0] if DEBUG else [i for i in range(N_FOLDS)] #Folds to be Trained\n    N_CLASSES = 1\n    in_channels = 1\n    \n    scheduler_name = 'CosineAnnealingWarmRestarts'\n    # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'GradualWarmupSchedulerV2']\n    scheduler_update = 'epoch' #['batch','epoch']\n    criterion_name = 'BCEWithLogitsLoss'\n    # ['BCEWithLogitsLoss','ROC-STAR']\n    optimizer_name = 'AdamP' #['Adam','AdamW','AdamP','Ranger'] -> AdamP doesn't work on TPUs\n    LR_RAMPUP_EPOCHS = 1\n    LR_SUSTAIN_EPOCHS = 0\n    \n    FREEZE = False #If you fine tune after START_FREEZE epochs\n    START_FREEZE = 8\n    \n    #Training Params\n    BATCH_SIZE = 1024\n    \n    LR = 2e-4\n    LR_START =1e-4\n    LR_MIN = 8e-7\n    weight_decay = 0\n    eps = 1e-8\n    PATIENCE = 1\n      \n    #CosineAnnealingWarmRestarts\n    T_0 = EPOCHS\n    \n    #CosineAnnealingLR\n    T_max = EPOCHS\n    \n    NUM_WORKERS = 4\n    \n    model_print = False #If the model architecture is printed\n    tqdm = True #If training bar is shown\n    \n    #n_procs = number of replicas -> TPU\n    n_procs = 8 #You can set it to 1 and run a TPU as a GPU if you want\n    SEED = 42\n    saved_models = {}","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:04.409342Z","iopub.execute_input":"2021-07-05T00:07:04.409785Z","iopub.status.idle":"2021-07-05T00:07:04.422878Z","shell.execute_reply.started":"2021-07-05T00:07:04.409707Z","shell.execute_reply":"2021-07-05T00:07:04.421529Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Installations","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Required Installations\n# ====================================================\n!pip install timm\n\nif CFG.device == 'TPU':\n    import os\n    !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n    !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev\n    os.system('export XLA_USE_BF16=1')\n    os.system('export XLA_TENSOR_ALLOCATOR_MAXSIZE=100000000')\n    import torch_xla.core.xla_model as xm\n    import torch_xla.distributed.parallel_loader as pl\n    import torch_xla.distributed.xla_multiprocessing as xmp\n    import ignite.distributed as idist\n    #CFG.LR = CFG.LR * CFG.n_procs\n    CFG.BATCH_SIZE = CFG.BATCH_SIZE / CFG.n_procs\n    \nif CFG.optimizer_name == 'Ranger':\n    !pip install --quiet '../input/pytorch-ranger'\nelif CFG.optimizer_name == 'AdamP':\n    !pip install adamp","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2021-07-05T00:07:04.424461Z","iopub.execute_input":"2021-07-05T00:07:04.424971Z","iopub.status.idle":"2021-07-05T00:07:21.131930Z","shell.execute_reply.started":"2021-07-05T00:07:04.424933Z","shell.execute_reply":"2021-07-05T00:07:21.130981Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Library","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Library\n# ====================================================\n\nimport random\nimport math\nimport time\n\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom matplotlib.pyplot import imread\nimport numpy as np\nimport cv2\nfrom sklearn.model_selection import GroupKFold, StratifiedKFold,KFold\nfrom sklearn.metrics import accuracy_score,roc_auc_score\n\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom torchvision import models as tvmodels\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, utils\nimport torch.nn.functional as F\nfrom tqdm import tqdm\nimport timm\n\n\n\nimport albumentations as A\nfrom albumentations import Compose\nfrom albumentations.pytorch import ToTensorV2\nimport seaborn as sns\n\n\n\nfrom PIL import Image, ImageOps, ImageEnhance, ImageChops\n\n\nif CFG.optimizer_name == 'Ranger':\n    from pytorch_ranger import Ranger\nelif CFG.optimizer_name == 'AdamP':\n    from adamp import AdamP\n    \nif CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n    from warmup_scheduler import GradualWarmupScheduler\n\n\n# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport librosa","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2021-07-05T00:07:21.135448Z","iopub.execute_input":"2021-07-05T00:07:21.135706Z","iopub.status.idle":"2021-07-05T00:07:26.135036Z","shell.execute_reply.started":"2021-07-05T00:07:21.135678Z","shell.execute_reply":"2021-07-05T00:07:26.134067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\n# credit to https://www.kaggle.com/thedrcat/g2net-fastai-resnet34-starter-mel\ndef id2path(id, is_test):\n    a, b, c = id[0], id[1], id[2]\n    if is_test: return f'../input/g2net-gravitational-wave-detection/test/{a}/{b}/{c}/{id}.npy'\n    return f'../input/g2net-gravitational-wave-detection/train/{a}/{b}/{c}/{id}.npy'\n\ndef get_score(y_true, y_pred):\n    score = roc_auc_score(y_true, y_pred)\n    return score\n\ndef seed_everything(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\n    \nif CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n    class GradualWarmupSchedulerV2(GradualWarmupScheduler):\n        def __init__(self, optimizer = None, multiplier = CFG.LR/CFG.LR_START, total_epoch = CFG.LR_RAMPUP_EPOCHS, after_scheduler=None):\n            super(GradualWarmupSchedulerV2, self).__init__(optimizer, multiplier, total_epoch, after_scheduler)\n            self.after_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = CFG.T_0 - CFG.LR_RAMPUP_EPOCHS, T_mult=1, eta_min=CFG.LR_MIN, last_epoch=-1)\n        def get_lr(self):\n            if self.last_epoch > self.total_epoch:\n                if self.after_scheduler:\n                    if not self.finished:\n                        self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]\n                        self.finished = True\n                    return self.after_scheduler.get_lr()\n                return [base_lr * self.multiplier for base_lr in self.base_lrs]\n            if self.multiplier == 1.0:\n                return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]\n            else:\n                return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]\n  \n    \ndef GetScheduler(scheduler_name,optimizer,batches):\n    #['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'GradualWarmupSchedulerV2']\n    if scheduler_name == 'OneCycleLR':\n        return torch.optim.lr_scheduler.OneCycleLR(optimizer,max_lr = 1e-2,epochs = CFG.EPOCHS,steps_per_epoch = batches+1,pct_start = 0.1)\n    elif scheduler_name == 'CosineAnnealingWarmRestarts':\n        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = CFG.T_0, T_mult=1, eta_min=CFG.LR_MIN, last_epoch=-1)\n    elif scheduler_name == 'CosineAnnealingLR':\n        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = CFG.T_max, eta_min=0, last_epoch=-1)\n    elif scheduler_name == 'ReduceLROnPlateau':\n        return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,factor=0.1, patience=1, threshold=0.0001, cooldown=0, min_lr=CFG.LR_MIN, eps=CFG.eps)\n    elif scheduler_name == 'GradualWarmupSchedulerV2':\n        return GradualWarmupSchedulerV2(optimizer=optimizer)\n    \ndef GetOptimizer(optimizer_name,parameters):\n    #['Adam','Ranger']\n    if optimizer_name == 'Adam':\n        if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n            return torch.optim.Adam(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay, amsgrad=False)\n        else:\n            return torch.optim.Adam(parameters, lr=CFG.LR, weight_decay=CFG.weight_decay, amsgrad=False)\n    elif optimizer_name == 'AdamW':\n        if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n            return torch.optim.AdamW(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay, amsgrad=False)\n        else:\n            return torch.optim.Adam(parameters, lr=CFG.LR, weight_decay=CFG.weight_decay, amsgrad=False)\n    elif optimizer_name == 'AdamP':\n        if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n            return AdamP(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay)\n        else:\n            return AdamP(parameters, lr=CFG.LR, weight_decay=CFG.weight_decay)\n    elif optimizer_name == 'Ranger':\n        return Ranger(parameters,lr = CFG.LR,alpha = 0.5, k = 6,N_sma_threshhold = 5,betas = (0.95,0.999),eps=CFG.eps,weight_decay=CFG.weight_decay)\n\ndef print_scheduler(scheduler = None,scheduler_update = CFG.scheduler_update,optimizer = None, batches = -1, epochs = -1, model = None):\n    lrs = []\n    if scheduler_update == 'epoch':\n        for epoch in range(epochs):\n            scheduler.step(epoch)\n            lrs.append(optimizer.param_groups[0][\"lr\"])\n        plt.figure(figsize=(15,4))\n        plt.plot(lrs)\n    elif scheduler_update == 'batch':\n        for epoch in range(epochs):\n            for batch in range(batches):\n                scheduler.step()\n                lrs.append(optimizer.param_groups[0][\"lr\"])\n        plt.figure(figsize=(15,4))\n        plt.plot(lrs)\n    \nSEED = CFG.SEED\nseed_everything(SEED)  \nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-07-05T00:07:26.136850Z","iopub.execute_input":"2021-07-05T00:07:26.137117Z","iopub.status.idle":"2021-07-05T00:07:26.205567Z","shell.execute_reply.started":"2021-07-05T00:07:26.137088Z","shell.execute_reply":"2021-07-05T00:07:26.204689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Transforms","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Transformations\n# ====================================================\n# credit to https://www.kaggle.com/thedrcat/g2net-fastai-resnet34-starter-mel\ndef sample_to_mel(_id, is_test):\n    x = np.load(id2path(_id, is_test))\n    spectrogram = []\n    for i in range(3):\n        mel = librosa.feature.melspectrogram(x[i]/x[i].max(), \n                                                sr=2048,\n                                                n_mels=16,\n                                               hop_length=16)\n        mel = mel[:,:256]\n        mel = librosa.power_to_db(mel).astype(np.float32)\n        mel = mel.reshape(64,64)\n        spectrogram.append(mel)\n    spectrogram = np.stack(spectrogram)\n    return spectrogram","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:26.206906Z","iopub.execute_input":"2021-07-05T00:07:26.207236Z","iopub.status.idle":"2021-07-05T00:07:26.216021Z","shell.execute_reply.started":"2021-07-05T00:07:26.207199Z","shell.execute_reply":"2021-07-05T00:07:26.215120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Datasets","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Datasets\n# ====================================================\ndef retrieve_df(df,name,idx):\n    series = df[name].iloc[idx]\n    series.reset_index(drop=True,inplace=True)\n    \n    return series\n\nclass G2NetDataset(torch.utils.data.Dataset):\n    def __init__(self, features, target, is_test=False,transformed = True,file_names = []):\n        self.features,self.target,self.is_test,self.transformed,self.file_names = features,target,is_test,transformed,file_names\n        \n    def __getitem__(self, i):\n        if not self.transformed:\n            img = self.features.loc[i]\n            img = sample_to_mel(img, self.is_test) if self.transform else img\n            if self.is_test:\n                return (torch.tensor(img, dtype=torch.float), torch.tensor(0, dtype=torch.float))\n            else:\n                tgt = self.target.loc[i]\n                return (torch.tensor(img, dtype=torch.float), torch.tensor(tgt, dtype=torch.float))\n        else:\n            #adapted from https://www.kaggle.com/yasufuminakama/g2net-efficientnet-b0-baseline-training\n            file_path = self.file_names[i]\n            image = np.load(file_path)\n            image = image[np.newaxis,:,:]\n            image = torch.from_numpy(image).float()\n            tgt = self.target.loc[i]\n            return (image, torch.tensor(tgt, dtype=torch.float))\n    \n    def __len__(self): return len(self.target)\n    ","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:26.217265Z","iopub.execute_input":"2021-07-05T00:07:26.217864Z","iopub.status.idle":"2021-07-05T00:07:26.228537Z","shell.execute_reply.started":"2021-07-05T00:07:26.217825Z","shell.execute_reply":"2021-07-05T00:07:26.227592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CV Split","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# CV Split\n# ====================================================\n#adapted from https://www.kaggle.com/yasufuminakama/g2net-efficientnet-b0-baseline-training\ntrain_df = pd.read_csv(\"../input/g2net-gravitational-wave-detection/training_labels.csv\")\ntest_df = pd.read_csv('../input/g2net-gravitational-wave-detection/sample_submission.csv')\n\ndef get_train_file_path(image_id):\n    return \"../input/g2net-n-mels-128-train-images/{}.npy\".format(image_id)\n\ndef get_test_file_path(image_id):\n    return \"../input/g2net-n-mels-128-test-images/{}.npy\".format(image_id)\n\ntrain_df['file_path'] = train_df['id'].apply(get_train_file_path)\ntest_df['file_path'] = test_df['id'].apply(get_test_file_path)\n\n#modified from https://www.kaggle.com/thedrcat/g2net-fastai-resnet34-starter-mel\nif CFG.DEBUG:\n    train_df = train_df.sample(frac=0.01).reset_index(drop=True)\nskf = StratifiedKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\nskf.get_n_splits(np.arange(train_df.shape[0]), train_df['target'])\nfolds = [(idxT,idxV) for i,(idxT,idxV) in enumerate(skf.split(np.arange(train_df.shape[0]), train_df['target']))]\n    \nsns.countplot(data=train_df, x=\"target\")","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:26.231192Z","iopub.execute_input":"2021-07-05T00:07:26.231618Z","iopub.status.idle":"2021-07-05T00:07:27.358125Z","shell.execute_reply.started":"2021-07-05T00:07:26.231580Z","shell.execute_reply":"2021-07-05T00:07:27.357363Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Model\n# ====================================================\nclass G2Net(nn.Module):\n    def __init__(self, model_name=CFG.MODEL_NAME, pretrained=CFG.pretrained,in_chans = CFG.in_channels):\n        super().__init__()\n        self.model_name = model_name\n        if model_name == 'deit_base_patch16_224' or model_name == 'deit_base_patch16_384':\n            self.model = torch.hub.load('facebookresearch/deit:main', model_name, pretrained=pretrained, in_chans=in_chans)\n        else:\n            self.model = timm.create_model(model_name, pretrained=pretrained, in_chans=in_chans)\n        if 'efficientnet' in model_name:\n            self.n_features = self.model.classifier.in_features\n            self.model.classifier = nn.Linear(self.n_features, CFG.N_CLASSES)\n        elif model_name == 'vit_large_patch16_384' or model_name == 'deit_base_patch16_224' or model_name == 'deit_base_patch16_384':\n            self.n_features = self.model.head.in_features\n            self.model.head = nn.Linear(self.n_features, CFG.N_CLASSES)\n        elif 'resnext' in model_name:\n            self.n_features = self.model.fc.in_features\n            self.model.fc = nn.Linear(self.n_features, CFG.N_CLASSES)\n        \n    def forward(self, x):\n        return self.model(x)\n    \n    def freeze(self):\n        # To freeze the residual layers\n        for param in self.model.parameters():\n            param.requires_grad = False\n            \n        if 'efficientnet' in self.model_name:\n            for param in self.model.classifier.parameters():\n                param.requires_grad = True\n        elif self.model_name == 'vit_large_patch16_384' or 'deit_base_patch16_224':\n            for param in self.model.head.parameters():\n                param.requires_grad = True\n        elif 'resnext' in self.model_name:\n            for param in self.model.fc.parameters():\n                param.requires_grad = True\n            \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.model.parameters():\n            param.requires_grad = True\nmodel = G2Net()","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:27.359553Z","iopub.execute_input":"2021-07-05T00:07:27.359872Z","iopub.status.idle":"2021-07-05T00:07:31.899069Z","shell.execute_reply.started":"2021-07-05T00:07:27.359837Z","shell.execute_reply":"2021-07-05T00:07:31.898233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training Loop","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Training Loop\n# ====================================================\ndef train_one_epoch(model,optimizer,scheduler,scaler,train_loader,criterion,batches,epoch,DEVICE):   \n    tr_loss = 0.0\n    scores = 0.0\n    trn_epoch_result = dict()\n    model.train()\n    if CFG.tqdm:\n        progress = tqdm(enumerate(train_loader), desc=\"Loss: \", total=len(train_loader))\n    else:\n        progress = enumerate(train_loader)\n    for i, (images,labels) in progress:\n        images = images.to(DEVICE)\n        labels = labels.to(DEVICE)\n        \n        optimizer.zero_grad()\n        \n        if CFG.device == 'TPU':\n            logits = model(images).reshape((-1))\n            loss = criterion(logits, labels)\n            loss.backward()\n            xm.optimizer_step(optimizer)\n        else:\n            with autocast():\n                logits = model(images).reshape((-1))\n                loss = criterion(logits, labels)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n\n        \n        if CFG.scheduler_update == 'batch':\n            if not CFG.scheduler_name == 'OneCycleLR':\n                scheduler.step(epoch + i/len(train_loader))\n            else:\n                scheduler.step()\n\n        tr_loss += loss.detach().item()\n        \n        if CFG.tqdm:\n            trn_epoch_result['Epoch'] = epoch\n            trn_epoch_result['train_loss'] = round(tr_loss/(i+1), 4)\n            trn_epoch_result['LR'] = round(optimizer.param_groups[0][\"lr\"],7)\n\n            progress.set_description(str(trn_epoch_result))\n        else:\n            print(tr_loss/(i+1))\n    if CFG.scheduler_update == 'epoch':\n            scheduler.step(epoch+1)\n        \ndef val_one_epoch(model,DEVICE,loader,val_criterion,epoch,get_output = False):\n    val_loss = 0.0\n    scores = 0.0\n    model.eval()\n    val_progress = tqdm(enumerate(loader), desc=\"Loss: \", total=len(loader))\n    with torch.no_grad():\n        for i, (images,labels) in val_progress:\n            images = images.to(DEVICE)\n            labels = labels.to(DEVICE)\n\n            logits = model(images).reshape((-1))\n            val_loss_value = val_criterion(logits,labels)\n            val_loss += val_loss_value.detach().item()\n\n            scores += get_score(labels.to('cpu').numpy(),logits.sigmoid().to('cpu').numpy())\n\n            val_epoch_result = dict()\n            val_epoch_result['Epoch'] = epoch\n            val_epoch_result['val_loss'] = round(val_loss/(i+1), 4)\n\n            val_epoch_result['val_acc'] = round(scores/(i+1), 4)\n            val_progress.set_description(str(val_epoch_result))\n    if get_output:\n        return val_loss/len(loader),scores/len(loader)\n        \ndef model_train():\n    if CFG.device == 'TPU':\n        DEVICE = xm.xla_device()\n    else:\n        DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        \n    for fold,(idxT, idxV) in enumerate(folds):\n        if fold not in CFG.TRAIN_FOLDS:\n            continue\n        #if xm.is_master_ordinal():\n        #    xm.master_print(fold)\n        #______INSTANTIATE TRAINING DATASETS_____\n        x_train = retrieve_df(train_df,'id',idxT)\n        y_train = retrieve_df(train_df,'target',idxT)\n        x_val = retrieve_df(train_df,'id',idxV)\n        y_val = retrieve_df(train_df,'target',idxV)\n        train_set = G2NetDataset(x_train,y_train, is_test=False,transformed = True,file_names = retrieve_df(train_df,'file_path',idxT))\n        val_set = G2NetDataset(x_val,y_val, is_test=False,transformed = True,file_names = retrieve_df(train_df,'file_path',idxV))\n        if CFG.device == 'TPU':\n            train_sampler = torch.utils.data.distributed.DistributedSampler(\n                train_set,\n                num_replicas=xm.xrt_world_size(),\n                rank=xm.get_ordinal(),\n                shuffle=True)\n            train_loader = DataLoader(train_set, batch_size=CFG.BATCH_SIZE, sampler=train_sampler,drop_last=True, num_workers=CFG.NUM_WORKERS)\n        \n            val_sampler = torch.utils.data.distributed.DistributedSampler(\n                val_set,\n                num_replicas=xm.xrt_world_size(),\n                rank=xm.get_ordinal(),\n                shuffle=False)\n            val_loader = DataLoader(val_set, batch_size=CFG.BATCH_SIZE, sampler=val_sampler,drop_last=True, num_workers=CFG.NUM_WORKERS)\n            scaler = None\n        else:\n            print(f\"Start of Fold {fold}\")\n            train_loader = DataLoader(train_set, batch_size=CFG.BATCH_SIZE, shuffle=True,drop_last=True, num_workers=CFG.NUM_WORKERS,pin_memory = True)\n            val_loader = DataLoader(val_set, batch_size=CFG.BATCH_SIZE, shuffle=False,drop_last=True, num_workers=CFG.NUM_WORKERS,pin_memory = True)\n            scaler = GradScaler()\n            \n        batches = len(train_loader)\n        val_batches = len(val_loader)\n\n        #INSTANTIATE FOLD MODEL\n        if CFG.model is None:\n            if CFG.device == 'TPU':\n                if xm.is_master_ordinal(local=True):\n                    CFG.model = G2Net(model_name=CFG.MODEL_NAME, pretrained=True)\n                xm.rendezvous('ModelDone')\n            else:\n                CFG.model = G2Net(model_name=CFG.MODEL_NAME, pretrained=True)\n        model = CFG.model.to(DEVICE)\n\n        criterion = CFG.criterion.to(DEVICE)\n        val_criterion = CFG.val_criterion.to(DEVICE)\n\n        optimizer = GetOptimizer(CFG.optimizer_name, model.parameters())\n        scheduler = GetScheduler(CFG.scheduler_name, optimizer,batches)\n        \n        saved_model = None\n        best_val_acc = 0.0\n        best_val_loss = 1e3\n        fold_patience = 0.0\n        for epoch in range(CFG.EPOCHS):\n            if epoch >= CFG.START_FREEZE and CFG.FREEZE:\n                print('Model Frozen -> Train Classifier Only')\n                info = torch.load(saved_model,map_location = torch.device(DEVICE))\n                model.load_state_dict(info)\n                model.freeze()\n                \n                CFG.FREEZE = False\n            #______TRAINING______\n            if CFG.device == 'TPU':\n                para_train_loader = pl.ParallelLoader(train_loader, [DEVICE])\n                train_one_epoch(model,optimizer,scheduler,scaler,para_train_loader.per_device_loader(DEVICE),criterion,batches,epoch,DEVICE)\n                del para_train_loader\n                gc.collect()\n            else:\n                train_one_epoch(model,optimizer,scheduler,scaler,train_loader,criterion,batches,epoch,DEVICE)\n            \n            #______VALIDATION_______\n            if CFG.device == 'TPU':\n                para_val_loader = pl.ParallelLoader(val_loader, [DEVICE])\n                val_loss, val_acc = val_one_epoch(model,DEVICE,para_val_loader.per_device_loader(DEVICE),val_criterion,epoch,get_output = True)\n                del para_val_loader\n                gc.collect()\n                val_loss = np.sum(idist.all_gather(torch.tensor(val_loss)).to('cpu').numpy())/CFG.n_procs\n                val_acc = np.sum(idist.all_gather(torch.tensor(val_acc)).to('cpu').numpy())/CFG.n_procs\n                xm.master_print(f'Fold Ended at {round(val_acc, 4)} val accuracy')\n            else:\n                val_loss, val_acc = val_one_epoch(model,DEVICE,val_loader,val_criterion,epoch,get_output = True)\n            \n            if val_acc > best_val_acc:\n                fold_patience = 0\n                best_val_loss = val_loss/val_batches\n                best_val_acc = val_acc\n                if CFG.device == 'TPU':\n                    xm.save(CFG.model.state_dict(),\n                                f'{CFG.MODEL_NAME}_f{fold}_TPU_b{round(best_val_acc, 4)}.pth')\n                    if saved_model is not None:\n                        try:\n                            os.remove(\"./\"+saved_model)\n                        except:\n                            a = 1\n                    xm.master_print(f'Model Saved at {round(best_val_acc, 5)} accuracy')\n                    saved_model = f'{CFG.MODEL_NAME}_f{fold}_TPU_b{round(best_val_acc, 4)}.pth'\n                    CFG.saved_models[fold] = round(best_val_acc, 4)\n                else:\n                    torch.save(model.state_dict(),\n                            f'{CFG.MODEL_NAME}_f{fold}_b{round(best_val_acc, 4)}.pth')\n                    if saved_model is not None:\n                        try:\n                            os.remove(\"./\"+saved_model)\n                        except:\n                            a = 1\n                    saved_model = f'{CFG.MODEL_NAME}_f{fold}_b{round(best_val_acc, 4)}.pth'\n                    CFG.saved_models[fold] = round(best_val_acc, 4)\n                    print(f'Model Saved at {round(best_val_acc, 5)} accuracy')\n            else:\n                fold_patience += 1\n                if fold_patience >= CFG.PATIENCE:\n                    print(f'Early stopping due to model not improving for {CFG.PATIENCE} epochs')\n                    CFG.model = None\n                    break\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n        CFG.model = None\n                \ndef _map_fn(index,flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = model_train()","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:31.900498Z","iopub.execute_input":"2021-07-05T00:07:31.900819Z","iopub.status.idle":"2021-07-05T00:07:31.938204Z","shell.execute_reply.started":"2021-07-05T00:07:31.900784Z","shell.execute_reply":"2021-07-05T00:07:31.937456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run","metadata":{}},{"cell_type":"code","source":"CFG.model = model\nCFG.criterion = nn.BCEWithLogitsLoss()\nCFG.val_criterion = nn.BCEWithLogitsLoss()\nif CFG.device == 'TPU':\n    FLAGS = {}\n    xmp.spawn(_map_fn, args=(FLAGS,), nprocs=CFG.n_procs, start_method='fork')\nelse:\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = model_train()","metadata":{"execution":{"iopub.status.busy":"2021-07-05T00:07:31.941217Z","iopub.execute_input":"2021-07-05T00:07:31.941715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference - Ensemble","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Inference\n# ====================================================\ntest_df['file_path'] = test_df['id'].apply(get_test_file_path)\nx_test = test_df['id']\ny_test = test_df['target']\nfolds_preds = np.zeros((len(CFG.TRAIN_FOLDS),len(test_df)))\nfor fold in CFG.TRAIN_FOLDS:\n    if CFG.device == 'TPU':\n        info = torch.load(f'{CFG.MODEL_NAME}_f{fold}_TPU_b{CFG.saved_models[fold]}.pth',map_location = torch.device(DEVICE))\n    else:\n        info = torch.load(f'{CFG.MODEL_NAME}_f{fold}_b{CFG.saved_models[fold]}.pth',map_location = torch.device(DEVICE))\n    model.load_state_dict(info)\n    test_set = G2NetDataset(x_test,y_test, is_test=True,transformed = True,file_names = test_df['file_path'])\n    test_loader = DataLoader(test_set, batch_size=CFG.BATCH_SIZE, shuffle=False,drop_last=False, num_workers=CFG.NUM_WORKERS,pin_memory = True)\n    test_progress = tqdm(enumerate(test_loader), desc=\"Loss: \", total=len(test_loader))\n    model.eval()\n    preds_arr = []\n    with torch.no_grad():\n        for i, (images,labels) in test_progress:\n            images = images.to(DEVICE)\n\n            logits = model(images).reshape((-1))\n            preds = logits.sigmoid().to('cpu').numpy()\n            preds_arr.append(preds)\n\n    folds_preds[fold,:] = np.concatenate(preds_arr)\nfolds_preds = np.mean(folds_preds,axis = 0)\ntest_df['target'] = folds_preds\ntest_df.drop('file_path',axis = 1,inplace=True)\ntest_df.to_csv('output.csv', index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}