{"cells":[{"metadata":{},"cell_type":"markdown","source":"**Hello kagglers,**\n\ni hope you are doing well in this pandemic time. This is Going to be my first Kaggle Notebook of this year and i've decided to give Pytorch XLA a try.last time i tried Pytorch XLA here [Pytorch TPU Transfer Learning Baseline](https://www.kaggle.com/mobassir/pytorch-tpu-transfer-learning-baseline?scriptVersionId=34030523) exactly 8 months ago on a image steganalysis problem.\n\nIn this Notebook i will try to convert Awesome GPU Baseline kernel \n[**Pytorch Efficientnet Baseline [Train] AMP+Aug** ](https://www.kaggle.com/khyeh0719/pytorch-efficientnet-baseline-train-amp-aug) of Mr. [Kun Hao Yeh](https://www.kaggle.com/khyeh0719) into TPU baseline, i hope it will help those who don't have personal GPU,love pytorch like me,want to do more experiments by investing less time(compared to GPU training) and also for beginners like me who want to learn how to make TPU baseline for a given data problem.\n"},{"metadata":{},"cell_type":"markdown","source":"**in this kernel [Faster torch tpu baseline for leaf disease](https://www.kaggle.com/mobassir/faster-torch-tpu-baseline-for-leaf-disease) i tried to convert qishen ha san's magnificent GPU baseline into tpu for this competition,you can give that a try as well**"},{"metadata":{},"cell_type":"markdown","source":"# ChangeLOG\n\n* **v1 :** replacing tf_efficientnet_b4_ns with vit_large_patch16_384 model, reduced image size down to 384x384x3, batch_size 8x8 = 64, 1 epoch took ~13 minutes so going for 15 epochs and training fold 0 only and also saving best weight file only(keeping everything else as it was in the GPU baseline)\n* **v2 :** version 1 commit didn't finish(don't know why),just training for 1 epoch now\n* **v3 :** vit_base_patch16_384\n* **v4 :** training for 10 epochs(stopped)\n* **v5 :** reducing batch size and workers and training for 10 epochs (failed)\n* **v6 :** i faced this problem in version 5 [Tpu kernel commit not finishing but worked well during interactive session](https://www.kaggle.com/discussion/210032) , doing xser.save() for saving weights and trying for single fold,5 epochs\n\n* **v7 :** after 1 epoch training stopped,don't know why(probably because of xser), going for 10 epoch and without saving weight file\n* **v8 :** running for 12 epochs and saving weight after each and every epoch(not best weight) - (failed)\n* **v9 :** trying to save weight of last epoch only\n* **v10 :** trying adamW with GradualWarmupScheduler for 10 epoch\n* **v11 :** increasing batch size,using batch size 10x8 = 80 for both train and valid,removing all unnecessary .item() call and before transfering things to cpu, doing mark_step,which will get faster performance once compilations stabilize.training for 12 epochs\n* **v12 :** 15 epochs,warmup_factor = 10,increasing batch size more,going for 12x8 = 96 batch size for train and valid\n* **v13 :** forgot to select tpu in version 12,fixing that\n* **v14 :** running for 10 epoch without scheduler\n* **v15 :**  in [**version 9 of ths notebook**](https://www.kaggle.com/mobassir/faster-torch-tpu-baseline-for-leaf-disease?scriptVersionId=51758471&select=vit_large_patch16_384_best_fold.pth) finally i was able to save best weight of fold 0,trying to save best weight here in version 15 as well,if it works then we can try 5 fold training and save best weight of each fold for further inference.\n* **v16 :**  GradualWarmupScheduler with **bi tempered logistic loss** (failed)\n* **v17 :** reducing batch size and running again\n* **v18 :** using CosineAnnealingWarmRestarts scheduler\n* **v19 :**   fold 1 training\n* **v20 :**   fold 2 training\n* **v21 :**   fold 3 training,12 epochs(stopped)\n* **v22 :**   fold 3 training,12 epochs\n* **v23 :**   fold 4 training,12 epochs (failed)\n* **v24 :**   fold 4 training,10 epochs\n* **v25 :** testing on tpu pod\n"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"'''\n\n!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\n    \n''' ","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"#!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n#!python pytorch-xla-env-setup.py --version 1.7\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# install latest version of torch xla nightly"},{"metadata":{"trusted":true,"_kg_hide-input":false,"_kg_hide-output":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version \"nightly\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# install gradual warmup scheduler"},{"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', #'../input/efficientnet-pytorch-07/efficientnet_pytorch-0.7.0'\n    '../input/image-fmix/FMix-master'\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 AdamW\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\n#from efficientnet_pytorch import EfficientNet\nfrom scipy.ndimage.interpolation import zoom\n\n\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# For parallelization in TPUs"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# check version of torch"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nprint(torch.__version__)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Configuration Dictionary"},{"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': 2,\n    'valid_bs': 2,\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': 2,\n    'accum_iter': 2, # suppoprt to do 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":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Class Distribution"},{"metadata":{"trusted":true},"cell_type":"code","source":"train.label.value_counts()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# visualizing a single image from train_images folder"},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef get_img(path):\n    im_bgr = cv2.imread(path)\n    im_rgb = im_bgr[:, :, ::-1]\n    return im_rgb\n\nimg = get_img('../input/cassava-leaf-disease-classification/train_images/4083726805.jpg')\nplt.imshow(img)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset Class\n\nCassavaDataset() class is our Custom dataset class,we can pass several parameters there,for example choosing do_fmix=True will let us apply [fmix](https://arxiv.org/abs/2002.12047) in our dataset,if we choose do_cutmix=True while calling this dataset class,means we will be applying [cutmix](https://www.kaggle.com/c/bengaliai-cv19/discussion/126504) on our dataset"},{"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\nwe will call get_train_transforms() function while loading training dataset and get_valid_transforms() function while loading validation dataset. get_valid_transforms() helps us do centercropping.resize image using the image size that we defined in CFG dictionary and normalize pixel values using the mean and std values that were used while training models on imagenet dataset.similarly,get_train_transforms() function applies several augmentations like horizontalflip,verticalflip,coarsedropout etc on our training dataset"},{"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\n* inside prepare_dataloader() we are calling CassavaDataset() class which prepares training and validation dataset for us,while using that dataset class we are doing  **transforms=get_train_transforms()** for training dataset preparation and  **transforms=get_valid_transforms()** for validation dataset preparation which will let us apply several augmentations on training and validation data differently as discussed above\n\n* unfortunately torch xla can't use our train_ds and valid_ds (that we prepared using CassavaDataset() class) directly for 8 core training and that's why we use **DistributedSampler()** for preparing **train_sampler** and **valid_sampler** that will appropriately distribute the dataset across the 8 cores.\n\n* using **train_sampler** and **valid_sampler** we finally create our **train_loader and valid_loader** that we will use for 8 core training on pytorch tpu"},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_dataloader(df, trn_idx, val_idx, data_root='../input/cassava-leaf-disease-classification/train_images/'):\n    \n    from catalyst.data.sampler import BalanceClassSampler\n    \n    train_ = df.loc[trn_idx,:].reset_index(drop=True)\n    valid_ = df.loc[val_idx,:].reset_index(drop=True)\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":{},"cell_type":"markdown","source":"[**bi tempered logistic loss**](https://arxiv.org/pdf/1906.03361.pdf)"},{"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()","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    for step, (imgs, image_labels) in pbar:\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\n          \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       \n    \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":"# Understanding Model"},{"metadata":{},"cell_type":"markdown","source":"in CFG dictionary above we used **'model_arch': 'vit_base_patch16_384'** that means we will be using the **vit_base_patch16_384** pretrained model from this excellent library called [Pytorch Image Models](https://github.com/rwightman/pytorch-image-models/tree/master/timm)\n\n[Visual transformers(VTs)](https://analyticsindiamag.com/hands-on-vision-transformers-with-pytorch/) are in recent research and moving the barrier to outperform the CNN models for several vision tasks. CNN architectures give equal weightage to all the pixels and thus have an issue of learning the essential features of an image. In the paper by Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, Peter Vajda: [“Visual Transformers: Token-based Image Representation and Processing for Computer Vision”](https://arxiv.org/abs/2006.03677), they have shown in place of convolutions, VTs apply transformers to relate semantic concepts in token-space directly.\n\n![](https://lh3.googleusercontent.com/rmkBtzSx7YZj3otixM57iatZwUd6z1slbCMNS7Kehi3gM_gdIo96IY775kCkREaFFut6qf0rO1IvTHjm6_WB9xNrb3omV-aFM7RV99gqmUew1CZDC4MipyfP-c3rLJNIoKzALhrC)"},{"metadata":{},"cell_type":"markdown","source":"[**How it works?**](https://analyticsindiamag.com/hands-on-vision-transformers-with-pytorch/)\n\nif we read the recent paper, Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby: [“An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale”](https://arxiv.org/abs/2010.11929) we will get to know that the SOTA model ViT breaks an input image of 16×16 to a  sequence of patches, just like a series of word embeddings generated by an NLP Transformers. Each patch gets flattened into a single vector in a series of interconnected channels of all pixels in a patch, then projects it to desired input dimension. Because transformers operate in self-attention mode, and they do not necessarily depend on the structure of the input elements, which in turns helps the architecture to learn and relate sparsely-distributed information more efficiently. In Vit, the relationship between the patches in an image is not known and thus allows it to learn more relevant features from the training data and encode in positional embedding in ViT. \n\n![](https://lh5.googleusercontent.com/xk2mQd9H3w4uS482ursxhbDZhr7UXxJ9RMZ7VVjErBMuhbsB1QfSary9pOWU4P5EeZHmB05R8KalB5GXx__eCiN2AQ5qRhXY4vHwYe2zoFqIO0XkpHHXIE8VP99lpcgW5HtjPKKx)"},{"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":"\nmodel = CassvaImgClassifier(CFG['model_arch'], train.label.nunique(), pretrained=True)\n","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":"\ndef train_model():\n    folds = StratifiedKFold(n_splits=CFG['fold_num'], shuffle=True, random_state=CFG['seed']).split(np.arange(train.shape[0]), train.label.values) \n    device = xm.xla_device()\n    model.to(device)\n        \n    for fold, (trn_idx, val_idx) in enumerate(folds):\n        # we'll train fold 0 first\n        if fold != 4:\n            continue\n\n        train_loader, val_loader = prepare_dataloader(train, trn_idx, val_idx, data_root='../input/cassava-leaf-disease-classification/train_images/')\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        #scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, CFG['epochs']-CFG['warmup_epo'])\n        #scheduler = GradualWarmupScheduler(optimizer, multiplier=CFG['warmup_factor'], total_epoch=CFG['warmup_epo'], after_scheduler=scheduler_cosine)\n\n        #loss = nn.CrossEntropyLoss() \n        \n        \n        #save best epoch of each fold\n        best_accuracy = 0.0\n        \n        for epoch in range(CFG['epochs']):\n            gc.collect()\n            para_loader = pl.ParallelLoader(train_loader, [device])\n            #scheduler.step(epoch-1)\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            \n            #with torch.no_grad():\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            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'.format(CFG['model_arch'], fold,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                #xser.save(model.state_dict(),'{}_fold_{}_best_epoch'.format(CFG['model_arch'], fold))\n                best_accuracy = cur_accuracy\n                \n        #xm.save(model.state_dict(),'{}_fold_{}'.format(CFG['model_arch'], fold))         \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\ndef _mp_fn(rank, flags):\n    global acc_list\n    torch.set_default_tensor_type('torch.FloatTensor')\n    res = train_model()\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","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":{},"cell_type":"markdown","source":"Thank You,\n\nHappy **KAGGLING**"},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}