{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"# Installing Pytorch-XLA and other dependencies\n!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 --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nos.environ['XLA_USE_BF16'] = \"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = \"100000000\"\n\nimport gc\nimport random\nimport timeit\nimport time\nimport datetime\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom typing import List\nfrom tqdm.notebook import tqdm\nfrom PIL import Image\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import roc_auc_score\n\n# pytorch imports\nimport torch\nimport torch.optim as optim\nfrom torchvision import transforms\nfrom torchvision.utils import make_grid\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\n\nimport timm\n\n# for TPU\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\n\nplt.style.use('bmh')\nplt.rcParams['figure.figsize'] = [20, 13]\nSEED = 421","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=1234):\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\nseed_everything(SEED)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_DIR =  \"../input/ranzcr-clip-catheter-line-classification/\"\nFOLDS_DIR = \"../input/ranzcr-folds/\"\n\n\nMODEL_NAME = \"efficientnet_b7\"\n# MODEL_NAME = \"resnet200d\"\n\nEPOCHS = 30\nINIT_LR = 1e-3\nBATCH_SIZE = 16\nIMAGE_SIZE = 600\nPRINT_EVERY = 50 # how often to print the losses/metric scores\nSTOP_TRAINING_AFTER = 5 # number of no-improvement epochs to wait before stopping training. ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The folds are taken from Abhishek's [kernel](https://www.kaggle.com/abhishek/ranzcr-tez-training-efficientnet-5)"},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_df = pd.read_csv(DATA_DIR + 'sample_submission.csv')\nfolds_df = pd.read_csv(FOLDS_DIR + \"train_folds.csv\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CLASSES = [col for col in folds_df.columns if col not in ['StudyInstanceUID', 'PatientID', 'kfold']]\nCLASSES","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"folds_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# we will perform validation on fold 0 and train the model on the remaining 4 folds\nvalid_fold = 0\n\ntrain_df = folds_df[folds_df.kfold != valid_fold].reset_index(drop=True)\nvalid_df = folds_df[folds_df.kfold == valid_fold].reset_index(drop=True)\n\n\nprint (\"NUMBER OF SAMPLES IN:\\n\")\nprint (f\"Training: {train_df.shape[0]}\")\nprint (f\"Validation: {valid_df.shape[0]}\")\nprint (f\"Testing: {sub_df.shape[0]}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Utility function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_input_images(imgs: torch.Tensor, title_string: str, nrow: int = 4) -> None:\n    image_grid = make_grid(imgs, nrow=nrow, padding=10, pad_value=1)\n    \n    # transform from CHW -> HWC\n    plt.imshow(image_grid.permute(1, 2, 0), cmap=plt.cm.bone)\n    plt.title(title_string)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset class"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RanczrDataset(Dataset):\n    def __init__(self, df, data_dir, transform=None):\n        super().__init__()\n        self.df = df\n        self.data_dir = data_dir\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, index):\n        img_name = self.df['StudyInstanceUID'][index]\n        targets = self.df.loc[index, CLASSES].values.astype(np.uint8)\n        \n        targets = torch.from_numpy(targets)\n        \n        img_path = os.path.join(self.data_dir, img_name+\".jpg\")\n        \n        image = np.array(Image.open(img_path))\n        \n        if self.transform:\n            image = self.transform(image=image)\n        \n        return image, targets","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Transformations"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_transform = A.Compose([\n    A.Resize(IMAGE_SIZE, IMAGE_SIZE, always_apply=True),\n    A.CLAHE(p=0.5),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.Normalize(mean=[0.485], std=[0.229], max_pixel_value=255.0, p=1.0),\n    ToTensorV2(),\n])\n\nvalid_transform = A.Compose([\n    A.Resize(IMAGE_SIZE, IMAGE_SIZE, always_apply=True),\n    A.Normalize(mean=[0.485], std=[0.229], max_pixel_value=255.0, p=1.0),\n    ToTensorV2(),\n])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create Datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Creating Pytorch Datasets\ntrain_dataset = RanczrDataset(df=train_df, data_dir=DATA_DIR+\"train/\", transform=train_transform)\nvalidation_dataset = RanczrDataset(df=valid_df, data_dir=DATA_DIR+\"train/\", transform=valid_transform)\n\n# Create Dataloaders\ntrain_loader = DataLoader(dataset=train_dataset, batch_size=16, drop_last=True, num_workers=0)\n\ninputs, targets = next(iter(train_loader))\n\n# plot input images\nplot_input_images(inputs['image'], title_string='batch images', nrow=8)\n\ndel train_loader\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_loader = DataLoader(dataset=validation_dataset, batch_size=16, drop_last=True, num_workers=0)\n\ninputs, targets = next(iter(valid_loader))\n\n# plot input images\nplot_input_images(inputs['image'], title_string='validation batch images', nrow=8)\n\ndel valid_loader\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RanczrModel(nn.Module):\n    def __init__(self):\n        super(RanczrModel, self).__init__()\n        \n        self.effnet = timm.create_model(MODEL_NAME, pretrained=True, in_chans=1).as_sequential()[:-2]\n        \n        self.dropout = nn.Dropout(p=0.5)\n        self.dense = nn.Linear(2560, len(CLASSES))\n        \n        \n    def forward(self, images):\n        pooled_features= self.effnet(images)\n        \n        outputs = self.dense(self.dropout(pooled_features))\n        \n        return outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"effnet = RanczrModel()\n\n# we want to train the models from scratch hence unfreezing all the weights\nfor param in effnet.parameters():\n    param.requires_grad = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"effnet","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# objective function\ndef loss_func(predictions, targets):\n    return nn.BCEWithLogitsLoss()(predictions, targets)\n\n\n# reduction function for xla\ndef reduce_func(vals):\n    # averaging the loss over all TPU cores\n    return sum(vals) / len(vals)\n\n\n# evaluation metric: multilabel AUC\ndef multilabel_auc(targets: torch.Tensor, predictions: torch.Tensor) -> float:\n#     xm.master_print (f'Shapes=> {targets.shape}\\t{predictions.shape}')\n    auc = 0\n    for j in range(targets.shape[1]):\n#         xm.master_print(f'Target class counts=> 1: {(targets[:, j]==1).sum()} 0: {(targets[:, j]==0).sum()}')\n        try:\n            auc += roc_auc_score(targets[:, j], predictions[:, j])\n        except ValueError:\n            # code will reach at this point when there is one one target class i.e all are 1s or 0s\n            assert torch.unique(targets[:, j]).shape[0] == 1, \"There should be only one unique value present in the target tensor\"\n            target_val = torch.unique(targets[:, j])\n            \n            # we will add one element to the target tensor with a value not present in the target tensor\n            target_tensor = torch.cat([targets[:, j], 1-target_val])\n            \n            # for the prediction column we can take the mean of the predictions and use that as the prediction value for the new target we just added\n            mean_prediction = torch.mean(predictions[:, j], dim=0, keepdim=True)\n            \n            predictions_tensor = torch.cat([predictions[:, j], mean_prediction])\n            \n            auc += roc_auc_score(target_tensor, predictions_tensor)\n    # calculate average over all the classes\n    return auc / len(CLASSES)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Train and validation methods"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop_fn(data_loader, model, optimizer, device):\n    # setting the model on train mode\n    model.train()\n    \n    train_loss_history = [] # this will contain the train loss for each step in the current epoch. These values will later be averaged over the entire epoch to calculate mean_train_loss\n\n\n    for batch_ix, (inputs, targets) in enumerate(data_loader):\n        images, targets = inputs['image'].to(device), targets.float().to(device)\n        \n        assert images is not None, \"input images are None\"\n        assert targets is not None, \"targets are None\"\n        # zeroing the gradients\n        optimizer.zero_grad()\n        \n        # forward pass\n        predictions = model(images)\n        \n        # calculate loss\n        train_loss = loss_func(predictions, targets)\n            \n        # perform mean of training loss over all the 8 TPU cores and return the value in train_loss_reduce variable; this value will be same for all the cores\n        train_loss_reduce = xm.mesh_reduce('train_loss_reduce', train_loss.item(), reduce_func)\n        \n        train_loss_history.append(train_loss_reduce)\n        \n        # print loss averaged over all TPU cores after `print_every` steps\n        if batch_ix % PRINT_EVERY == 0:    \n            # master_print will only print once (not from all 8 cores)\n            xm.master_print(f\"batch_ix={batch_ix}, train_loss={train_loss_reduce:.4f}\")\n        \n        # backward pass\n        train_loss.backward()\n        \n        # gradient clipping\n#         grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)\n        \n        # DEBUGGING        \n        if model.effnet[0].weight.isnan().any():\n            sys.exit(\"NaN encountered!!!\")\n        # Use Pytorch XLA optimizer for applying the weight updates\n        xm.optimizer_step(optimizer)\n        \n\n    return train_loss_history\n\ndef valid_loop_fn(data_loader, model, device):\n    # setting the model in eval() model for validation step\n    model.eval()\n    \n    val_loss_history = [] # this will contain the validation loss for each step in the current epoch. These values will later be averaged over the entire epoch to calculate mean_val_loss \n    \n    all_targets = None\n    all_predictions = None\n    for batch_ix, (inputs, targets) in enumerate(data_loader):\n        images, targets = inputs['image'].to(device), targets.float().to(device)\n        \n        # make prediction on a hold-out validation set\n        predictions = model(images)\n        \n        val_loss = loss_func(predictions.detach(), targets)\n        val_loss_reduce = xm.mesh_reduce('val_loss_reduce', val_loss.item(), reduce_func)\n        val_loss_history.append(val_loss_reduce)\n        \n        if all_predictions is None:\n            all_predictions = predictions.detach().cpu()\n            all_targets = targets.cpu()\n        else:\n            all_predictions = torch.cat((all_predictions, predictions.detach().cpu()), dim=0)\n            all_targets = torch.cat((all_targets, targets.cpu()), dim=0)\n        \n    val_auc = multilabel_auc(all_targets, all_predictions)\n    val_auc_reduce = xm.mesh_reduce('val_auc_reduce', val_auc, reduce_func)\n    \n    return val_loss_history, val_auc_reduce\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n    # this is the main function that calls the above functions. This function will be spawned by Pytorch XLA multiprocessing.\n    # this function will be run on each of the 8 cores\n    \n    # We need to define a data sampler to appropriately distribute the data across 8 cores\n    train_sampler = DistributedSampler(\n                    train_dataset,\n                    num_replicas=xm.xrt_world_size(), # tell Pytorch how many devices (TPU cores) we are using for training\n                    rank=xm.get_ordinal(), # tell Pytorch which device (core) we are on currently,\n                    shuffle=True, # sampler will shuffle the indices\n                )\n\n    train_data_loader = DataLoader(\n                        train_dataset,\n                        batch_size=BATCH_SIZE,\n                        sampler=train_sampler,\n                        drop_last=True,\n                        num_workers=0, # We will use only the main process to load the data hence saving up a lot of VM's memory \n    )\n    \n    valid_sampler = DistributedSampler(\n                    validation_dataset,\n                    num_replicas=xm.xrt_world_size(),\n                    rank=xm.get_ordinal(),\n                    shuffle=False,\n                )\n    \n    valid_data_loader = DataLoader(\n                        validation_dataset,\n                        batch_size=4,\n                        sampler=valid_sampler,\n                        drop_last=False,\n                        num_workers=0,\n                )\n\n    device = xm.xla_device() # our device (single TPU core)\n    model = effnet.to(device) # put model on single TPU core\n    \n    xm.master_print(\"Model loading on TPU completed\")\n    # Whatever needs to printed only once for all 8 TPU cores, should be printed using `master_print` function\n    xm.master_print(\"Training initiated\")\n    \n#     scaled_lr = LR * xm.xrt_world_size()# scale the learning rate as per the number of devices\n    # calculate the total number of training steps\n    num_train_steps = int(len(train_dataset) / BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n    \n    # define the optimizer; CHECK THE HYPERPARAMETERS OF THIS OPTIMIZER AS WELL\n    optimizer = optim.Adam(model.parameters(), lr=INIT_LR)\n    \n    # learning rate scheduler; CHECK THE HYPERPARAMETERS OF THIS SCHEDULER AS WELL\n#     scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, EPOCHS, eta_min=1e-7)\n    \n#     scheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=2, mode='max', verbose=True)\n    \n    xm.master_print(f\"num_training_steps = {num_train_steps}, world_size={xm.xrt_world_size()}\")\n    \n    max_mean_auc_score = -np.Inf\n\n    epoch_since_last_improvement = 0\n    stop_training = False\n\n    # initiate training\n    for epoch in range(1, EPOCHS+1):\n        xm.master_print(f\"EPOCH: {epoch}/{EPOCHS}\")\n        epoch_start_tick = timeit.default_timer() # record time at the start of the epoch\n        gc.collect()\n        # we will encapsulate our dataloader with Pytorch XLA's ParallelLoader for TPU-core-specific dataloading\n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        xm.master_print(\"Training Parallel loader created...\\nTraining now\")\n        gc.collect()\n        \n        # call training function\n        train_loss_hist = train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device)\n        del para_loader\n        gc.collect()        \n        \n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        xm.master_print(\"Validation Parallel loader created...\\nValidating now\")\n        gc.collect()\n        \n        # call evaluation function\n        val_loss_hist, epoch_mean_auc_score = valid_loop_fn(para_loader.per_device_loader(device), model, device)\n        del para_loader\n        \n        gc.collect()\n        \n        # print train loss, validation loss and validation AUC at the end of each epoch\n        xm.master_print(f\"EPOCH: {epoch}==> mean_train_loss: {np.mean(train_loss_hist):.4f}, mean_valid_loss: {np.mean(val_loss_hist):.4f}, mean_valid_auc: {epoch_mean_auc_score:.4f}\")\n        gc.collect()\n        \n        elapsed = timeit.default_timer() - epoch_start_tick\n        elapsed = str(datetime.timedelta(seconds=elapsed)).split('.')[0]\n        \n        xm.master_print(f\"Time taken for epoch: {elapsed}\\n\")\n        \n#         scheduler_warmup.step(epoch-1)\n        \n        # decrease the learning rate if the epoch mean auc score does not improve as per the schedule specified in the scheduler\n        if scheduler:\n            scheduler.step(epoch_mean_auc_score)\n        \n        # model saving code\n        if epoch_mean_auc_score >= max_mean_auc_score:\n            xm.master_print(f\"Validation AUC score increased ({max_mean_auc_score:.6f} -> {epoch_mean_auc_score:.6f}) at EPOCH {epoch}\\n  Saving model ...\\n\")\n            xm.save(model.state_dict(), 'timm-effnet-b7-res600-final.pt')\n            max_mean_auc_score = epoch_mean_auc_score\n            epoch_since_last_improvement = 0\n\n        # check if validation auc didn't improve\n        if epoch_mean_auc_score < max_mean_auc_score:\n            epoch_since_last_improvement+=1\n            xm.master_print(f'{epoch_since_last_improvement} epochs have finished since the last improvement in val AUC')\n\n            if epoch_since_last_improvement > STOP_TRAINING_AFTER:\n                xm.master_print('Stopping training prematurely due to no improvement in val AUC')\n                stop_training = True\n        if stop_training:\n            break\n    xm.master_print(f'Best Validation AUC is {max_mean_auc_score:.6f}')\n    gc.collect()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# start training process; we need to spawn the training processes on each of the TPU cores.\ndef _map_fn(rank, flags):\n    a = _run()\n\nFLAGS = {}\nstart_time = timeit.default_timer()\nxmp.spawn(_map_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}