{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Install PyTorch XLA"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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 nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Import PyTorch XLA"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.distributed.parallel_loader as pl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nfrom PIL import Image\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold\nfrom sklearn.metrics import roc_auc_score\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport os\nimport torch\nimport torchvision\nimport cv2\nimport numpy as np\nimport time\nimport warnings\nimport datetime","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/melanoma-merged-external-data-512x512-jpeg/folds_13062020.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MelanomaDataset(torch.utils.data.Dataset):\n    def __init__(self, path_files, path_csv=None, pd_loaded=None, extension=\".jpg\", transforms=None):\n        self.path_files = path_files\n        self.ext = extension\n        self.transforms = transforms\n        \n        if pd_loaded is not None:\n            self.df = pd_loaded\n        elif path_csv is None and pd_loaded is None:\n            raise Exception(\"Both parameters path_csv and pd_loaded can't be none\")\n        else:\n            self.df = pd.read_csv(path_csv)\n        \n    def __getitem__(self, index):\n        img_path = os.path.join(self.path_files, self.df.iloc[index]['image_id'] + self.ext)\n        img_arr = cv2.imread(img_path)\n        img_arr_rgb = cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)\n        \n        if self.df.iloc[index]['target'] == 0:\n            img_label = [1,0]\n        else:\n            img_label = [0,1]\n    \n        if self.transforms:\n            sample = {'image':img_arr_rgb, 'label':img_label}\n            sample = self.transforms(**sample)\n            img_tens = sample['image']\n            img_label = sample['label']\n        else:\n            img_tens = torchvision.transforms.ToTensor()(img_arr_rgb_no_hair)\n            \n        img_label = torch.tensor(img_label)\n        \n        return img_tens, img_label\n    \n    def __len__(self):\n        return len(self.df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# use albumentations to transform the images\ntrain_transform = A.Compose(\n    [A.HorizontalFlip(p=0.5),\n     A.VerticalFlip(p=0.5),\n     A.RandomRotate90(p=0.5),\n     A.RandomBrightnessContrast(32./255, 0.5),\n     A.CoarseDropout(max_holes=15, max_height=15, max_width=15),\n     A.Normalize(),\n     ToTensorV2()])\n\nvalid_transform = A.Compose(\n    [A.Normalize(),\n     ToTensorV2()])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Thank you @rwightman (Ross Wightman) for that awesome libary\npip install timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_model():\n    model = timm.create_model(\"tf_efficientnet_b4\", pretrained=True)\n    # two classes only\n    num_classes = 2\n    model.classifier = torch.nn.Linear(model.classifier.in_features, num_classes)\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training function"},{"metadata":{},"cell_type":"markdown","source":"The function that will be executed by each core is _mp_fn:\n* the arguments are passed through a dict (flags in my case)\n* make sure to not execute xm.xla_device() anywhere else than in _mp_fn or it will cause an error\n* make sure to not execute xm.xrt_world_size() anywhere else than in _mp_fn or it will cause an error"},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(index, flags):\n    # set seed make to make sure all same trained\n    torch.manual_seed(flags['seed'])\n    # to use TPU\n    device = xm.xla_device()\n    \n    best_score = None\n    patience_counter = 0\n    \n    if flags['es_criterion'] == 'roc_score':\n        mode = 'max'\n    else: \n        mode = 'min'\n        \n    # samplers so each core gets its own set of data    \n    train_sampler = torch.utils.data.distributed.DistributedSampler(flags['train_ds'],num_replicas=xm.xrt_world_size(),rank=xm.get_ordinal(),shuffle=True)\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(flags['valid_ds'],num_replicas=xm.xrt_world_size(),rank=xm.get_ordinal(),shuffle=False)\n    \n            \n    train_dl = torch.utils.data.DataLoader(dataset=flags['train_ds'], batch_size=flags['batch_size'], sampler=train_sampler, num_workers=flags['num_workers'])\n    valid_dl = torch.utils.data.DataLoader(dataset=flags['valid_ds'], batch_size=flags['batch_size'], sampler=valid_sampler, num_workers=flags['num_workers'])\n    \n    model = create_model().to(device)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=flags['lr'])\n    loss_criterion = torch.nn.BCEWithLogitsLoss()\n\n    \n    if flags['use_lr_scheduler']:\n        if xm.is_master_ordinal():\n            lr_scheduler = ReduceLROnPlateau(optimizer=optimizer, mode=mode, patience=1, verbose=True, factor=0.2)\n        else:\n            lr_scheduler = ReduceLROnPlateau(optimizer=optimizer, mode=mode, patience=1, verbose=False, factor=0.2)\n\n    \n    for epoch in range(flags['num_epochs']):\n        train_size = 0\n        train_loss = 0.0\n        train_acc = 0.0\n\n        valid_size = 0\n        valid_loss = 0.0\n        valid_acc = 0.0\n        \n        # Training\n        start_time = time.time()\n        model.train()\n        # only master prints\n        xm.master_print('=' * 20, 'Training - begin for epoch', epoch+1, '=' * 20)\n        para_loader_train = pl.ParallelLoader(train_dl, [device]).per_device_loader(device)\n        for images, labels in para_loader_train:\n            \n            outputs = model(images)\n            labels = torch.as_tensor(data=labels, dtype=torch.float32)\n            loss = loss_criterion(outputs, labels)\n            \n            optimizer.zero_grad()\n            loss.backward()\n            \n            # stepping with TPU \n            # note for multiple cores barrier=True is not needed!\n            xm.optimizer_step(optimizer)\n            \n            # Compute the total loss for the batch and add it to train_loss\n            batch_size = images.size(0)\n            train_size += batch_size\n            train_loss += loss.item() * batch_size\n\n            # Compute the accuracy\n            ret, predictions = torch.max(outputs.data, 1)\n            ret_, predictions_ = torch.max(labels.data, 1)\n            correct_counts = predictions.eq(predictions_.data.view_as(predictions))\n\n            # Convert correct_counts to float and then compute the mean\n            acc = torch.mean(correct_counts.type(torch.FloatTensor))\n\n            # Compute total accuracy in the whole batch and add to train_acc\n            train_acc += acc.item() * batch_size\n\n            # Find average training loss and training accuracy\n            avg_train_loss = train_loss / train_size\n            avg_train_acc = train_acc / train_size\n            \n        print(f'Core: {xm.get_ordinal()} Train: Epoch {epoch+1} - loss: {avg_train_loss:.4f} - acc: {avg_train_acc:.4f} - Time: {str(datetime.timedelta(seconds=time.time() - start_time))[:7]}')\n        \n        if flags['validation_on']:\n            # Validation\n            start_time = time.time()\n            model.eval()\n            para_loader_valid = pl.ParallelLoader(valid_dl, [device]).per_device_loader(device)\n            sigmoid = torch.nn.Sigmoid()\n            val_preds = []\n            val_targets = []\n\n            for images, labels in para_loader_valid:\n                labels = torch.as_tensor(data=labels, dtype=torch.float32)\n                with torch.no_grad():\n                    outputs = model(images)\n                    outputs_sig = sigmoid(outputs)\n                \n                # BCEWithLogitsLoss will convert outputs to Sigmoid\n                loss = loss_criterion(outputs, labels)\n\n                # Compute the total loss for the batch and add it to train_loss\n                batch_size = images.size(0)\n                valid_size += batch_size\n                valid_loss += loss.item() * batch_size\n\n                # Compute the accuracy\n                ret, predictions = torch.max(outputs_sig.data, 1)\n                ret_, predictions_ = torch.max(labels.data, 1)\n                correct_counts = predictions.eq(predictions_.data.view_as(predictions))\n\n                # Convert correct_counts to float and then compute the mean\n                acc = torch.mean(correct_counts.type(torch.FloatTensor))\n\n                # Compute total accuracy in the whole batch and add to train_acc\n                valid_acc += acc.item() * batch_size\n\n                # Find average training loss and training accuracy\n                avg_valid_loss = valid_loss / valid_size\n                avg_valid_acc = valid_acc / valid_size\n\n                val_preds.append(outputs_sig.detach().cpu().numpy())\n                val_targets.append(labels.detach().cpu().numpy())\n\n            val_preds = np.concatenate(val_preds)\n            val_targets = np.concatenate(val_targets)\n            # calculate roc-score\n            roc_score =  roc_auc_score(val_targets, val_preds)\n            # each process sends its roc-score to all other processes and build a mean with it\n            # this causes a rendevouz where all processes are waiting for each other till they reach that point together\n            avg_roc_score = xm.mesh_reduce('mean-roc-score', roc_score, np.mean)\n            xm.master_print(f'Valid: Epoch {epoch+1} - master-loss: {avg_valid_loss:.4f} - master-acc: {avg_valid_acc:.4f} - roc: {avg_roc_score:.4f} - Time: {str(datetime.timedelta(seconds=time.time() - start_time))[:7]}')\n            \n            if flags['es_criterion'] == 'val_loss':\n                crit_var = avg_valid_loss\n            elif flags['es_criterion'] == 'roc_score':\n                crit_var = avg_roc_score\n\n            if flags['use_lr_scheduler']:\n                lr_scheduler.step(crit_var)\n\n            # early stopping\n            if best_score is None:\n                best_score = crit_var\n                # saving the model with the master-process only\n                # this causes a rendevouz where all processes are waiting for each other till they reach that point together\n                xm.save(model.state_dict(), f\"trained_melanoma_weights_{flags['id']+1}\")\n            elif best_score > crit_var and flags['es_criterion'] == 'val_loss':\n                patience_counter = 0\n                xm.master_print(f'Loss reduced from {best_score} ----> {crit_var}')\n                best_score = crit_var\n                xm.save(model.state_dict(), f\"trained_melanoma_weights_{flags['id']+1}\")\n                xm.master_print(\"Saving model...\")\n            elif best_score < crit_var and flags['es_criterion'] == 'roc_score':\n                patience_counter = 0\n                xm.master_print(f'ROC Score increased from {best_score} ----> {crit_var}')\n                best_score = crit_var\n                xm.save(model.state_dict(), f\"trained_melanoma_weights_{flags['id']+1}\")\n                xm.master_print(\"Saving model...\")\n            else:\n                if flags['patience'] is not None and patience_counter == flags['patience']:\n                    xm.master_print(f'Early Stopping at Epoch {epoch+1}')\n                    break\n                patience_counter += 1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Call training with all 8 cores"},{"metadata":{},"cell_type":"markdown","source":"Training Foldwise, for each model I run the notebook all over again so I can train each model for 3 hours instead of pushing all models to train within 3 hours. id_ has to be 0 <= id_ < num_folds"},{"metadata":{"trusted":true},"cell_type":"code","source":"flags={}\nflags['num_epochs'] = 8\nflags['es_criterion'] = 'roc_score'\nflags['validation_on'] = True\nflags['use_lr_scheduler'] = True\nflags['patience'] = None\nflags['seed'] = 1234\nflags['batch_size'] = 5\nflags['num_workers'] = 8\nflags['num_folds'] = 5\nflags['lr'] = 0.0001\n\n# GroupKFold\ngkf = GroupKFold(n_splits=flags['num_folds'])\n\n# Foldwise Model Training, increment of id_ and rerun will create new model\n# Model 1\nid_ = 0\ntrain_index, test_index = list(gkf.split(np.zeros(len(df)), df['target'], list(df['patient_id'])))[id_]\nprint('-' * 20, 'Training Model', id_+1, '-' * 20)\ntrain_ds = MelanomaDataset(\"/kaggle/input/melanoma-external-data-jpeg-384x384/melanoma/\", pd_loaded=df.iloc[train_index], transforms=train_transform)\nvalid_ds = MelanomaDataset(\"/kaggle/input/melanoma-external-data-jpeg-384x384/melanoma/\", pd_loaded=df.iloc[test_index],  transforms=valid_transform)\nflags['id'] = id_\nflags['train_ds'] = train_ds\nflags['valid_ds'] = valid_ds\n# using 8 cores running through _mp_fn()\nif __name__ == '__main__':\n    xmp.spawn(_mp_fn, args=(flags,), nprocs=8, start_method='fork')","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}