{"cells":[{"metadata":{},"cell_type":"markdown","source":"### A Pytorch Template for all users (beginners/advanced)\n\nThe goal of this notebook is to allow users to use a script that runs at the command line and allows the user to try out different options for experimentation. \n\nPytorch AMP is used for training and augmentations are via `Albumentations` (https://github.com/albumentations-team/albumentations) \n\n### Specifically, the features supported are:\n\n1) Get list of models via `--find_model` eg. `--find_model se*resnet*50`\n\n2) Testing a sample of data via `--n_samples` to make sure main code works fine for changes\n\n3) Testing augmentation visually via `--test_loader` and changing the augments in the `get_sample_transforms` function. I noticed that a lot of people dont really look at the details of the augmentations and this helps quickly look at what they are really doing\n\n4) Providing a scheduler or loss function via `--scheduler <scheduler name>` / `--loss_fn <loss function name>` Supported loss functions are \n*  CrossEntropyLoss\n*  SmoothCrossEntropyLoss\n*  SymmetricCrossEntropy\n*  BCEWithLogitsLoss\n*  TruncatedLoss\n*  TaylorCrossEntropyLoss\n*  BiTemperedLoss\n*  FocalCosineLoss\n\n5) Optionally get noisy labels using `CleanLab` (https://github.com/cgnorthcutt/cleanlab)\n\n6) Save best model for each epoch and each fold\n\n7) Support for upsampling/downsampling \n\n8) Optional Cutmix (per batch with 25% probability) \n\n9) Optional SVM head on CNN model \n\n### The script saves all states for reproducibility in the future - eg. List. ofcommand line arguments, list of config options, augmentations performed, scores etc"},{"metadata":{},"cell_type":"markdown","source":"### With the above functionality, users should be able to try out different architectures and parameters for experimentation. "},{"metadata":{},"cell_type":"markdown","source":"### Sample Command line usage: \n`python3 pytorch_amp.py --train --model tf_efficientnet_b4_ns --model_dir effnetb4_2019_epochs --img_size 512 --batch_size 16 --loss_fn BiTemperedLoss --augment` "},{"metadata":{},"cell_type":"markdown","source":"\n### Credits:\n\nThis kernel would not have been possible without the contributions of various authors below. \n\nLoss functions:\nhttps://www.kaggle.com/piantic/train-cassava-starter-using-various-loss-funcs\n\nPytorch AMP: \nhttps://www.kaggle.com/khyeh0719/pytorch-efficientnet-baseline-train-amp-aug\n\nCutmix:\nhttps://www.kaggle.com/ar2017/pytorch-efficientnet-train-aug-cutmix-fmix\n\nCleanlab Tutorial: \nhttps://www.kaggle.com/telljoy/noisy-label-eda-with-cleanlab\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nfrom pathlib import Path\nfrom sklearn.utils import resample\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import f1_score, accuracy_score\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport time\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\nfrom transformers import get_linear_schedule_with_warmup\nfrom torch.nn import functional as F\nimport random\nfrom sklearn.preprocessing import LabelEncoder\nimport cv2\nfrom albumentations import (\n    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout,\n    ShiftScaleRotate, CenterCrop, Resize, RandomCrop\n)\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom torch.nn.modules.loss import _WeightedLoss\nfrom scipy.special import softmax\nimport argparse\nimport matplotlib.pyplot as plt\nfrom torch.cuda.amp import autocast, GradScaler\n\nfrom cleanlab.pruning import get_noise_indices\nfrom cleanlab.classification import LearningWithNoisyLabels\nfrom sklearn.base import BaseEstimator\nimport sklearn\nfrom sklearn import svm\nimport pickle\n\n\nSEED = 1234\n\n# Base data path\nDATA_PATH = \"../input/cassava-leaf-disease-classification/\"\n# Where the images/audio are stored\nFILE_PATH = \"../input/cassava-leaf-disease-classification/train_images/\"\nTEST_FILE_PATH = \"../input/cassava-leaf-disease-classification/test_images/\"\n\nTRAIN_DF_FILE = 'train_2019_2020_clean.csv'\nTARGET_COL = 'label'\nID_COL = 'image_id'\n\n# Directory to store experiment results\nBASE_MODEL_FOLDER = \"./models/\"\n\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\nNUM_WORKERS = 4\nSMOOTHING = 0.05\n\n\n# model_names = timm.list_models('*inception*')\n# print(model_names)\n\ndef detect_leaf(img):\n    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n    # find the brown color\n    mask_brown = cv2.inRange(hsv, (8, 60, 20), (30, 255, 200))\n    # find the yellow and green color in the leaf\n    # 21, 86\n    mask_yellow_green = cv2.inRange(hsv, (10, 39, 64), (86, 255, 255))\n    # find any of the three colors(green or brown or yellow) in the image\n    mask = cv2.bitwise_or(mask_yellow_green, mask_brown)\n    # Bitwise-AND mask and original image\n    res = cv2.bitwise_and(img, img, mask=mask)\n    return res\n\n\ndef rand_bbox(size, lam):\n    W = size[2]\n    H = size[3]\n    cut_rat = np.sqrt(1. - lam)\n    cut_w = np.int(W * cut_rat)\n    cut_h = np.int(H * cut_rat)\n\n    # uniform\n    cx = np.random.randint(W)\n    cy = np.random.randint(H)\n\n    bbx1 = np.clip(cx - cut_w // 2, 0, W)\n    bby1 = np.clip(cy - cut_h // 2, 0, H)\n    bbx2 = np.clip(cx + cut_w // 2, 0, W)\n    bby2 = np.clip(cy + cut_h // 2, 0, H)\n    return bbx1, bby1, bbx2, bby2\n\n\ndef cutmix(data, target, alpha):\n    indices = torch.randperm(data.size(0))\n    shuffled_data = data[indices]\n    shuffled_target = target[indices]\n\n    lam = np.clip(np.random.beta(alpha, alpha), 0.3, 0.4)\n    bbx1, bby1, bbx2, bby2 = rand_bbox(data.size(), lam)\n    new_data = data.clone()\n    new_data[:, :, bby1:bby2, bbx1:bbx2] = data[indices, :, bby1:bby2, bbx1:bbx2]\n    # adjust lambda to exactly match pixel ratio\n    lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (data.size()[-1] * data.size()[-2]))\n    targets = (target, shuffled_target, lam)\n\n    return new_data, targets\n\n\n# CutMix\nclass CutMixCollator:\n    def __call__(self, batch):\n        # batch = torch.utils.data.dataloader.default_collate(batch)\n        batch = cutmix(batch)\n        return batch\n\n\nclass CutMixCriterion(nn.Module):\n    def __init__(self, criterion):\n        super(CutMixCriterion, self).__init__()\n        self.criterion = criterion\n\n    def forward(self, preds, targets):\n        targets1 = targets[0].to(device)\n        targets2 = targets[1].to(device)\n        lam = targets[2]\n        return lam * self.criterion.forward(\n            preds, targets1) + (1 - lam) * self.criterion.forward(preds, targets2)\n\n\n# Focal cosine loss\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,\n                                              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\n\n# BiTemperedLoss\n# Code taken from https://github.com/fhopfmueller/bi-tempered-loss-pytorch/blob/master/bi_tempered_loss_pytorch.py\n\ndef 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\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\n\ndef compute_normalization_fixed_point(activations, t, num_iters):\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\n\ndef compute_normalization_binary_search(activations, t, num_iters):\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\n\nclass ComputeNormalization(torch.autograd.Function):\n    \"\"\"\n    Class implementing custom backward pass for compute_normalization. See compute_normalization.\n    \"\"\"\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\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\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\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    \"\"\"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\n\ndef bi_tempered_logistic_loss(activations,\n                              labels,\n                              t1,\n                              t2,\n                              label_smoothing=0.0,\n                              num_iters=5,\n                              reduction='mean'):\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\n\nclass BiTemperedLogisticLoss(nn.Module):\n    def __init__(self, t1, t2, smoothing=0.0):\n        super(BiTemperedLogisticLoss, self).__init__()\n        self.t1 = t1\n        self.t2 = t2\n        self.smoothing = smoothing\n\n    def forward(self, logit_label, truth_label):\n        loss_label = bi_tempered_logistic_loss(\n            logit_label, truth_label,\n            t1=self.t1, t2=self.t2,\n            label_smoothing=self.smoothing,\n            reduction='none'\n        )\n\n        loss_label = loss_label.mean()\n        return loss_label\n\n\n## End bitempered loss\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\n    def forward(self, pred, target):\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\n\nclass TaylorSoftmax(nn.Module):\n    '''\n    This is the autograd version\n    '''\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\n\nclass TaylorCrossEntropyLoss(nn.Module):\n    def __init__(self, n=2, ignore_index=-1, reduction='mean', smoothing=0.05):\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        self.lab_smooth = LabelSmoothingLoss(Config.img_size, smoothing=smoothing)\n\n    def forward(self, logits, labels):\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\n\n\n# https://github.com/AlanChou/Truncated-Loss/blob/master/TruncatedLoss.py\nclass TruncatedLoss(nn.Module):\n\n    def __init__(self, q=0.7, k=0.5, trainset_size=50000):\n        super(TruncatedLoss, self).__init__()\n        self.q = q\n        self.k = k\n        self.weight = torch.nn.Parameter(data=torch.ones(trainset_size, 1), requires_grad=False)\n\n    def forward(self, logits, targets, indexes):\n        p = F.softmax(logits, dim=1)\n        Yg = torch.gather(p, 1, torch.unsqueeze(targets, 1))\n\n        loss = ((1 - (Yg ** self.q)) / self.q) * self.weight[indexes] - ((1 - (self.k ** self.q)) / self.q) * \\\n               self.weight[indexes]\n        loss = torch.mean(loss)\n\n        return loss\n\n    def update_weight(self, logits, targets, indexes):\n        p = F.softmax(logits, dim=1)\n        Yg = torch.gather(p, 1, torch.unsqueeze(targets, 1))\n        Lq = ((1 - (Yg ** self.q)) / self.q)\n        Lqk = np.repeat(((1 - (self.k ** self.q)) / self.q), targets.size(0))\n        Lqk = torch.from_numpy(Lqk).type(torch.cuda.FloatTensor)\n        Lqk = torch.unsqueeze(Lqk, 1)\n\n        condition = torch.gt(Lqk, Lq)\n        self.weight[indexes] = condition.type(torch.cuda.FloatTensor)\n\n\n# source: https://www.kaggle.com/c/siim-isic-melanoma-classification/discussion/173733\nclass SmoothCrossEntropyLoss(_WeightedLoss):\n    def __init__(self, weight=None, reduction='mean', smoothing=SMOOTHING):\n        super().__init__(weight=weight, reduction=reduction)\n        self.smoothing = smoothing\n        self.weight = weight\n        self.reduction = reduction\n\n    @staticmethod\n    def _smooth_one_hot(targets: torch.Tensor, n_classes: int, smoothing=SMOOTHING):\n        assert 0 <= smoothing < 1\n        with torch.no_grad():\n            targets = torch.empty(size=(targets.size(0), n_classes),\n                                  device=targets.device) \\\n                .fill_(smoothing / (n_classes - 1)) \\\n                .scatter_(1, targets.data.unsqueeze(1), 1. - smoothing)\n        return targets\n\n    def forward(self, inputs, targets):\n        targets = SmoothCrossEntropyLoss._smooth_one_hot(targets, inputs.size(-1),\n                                                         self.smoothing)\n        lsm = F.log_softmax(inputs, -1)\n\n        if self.weight is not None:\n            lsm = lsm * self.weight.unsqueeze(0)\n\n        loss = -(targets * lsm).sum(-1)\n\n        if self.reduction == 'sum':\n            loss = loss.sum()\n        elif self.reduction == 'mean':\n            loss = loss.mean()\n\n        return loss\n\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\ndef mono_to_color(X, eps=1e-6, mean=None, std=None):\n    X = np.stack([X, X, X], axis=-1)\n\n    # Standardize\n    mean = mean or X.mean()\n    std = std or X.std()\n    X = (X - mean) / (std + eps)\n\n    # Normalize to [0, 255]\n    _min, _max = X.min(), X.max()\n\n    if (_max - _min) > eps:\n        V = np.clip(X, _min, _max)\n        V = 255 * (V - _min) / (_max - _min)\n        V = V.astype(np.uint8)\n    else:\n        V = np.zeros_like(X, dtype=np.uint8)\n\n    return V\n\n\ndef resize(image, size=None):\n    if size is not None:\n        h, w, _ = image.shape\n        new_w, new_h = int(w * size / h), size\n        image = cv2.resize(image, (new_w, new_h))\n\n    return image\n\n\ndef normalize(image, mean=None, std=None):\n    image = image / 255.0\n    if mean is not None and std is not None:\n        image = (image - mean) / std\n    # return np.moveaxis(image, 2, 0).astype(np.float32)\n    return image.astype(np.float32)\n\n\n# Data loader\nclass CustomDataset(Dataset):\n    def __init__(self, df, file_path, train=True, transforms=None):\n        self.train = train\n        self.df = df\n        self.file_path = file_path\n        self.filename = df[ID_COL].values\n        self.transforms = transforms\n        self.labels = df[TARGET_COL].values\n\n    def __len__(self):\n        return len(self.filename)\n\n    def __getitem__(self, idx: int):\n        # Return audio and sampling rate\n        file = self.file_path + self.filename[idx]\n        if not os.path.exists(file):\n            print(file)\n        image = cv2.imread(file)\n        orig_image = image.copy()\n\n        # Special user defined mask\n        if Config.cv_mask:\n            image = detect_leaf(image)\n        else:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        # Normalize if no augmentation. Test set also needs this\n        if not self.transforms:\n            image = cv2.resize(image, dsize=(Config.img_size, Config.img_size), interpolation=cv2.INTER_LINEAR)\n            image = normalize(image, mean=None, std=None)\n        # Augment\n        else:\n            image = cv2.resize(image, dsize=(Config.img_size, Config.img_size), interpolation=cv2.INTER_LINEAR)\n            image = self.transforms(image=image)['image']\n\n        # Uncomment to switch channel to first dimension if image shape is H x W x C\n        # image = np.transpose(image, axes=[2,0,1])\n\n        # Return image and raw data\n        if self.train:\n            return image, ONE_HOT[self.df[TARGET_COL][idx]]\n        elif Config.test_loader:\n            return image, orig_image\n        else:\n            return image\n\n\n# Test data loader\n# audio_data = AudioDataset(df=train_df, params=AudioParams, audio_path=AUDIO_PATH)\n# audio_data[0]\n# Torch utils\ndef seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results\n\n    Arguments:\n        seed {int} -- Number of the seed\n    \"\"\"\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  # False\n\n\ndef save_model_params(val_scores, fold, cp_folder, use_svm=False):\n    val_scores = np.array(val_scores).mean()\n    params_file = os.path.join(cp_folder, 'params.txt')\n    with open(params_file, 'a+') as f:\n        if use_svm:\n            f.write(f'Fold: {fold} Mean Val SVM Score: {val_scores}\\n')\n        else:\n            f.write(f'Fold: {fold} Mean Val Score: {val_scores}\\n')\n\n\ndef save_model_weights(model, filename, verbose=1, cp_folder=\"\"):\n    \"\"\"\n    Saves the weights of a PyTorch model\n\n    Arguments:\n        model {torch module} -- Model to save the weights of\n        filename {str} -- Name of the checkpoint\n\n    Keyword Arguments:\n        verbose {int} -- Whether to display infos (default: {1})\n        cp_folder {str} -- Folder to save to (default: {''})\n    \"\"\"\n    if verbose:\n        print(f\"\\n -> Saving weights to {os.path.join(cp_folder, filename)}\\n\")\n    torch.save(model.state_dict(), os.path.join(cp_folder, filename))\n\n\ndef load_model_weights(model, filename, verbose=1, cp_folder=\"\"):\n    \"\"\"\n    Loads the weights of a PyTorch model. The exception handles cpu/gpu incompatibilities\n\n    Arguments:\n        model {torch module} -- Model to load the weights to\n        filename {str} -- Name of the checkpoint\n\n    Keyword Arguments:\n        verbose {int} -- Whether to display infos (default: {1})\n        cp_folder {str} -- Folder to load from (default: {''})\n\n    Returns:\n        torch module -- Model with loaded weights\n    \"\"\"\n    if verbose:\n        print(f\"\\n -> Loading weights from {os.path.join(cp_folder, filename)}\\n\")\n    try:\n        model.load_state_dict(os.path.join(cp_folder, filename), strict=True)\n    except BaseException:\n        model.load_state_dict(\n            torch.load(os.path.join(cp_folder, filename), map_location=\"cpu\"),\n            strict=True,\n        )\n    return model\n\n\ndef count_parameters(model, all=False):\n    \"\"\"\n    Count the parameters of a model\n\n    Arguments:\n        model {torch module} -- Model to count the parameters of\n\n    Keyword Arguments:\n        all {bool} -- Whether to include not trainable parameters in the sum (default: {False})\n\n    Returns:\n        int -- Number of parameters\n    \"\"\"\n    if all:\n        return sum(p.numel() for p in model.parameters())\n    else:\n        return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\n\ndef get_metric(truth, pred, avg=\"micro\", metrics=[\"f1\"]):\n    if len(truth.shape) == 1:\n        truth = ONE_HOT[truth]\n    pred_class = np.argmax(pred, axis=1)\n    truth_class = np.argmax(truth, axis=1)\n    results = []\n    for metric in metrics:\n        if metric == 'f1':\n            results.append(sklearn.metrics.f1_score(truth_class, pred_class, average=avg))\n        elif metric == 'acc':\n            results.append(sklearn.metrics.accuracy_score(truth_class, pred_class))\n        elif metric == 'mae':\n            results.append(sklearn.metrics.mean_absolute_error(truth_class, pred_class))\n    return results\n\n\ndef smooth_label(y, alpha=0.01):\n    y = y * (1 - alpha)\n    y[y == 0] = alpha\n    return y\n\n\ndef get_model(name, num_classes=1):\n    model = timm.create_model(name, pretrained=True)\n    if 'resne' in name or 'inception' in name:\n        nb_ft = model.fc.in_features\n        del model.fc\n        model.fc = nn.Linear(nb_ft, num_classes)\n    elif 'vit' in name:\n        nb_ft = model.head.in_features\n        del model.head\n        model.head = nn.Linear(nb_ft, num_classes)\n    else:\n        #for param in model.parameters():\n        #    param.requires_grad = False\n        nb_ft = model.classifier.in_features\n        del model.classifier\n        model.classifier = nn.Linear(nb_ft, num_classes)\n    return model\n\n\ndef train_one_epoch(epoch, model, scaler, criterion, optimizer, train_loader, device, scheduler=None,\n                    schd_batch_update=False):\n    model.train()\n\n    t = time.time()\n    running_loss = None\n    base_criterion = criterion\n    image_preds_all = []\n    image_labels_all = []\n\n    for step, (images, image_labels) in enumerate(train_loader):\n\n        mix_decision = np.random.rand()\n\n        # Cutmix to modify the images\n        if Config.cutmix and mix_decision < 0.25:\n            images, image_labels_tuple = cutmix(images, image_labels, 1.)\n            image_labels = image_labels_tuple[0]\n\n        images = images.to(device)\n        # Probablities\n        image_labels = image_labels.to(device).long()\n        # Get label of max prob\n        image_labels = torch.max(image_labels, 1)[1]\n\n        with autocast():\n            image_preds = model(images)\n\n            # For SVM aggregate preds and labels to fit\n            if Config.use_svm:\n                image_preds_all.append(image_preds.cpu().detach().numpy())\n                image_labels_all.append(image_labels.cpu().detach().numpy())\n\n            # Change criterion if cutmix is enabled for this batch\n            if Config.cutmix and mix_decision < 0.25:\n                criterion = CutMixCriterion(base_criterion).to(device)\n                loss = criterion(image_preds, image_labels_tuple)\n\n            else:\n                criterion = base_criterion\n                loss = criterion(image_preds, image_labels)\n\n        scaler.scale(loss).backward()\n\n        if running_loss is None:\n            running_loss = loss.item()\n        else:\n            running_loss = running_loss * .99 + loss.item() * .01\n\n        if ((step + 1) % Config.accum_iter == 0) or ((step + 1) == len(train_loader)):\n            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)\n\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n\n            if scheduler is not None and schd_batch_update:\n                scheduler.step()\n\n    if Config.use_svm:\n        image_preds_all = np.concatenate(image_preds_all)\n        image_labels_all = np.concatenate(image_labels_all)\n        svm_clf.fit(image_preds_all, image_labels_all)\n\n    train_time = time.time() - t\n    if ((step + 1) % Config.verbose == 0) or ((step + 1) == len(train_loader)):\n        print(f'Epoch {epoch + 1}: Train loss: {running_loss:.4f} Time: {train_time:.3f} secs')\n\n    if scheduler is not None and not schd_batch_update:\n        scheduler.step()\n\n\ndef valid_one_epoch(epoch, model, criterion, val_loader, device, scheduler=None, schd_loss_update=False):\n    model.eval()\n\n    loss_sum = 0\n    sample_num = 0\n    image_preds_all = []\n    image_targets_all = []\n    image_preds_svm_all = []\n    svm_metric_score = 0\n\n    for step, (images, image_labels) in enumerate(val_loader):\n        image_labels = image_labels.to(device).long()\n        image_labels = torch.max(image_labels, 1)[1]\n        images = images.to(device)\n\n        image_preds = model(images)  # output = model(input)\n        # print(image_preds.shape, exam_pred.shape)\n        image_preds_all += [torch.argmax(image_preds, 1).detach().cpu().numpy()]\n        image_targets_all += [image_labels.detach().cpu().numpy()]\n\n        loss = criterion(image_preds, image_labels)\n\n        loss_sum += loss.item() * image_preds.shape[0]\n        sample_num += image_preds.shape[0]\n\n        if Config.use_svm:\n            svm_preds = svm_clf.predict(image_preds.cpu().detach().numpy())\n            image_preds_svm_all += [svm_preds]\n\n    if ((step + 1) % Config.verbose == 0) or ((step + 1) == len(val_loader)):\n        print(f'Epoch {epoch + 1}: Val loss: {loss_sum / sample_num:.4f}')\n\n    image_preds_all = np.concatenate(image_preds_all)\n    image_targets_all = np.concatenate(image_targets_all)\n    if Config.use_svm:\n        image_preds_svm_all = np.concatenate(image_preds_svm_all)\n\n    metric_score = (image_preds_all == image_targets_all).mean()\n    print('Validation multi-class accuracy = {:.4f}'.format(metric_score))\n\n    if Config.use_svm:\n        svm_metric_score = (image_preds_svm_all == image_targets_all).mean()\n        print('SVM validation accuracy = {:.5f}'.format(svm_metric_score))\n\n    if scheduler is not None:\n        if schd_loss_update:\n            scheduler.step(loss_sum / sample_num)\n        else:\n            scheduler.step()\n\n    return metric_score, svm_metric_score\n\n\ndef predict(model, dataset, batch_size=64, infer=False):\n    \"\"\"\n    Usual torch predict function\n\n    Arguments:\n        model {torch model} -- Model to predict with\n        dataset {torch dataset} -- Dataset to predict with on\n\n    Keyword Arguments:\n        batch_size {int} -- Batch size (default: {32})\n\n    Returns:\n        numpy array -- Predictions\n    \"\"\"\n    model.eval()\n    preds = np.empty((0, NUM_CLASSES))\n\n    loader = DataLoader(\n        dataset, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS\n    )\n    with torch.no_grad():\n        if not infer:\n            for x, _ in loader:\n                # x = x.type(torch.LongTensor)\n                y_pred = model(x.detach())\n                preds = np.concatenate([preds, y_pred.cpu().numpy()])\n                # preds = np.concatenate([preds, torch.sigmoid(y_pred).cpu().numpy()])\n        else:\n            for x in loader:\n                y_pred = model(x.to(device).detach())\n                preds = np.concatenate([preds, y_pred.cpu().numpy()])\n                # preds = np.concatenate([preds, torch.sigmoid(y_pred).cpu().numpy()])\n    return preds\n\n\nclass ImageClassifier(nn.Module):\n    def __init__(self, model, n_class, infer=False):\n        super().__init__()\n        if not infer:\n            self.model = get_model(model, n_class)\n        else:\n            self.model = model\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\n\ndef train_model(config, df_train, df_val, fold):\n    print(f\"{len(df_train)} training samples \")\n    print(f\"{len(df_val)} validation samples \")\n    seed_everything(config.seed)\n\n    # Pretrained model\n    model = get_model(config.selected_model, num_classes=NUM_CLASSES).to(device)\n\n    if torch.cuda.device_count() > 1:\n        print(f\"Found {torch.cuda.device_count()} GPUs. Using DataParallel\")\n        model = nn.DataParallel(model)\n\n    model.zero_grad()\n\n    df_train = df_train.reset_index().drop('index', axis=1)\n    df_val = df_val.reset_index().drop('index', axis=1)\n\n    # Upsample train set to prevent leakage in val set\n    if args.upsample:\n        majority_class = 3\n        unique_classes = df_train[TARGET_COL].unique()\n        # Initialize with majority class\n        final_df = df_train[df_train[TARGET_COL] == majority_class]\n        # number of samples in majority class\n        samples = int(args.upsample * final_df.shape[0])\n        print(f\"Before sampling:\")\n        print(df_train[TARGET_COL].value_counts())\n        print(f\"Sampling to match {samples} samples\")\n        for class_val in unique_classes:\n            if class_val == majority_class:\n                continue\n            df_train_min = df_train[df_train[TARGET_COL] == class_val]\n            upsample_df = resample(df_train_min,\n                                   replace=True,  # sample without replacement\n                                   n_samples=samples,  # to match minority class\n                                   random_state=123)\n            final_df = pd.concat([final_df, upsample_df])\n        df_train = final_df\n        print('After upsampling:')\n        print(df_train[TARGET_COL].value_counts())\n        df_train = df_train.reset_index().drop('index', axis=1)\n\n    epochs = config.epochs\n\n    # Datasets\n    if Config.augment:\n        train_dataset = CustomDataset(df=df_train, file_path=FILE_PATH, train=True,\n                                      transforms=get_train_transforms())\n        val_dataset = CustomDataset(df=df_val, file_path=FILE_PATH, train=True, transforms=get_valid_transforms())\n    else:\n        train_dataset = CustomDataset(df=df_train, file_path=FILE_PATH, train=True)\n        val_dataset = CustomDataset(df=df_val, file_path=FILE_PATH, train=True)\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config.batch_size,\n        shuffle=True,\n        drop_last=False,\n        pin_memory=False,\n        num_workers=NUM_WORKERS,\n    )\n    val_loader = DataLoader(\n        val_dataset, batch_size=config.val_batch_size, shuffle=False, pin_memory=False, num_workers=NUM_WORKERS\n    )\n\n    n_parameters = count_parameters(model)\n    print(f\"Trainable parameters: {n_parameters}\")\n\n    model = ImageClassifier(config.selected_model, NUM_CLASSES).to(device)\n    scaler = GradScaler()\n    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)\n\n    # Loss function\n    if Config.loss_fn == 'CrossEntropyLoss':\n        criterion = nn.CrossEntropyLoss().to(device)\n    elif Config.loss_fn == 'SmoothCrossEntropyLoss':\n        criterion = SmoothCrossEntropyLoss().to(device)\n    elif Config.loss_fn == 'SymmetricCrossEntropy':\n        criterion = SymmetricCrossEntropy().to(device)\n    elif Config.loss_fn == 'BCEWithLogitsLoss':\n        criterion = nn.BCEWithLogitsLoss(reduction=\"mean\").to(device)\n    elif Config.loss_fn == 'TruncatedLoss':\n        criterion = TruncatedLoss(trainset_size=len(train_dataset)).to(device)\n    elif Config.loss_fn == 'MAE':\n        criterion = nn.L1Loss(reduction='mean').to(device)\n    elif Config.loss_fn == 'TaylorCrossEntropyLoss':\n        criterion = TaylorCrossEntropyLoss(smoothing=Config.label_smoothing).to(device)\n    elif Config.loss_fn == 'BiTemperedLoss':\n        criterion = BiTemperedLogisticLoss(t1=Config.t1, t2=Config.t2, smoothing=Config.label_smoothing).to(device)\n    elif Config.loss_fn == 'FocalCosineLoss':\n        criterion = FocalCosineLoss().to(device)\n\n    # Schedulers\n    num_warmup_steps = int(config.warmup_prop * config.epochs * len(train_loader))\n    num_training_steps = int(config.epochs * len(train_loader))\n\n    if Config.scheduler == 'CosineAnnealingLR':\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=epochs, T_mult=1, eta_min=1e-6,\n                                                                         last_epoch=-1)\n    elif Config.scheduler == 'ReduceLROnPlateau':\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            patience=1,\n            factor=0.25,\n            min_lr=1e-6,\n            verbose=True,\n            mode=\"max\"\n        )\n    # Reduce LR every step size epochs by gamma\n    elif Config.scheduler == 'StepLR':\n        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)\n    else:\n        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n\n    val_scores = []\n    val_svm_scores = []\n\n    \"\"\"\n    # Main training loop over epochs\n    \"\"\"\n    print(f\"Fold: {fold + 1}\")\n    best_epoch_score = 0\n    save_filename = ''\n    for epoch in range(epochs):\n        train_one_epoch(epoch, model, scaler, criterion, optimizer, train_loader, device, scheduler=scheduler,\n                        schd_batch_update=False)\n\n        with torch.no_grad():\n            acc, svm_acc = valid_one_epoch(epoch, model, criterion, val_loader, device, scheduler=None,\n                                           schd_loss_update=False)\n            val_scores.append(acc)\n            val_svm_scores.append(svm_acc)\n\n        # Save model if this is the best epoch so far\n        if acc > best_epoch_score:\n            best_epoch_score = acc\n            if config.save:\n                # Delete previous checkpoint if exists\n                if os.path.exists(f'{SAVE_MODEL_FOLDER}/{save_filename}'):\n                    os.system(f'rm {SAVE_MODEL_FOLDER}/{save_filename}')\n                acc_str = str(acc).replace('.','_')\n                save_filename = f'{config.selected_model}_{config.name}_fold{fold}_epoch{epoch}_{acc_str}.pt'\n                save_model_weights(\n                    model,\n                    save_filename,\n                    cp_folder=SAVE_MODEL_FOLDER,\n                )\n\n\n    # Save model per Fold\n    if config.save:\n        save_model_weights(\n            model,\n            f\"{config.selected_model}_{config.name}_fold{fold}.pt\",\n            cp_folder=SAVE_MODEL_FOLDER,\n        )\n        save_model_params(val_scores, fold, SAVE_MODEL_FOLDER)\n\n        if Config.use_svm:\n            save_model_params(val_svm_scores, fold, SAVE_MODEL_FOLDER, use_svm=True)\n            pickle.dump(svm_clf, open(f'{SAVE_MODEL_FOLDER}/SVM_model.pkl', 'wb'))\n\n    # torch.save(model.cnn_model.state_dict(),'{}/cnn_model_fold_{}_{}'.format(CFG['model_path'], fold, CFG['tag']))\n    del model, optimizer, train_loader, val_loader, scaler, scheduler\n    torch.cuda.empty_cache()\n\n\ndef k_fold(config, df):\n    skf = StratifiedKFold(n_splits=config.k, random_state=config.random_state)\n    splits = list(skf.split(X=df, y=df[TARGET_COL]))\n\n    # Write all config values to save model dir for replication\n    params_file = os.path.join(SAVE_MODEL_FOLDER, 'params.txt')\n    with open(params_file, 'w') as f:\n        f.write('Arguments:\\n')\n        attrs = vars(args)\n        f.write(', '.join(\"%s: %s\" % item for item in attrs.items()))\n        f.write('\\n\\n')\n        f.write('Config:\\n')\n        attrs = vars(Config)\n        f.write(', '.join(\"%s: %s\" % item for item in attrs.items()))\n        f.write('\\n\\n')\n        f.write(f'Train transforms: {get_train_transforms()}\\n\\n')\n        f.write(f'Valid transforms: {get_valid_transforms()}\\n\\n')\n\n    for i, (train_idx, val_idx) in enumerate(splits):\n        print(f\"\\n-------------   Fold {i + 1} / {config.k}  -------------\\n\")\n\n        df_train = df.iloc[train_idx].copy()\n        df_val = df.iloc[val_idx].copy()\n\n        train_model(config, df_train, df_val, i)\n\n\ndef get_scores(folds=[], base_model_path=''):\n    test_df = pd.read_csv(DATA_PATH + '/sample_submission.csv')\n\n    # local\n    # model = ImageClassifier(Config.selected_model, NUM_CLASSES).to(device)\n\n    # kaggle- no internet.\n    # base_model = get_model(Config.selected_model, num_classes=NUM_CLASSES).to(device)\n    # torch.save(base_model, 'effnetb4')\n    base_model = torch.load(base_model_path)\n    model = ImageClassifier(base_model, NUM_CLASSES, infer=True).to(device)\n\n    avg_preds = []\n\n    if Config.tta:\n        for tta in range(Config.tta):\n            test_dataset = CustomDataset(df=test_df, file_path=TEST_FILE_PATH, train=False,\n                                         transforms=get_test_transforms())\n\n            for fold in folds:\n                model = load_model_weights(model, fold, verbose=1, cp_folder=SAVE_MODEL_FOLDER)\n                preds = predict(model, test_dataset, infer=True)\n                avg_preds.append(preds)\n                # avg_preds.append(np.argmax(preds, axis=1))\n                # preds = np.argmax(preds, axis=1)\n\n        avg_preds = np.mean(avg_preds, axis=0)\n    else:\n        test_dataset = CustomDataset(df=test_df, file_path=TEST_FILE_PATH, train=False,\n                                     transforms=get_test_transforms())\n\n        for fold in folds:\n            model = load_model_weights(model, fold, verbose=1, cp_folder=SAVE_MODEL_FOLDER)\n            preds = predict(model, test_dataset, infer=True)\n            avg_preds.append(np.argmax(preds, axis=1))\n            # preds = np.argmax(preds, axis=1)\n\n        avg_preds = np.mean(avg_preds, axis=0)\n\n    return avg_preds\n\n\ndef get_train_transforms():\n    return Compose([\n        #Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], p=1.0),\n        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n        RandomResizedCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        # RandomCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        HorizontalFlip(p=0.5),\n        VerticalFlip(p=0.5),\n        #CoarseDropout(p=0.5),\n        RandomRotate90(p=0.5),\n        #Transpose(p=0.5),\n        #ShiftScaleRotate(p=0.5),\n        #HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n        #RandomBrightnessContrast(brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n        ToTensorV2(p=1.0),\n    ], p=1.)\n\n\ndef get_valid_transforms():\n    return Compose([\n        #Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], p=1.0),\n        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n        # CenterCrop(Config.img_size, Config.img_size, p=1.),\n        #RandomResizedCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        # RandomCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        # RandomRotate90(p=0.5),\n        # HorizontalFlip(p=0.5),\n        ToTensorV2(p=1.0),\n    ], p=1.)\n\n\ndef get_test_transforms():\n    return Compose([\n        #Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], p=1.0),\n        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n        RandomResizedCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        #RandomCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        # RandomRotate90(p=0.5),\n        # HorizontalFlip(p=0.5),\n        ToTensorV2(p=1.0),\n    ], p=1.)\n\n\ndef get_sample_transforms():\n    return Compose([\n        # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n        RandomResizedCrop(Config.aug_img_size, Config.aug_img_size, p=1),\n        HorizontalFlip(p=0.5),\n        VerticalFlip(p=0.5),\n        # RandomRotate90(p=0.5),\n        # Cutout(p=0.5),\n        # Transpose(p=0.5),\n        # ShiftScaleRotate(p=0.5),\n        # HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n        # RandomBrightnessContrast(brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n        CoarseDropout(p=0.5),\n        ToTensorV2(p=1.0),\n    ], p=1.)\n\n\n# Test the data loader and augments\ndef test_loader(df, samples):\n    train_df = df.sample(n=samples)\n    dataset = CustomDataset(df=train_df, file_path=FILE_PATH, train=False, transforms=get_sample_transforms())\n    data_loader = DataLoader(\n        dataset,\n        batch_size=samples,\n        shuffle=True,\n        drop_last=False,\n        pin_memory=False,\n        num_workers=NUM_WORKERS,\n    )\n    for images, orig_images in (data_loader):\n        for i in range(len(images)):\n            plt.figure(figsize=(20, 16))\n            img = images[i]\n            orig_image = orig_images[i]\n            img = img.permute(2, 1, 0)\n            plt.subplot(1, 2, 1)\n            plt.imshow(orig_image)\n            plt.subplot(1, 2, 2)\n            plt.imshow(img)\n            plt.show()\n\n\ndef fit_noisy(\n        model,\n        train_dataset,\n        epochs=50,\n        batch_size=32,\n        val_batch_size=32,\n        warmup_prop=0.1,\n        lr=1e-3,\n        verbose=1,\n        verbose_eval=1\n):\n    \"\"\"\n    Usual torch fit function\n\n    Arguments:\n        model {torch model} -- Model to train\n        train_dataset {torch dataset} -- Dataset to train with\n        val_dataset {torch dataset} -- Dataset to validate with\n\n    Keyword Arguments:\n        epochs {int} -- Number of epochs (default: {50})\n        batch_size {int} -- Training batch size (default: {32})\n        val_bs {int} -- Validation batch size (default: {32})\n        warmup_prop {float} -- Warmup proportion (default: {0.1})\n        lr {float} -- Start (or maximum) learning rate (default: {1e-3})\n        alpha {float} -- alpha value for mixup (default: {0.4})\n        mixup_proba {float} -- Probability to apply mixup (default: {0.})\n        verbose {int} -- Period (in epochs) to display logs at (default: {1})\n        verbose_eval {int} -- Period (in epochs) to perform evaluation at (default: {1})\n\n    Returns:\n        numpy array -- Predictions at the last epoch\n    \"\"\"\n\n    # model.set_callbacks(callbacks)\n\n    optimizer = Adam(model.parameters(), lr=lr)\n\n    # Loss function\n    if Config.loss_fn == 'CrossEntropyLoss':\n        criterion = nn.CrossEntropyLoss().to(device)\n    elif Config.loss_fn == 'SmoothCrossEntropyLoss':\n        criterion = SmoothCrossEntropyLoss().to(device)\n    elif Config.loss_fn == 'BCEWithLogitsLoss':\n        criterion = nn.BCEWithLogitsLoss(reduction=\"mean\").to(device)\n    elif Config.loss_fn == 'TruncatedLoss':\n        criterion = TruncatedLoss(trainset_size=len(train_dataset)).to(device)\n    elif Config.loss_fn == 'MAE':\n        criterion = nn.L1Loss(reduction='mean').to(device)\n    elif Config.loss_fn == 'TaylorCrossEntropyLoss':\n        criterion = TaylorCrossEntropyLoss(smoothing=Config.label_smoothing).to(device)\n    elif Config.loss_fn == 'BiTemperedLoss':\n        criterion = BiTemperedLogisticLoss(t1=Config.t1, t2=Config.t2, smoothing=Config.label_smoothing).to(device)\n    elif Config.loss_fn == 'FocalCosineLoss':\n        criterion = FocalCosineLoss().to(device)\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=False,\n        pin_memory=False,\n        num_workers=NUM_WORKERS,\n    )\n\n    num_warmup_steps = int(warmup_prop * epochs * len(train_loader))\n    num_training_steps = int(epochs * len(train_loader))\n\n    if Config.scheduler == 'CosineAnnealingLR':\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=epochs, T_mult=1, eta_min=1e-6,\n                                                                         last_epoch=-1)\n    elif Config.scheduler == 'ReduceLROnPlateau':\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            patience=1,\n            factor=0.25,\n            min_lr=1e-6,\n            verbose=True,\n            mode=\"max\"\n        )\n    # Reduce LR every step size epochs by gamma\n    elif Config.scheduler == 'StepLR':\n        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)\n    else:\n        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n\n    for epoch in range(epochs):\n        model.train()\n        start_time = time.time()\n        optimizer.zero_grad()\n\n        avg_loss = 0\n        for step, (x, y_batch) in enumerate(train_loader):\n            y_pred = model(x.to(device))\n            y_pred_labels = torch.max(y_pred, 1)[1]\n            y_true = y_batch.to(device).long()\n            y_true_labels = torch.max(y_true, 1)[1]\n            if Config.loss_fn == 'CrossEntropyLoss':\n                loss = criterion(y_pred, y_true_labels)\n            elif Config.loss_fn == 'BiTemperedLoss':\n                loss = criterion(y_pred, y_true_labels)\n            else:\n                loss = criterion(y_pred, y_true_labels)\n\n            loss.backward()\n            avg_loss += loss.item() / len(train_loader)\n            optimizer.step()\n            # Zero out gradients for next loop\n            optimizer.zero_grad()\n\n        # Update learning rate\n        if Config.scheduler == 'CosineAnnealingLR':\n            scheduler.step(epoch)\n        else:\n            scheduler.step()\n\n        # Calculate learning rate\n        lr = scheduler.get_last_lr()[0]\n\n        elapsed_time = time.time() - start_time\n        if (epoch + 1) % verbose == 0:\n            elapsed_time = elapsed_time * verbose\n            print(\n                f\"Epoch {epoch + 1}/{epochs} \\t lr={lr:.1e} \\t t={elapsed_time:.0f}s  \\t loss={avg_loss:.4f}\\n\",\n                end=\"\",\n            )\n\n    torch.cuda.empty_cache()\n\n\ndef get_scores_train(folds=[], base_model=''):\n    train_df_path = os.path.join(DATA_PATH, TRAIN_DF_FILE)\n    test_df = pd.read_csv(train_df_path)\n\n    # local\n    model = get_model(Config.selected_model, num_classes=NUM_CLASSES).to(device)\n    avg_preds = []\n\n    if Config.tta:\n        for tta in range(Config.tta):\n            test_dataset = CustomDataset(df=test_df, file_path=FILE_PATH, train=False,\n                                         transforms=get_test_transforms())\n\n            for fold in folds:\n                model = load_model_weights(model, fold, verbose=1, cp_folder=base_model)\n                preds = predict(model, test_dataset, infer=True)\n                avg_preds.append(preds)\n                # avg_preds.append(np.argmax(preds, axis=1))\n                # preds = np.argmax(preds, axis=1)\n\n        avg_preds = np.mean(avg_preds, axis=0)\n    else:\n        test_dataset = CustomDataset(df=test_df, file_path=FILE_PATH, train=False,\n                                     transforms=get_test_transforms())\n\n        for fold in folds:\n            model = load_model_weights(model, fold, verbose=1, cp_folder=base_model)\n            preds = predict(model, test_dataset, infer=True)\n            avg_preds.append(preds)\n            # preds = np.argmax(preds, axis=1)\n\n        avg_preds = np.mean(avg_preds, axis=0)\n\n    preds_df = pd.DataFrame()\n    preds_df[ID_COL] = test_df[ID_COL]\n    preds_df[TARGET_COL] = test_df[TARGET_COL]\n    preds_df['pred'] = avg_preds.tolist()\n\n    return preds_df, avg_preds\n\n\nCV_FOLD = 0\n\n\nclass Classifier(BaseEstimator):\n    def __init__(self, model, config):\n        self.model = model\n        self.config = config\n        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n        self.model.to(self.device)\n        self.best_model = None\n\n    def fit(self, img_idx, labels):\n        global CV_FOLD\n        print(f\"CV Fold: {CV_FOLD + 1}\")\n        CV_FOLD += 1\n        df_train = train_df.iloc[img_idx].copy()\n        df_train = df_train.reset_index().drop('index', axis=1)\n        train_dataset = CustomDataset(df=df_train, file_path=FILE_PATH, train=True, transforms=get_train_transforms())\n\n        fit_noisy(\n            self.model,\n            train_dataset,\n            epochs=self.config.epochs,\n            batch_size=self.config.batch_size,\n            val_batch_size=self.config.val_batch_size,\n            lr=self.config.lr,\n            warmup_prop=self.config.warmup_prop,\n            verbose_eval=self.config.verbose_eval\n        )\n\n    def predict_proba(self, img_idx, phase=\"train\"):\n        df_val = train_df.iloc[img_idx].copy()\n        df_val = df_val.reset_index().drop('index', axis=1)\n        test_dataset = CustomDataset(df=df_val, file_path=FILE_PATH, train=False, transforms=get_test_transforms())\n        prob = predict(self.model, test_dataset, infer=True)\n        return prob\n\n    def predict(self, img_idx, phase=\"train\"):\n        prob = self.predict_proba(img_idx, phase=phase)\n        preds = np.argmax(prob, axis=1)\n        return preds\n\n    def score(self, img_idx, label, phase=\"train\"):\n        preds = self.predict(img_idx, phase=phase)\n        return accuracy_score(label, preds)\n\n\nclass Config:\n    selected_model = 'resnet50'\n    scheduler = 'CosineAnnealingLR'\n    loss_fn = 'TaylorCrossEntropyLoss'\n    augment = False\n    tta = 3\n    # Value to resize to\n    img_size = 512\n    aug_img_size = 512\n    # Hyper params\n    batch_size = 16\n    val_batch_size = 16\n    epochs = 10\n    lr = 1e-4\n    weight_decay = 1e-6\n\n    # General\n    seed = 1234\n    # Iterations to show loss\n    verbose = 1\n    # Iterations to accumulate grads\n    accum_iter = 4\n    # verbose_eval = 31\n    save = True\n    # k-fold\n    k = 5\n    random_state = None\n    # Model\n    use_msd = False\n    use_conf = False\n\n    warmup_prop = 0.05\n    # For bitempered loss\n    t1 = 0.8\n    t2 = 1.4\n    # Label smoothing\n    label_smoothing = 0.2\n    alpha = 5\n    name = \"extra\"\n    cv_mask = False\n    test_loader = False\n    cutmix = False\n    use_svm = False\n\n\ntrain_df_path = os.path.join(DATA_PATH, TRAIN_DF_FILE)\ntrain_df = pd.read_csv(train_df_path)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Command line arguments"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nparser = argparse.ArgumentParser(description='Process pytorch params.')\nparser.add_argument('-model', '--model', type=str, help='Pytorch (timm) model name')\nparser.add_argument('-model_dir', '--model_dir', type=str, help='Model save dir name')\nparser.add_argument('--folds', type=int, help='Number of folds')\nparser.add_argument('-epochs', '--epochs', type=int, help='Number of epochs')\nparser.add_argument('-img_size', '--img_size', type=int, help='Image size to resize')\nparser.add_argument('-batch_size', '--batch_size', type=int, help='batch size')\nparser.add_argument('-lr', '--lr', type=float, help='Learning rate')\nparser.add_argument('--augment', action='store_true', help='Augment data')\nparser.add_argument('--loss_fn', type=str, help='Loss function')\nparser.add_argument('--scheduler', type=str, help='Scheduler')\nparser.add_argument('--n_samples', type=int, help='Number of samples of train data')\nparser.add_argument('--train', action='store_true', help='Pytorch train mode')\nparser.add_argument('--test', action='store_true', help='Pytorch test mode')\nparser.add_argument('--find_model', type=str, help='Regex for timm model search')\nparser.add_argument('--cv_mask', action='store_true', help='Pytorch train mode')\nparser.add_argument('--stats', action='store_true', help='Pytorch train mode')\nparser.add_argument('--test_loader', action='store_true', help='Pytorch train mode')\nparser.add_argument('--downsample', action='store_true', help='Downsample')\nparser.add_argument('--upsample', type=float, help='Upsample pct e.g value of 0.5 is 50% of majority class')\nparser.add_argument('--get_noise_indices', action='store_true', help='Get noise indices from cleanlab')\nparser.add_argument('--train_noisy_clf', action='store_true', help='Train a noisy classifier (Cleanlab)')\nparser.add_argument('--clean_data', type=str, help='Remove user specified indices from train')\nparser.add_argument('--clean_samples', type=int, help='Remove user specified indices from train')\nparser.add_argument('--train_df', type=str, help='Remove user specified indices from train')\nparser.add_argument('--cutmix', action='store_true', help='Cutmix for images')\nparser.add_argument('--use_svm', action='store_true', help='Use SVM on top of model')\n\nargs = parser.parse_args()\n\n# User defined train df\nif args.train_df:\n    train_df = pd.read_csv(args.train_df)\n\nNUM_CLASSES = train_df[TARGET_COL].nunique()\nprint(f\"Number of classes: {NUM_CLASSES}\")\n\n# %% [code]\nONE_HOT = np.eye(NUM_CLASSES)\nle = LabelEncoder()\ntrain_df[TARGET_COL + '_encoded'] = le.fit_transform(train_df[TARGET_COL])\n\n# Parse arguments\nConfig.selected_model = args.model\nConfig.cv_mask = args.cv_mask\nConfig.cutmix = args.cutmix\n\nif args.img_size:\n    Config.img_size = args.img_size\nif args.batch_size:\n    Config.batch_size = args.batch_size\nif args.folds:\n    Config.k = args.folds\nif args.epochs:\n    Config.epochs = args.epochs\nif args.augment:\n    Config.augment = args.augment\nif args.loss_fn:\n    Config.loss_fn = args.loss_fn\nif args.scheduler:\n    Config.scheduler = args.scheduler\nif args.lr:\n    Config.lr = args.lr\nif args.use_svm:\n    Config.use_svm = args.use_svm\n    svm_clf = svm.LinearSVC(C=1, verbose=0, max_iter=100000, loss='squared_hinge', penalty='l2', dual=True)\n\nprint(f'Device:{device}')\nprint(f'Arguments: {args}\\n')\n\nif args.stats:\n    print(f'Train shape: {train_df.shape}')\n    print(f'Target dist:\\n{train_df[TARGET_COL].value_counts()}')\n\nif args.find_model:\n    model_names = timm.list_models(f'*{args.find_model}*')\n    print(model_names)\n\n# Try a subset of data\nif args.n_samples:\n    train_df = train_df.sample(n=args.n_samples)\n\n# FIXME - move this to inside fold/epoch loop\nif args.downsample:\n    majority_class = 3\n    train_df_min = train_df[train_df[TARGET_COL] != majority_class]\n    samples = int(np.ceil(train_df_min[TARGET_COL].value_counts().max()))\n    train_df_maj = train_df[train_df[TARGET_COL] == majority_class]\n    df_majority_downsampled = resample(train_df_maj,\n                                       replace=False,  # sample without replacement\n                                       n_samples=samples,  # to match minority class\n                                       random_state=123)\n    train_df = pd.concat([train_df_min, df_majority_downsampled])\n    # print(train_df[TARGET_COL].value_counts())\n\nif args.model_dir:\n    model_dir = args.model_dir\n    SAVE_MODEL_FOLDER = BASE_MODEL_FOLDER + model_dir\n    if not os.path.exists(SAVE_MODEL_FOLDER):\n        os.makedirs(SAVE_MODEL_FOLDER)\n\nif args.train:\n    k_fold(Config, train_df)\n\nif args.test_loader:\n    Config.test_loader = True\n    test_loader(train_df, 4)\n\nif args.test:\n    SAVE_MODEL_FOLDER = '../input/image-pretrained-models/'\n    test_df = pd.read_csv(DATA_PATH + '/sample_submission.csv')\n    ensemble = False\n    if ensemble:\n        preds1 = get_scores(['tf_efficientnet_b4_ns_extra_2.pt'], '../input/base-pretrained/effnetb4')\n        preds2 = get_scores(['seresnet50_extra_2.pt'], '../input/base-pretrained/seresnet50')\n        preds = 0.5 * preds1 + 0.5 * preds2\n    else:\n        # preds = get_scores(['tf_efficientnet_b4_ns_extra_2_512_1_4.pt'], '../input/base-pretrained/effnetb4')\n        preds = get_scores(['seresnet50_extra_0.pt'], '../input/base-pretrained/seresnet50')\n\n    submission_df = pd.DataFrame()\n    submission_df[ID_COL] = test_df[ID_COL]\n    submission_df[TARGET_COL] = softmax(preds).argmax(1)\n    submission_df[TARGET_COL] = submission_df[TARGET_COL].astype(int)\n\n    print(submission_df.head())\n    submission_df.to_csv('submission.csv', index=False)\n\nif args.get_noise_indices:\n    # Prediction\n    preds_df, psx = get_scores_train(['./models/seresnet50_1_3_test/seresnet50_extra_2.pt'])\n\n    labels = preds_df[TARGET_COL].values\n\n    ordered_label_errors = get_noise_indices(\n        s=labels,\n        psx=psx,\n        sorted_index_method='normalized_margin',  # Orders label errors\n    )\n\n    print(f'Label errors found={len(ordered_label_errors)}')\n\n    errors_df = preds_df.iloc[ordered_label_errors, :].reset_index()\n\n    print(errors_df.shape)\n    print(errors_df.head())\n    errors_df.to_csv('errors.csv', index=False)\n\nif args.train_noisy_clf:\n    train_image_id = train_df[ID_COL].values\n    train_label = train_df[TARGET_COL].values\n    val_preds = np.zeros((train_label.shape[0], 5))\n    kfold = StratifiedKFold(n_splits=Config.k, random_state=None)\n    seed_everything(0)\n\n    val_scores = []\n\n    for fold, (train_idx, val_idx) in enumerate(kfold.split(train_image_id, train_label)):\n        CV_FOLD = 0\n        print(f'Fold: {fold + 1}')\n        X_train, y_train = train_image_id[train_idx], train_label[train_idx]\n        X_val, y_val = train_image_id[val_idx], train_label[val_idx]\n\n        base_model = get_model(Config.selected_model, num_classes=NUM_CLASSES).to(device)\n        model = Classifier(base_model, Config)\n        lnl = LearningWithNoisyLabels(clf=model, seed=0, n_jobs=os.cpu_count(), cv_n_folds=5)\n        clf = lnl.fit(train_idx, y_train)\n\n        val_preds[val_idx, :] = clf.predict_proba(val_idx)\n        acc = accuracy_score(y_val, np.argmax(val_preds[val_idx, :], axis=1))\n        val_scores.append(acc)\n        print(f'Accuracy:{acc}')\n        # Save model per Fold\n        if Config.save:\n            save_model_weights(\n                clf.model,\n                f\"{Config.selected_model}_{Config.name}_{fold}.pt\",\n                cp_folder=SAVE_MODEL_FOLDER,\n            )\n            save_model_params(val_scores, fold, SAVE_MODEL_FOLDER)\n\nif args.clean_data:\n    errors_df = pd.read_csv(args.clean_data)\n    errors_df = errors_df.iloc[0:args.clean_samples - 1, :]\n    error_indices = errors_df['index'].values\n    print(f\"Input shape: {train_df.shape}\")\n    train_df = train_df.loc[~train_df.index.isin(error_indices)]\n    print(f'Removing {args.clean_samples} samples from the dataset')\n    print(f\"Output shape: {train_df.shape}\")\n    train_df.to_csv('mod_train_df.csv', index=False)\n","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}