{"cells":[{"metadata":{},"cell_type":"markdown","source":"This is the training notebook of the VIT 2020 model which was used for the our final submission which scored 91.3% on public and private leaderboard of the Cassava Leaf Disease Classification 2020 competition (1st place).\n\nSpecial thanks go to [Mobassir](https://www.kaggle.com/mobassir) who has posted his notebook on VIT TPU training, the basis of this notebook."},{"metadata":{},"cell_type":"markdown","source":"# Installation of TPU dependencies"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!python3.7 -m pip install --upgrade pip\n!pip install   torch==1.7.0\n!pip install  torchvision \n!pip install  cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# imports "},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import os\nimport torch\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.utils.serialization as xser\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\nfrom warmup_scheduler import GradualWarmupScheduler\nimport sys; \npackage_paths = [\n    '../input/pytorch-image-models/pytorch-image-models-master',\n    '../input/image-fmix/FMix-master'\n]\n\n\nfor pth in package_paths:\n    sys.path.append(pth)\n    \nimport warnings\nimport pandas as pd\nimport numpy as np\nimport torch.nn as nn\nfrom sklearn.model_selection import train_test_split\nfrom sklearn import metrics\nfrom transformers import get_linear_schedule_with_warmup\nimport time\nimport torchvision\nimport torch.nn as nn\nfrom tqdm import tqdm_notebook as tqdm\nfrom PIL import Image, ImageFile\nfrom torch.utils.data import Dataset\nimport torch.optim as optim\nfrom torchvision import transforms\nfrom torch.optim import lr_scheduler\nimport sys\nimport gc\nimport os\nimport random\nimport skimage.io\nfrom PIL import Image\nimport scipy as sp\nimport sklearn.metrics\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold\nfrom functools import partial\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.models as models\nfrom albumentations import Compose, Normalize, HorizontalFlip, VerticalFlip\nfrom albumentations.pytorch import ToTensorV2\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, ShiftScaleRotate, CenterCrop, Resize\n)\nfrom contextlib import contextmanager\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nfrom fmix import sample_mask, make_low_freq_image, binarise_mask\nfrom glob import glob\nfrom sklearn.model_selection import GroupKFold, StratifiedKFold\nimport cv2\nfrom skimage import io\nimport torch\nfrom torch import nn\nimport os\nfrom datetime import datetime\nimport time\nimport random\n\nimport torchvision\nfrom torchvision import transforms\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\n\nimport albumentations\n\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.nn.modules.loss import _WeightedLoss\nimport torch.nn.functional as F\n\nimport timm\n\nimport sklearn\n\nimport joblib\nfrom sklearn.metrics import roc_auc_score, log_loss\nfrom sklearn import metrics\nimport cv2\nimport pydicom\nfrom scipy.ndimage.interpolation import zoom\n\n\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"\n\nwarnings.filterwarnings(\"ignore\")\n\nprint(torch.__version__)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Configuration of parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG = {\n    'fold_num': 5,\n    'seed': 719,\n    'model_arch': 'vit_base_patch16_384',\n    'img_size': 384,\n    'epochs': 10,\n    'train_bs': 10,\n    'valid_bs': 10,\n    'T_0': 10,\n    'lr': 1e-4,\n    'min_lr': 1e-4, \n    'smoothing' : 0.06,\n    't1' : 0.8,\n    't2' : 1.4,\n    'warmup_factor' : 7,\n    'warmup_epo' : 1,\n    'num_workers': 4,#8\n    'accum_iter': 2, # batch accumulation for backprop with effectively larger batch size\n    'verbose_step': 1\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Reading train.csv file"},{"metadata":{},"cell_type":"markdown","source":"We load here meta data necessary for the out of fold training and prediction process."},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\n\ntrain = pd.read_csv('../input/cassava-leaf-disease-metadata-vit/train_j.csv')\nval = pd.read_csv('../input/cassava-leaf-disease-metadata-vit/val_j.csv')\n\ntrain = train.rename(columns={'target': 'label', 'img_name' : \"image_id\"})\nval = val.rename(columns={'target': 'label', 'img_name': \"image_id\"})\n\n\ndef get_img(path):\n    im_bgr = cv2.imread(path)\n    im_rgb = im_bgr[:, :, ::-1]\n    return im_rgb","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset Class\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rand_bbox(size, lam):\n    W = size[0]\n    H = size[1]\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\nclass CassavaDataset(Dataset):\n    def __init__(self, df, data_root, \n                 transforms=None, \n                 output_label=True, \n                 one_hot_label=False,\n                 do_fmix=False, \n                 fmix_params={\n                     'alpha': 1., \n                     'decay_power': 3., \n                     'shape': (CFG['img_size'], CFG['img_size']),\n                     'max_soft': True, \n                     'reformulate': False\n                 },\n                 do_cutmix=False,\n                 cutmix_params={\n                     'alpha': 1,\n                 }\n                ):\n        \n        super().__init__()\n        self.df = df.reset_index(drop=True).copy()\n        self.transforms = transforms\n        self.data_root = data_root\n        self.do_fmix = do_fmix\n        self.fmix_params = fmix_params\n        self.do_cutmix = do_cutmix\n        self.cutmix_params = cutmix_params\n        \n        self.output_label = output_label\n        self.one_hot_label = one_hot_label\n        \n        if output_label == True:\n            self.labels = self.df['label'].values\n\n            if one_hot_label is True:\n                self.labels = np.eye(self.df['label'].max()+1)[self.labels]\n\n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, index: int):\n        \n        # get labels\n        if self.output_label:\n            target = self.labels[index]\n          \n        img  = get_img(\"{}/{}\".format(self.data_root, self.df.loc[index]['image_id']))\n\n        if self.transforms:\n            img = self.transforms(image=img)['image']\n        \n        if self.do_fmix and np.random.uniform(0., 1., size=1)[0] > 0.5:\n            with torch.no_grad():\n                lam = np.clip(np.random.beta(self.fmix_params['alpha'], self.fmix_params['alpha']),0.6,0.7)\n                \n                # Make mask, get mean / std\n                mask = make_low_freq_image(self.fmix_params['decay_power'], self.fmix_params['shape'])\n                mask = binarise_mask(mask, lam, self.fmix_params['shape'], self.fmix_params['max_soft'])\n    \n                fmix_ix = np.random.choice(self.df.index, size=1)[0]\n                fmix_img  = get_img(\"{}/{}\".format(self.data_root, self.df.iloc[fmix_ix]['image_id']))\n\n                if self.transforms:\n                    fmix_img = self.transforms(image=fmix_img)['image']\n\n                mask_torch = torch.from_numpy(mask)\n                \n                # mix image\n                img = mask_torch*img+(1.-mask_torch)*fmix_img\n\n                rate = mask.sum()/CFG['img_size']/CFG['img_size']\n                target = rate*target + (1.-rate)*self.labels[fmix_ix]\n       \n        if self.do_cutmix and np.random.uniform(0., 1., size=1)[0] > 0.5:\n            with torch.no_grad():\n                cmix_ix = np.random.choice(self.df.index, size=1)[0]\n                cmix_img  = get_img(\"{}/{}\".format(self.data_root, self.df.iloc[cmix_ix]['image_id']))\n                if self.transforms:\n                    cmix_img = self.transforms(image=cmix_img)['image']\n                    \n                lam = np.clip(np.random.beta(self.cutmix_params['alpha'], self.cutmix_params['alpha']),0.3,0.4)\n                bbx1, bby1, bbx2, bby2 = rand_bbox((CFG['img_size'], CFG['img_size']), lam)\n\n                img[:, bbx1:bbx2, bby1:bby2] = cmix_img[:, bbx1:bbx2, bby1:bby2]\n\n                rate = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (CFG['img_size'] * CFG['img_size']))\n                target = rate*target + (1.-rate)*self.labels[cmix_ix]\n                      \n        if self.output_label == True:\n            return img, target\n        else:\n            return img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentation\n\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return Compose([\n            RandomResizedCrop(CFG['img_size'], CFG['img_size']),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(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            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            CoarseDropout(p=0.5),\n            Cutout(p=0.5),\n            ToTensorV2(p=1.0),\n        ], p=1.)\n  \n        \ndef get_valid_transforms():\n    return Compose([\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            CenterCrop(CFG['img_size'], CFG['img_size'], p=1.),\n            Resize(CFG['img_size'], CFG['img_size']),\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            ToTensorV2(p=1.0),\n        ], p=1.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataloader\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_dataloader(train_, valid_, data_root='../input/cassava-leaf-disease-classification/train_images/'):\n    \n    from catalyst.data.sampler import BalanceClassSampler\n    \n    train_ds = CassavaDataset(train_, data_root, transforms=get_train_transforms(), output_label=True, one_hot_label=False, do_fmix=0, do_cutmix=0)\n    valid_ds = CassavaDataset(valid_, data_root, transforms=get_valid_transforms(), output_label=True)\n    \n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_ds,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True)\n\n\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n        valid_ds,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False,\n        )\n    \n    \n    train_loader = torch.utils.data.DataLoader(\n        dataset=train_ds,\n        batch_size=CFG['train_bs'],\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=CFG['num_workers'],\n    )\n   \n    valid_loader = torch.utils.data.DataLoader(\n        dataset=valid_ds,\n        batch_size=CFG['valid_bs'],\n        sampler=valid_sampler,\n        drop_last=True,\n        num_workers=CFG['num_workers'],\n    )\n    return train_loader, valid_loader\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"# 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\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,\n        t2,\n        label_smoothing=0.0,\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    \n    \n\ndef get_probs(activations,\n        labels,\n        t1,\n        t2,\n        label_smoothing=0.0,\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    return probabilities","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training and Validation\n\n1. when we will call train_one_epoch() function, we will be training our model **for 1 epoch** every time using the train_loader that we have prepared above for 8 core tpu training,**from *train_loader* we are using batch images and labels and later feeding them to our model,calculating loss,doing backpropagation,doing gradient accumulation,taking optimizer step and scheduler step**\n\n2. valid_one_epoch() function uses our validation dataloader that we have created above for 8 core training on tpu.it calculates validation loss and validation accuracy using our validation dataset and returns validation accuracy, while calculatin validation accuracy we should do this : **accuracy = xm.mesh_reduce('test_accuracy', acc, np.mean)**,if you don't do this then you won't be able to save best weight file for each fold and xm.save() will hang for forever(i made this silly mistake few days ago while working on this kernel,so keep this in mind)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(epoch, model, optimizer, train_loader, device, scheduler=None, schd_batch_update=False):\n    model.train()\n\n    t = time.time()\n\n    pbar = tqdm(enumerate(train_loader), total=len(train_loader))\n    z = 0\n    for step, (imgs, image_labels) in pbar:\n        z = z + 1\n        if z % 20 == 0:\n            gc.collect()\n        \n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n\n        with autocast():\n            image_preds = model(imgs)   #output = model(input)\n\n            #loss = loss_fn(image_preds, image_labels)\n            loss = bi_tempered_logistic_loss(image_preds, image_labels, t1=CFG['t1'], t2=CFG['t2'], label_smoothing=CFG['smoothing'])\n            \n            #scaler.scale(loss).backward()\n            \n            loss.backward()\n\n            if ((step + 1) %  CFG['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                xm.optimizer_step(optimizer)\n                #scaler.update()\n                optimizer.zero_grad() \n                \n                if scheduler is not None and schd_batch_update:\n                    scheduler.step()\n    pbar.close()\n    if scheduler is not None and not schd_batch_update:\n        scheduler.step()\n        \ndef valid_one_epoch(epoch, model, val_loader, device):\n    model.eval()\n\n    t = time.time()\n    loss_sum = 0\n    sample_num = 0\n    image_preds_all = []\n    image_targets_all = []\n\n    pbar = tqdm(enumerate(val_loader), total=len(val_loader))\n    for step, (imgs, image_labels) in pbar:\n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n        \n        image_preds = model(imgs)   #output = model(input)\n        xm.mark_step()\n        image_preds_all += [torch.argmax(image_preds, 1).detach().cpu().numpy()]\n        image_targets_all += [image_labels.detach().cpu().numpy()]\n        \n        #loss = loss_fn(image_preds, image_labels)\n        loss = bi_tempered_logistic_loss(image_preds, image_labels, t1=CFG['t1'], t2=CFG['t2'], label_smoothing=CFG['smoothing'])\n        \n        #loss_sum += loss*image_labels.shape[0]\n        #sample_num += image_labels.shape[0]  \n\n    pbar.close()\n    image_preds_all = np.concatenate(image_preds_all)\n    image_targets_all = np.concatenate(image_targets_all)\n    acc = (image_preds_all==image_targets_all).mean()\n    #LOGGER.debug('validation multi-class accuracy = {:.4f}'.format(acc))\n    accuracy = xm.mesh_reduce('test_accuracy', acc, np.mean)\n    xm.master_print(\"Validation Accuracy = \",accuracy)\n   \n    return loss,accuracy","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**CassvaImgClassifier()** class below is used for loading pretrained **vit_base_patch16_384** model from [timm](https://github.com/rwightman/pytorch-image-models/tree/master/timm)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassvaImgClassifier(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        #print(self.model)\n        n_features = self.model.head.in_features\n        self.model.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loading the Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = CassvaImgClassifier(CFG['model_arch'], train.label.nunique(), pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train Model\n\n**In order to train the model, we need to spawn the training processes on each of the TPU cores.we will call train_model() function during spawn call**\n\nwe use StratifiedKFold for reducing overfitting,calling prepare_dataloader() function for preparing train and valid loader,Adam optimizer,CosineAnnealingWarmRestarts scheduler,CrossEntropyLoss as our loss function and inside epoch loop saving best weight of each fold and doing **train_one_epoch() and valid_one_epoch()** as discussed above"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_model(folds = range(0, 5)):\n    device = xm.xla_device()\n    model.to(device)\n        \n    \n    for fold in folds:\n        trn_idx = train[train['fold'] == fold]\n        val_idx = val[val['fold'] == fold]\n        trn_idx = pd.concat([trn_idx, val_idx])\n\n        lr =  CFG['lr']* xm.xrt_world_size()\n        optimizer = torch.optim.Adam(model.parameters(), lr=lr/CFG['warmup_factor'])\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=CFG['T_0'], T_mult=1, eta_min=CFG['min_lr'], last_epoch=-1)\n\n        best_accuracy = 0.0\n        \n        for epoch in range(CFG['epochs']):\n            gc.collect()\n            train_loader, val_loader = prepare_dataloader(trn_idx, val_idx, data_root='../input/cassava-leaf-disease-classification/train_images/')\n            para_loader = pl.ParallelLoader(train_loader, [device])\n            train_one_epoch(epoch, model, optimizer, para_loader.per_device_loader(device), device, scheduler=scheduler, schd_batch_update=False)\n    \n            del para_loader\n            del(train_loader)\n            gc.collect()\n            para_loader = pl.ParallelLoader(val_loader, [device])\n            val_loss,cur_accuracy = valid_one_epoch(epoch, model, para_loader.per_device_loader(device), device)\n            \n            del para_loader\n            del(val_loader)\n            gc.collect()\n\n            content = time.ctime() + ' ' + f'FOLD -> {fold} --> Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, val_loss : {(val_loss):.5f},Validation_Accuracy: {(cur_accuracy):.5f}'\n\n            with open(f'log.txt', 'a') as appender:\n                appender.write(content + '\\n')\n            \n            xm.save(model.state_dict(),'{}_fold_{}_best_epoch_{}_Validation_Accuracy.h5'.format(CFG['model_arch'], epoch,cur_accuracy))\n            \n             \n            if cur_accuracy >= best_accuracy:\n                xm.save(model.state_dict(),'{}_fold_{}_best_epoch'.format(CFG['model_arch'], fold))\n                best_accuracy = cur_accuracy\n                    \n        \n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Start training processes"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"%%time\n\n#important, we specify here the fold to train (one out of [0,4])\nfolds_to_train  = [1]\n\ndef _mp_fn(rank, flags):\n    global acc_list\n    torch.set_default_tensor_type('torch.FloatTensor')\n    res = train_model(folds=folds_to_train)\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()\ntorch.cuda.empty_cache()\ntry:\n    del(model)\n    \nexcept:\n    pass\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Full Training Log"},{"metadata":{"trusted":true},"cell_type":"code","source":"f = open(f'./log.txt', \"r\")\n\nprint(f.read())","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}