{"cells":[{"metadata":{},"cell_type":"markdown","source":"![](https://storage.googleapis.com/kaggle-competitions/kaggle/13836/logos/header.png?t=2020-10-01-17-22-54)"},{"metadata":{},"cell_type":"markdown","source":"Motivation: I wanted to implement a notebook where it would be easy to test out different ideas that occured along the competition, and for it to be easy enough as just editing the config (first cell) and transformations (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   : ['CrossEntropyLoss', LabelSmoothing', 'FocalLoss','FocalCosineLoss', 'SymmetricCrossEntropyLoss', 'BiTemperedLoss', 'TaylorCrossEntropyLoss', 'TaylorSmoothedLoss']\n* Augmentations: HeavyAug+Cutout(Train),LightAug (only rotations), HeavyAug, AutoAugment(imagenet)\n\nAny feedback is appreciated!\nBe sure to also checkout:\n* Inference: https://www.kaggle.com/capiru/cassavanet-inference-tta-easy-submission"},{"metadata":{},"cell_type":"markdown","source":"Credits:\n\n* AdamP: https://www.kaggle.com/seriousran/adam-vs-adamp-iclr-2021\n* Ranger:\n* Loss Functions: https://www.kaggle.com/piantic/train-cassava-starter-using-various-loss-funcs/data\n* TaylorSmoothed: https://www.kaggle.com/yerramvarun/cassava-taylorce-loss-label-smoothing-combo\n* GradualWarmup: https://www.kaggle.com/mobassir/faster-pytorch-tpu-baseline-for-cld-cv-0-9#Train-and-eval-function\n* FineTuning: https://www.kaggle.com/piantic/how-to-finetuning-models-pytorch-xla-tpu/data#Dataset"},{"metadata":{},"cell_type":"markdown","source":"Best CV results for each fold so far (seed 42):\n* Fold 0: 90.1\n* Fold 1: 90.1\n* Fold 2: 89.4\n* Fold 3: 89.3\n* Fold 4: 89.2\n\nLB: 0.901 (5 folds 4x TTA)\n\nCurrent goal is improving other folds. Any suggestions greatly appreciated :)\nInference is coming soon"},{"metadata":{},"cell_type":"markdown","source":"* v1 - Initial Release / Debug Run\n* v2 - Fold 4, 10 Epoch, EffnetB4, TaylorSmoothed\n* v3/4 - No longer saves when loss is decreased, only when we improve accuracy\n* v5 - Tried running with heavyaugs"},{"metadata":{},"cell_type":"markdown","source":"# Run Config"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# ====================================================\n# Configurations\n# ====================================================\nimport os\nclass CFG:\n    DEBUG = False\n    \n    #Model Params\n    device = 'TPU' #['CPU','GPU','TPU']\n    N_FOLDS = 5\n    MODEL_NAME = 'tf_efficientnet_b4_ns' # Recommended : ['deit_base_patch16_384','vit_large_patch16_384','tf_efficientnet_b4_ns','resnext50_32x4d']\n    pretrained = True   \n    N_CLASSES = 5\n    #TRAIN_FOLDS = [0,1,2,3,4]\n    TRAIN_FOLDS = [1] #Folds to be Trained\n    \n    scheduler_name = 'CosineAnnealingWarmRestarts'\n    # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'GradualWarmupSchedulerV2']\n    scheduler_update = 'epoch' #['batch','epoch']\n    criterion_name = 'TaylorSmoothedLoss'\n    # ['CrossEntropyLoss', LabelSmoothing', 'FocalLoss','FocalCosineLoss', 'SymmetricCrossEntropyLoss', 'BiTemperedLoss', 'TaylorCrossEntropyLoss', 'TaylorSmoothedLoss']\n    optimizer_name = 'AdamW' #['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    #Image Size\n    HEIGHT = 512 #If VIT or deit is chosen as model: need 384 x 384\n    WIDTH = 512\n    CHANNELS = 3\n    TRAIN_AUG_TYPE = 'train' #['train','lightaug','heavyaug','autoaugment']\n    VALID_AUG_TYPE = 'valid' #['valid']\n    \n    #Training Params\n    BATCH_SIZE = 64 # PER REPLICA FOR TPUS\n    #RECOMMENDED : effnet = 16 ; resnext = 8 ; vit = 4 ; deit = 4\n    EPOCHS = 7# more is definitely plausible and recommended around 10\n    LR = 1e-4\n    LR_START =1e-4\n    LR_MIN = 8e-7\n    weight_decay = 0\n    eps = 1e-8\n    PATIENCE = 4\n    \n    #BiTemperedLoss\n    T1 = 0.2\n    T2 = 1.1\n    LABEL_SMOOTH = 0.20\n    \n    #CosineAnnealingWarmRestarts\n    T_0 = 7\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    IMG_MEAN = [0.485, 0.456, 0.406] #Mean for normalization Transform cassava = [0.4303, 0.4967, 0.3134] imgnet = [0.485, 0.456, 0.406]\n    IMG_STD = [0.229, 0.224, 0.225] #STD for normalization Transform cassava = [0.2142, 0.2191, 0.1954] imgnet = [0.229, 0.224, 0.225]\n    \n    USE_2019 = True #Use 2019 images?\n    \n    #n_procs = number of replicas -> TPU\n    n_procs = 1 #You can set it to 1 and run a TPU as a GPU if you want\n    BGR = False #Alternate method for loading images -> set to true is a bit slower \n    SEED = 42","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Required Installations\n# ====================================================\n!pip install --quiet timm==0.3.2\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\n    \nif CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n    !pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Library"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Library\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\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\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import KFold,StratifiedKFold\n\nimport albumentations as A\nfrom albumentations import Compose\nfrom albumentations.pytorch import ToTensorV2\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\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss Functions"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def log_t(u, t):\n    \"\"\"Compute log_t for `u'.\"\"\"\n    if t==1.0:\n        return u.log()\n    else:\n        return (u.pow(1.0 - t) - 1.0) / (1.0 - t)\n\ndef exp_t(u, t):\n    \"\"\"Compute exp_t for `u'.\"\"\"\n    if t==1:\n        return u.exp()\n    else:\n        return (1.0 + (1.0-t)*u).relu().pow(1.0 / (1.0 - t))\n\ndef compute_normalization_fixed_point(activations, t, num_iters):\n\n    \"\"\"Returns the normalization value for each example (t > 1.0).\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (> 1.0 for tail heaviness).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same shape as activation with the last dimension being 1.\n    \"\"\"\n    mu, _ = torch.max(activations, -1, keepdim=True)\n    normalized_activations_step_0 = activations - mu\n\n    normalized_activations = normalized_activations_step_0\n\n    for _ in range(num_iters):\n        logt_partition = torch.sum(\n                exp_t(normalized_activations, t), -1, keepdim=True)\n        normalized_activations = normalized_activations_step_0 * \\\n                logt_partition.pow(1.0-t)\n\n    logt_partition = torch.sum(\n            exp_t(normalized_activations, t), -1, keepdim=True)\n    normalization_constants = - log_t(1.0 / logt_partition, t) + mu\n\n    return normalization_constants\n\ndef compute_normalization_binary_search(activations, t, num_iters):\n\n    \"\"\"Returns the normalization value for each example (t < 1.0).\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (< 1.0 for finite support).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same rank as activation with the last dimension being 1.\n    \"\"\"\n\n    mu, _ = torch.max(activations, -1, keepdim=True)\n    normalized_activations = activations - mu\n\n    effective_dim = \\\n        torch.sum(\n                (normalized_activations > -1.0 / (1.0-t)).to(torch.int32),\n            dim=-1, keepdim=True).to(activations.dtype)\n\n    shape_partition = activations.shape[:-1] + (1,)\n    lower = torch.zeros(shape_partition, dtype=activations.dtype, device=activations.device)\n    upper = -log_t(1.0/effective_dim, t) * torch.ones_like(lower)\n\n    for _ in range(num_iters):\n        logt_partition = (upper + lower)/2.0\n        sum_probs = torch.sum(\n                exp_t(normalized_activations - logt_partition, t),\n                dim=-1, keepdim=True)\n        update = (sum_probs < 1.0).to(activations.dtype)\n        lower = torch.reshape(\n                lower * update + (1.0-update) * logt_partition,\n                shape_partition)\n        upper = torch.reshape(\n                upper * (1.0 - update) + update * logt_partition,\n                shape_partition)\n\n    logt_partition = (upper + lower)/2.0\n    return logt_partition + mu\n\nclass ComputeNormalization(torch.autograd.Function):\n    \"\"\"\n    Class implementing custom backward pass for compute_normalization. See compute_normalization.\n    \"\"\"\n    @staticmethod\n    def forward(ctx, activations, t, num_iters):\n        if t < 1.0:\n            normalization_constants = compute_normalization_binary_search(activations, t, num_iters)\n        else:\n            normalization_constants = compute_normalization_fixed_point(activations, t, num_iters)\n\n        ctx.save_for_backward(activations, normalization_constants)\n        ctx.t=t\n        return normalization_constants\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        activations, normalization_constants = ctx.saved_tensors\n        t = ctx.t\n        normalized_activations = activations - normalization_constants \n        probabilities = exp_t(normalized_activations, t)\n        escorts = probabilities.pow(t)\n        escorts = escorts / escorts.sum(dim=-1, keepdim=True)\n        grad_input = escorts * grad_output\n        \n        return grad_input, None, None\n\ndef compute_normalization(activations, t, num_iters=5):\n    \"\"\"Returns the normalization value for each example. \n    Backward pass is implemented.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      num_iters: Number of iterations to run the method.\n    Return: A tensor of same rank as activation with the last dimension being 1.\n    \"\"\"\n    return ComputeNormalization.apply(activations, t, num_iters)\n\ndef tempered_sigmoid(activations, t, num_iters = 5):\n    \"\"\"Tempered sigmoid function.\n    Args:\n      activations: Activations for the positive class for binary classification.\n      t: Temperature tensor > 0.0.\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A probabilities tensor.\n    \"\"\"\n    internal_activations = torch.stack([activations,\n        torch.zeros_like(activations)],\n        dim=-1)\n    internal_probabilities = tempered_softmax(internal_activations, t, num_iters)\n    return internal_probabilities[..., 0]\n\n\ndef tempered_softmax(activations, t, num_iters=5):\n    \"\"\"Tempered softmax function.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      t: Temperature > 1.0.\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A probabilities tensor.\n    \"\"\"\n    if t == 1.0:\n        return activations.softmax(dim=-1)\n\n    normalization_constants = compute_normalization(activations, t, num_iters)\n    return exp_t(activations - normalization_constants, t)\n\ndef bi_tempered_binary_logistic_loss(activations,\n        labels,\n        t1,\n        t2,\n        label_smoothing = 0.0,\n        num_iters=5,\n        reduction='mean'):\n\n    \"\"\"Bi-Tempered binary logistic loss.\n    Args:\n      activations: A tensor containing activations for class 1.\n      labels: A tensor with shape as activations, containing probabilities for class 1\n      t1: Temperature 1 (< 1.0 for boundedness).\n      t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      label_smoothing: Label smoothing\n      num_iters: Number of iterations to run the method.\n    Returns:\n      A loss tensor.\n    \"\"\"\n    internal_activations = torch.stack([activations,\n        torch.zeros_like(activations)],\n        dim=-1)\n    internal_labels = torch.stack([labels.to(activations.dtype),\n        1.0 - labels.to(activations.dtype)],\n        dim=-1)\n    return bi_tempered_logistic_loss(internal_activations, \n            internal_labels,\n            t1,\n            t2,\n            label_smoothing = label_smoothing,\n            num_iters = num_iters,\n            reduction = reduction)\n\ndef bi_tempered_logistic_loss(activations,\n        labels,\n        t1 = CFG.T1,\n        t2 = CFG.T2,\n        label_smoothing=CFG.LABEL_SMOOTH,\n        num_iters=5,\n        reduction = 'mean'):\n\n    \"\"\"Bi-Tempered Logistic Loss.\n    Args:\n      activations: A multi-dimensional tensor with last dimension `num_classes`.\n      labels: A tensor with shape and dtype as activations (onehot), \n        or a long tensor of one dimension less than activations (pytorch standard)\n      t1: Temperature 1 (< 1.0 for boundedness).\n      t2: Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).\n      label_smoothing: Label smoothing parameter between [0, 1). Default 0.0.\n      num_iters: Number of iterations to run the method. Default 5.\n      reduction: ``'none'`` | ``'mean'`` | ``'sum'``. Default ``'mean'``.\n        ``'none'``: No reduction is applied, return shape is shape of\n        activations without the last dimension.\n        ``'mean'``: Loss is averaged over minibatch. Return shape (1,)\n        ``'sum'``: Loss is summed over minibatch. Return shape (1,)\n    Returns:\n      A loss tensor.\n    \"\"\"\n\n    if len(labels.shape)<len(activations.shape): #not one-hot\n        labels_onehot = torch.zeros_like(activations)\n        labels_onehot.scatter_(1, labels[..., None], 1)\n    else:\n        labels_onehot = labels\n\n    if label_smoothing > 0:\n        num_classes = labels_onehot.shape[-1]\n        labels_onehot = ( 1 - label_smoothing * num_classes / (num_classes - 1) ) \\\n                * labels_onehot + \\\n                label_smoothing / (num_classes - 1)\n\n    probabilities = tempered_softmax(activations, t2, num_iters)\n\n    loss_values = labels_onehot * log_t(labels_onehot + 1e-10, t1) \\\n            - labels_onehot * log_t(probabilities, t1) \\\n            - labels_onehot.pow(2.0 - t1) / (2.0 - t1) \\\n            + probabilities.pow(2.0 - t1) / (2.0 - t1)\n    loss_values = loss_values.sum(dim = -1) #sum over classes\n\n    if reduction == 'none':\n        return loss_values\n    if reduction == 'sum':\n        return loss_values.sum()\n    if reduction == 'mean':\n        return loss_values.mean()\n\nclass BiTemperedLogistic(nn.Module):\n    def __init__(self, T1 = CFG.T1, T2 = CFG.T2, LABEL_SMOOTH = CFG.LABEL_SMOOTH):\n        super().__init__()\n        self.T1 = T1\n        self.T2 = T2\n        self.LABEL_SMOOTH = LABEL_SMOOTH\n\n    def forward(self, logits,labels):\n        return bi_tempered_logistic_loss(logits, labels,t1 = self.T1,t2 = self.T2, label_smoothing = self.LABEL_SMOOTH)\n    \nclass SymmetricCrossEntropy(nn.Module):\n\n    def __init__(self, alpha=0.1, beta=1.0, num_classes= 5):\n        super(SymmetricCrossEntropy, self).__init__()\n        self.alpha = alpha\n        self.beta = beta\n        self.num_classes = num_classes\n\n    def forward(self, logits, targets, reduction='mean'):\n        onehot_targets = torch.eye(self.num_classes)[targets].cuda()\n        ce_loss = F.cross_entropy(logits, targets, reduction=reduction)\n        rce_loss = (-onehot_targets*logits.softmax(1).clamp(1e-7, 1.0).log()).sum(1)\n        if reduction == 'mean':\n            rce_loss = rce_loss.mean()\n        elif reduction == 'sum':\n            rce_loss = rce_loss.sum()\n        return self.alpha * ce_loss + self.beta * rce_loss\n    \nclass LabelSmoothingLoss(nn.Module): \n    def __init__(self, classes=5, smoothing=0.0, dim=-1): \n        super(LabelSmoothingLoss, self).__init__() \n        self.confidence = 1.0 - smoothing \n        self.smoothing = smoothing \n        self.cls = classes \n        self.dim = dim \n    def forward(self, pred, target): \n        if CFG.criterion_name == 'LabelSmoothingLoss':\n            pred = pred.log_softmax(dim=self.dim) \n        with torch.no_grad():\n            true_dist = torch.zeros_like(pred) \n            true_dist.fill_(self.smoothing / (self.cls - 1)) \n            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) \n        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))\n    \nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=1, gamma=2, reduce=True):\n        super(FocalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.reduce = reduce\n\n    def forward(self, inputs, targets):\n        BCE_loss = nn.CrossEntropyLoss()(inputs, targets)\n\n        pt = torch.exp(-BCE_loss)\n        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss\n\n        if self.reduce:\n            return torch.mean(F_loss)\n        else:\n            return F_loss\n        \nclass FocalCosineLoss(nn.Module):\n    def __init__(self, alpha=1, gamma=2, xent=.1):\n        super(FocalCosineLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n\n        self.xent = xent\n\n        self.y = torch.Tensor([1]).cuda()\n\n    def forward(self, input, target, reduction=\"mean\"):\n        cosine_loss = F.cosine_embedding_loss(input, F.one_hot(target, num_classes=input.size(-1)), self.y, reduction=reduction)\n\n        cent_loss = F.cross_entropy(F.normalize(input), target, reduce=False)\n        pt = torch.exp(-cent_loss)\n        focal_loss = self.alpha * (1-pt)**self.gamma * cent_loss\n\n        if reduction == \"mean\":\n            focal_loss = torch.mean(focal_loss)\n\n        return cosine_loss + self.xent * focal_loss\n    \nclass TaylorSoftmax(nn.Module):\n    '''\n    This is the autograd version\n    '''\n    def __init__(self, dim=1, n=2):\n        super(TaylorSoftmax, self).__init__()\n        assert n % 2 == 0\n        self.dim = dim\n        self.n = n\n\n    def forward(self, x):\n        '''\n        usage similar to nn.Softmax:\n            >>> mod = TaylorSoftmax(dim=1, n=4)\n            >>> inten = torch.randn(1, 32, 64, 64)\n            >>> out = mod(inten)\n        '''\n        fn = torch.ones_like(x)\n        denor = 1.\n        for i in range(1, self.n+1):\n            denor *= i\n            fn = fn + x.pow(i) / denor\n        out = fn / fn.sum(dim=self.dim, keepdims=True)\n        return out\n    \nclass TaylorCrossEntropyLoss(nn.Module):\n    '''\n    This is the autograd version\n    '''\n    def __init__(self, n=2, ignore_index=-1, reduction='mean'):\n        super(TaylorCrossEntropyLoss, self).__init__()\n        assert n % 2 == 0\n        self.taylor_softmax = TaylorSoftmax(dim=1, n=n)\n        self.reduction = reduction\n        self.ignore_index = ignore_index\n\n    def forward(self, logits, labels):\n        '''\n        usage similar to nn.CrossEntropyLoss:\n            >>> crit = TaylorCrossEntropyLoss(n=4)\n            >>> inten = torch.randn(1, 10, 64, 64)\n            >>> label = torch.randint(0, 10, (1, 64, 64))\n            >>> out = crit(inten, label)\n        '''\n        log_probs = self.taylor_softmax(logits).log()\n        loss = F.nll_loss(log_probs, labels, reduction=self.reduction,\n                ignore_index=self.ignore_index)\n        return loss\n    \nclass TaylorSmoothedLoss(nn.Module):\n\n    def __init__(self, n=2, ignore_index=-1, reduction='mean', smoothing=0.2):\n        super(TaylorSmoothedLoss, self).__init__()\n        assert n % 2 == 0\n        self.taylor_softmax = TaylorSoftmax(dim=1, n=n)\n        self.reduction = reduction\n        self.ignore_index = ignore_index\n        self.lab_smooth = LabelSmoothingLoss(CFG.N_CLASSES, smoothing=CFG.LABEL_SMOOTH)\n\n    def forward(self, logits, labels):\n\n        log_probs = self.taylor_softmax(logits).log()\n        #loss = F.nll_loss(log_probs, labels, reduction=self.reduction,\n        #        ignore_index=self.ignore_index)\n        loss = self.lab_smooth(log_probs, labels)\n        return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Augmentations\n# ====================================================\nAug_Norm = A.Normalize(mean=CFG.IMG_MEAN, std=CFG.IMG_STD, max_pixel_value=255.0, p=1.0)\nDrop_Rand = A.CoarseDropout(max_holes=12, max_height=int(0.11*CFG.HEIGHT), max_width=int(0.11*CFG.WIDTH),\n                            min_holes=1, min_height=int(0.03*CFG.HEIGHT), min_width=int(0.03*CFG.WIDTH),\n                            always_apply=False, p=0.5)\nRand_Crop = A.RandomCrop(height= CFG.HEIGHT, width = CFG.WIDTH,always_apply=True, p=1.0)\nResize_Crop = A.RandomResizedCrop(CFG.HEIGHT, CFG.WIDTH,p=1.0)\ntrain_transforms = Compose([\n            A.Transpose(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.5),\n            A.HueSaturationValue(\n                hue_shift_limit=0.2, \n                sat_shift_limit=0.2, \n                val_shift_limit=0.2, \n                p=0.5\n            ),\n            A.RandomBrightnessContrast(\n                    brightness_limit=(-0.1,0.1), \n                    contrast_limit=(-0.1, 0.1), \n                    p=0.5\n                ),\n            Resize_Crop,\n            Drop_Rand,           \n            Aug_Norm,   \n            ToTensorV2(p=1.0),\n        ], p=1.)\n\nlight_transforms = Compose([\n            A.Transpose(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.5),\n            Resize_Crop,\n            Aug_Norm,   \n            ToTensorV2(p=1.0),\n        ], p=1.)\n\nheavy_transforms = Compose([\n    A.HorizontalFlip(p=0.5),\n    \n    A.Resize(CFG.HEIGHT, CFG.WIDTH),\n    \n    A.Transpose(p=0.5),\n    A.VerticalFlip(p=0.5),\n    #A.augmentations.transforms.ColorJitter(brightness=0.10, contrast=0.2, saturation=0.2, hue=0.00, always_apply=False, p=0.5),\n    A.ShiftScaleRotate(p=0.5),\n    A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n    A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n    A.augmentations.transforms.RGBShift (r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, always_apply=False, p=0.5),\n    A.augmentations.transforms.ChannelDropout (channel_drop_range=(1, 1), fill_value=0, always_apply=False, p=0.5),\n    \n    A.augmentations.transforms.GridDistortion (num_steps=5, distort_limit=0.3, interpolation=1, border_mode=4, value=None, mask_value=None, always_apply=False, p=0.5),\n    A.CoarseDropout(p=0.5),\n    A.Cutout(p=0.5),\n    Aug_Norm,\n    ToTensorV2(p=1.0),])\n\nvalid_transforms = Compose([\n            A.CenterCrop(CFG.HEIGHT, CFG.WIDTH),\n            Aug_Norm,   \n            ToTensorV2(p=1.0),\n        ], p=1.)\n\ntest_aug = Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(p = 1.0),\n            #A.ColorJitter(brightness=0.1, contrast=0.2, saturation=0.2, hue=0.00, always_apply=False, p=1.0),\n            Rand_Crop,\n            Aug_Norm,\n            ToTensorV2(p=1.0)\n        ], p=1.)\n\nimage_net_post = Compose([\n            Resize_Crop,\n            Drop_Rand,\n            Aug_Norm,    \n            ToTensorV2(p=1.0)\n        ], p=1.)\n\nclass UnNormalize(object):\n    def __init__(self, mean, std):\n        self.mean = mean\n        self.std = std\n\n    def __call__(self, tensor):\n        \"\"\"\n        Args:\n            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.\n        Returns:\n            Tensor: Normalized image.\n        \"\"\"\n        for t, m, s in zip(tensor, self.mean, self.std):\n            t.mul_(s).add_(m)\n            # The normalize code -> t.sub_(m).div_(s)\n        return tensor","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\ndef retrieve_df(df,name,idx):\n    series = df[name].iloc[idx]\n    series.reset_index(drop=True,inplace=True)\n    return series\n\ndef accuracy_metric(input, targs):\n    return accuracy_score(targs.cpu(), input.cpu())\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#Choose Criterions for the Training Loop\nif CFG.criterion_name == 'BiTemperedLoss':\n    CFG.criterion = BiTemperedLogistic()\n    CFG.val_criterion = BiTemperedLogistic()\nelif CFG.criterion_name == 'SymmetricCrossEntropyLoss':\n    CFG.criterion = SymmetricCrossEntropy()\n    CFG.val_criterion = SymmetricCrossEntropy()\nelif CFG.criterion_name == 'CrossEntropyLoss':\n    CFG.criterion = nn.CrossEntropyLoss()\n    CFG.val_criterion = nn.CrossEntropyLoss()\nelif CFG.criterion_name == 'LabelSmoothingLoss':\n    CFG.criterion = LabelSmoothingLoss()\n    CFG.val_criterion = LabelSmoothingLoss()\nelif CFG.criterion_name == 'FocalLoss':\n    CFG.criterion = FocalLoss()\n    CFG.val_criterion = FocalLoss()\nelif CFG.criterion_name == 'FocalCosineLoss':\n    CFG.criterion = FocalCosineLoss()\n    CFG.val_criterion = FocalCosineLoss()\nelif CFG.criterion_name == 'TaylorCrossEntropyLoss':\n    CFG.criterion = TaylorCrossEntropyLoss()\n    CFG.val_criterion = TaylorCrossEntropyLoss()\nelif CFG.criterion_name == 'TaylorSmoothedLoss':\n    CFG.criterion = TaylorSmoothedLoss()\n    CFG.val_criterion = TaylorSmoothedLoss()\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\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Datasets\n# ====================================================\nclass GetData(Dataset):\n    def __init__(self, Dir, FNames, labels,Type):\n        self.dir = Dir\n        self.fnames = FNames\n        self.lbs = labels\n        self.type = Type\n        self.auto_augment = timm.data.auto_augment.auto_augment_transform('originalr',None)\n        \n    def __len__(self):\n        return len(self.fnames)\n    \n    def get_x(self,index):\n        if CFG.BGR:\n            x = cv2.imread(os.path.join(self.dir, self.fnames[index]))\n            x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)\n        else:\n            x = imread(os.path.join(self.dir, self.fnames[index]))\n        return x\n\n    def __getitem__(self, index):\n        if \"train\" in self.type:\n            x = self.get_x(index)\n            aug_data = train_transforms(image = x)\n            return aug_data['image'], self.lbs[index]\n        elif \"lightaug\" in self.type:\n            x = self.get_x(index)\n            aug_data = light_transforms(image = x)\n            return aug_data['image'], self.lbs[index]\n        elif \"heavyaug\" in self.type:\n            x = self.get_x(index)\n            aug_data = heavy_transforms(image = x)\n            return aug_data['image'], self.lbs[index]\n        elif \"autoaugment\" in self.type:\n            x = Image.open(os.path.join(self.dir, self.fnames[index]))\n            aug_image = self.auto_augment(x)\n            aug_data = image_net_post(image = np.asarray(aug_image,dtype = np.float32))\n            return aug_data['image'], self.lbs[index]\n        elif \"valid\" in self.type:\n            x = self.get_x(index)\n            aug_data = valid_transforms(image = x)\n            return aug_data['image'], self.lbs[index]\n        elif \"tr-tst\" in self.type:\n            x = self.get_x(index)\n            return x, self.lbs[index]\n        elif \"test\" in self.type:\n            x = self.get_x(index)\n            return x, self.fnames[index]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# CV SPLIT"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# CV Split\n# ====================================================\nDATA_PATH = '../input/cassava-leaf-disease-classification/'\nTRAIN_DIR = DATA_PATH + 'train_images/'\n\nDATA_PATH_2019 = '../input/cassava-leaf-disease-merged/'\nTRAIN_DIR_2019 = DATA_PATH_2019 + 'train/'\nTEST_DIR = DATA_PATH + 'test_images/'\n\n#This guarantees that no images from 2019 contaminate the validation split\nif CFG.USE_2019:\n    train_df_merged = pd.read_csv(DATA_PATH_2019 + 'merged.csv')\n    train_df = train_df_merged.loc[train_df_merged.source == 2020]\n    if CFG.DEBUG:\n        train_df = train_df.sample(500).reset_index(drop=True)\n    train_df_2019 = train_df_merged.loc[train_df_merged.source == 2019]\n    skf = StratifiedKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\n    skf.get_n_splits(np.arange(train_df.shape[0]), train_df['label'])\n    folds = [(idxT,idxV) for i,(idxT,idxV) in enumerate(skf.split(np.arange(train_df.shape[0]), train_df['label']))]\n    if not CFG.DEBUG:\n        folds_2019 = [np.concatenate((idxT,idxV)) for i,(idxT,idxV) in enumerate(skf.split(np.arange(train_df_2019.shape[0]), train_df_2019['label']))]\n        for i in range(CFG.N_FOLDS):\n            (idxT,idxV) = folds[i]\n            folds[i] = (np.concatenate((idxT,train_df_2019.iloc[folds_2019[i]].index)),idxV)\n            (idxT,idxV) = folds[i]\n            print(np.bincount(train_df_merged['label'].iloc[idxT]),np.bincount(train_df['label'].iloc[idxV]))\n    DATA_FOLD = TRAIN_DIR_2019\n    del train_df_2019\nelse:\n    train_df = pd.read_csv(DATA_PATH + 'train.csv')\n    if CFG.DEBUG:\n        train_df = train_df.sample(500).reset_index(drop=True)\n    skf = StratifiedKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\n    skf.get_n_splits(np.arange(train_df.shape[0]), train_df['label'])\n    folds = [(idxT,idxV) for i,(idxT,idxV) in enumerate(skf.split(np.arange(train_df.shape[0]), train_df['label']))]\n    for i in range(CFG.N_FOLDS):\n        (idxT,idxV) = folds[i]\n        print(np.bincount(train_df['label'].iloc[idxT]),np.bincount(train_df['label'].iloc[idxV]))\n    \n    train_df_merged = train_df\n    DATA_FOLD = TRAIN_DIR","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentation Visualization"},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = [i for i in range(16)]\nx_train = retrieve_df(train_df_merged,'image_id',idx)\ny_train = retrieve_df(train_df_merged,'label',idx)\ntrain_set = GetData(DATA_FOLD, x_train, y_train, Type = CFG.TRAIN_AUG_TYPE)\nplt.figure(figsize=(10, 10))\ninv_normalize = UnNormalize(mean = CFG.IMG_MEAN,std = CFG.IMG_STD)\nfor i,(image,label) in enumerate(train_set):\n    ax = plt.subplot(4, 4, i + 1)\n    aug_image = inv_normalize(image)\n    aug_image = np.transpose(aug_image.numpy(),[1,2,0])\n    plt.imshow(aug_image)\n    label = train_df['label'].iloc[i]\n    plt.title(label)\n    plt.axis(\"off\")\n    if i == 15:\n        break\ndel aug_image,label,image,ax","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ====================================================\n# Model\n# ====================================================\nclass CassavaNet(nn.Module):\n    def __init__(self, model_name=CFG.MODEL_NAME, pretrained=CFG.pretrained):\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)\n        else:\n            self.model = timm.create_model(model_name, pretrained=pretrained)\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = CassavaNet()\nif CFG.model_print:\n    print(model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train Loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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)\n            loss = criterion(logits, labels)\n            loss.backward()\n            xm.optimizer_step(optimizer)\n        else:\n            with autocast():\n                logits = model(images)\n                loss = criterion(logits, labels)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        \n        preds = F.softmax(logits).argmax(axis = 1)\n        scores += (preds==labels).sum().cpu().numpy()\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['train_acc'] = round(scores/(i+1)/CFG.BATCH_SIZE, 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)\n            val_loss_value = val_criterion(logits,labels)\n            val_loss += val_loss_value.detach().item()\n\n            preds = F.softmax(logits).argmax(axis = 1)\n            scores += (preds==labels).sum().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)/CFG.BATCH_SIZE, 4)\n            val_progress.set_description(str(val_epoch_result))\n    if get_output:\n        return val_loss/len(loader),scores/len(loader)/CFG.BATCH_SIZE\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_merged,'image_id',idxT)\n        y_train = retrieve_df(train_df_merged,'label',idxT)\n        x_val = retrieve_df(train_df_merged,'image_id',idxV)\n        y_val = retrieve_df(train_df_merged,'label',idxV)\n        train_set = GetData(DATA_FOLD, x_train, y_train, Type = CFG.TRAIN_AUG_TYPE)\n        val_set = GetData(DATA_FOLD, x_val, y_val, Type = CFG.VALID_AUG_TYPE)\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            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 = CassavaNet(model_name=CFG.MODEL_NAME, pretrained=True)\n                xm.rendezvous('ModelDone')\n            else:\n                CFG.model = CassavaNet(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                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                \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()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG.model = CassavaNet(model_name=CFG.MODEL_NAME, pretrained=True)\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()","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}