{"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/13836/logos/header.png?t=2020-10-01-17-22-54)","metadata":{}},{"cell_type":"markdown","source":"[Sharpness-Aware Minimization](http://arxiv.org/abs/2010.01412) used as optimizer as paper states that \"we find that SAM natively provides robustness to label noise on par with that provided by state-of-the-art procedures that specifically target learning with noisy labels\".\n![loss surface](https://raw.githubusercontent.com/davda54/sam/main/img/loss_landscape.png)\n","metadata":{}},{"cell_type":"code","source":"# !pip install torchcontrib\n# from torchcontrib.optim.swa import SWA","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --quiet timm","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run Config","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Configurations\n# ====================================================\nimport os, sys\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' # [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] #Folds to be Trained\n    \n    scheduler_name = 'CosineAnnealingWarmRestarts'\n    # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'GradualWarmupSchedulerV2']\n    scheduler_update = 'epoch' #['batch','epoch']\n    criterion_name = 'BiTemperedLoss'\n    # ['CrossEntropyLossOneHot', 'CrossEntropyLoss', 'LabelSmoothing', 'FocalLoss','FocalCosineLoss', 'SymmetricCrossEntropyLoss', 'BiTemperedLoss', 'TaylorCrossEntropyLoss', 'TaylorSmoothedLoss']\n    optimizer_name = 'SAM' #['Adam','AdamW','AdamP','Ranger', 'SAM'] \n    LR_RAMPUP_EPOCHS = 1\n    LR_SUSTAIN_EPOCHS = 0\n    \n    FREEZE = True \n    START_FREEZE = 8\n    \n    #Image Size\n    HEIGHT = 512\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 = 8 # PER REPLICA FOR TPUS\n    EPOCHS = 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 = 5\n    \n    #Knowledge distillation\n    GT_RATE = 0.7\n    ST_RATE = 0.3\n    \n    #BiTemperedLoss\n    T1 = 0.2\n    T2 = 1.1\n    LABEL_SMOOTH = 0.10\n    \n    #CosineAnnealingWarmRestarts\n    T_0 = 7\n    \n    #CosineAnnealingLR\n    T_max = EPOCHS\n    \n    NUM_WORKERS = 4\n    \n    model_print = False\n    tqdm = True \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 \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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# Required Installations\n# ====================================================\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    ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Library","metadata":{}},{"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, KFold\nfrom sklearn.metrics import classification_report, confusion_matrix\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\n\nimport albumentations as A\nfrom albumentations import Compose\nfrom albumentations.pytorch import ToTensorV2\n\nfrom PIL import Image, ImageOps, ImageEnhance, ImageChops\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\nsys.path.append('../input/pytorch-sam-optimizer')\nfrom sam import SAM    \n\n# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Functions","metadata":{}},{"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 CrossEntropyLossOneHot(nn.Module):\n    def __init__(self):\n        super(CrossEntropyLossOneHot, self).__init__()\n        self.log_softmax = nn.LogSoftmax(dim=-1)\n\n    def forward(self, preds, labels):\n        return torch.mean(torch.sum(-labels * self.log_softmax(preds), -1))    \n    \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","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Augmentations","metadata":{}},{"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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"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()\nelif CFG.criterion_name == 'CrossEntropyLossOneHot':\n    CFG.criterion = CrossEntropyLossOneHot()\n    CFG.val_criterion = CrossEntropyLossOneHot()\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 == 'SAM':\n        base_opt = torch.optim.Adam\n        return SAM(parameters, base_opt, lr=CFG.LR)\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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Datasets","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Datasets\n# ====================================================\nclass GetData(Dataset):\n    def __init__(self, Dir, FNames, labels, Type,  soft = None):\n        self.dir = Dir\n        self.fnames = FNames\n        if soft is not None:\n            self.lbs = torch.FloatTensor((soft.values).astype(np.float))\n        else:\n            self.lbs = labels\n            \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]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CV SPLIT","metadata":{}},{"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\n\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\n    soft_labels = pd.read_csv('../input/cassava-leaf-disease-soft-targets-09-model/soft_targets_2020.csv')\n    soft_labels_19 = pd.read_csv('../input/cassava-leaf-disease-soft-targets-09-model/soft_targets_2019.csv')\n    soft_labels = soft_labels.append(soft_labels_19, ignore_index=True)\n    \n    ground_truth = train_df_merged.join(pd.get_dummies(train_df['label'], prefix='p', prefix_sep=''))\n\n    ground_truth['p0'] = ground_truth['p0']*CFG.GT_RATE\n    ground_truth['p1'] = ground_truth['p1']*CFG.GT_RATE\n    ground_truth['p2'] = ground_truth['p2']*CFG.GT_RATE\n    ground_truth['p3'] = ground_truth['p3']*CFG.GT_RATE\n    ground_truth['p4'] = ground_truth['p4']*CFG.GT_RATE\n\n    soft_labels['p0'] = soft_labels['p0']*CFG.ST_RATE\n    soft_labels['p1'] = soft_labels['p1']*CFG.ST_RATE\n    soft_labels['p2'] = soft_labels['p2']*CFG.ST_RATE\n    soft_labels['p3'] = soft_labels['p3']*CFG.ST_RATE\n    soft_labels['p4'] = soft_labels['p4']*CFG.ST_RATE\n    \n    train_df_merged = ground_truth.set_index('image_id').add(soft_labels.set_index('image_id'), fill_value =0).reset_index()\n    \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    \n    \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    \n    soft_labels = pd.read_csv('../input/cassava-leaf-disease-soft-targets-09-model/soft_targets_2020.csv')\n    \n    ground_truth = train_df.join(pd.get_dummies(train_df['label'], prefix='p', prefix_sep=''))\n    ground_truth['p0'] = ground_truth['p0']*CFG.GT_RATE\n    ground_truth['p1'] = ground_truth['p1']*CFG.GT_RATE\n    ground_truth['p2'] = ground_truth['p2']*CFG.GT_RATE\n    ground_truth['p3'] = ground_truth['p3']*CFG.GT_RATE\n    ground_truth['p4'] = ground_truth['p4']*CFG.GT_RATE\n\n    soft_labels['p0'] = soft_labels['p0']*CFG.ST_RATE\n    soft_labels['p1'] = soft_labels['p1']*CFG.ST_RATE\n    soft_labels['p2'] = soft_labels['p2']*CFG.ST_RATE\n    soft_labels['p3'] = soft_labels['p3']*CFG.ST_RATE\n    soft_labels['p4'] = soft_labels['p4']*CFG.ST_RATE\n    \n    train_df = ground_truth.set_index('image_id').add(soft_labels.set_index('image_id'), fill_value =0).reset_index()\n    \n    \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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Augmentation Visualization","metadata":{}},{"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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"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        elif 'nfnet' in model_name:\n            self.n_features = self.model.head.fc.in_features\n            self.model.head.fc = nn.Linear(self.n_features, CFG.N_CLASSES)\n            \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        elif 'nfnet' in model_name:\n            for param in self.model.head.fc.parameters():\n                param.requires_grad = True\n            \n            \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.model.parameters():\n            param.requires_grad = True","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = CassavaNet()\nif CFG.model_print:\n    print(model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Loop","metadata":{}},{"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        \n        def closure():\n            loss_cl = criterion(logits, labels)\n            loss_cl.backward()\n            return loss_cl\n        \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\n            # first forward-backward step\n            loss.backward()\n            optimizer.first_step(zero_grad=True)\n\n            # second forward-backward step\n            criterion(model(images), labels).backward()\n            optimizer.second_step(zero_grad=True)\n        \n        preds = F.softmax(logits).argmax(axis = 1)\n        scores += (preds==labels.argmax(axis = 1)).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.argmax(axis = 1)).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        y_train_soft = retrieve_df(train_df_merged, ['p0', 'p1','p2', 'p3', 'p4'], idxT)\n        y_val_soft = retrieve_df(train_df_merged, ['p0', 'p1','p2', 'p3', 'p4'], idxV)\n\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, soft = y_train_soft, Type = CFG.TRAIN_AUG_TYPE)\n        val_set = GetData(DATA_FOLD, x_val, y_val, soft = y_val_soft, Type = CFG.VALID_AUG_TYPE)\n        \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        \n#         optimizer = SWA(optimizer, swa_start=5,\n#                     swa_freq=1)\n\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#             if epoch == 0 or epoch == CFG.EPOCHS - 1:\n#                 # Batchnorm update\n#                 optimizer.swap_swa_sgd()\n#                 optimizer.bn_update(train_loader, model, device='cuda')\n                \n#                 optimizer.swap_swa_sgd()\n            \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()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\ngc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}