{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Crazy Fast Baseline using Torch XLA by [@shonenkov](https://www.kaggle.com/shonenkov)\n\n\nHi everyone!\n\nThank you all, my friends, for reading my kernels about this competition:\n\n- [[Merge External Data]](https://www.kaggle.com/shonenkov/merge-external-data)\n- [[Training CV] Melanoma Starter](https://www.kaggle.com/shonenkov/training-cv-melanoma-starter)\n- [[Inference Single Model] Melanoma Starter](https://www.kaggle.com/shonenkov/inference-single-model-melanoma-starter)\n\nNow I would like to share with you Torch XLA baseline with crazy fast learning using TPU ;)\nDon't afraid! If you understand [[Training CV] Melanoma Starter](https://www.kaggle.com/shonenkov/training-cv-melanoma-starter), then TPU and this kernel are super easy for understanding !","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Main Idea\n\n[PyTorch XLA](https://github.com/pytorch/xla) is enabling PyTorch on Google TPU. So lets do it for Melanoma Competition! At the end of this kernel you will find link on Google Colab with prepared baseline!","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Dependencies","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":true,"_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 > /dev/null\n!python pytorch-xla-env-setup.py --version 20200610 --apt-packages libomp5 libopenblas-dev > /dev/null\n!pip install catalyst > /dev/null\n!pip install --no-deps albumentations --force-reinstall > /dev/null\n!pip install -q efficientnet_pytorch > /dev/null\n!pip install --no-deps timm > /dev/null\n!git clone https://github.com/4uiiurz1/pytorch-auto-augment > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import os\nos.environ['XLA_USE_BF16'] = \"1\"\n\nimport pandas as pd\nfrom sklearn.model_selection import GroupKFold\nimport cv2\nfrom skimage import io\nimport albumentations as A\nimport torch\nfrom datetime import datetime\nimport time\nimport random\nimport cv2\nimport pandas as pd\nimport numpy as np\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom torch.nn import functional as F\nfrom glob import glob\nimport sklearn\nfrom torch import nn\nimport warnings\n\nwarnings.filterwarnings(\"ignore\") \nwarnings.filterwarnings(\"ignore\", category=DeprecationWarning) \n\nSEED = 42\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ROOT_PATH = './'\nDATA_PATH = f'../input/melanoma-merged-external-data-512x512-jpeg'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_folds = pd.read_csv(f'{DATA_PATH}/folds_08062020.csv', index_col='image_id')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# [PyTorch Auto Augment](https://github.com/4uiiurz1/pytorch-auto-augment)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from albumentations.pytorch import ToTensor\nfrom torchvision import transforms\n\nimport sys\nsys.path.insert(0, './pytorch-auto-augment')\nfrom auto_augment import AutoAugment, Cutout\n\n\ndef get_train_transforms():\n    return transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomVerticalFlip(),\n        AutoAugment(),\n        Cutout(),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n    ])\n\ndef get_valid_transforms():\n    return transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n    ])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_ROOT_PATH = f'{DATA_PATH}/512x512-dataset-melanoma/512x512-dataset-melanoma'\n\ndef onehot(size, target):\n    vec = torch.zeros(size, dtype=torch.float32)\n    vec[target] = 1.\n    return vec\n\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, image_ids, labels, transforms=None):\n        super().__init__()\n        self.image_ids = image_ids\n        self.labels = labels\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{TRAIN_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.resize(image, (256, 256), cv2.INTER_AREA)\n        label = self.labels[index]\n        target = onehot(2, label)        \n        if self.transforms:\n            image = self.transforms(image)\n        return image, target\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def get_labels(self):\n        return list(self.labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fold_number = 0\n\ntrain_dataset = DatasetRetriever(\n    image_ids=df_folds[df_folds['fold'] != fold_number].index.values,\n    labels=df_folds[df_folds['fold'] != fold_number].target.values,\n    transforms=get_train_transforms(),\n)\n\ndf_val = df_folds[(df_folds['fold'] == fold_number) & (df_folds['source'] == 'ISIC20')]\n\nvalidation_dataset = DatasetRetriever(\n    image_ids=df_val.index.values,\n    labels=df_val.target.values,\n    transforms=get_valid_transforms(),\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from skimage import io\n\nimage, target = train_dataset[777]\nnumpy_image = image.permute(1,2,0).cpu().numpy()\n\nio.imshow(numpy_image);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Metrics","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn import metrics\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\nclass RocAucMeter(object):\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.y_true = np.array([0,1])\n        self.y_pred = np.array([0.5,0.5])\n        self.score = 0\n\n    def update(self, y_true, y_pred):\n        y_true = y_true.cpu().numpy().argmax(axis=1).clip(min=0, max=1).astype(int)\n        y_pred = nn.functional.softmax(y_pred, dim=1).data.cpu().numpy()[:,1]\n        self.y_true = np.hstack((self.y_true, y_true))\n        self.y_pred = np.hstack((self.y_pred, y_pred))\n        self.score = sklearn.metrics.roc_auc_score(self.y_true, self.y_pred)\n\n    @property\n    def avg(self):\n        return self.score\n\nclass APScoreMeter(RocAucMeter):\n    def __init__(self):\n        super(APScoreMeter, self).__init__()\n\n    def update(self, y_true, y_pred):\n        y_true = y_true.cpu().numpy().argmax(axis=1).clip(min=0, max=1).astype(int)\n        y_pred = nn.functional.softmax(y_pred, dim=1).data.cpu().numpy()[:,1]\n        self.y_true = np.hstack((self.y_true, y_true))\n        self.y_pred = np.hstack((self.y_pred, y_pred))\n        self.score = sklearn.metrics.average_precision_score(self.y_true, self.y_pred)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Losses","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class FocalLoss(nn.Module):\n    def __init__(self, alpha=1, gamma=2, logits=False, reduce=True):\n        super(FocalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.logits = logits\n        self.reduce = reduce\n\n    def forward(self, inputs, targets):\n        if self.logits:\n            BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets)\n        else:\n            BCE_loss = F.binary_cross_entropy(inputs, targets)\n\n        pt = torch.exp(-BCE_loss)\n        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss\n\n        if self.reduce:\n            return torch.mean(F_loss)\n        else:\n            return F_loss\n        \nclass LabelSmoothing(nn.Module):\n    def __init__(self, smoothing = 0.1):\n        super(LabelSmoothing, self).__init__()\n        self.confidence = 1.0 - smoothing\n        self.smoothing = smoothing\n\n    def forward(self, x, target):\n        if self.training:\n            x = x.float()\n            target = target.float()\n            logprobs = torch.nn.functional.log_softmax(x, dim = -1)\n\n            nll_loss = -logprobs * target\n            nll_loss = nll_loss.sum(-1)\n    \n            smooth_loss = -logprobs.mean(dim=-1)\n\n            loss = self.confidence * nll_loss + self.smoothing * smooth_loss\n\n            return loss.mean()\n        else:\n            return torch.nn.functional.cross_entropy(x, target)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TPU Fitter","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nfrom catalyst.data.sampler import DistributedSamplerWrapper, BalanceClassSampler\n\n\nclass TPUFitter:\n    \n    def __init__(self, model, device, config, sub_folder):\n\n        self.config = config\n        self.epoch = 0\n        self.best_score = 0\n        self.base_dir = f'./{config.folder}/{sub_folder}'\n        time.sleep(1)\n        if not os.path.exists(self.base_dir):\n            os.makedirs(self.base_dir)\n\n        self.log_path = f'{self.base_dir}/log.txt'\n\n        self.model = model\n        self.device = device\n\n        param_optimizer = list(self.model.named_parameters())\n        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n        optimizer_grouped_parameters = [\n            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n        ]\n\n        self.optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=config.lr*xm.xrt_world_size())\n        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n\n        self.criterion = FocalLoss(logits=True).to(self.device)\n        # self.criterion = LabelSmoothing().to(self.device)\n        xm.master_print(f'Fitter prepared. Device is {self.device}')\n\n        self.best_loss = 10**5\n\n    def fit(self, train_loader, validation_loader):\n        for e in range(self.config.n_epochs):\n            if self.config.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n\n            t = time.time()\n            para_loader = pl.ParallelLoader(train_loader, [self.device])\n            summary_loss, roc_auc_scores, ap_scores = self.train_one_epoch(para_loader.per_device_loader(self.device))\n            \n            self.log(f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, roc_auc: {roc_auc_scores.avg:.5f}, ap: {ap_scores.avg:.5f}, time: {(time.time() - t):.5f}')\n\n            t = time.time()\n            para_loader = pl.ParallelLoader(validation_loader, [self.device])\n            summary_loss, roc_auc_scores, ap_scores = self.validation(para_loader.per_device_loader(self.device))\n            self.log(f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, roc_auc: {roc_auc_scores.avg:.5f}, ap: {ap_scores.avg:.5f}, time: {(time.time() - t):.5f}')\n\n            if self.config.validation_scheduler:\n                self.scheduler.step(metrics=roc_auc_scores.avg)\n\n            if e > 20: # hardcode for keep memory\n                self.save(f'{self.base_dir}/checkpoint-{str(self.epoch).zfill(3)}epoch.bin')\n\n            self.epoch += 1\n    \n    def validation(self, val_loader):\n        self.model.eval()\n        summary_loss = AverageMeter()\n        roc_auc_scores = RocAucMeter()\n        ap_scores = APScoreMeter()\n\n        t = time.time()\n        for step, (images, targets) in enumerate(val_loader):\n            with torch.no_grad():\n                targets = targets.to(self.device, dtype=torch.float32)\n                batch_size = images.shape[0]\n                images = images.to(self.device, dtype=torch.float32)\n                outputs = self.model(images)\n                loss = self.criterion(outputs, targets)\n                roc_auc_scores.update(targets, outputs)\n                ap_scores.update(targets, outputs)\n                summary_loss.update(loss.detach().item(), batch_size)\n                \n        return summary_loss, roc_auc_scores, ap_scores\n         \n    def train_one_epoch(self, train_loader):\n        self.model.train()\n\n        summary_loss = AverageMeter()\n        roc_auc_scores = RocAucMeter()\n        ap_scores = APScoreMeter()\n        t = time.time()\n        for step, (images, targets) in enumerate(train_loader):\n            images = images.to(self.device, dtype=torch.float32)\n            targets = targets.to(self.device, dtype=torch.float32)\n\n            self.optimizer.zero_grad()\n\n            outputs = self.model(images)\n            loss = self.criterion(outputs, targets)\n\n            batch_size = images.size(0)\n\n            loss.backward()\n\n            roc_auc_scores.update(targets, outputs)\n            ap_scores.update(targets, outputs)\n            summary_loss.update(loss.detach().item(), batch_size)\n\n            xm.optimizer_step(self.optimizer)\n\n            if self.config.step_scheduler:\n                self.scheduler.step()\n\n        return summary_loss, roc_auc_scores, ap_scores\n\n    def save(self, path):\n        self.model.eval()        \n        xm.save(self.model.state_dict(), path)\n\n    def log(self, message):\n        if self.config.verbose:\n            xm.master_print(message)\n        with open(self.log_path, 'a+') as logger:\n            xm.master_print(f'{message}\\n', logger)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 1 # 4\n    batch_size = 16 \n    n_epochs = 2 # 40\n    lr = 0.000006\n\n    # -------------------\n    folder = 'resnext50d_32x4d-autoaugment-KFOLD' \n    verbose = True\n    verbose_step = 1\n    # -------------------\n\n    # --------------------\n    step_scheduler = False  # do scheduler.step after optimizer.step\n    validation_scheduler = True  # do scheduler.step after validation stage loss\n\n#     SchedulerClass = torch.optim.lr_scheduler.OneCycleLR\n#     scheduler_params = dict(\n#         max_lr=0.001,\n#         epochs=n_epochs,\n#         steps_per_epoch=int(len(train_dataset) / batch_size),\n#         pct_start=0.1,\n#         anneal_strategy='cos', \n#         final_div_factor=10**5\n#     )\n    \n    SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n    scheduler_params = dict(\n        mode='max',\n        factor=0.8,\n        patience=2,\n        verbose=False, \n        threshold=0.0001,\n        threshold_mode='abs',\n        cooldown=0, \n        min_lr=1e-8,\n        eps=1e-08\n    )\n    # --------------------","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Class Balance \"on fly\" with XLA\n\nRecently I have created kernel about it in another competition with demonstration how it works! Welcome for reading :)\n\n- [Class Balance with PyTorch/XLA](https://www.kaggle.com/shonenkov/class-balance-with-pytorch-xla)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    xm.set_rng_state(SEED)\n    device = xm.xla_device()\n    net.to(device)\n\n    train_sampler = DistributedSamplerWrapper(\n        sampler=BalanceClassSampler(labels=train_dataset.get_labels(), mode=\"downsampling\"),\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=train_sampler,\n        pin_memory=False,\n        drop_last=True,\n        num_workers=TrainGlobalConfig.num_workers,\n    )\n    validation_sampler = torch.utils.data.distributed.DistributedSampler(\n        validation_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False\n    )\n    validation_loader = torch.utils.data.DataLoader(\n        validation_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=validation_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=TrainGlobalConfig.num_workers\n    )\n\n    fitter = TPUFitter(model=net, device=device, config=TrainGlobalConfig, sub_folder=sub_folder)\n    if rank == 0:\n        time.sleep(1)\n    fitter.fit(train_loader, validation_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import timm\n\ndef get_net():\n    net = timm.create_model('resnext50d_32x4d', pretrained=True)\n    net.fc = nn.Linear(in_features=net.fc.in_features, out_features=2, bias=True)\n    return net","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n\n\nfor fold_number in range(1): # range(5)\n    sub_folder = f'fold{fold_number}'\n    net = get_net()\n\n    train_dataset = DatasetRetriever(\n        image_ids=df_folds[df_folds['fold'] != fold_number].index.values,\n        labels=df_folds[df_folds['fold'] != fold_number].target.values,\n        transforms=get_train_transforms(),\n    )\n\n    df_val = df_folds[(df_folds['fold'] == fold_number) & (df_folds['source'] == 'ISIC20')]\n    validation_dataset = DatasetRetriever(\n        image_ids=df_val.index.values,\n        labels=df_val.target.values,\n        transforms=get_valid_transforms(),\n    )\n\n    FLAGS={}\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Here not enough CPU and RAM for really crazy fast! \n\nSo I prepared Google Colab pipeline for you, my friends! Using PyTorch XLA you can increase speed of research in ~5-10 times! More experiments --> higher LB positions!\nI am using Colab Pro with TPU Runtime Type and HIGH-RAM mode ~36GB and 40 CPU. For this model `resnext50d_32x4d` 1 epoch has time: ~60s (train), ~35s (validation).\n\nYou should save copy on your drive and run research right now!\n- [Prepared Baseline on Colab!](https://drive.google.com/drive/folders/1VerwMHYuP6lrL-CPWJ1_aKYj5Fq2SNsX?usp=sharing)\n\nI have run one experiment for example. Inference you can find here:\n\n- [[Inference] Melanoma Crazy Fast](https://www.kaggle.com/shonenkov/inference-melanoma-crazy-fast)\n\n# Thank you for reading my kernels!","execution_count":null}],"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}