{"cells":[{"metadata":{},"cell_type":"markdown","source":"<center><img src=\"https://i.imgur.com/20rO1id.jpg\" width=\"500px\"></center>","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nIn this project called \"Prostate cANcer graDe Assessment (PANDA) Challenge\", I will a build machine learning model to diagnose Prostate Cancer from biopsy scans (images and masks). This problem is important because fast and accurate automated diagnosis can help reduce burden on doctors and let them focus on curing patients.\n\nI will show how one can build a **multitask model** to solve this problem. I will build a ResNet-based model, which takes a biopsy scan as input and predicts two quantities: the **ISUP grade and  Gleason score**. These are two different, bu related scales used to measure the severity of Prostate Cancer. Training a model on two different, but related tasks **can improve the model's performance on both tasks.\n\n\n<center><img src=\"https://i.imgur.com/piOxK6F.png\" width=\"750px\"></center>","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Acknowledgements\n\n1. [PyTorch XLA ~ by PyTorch](https://pytorch.org/xla/release/1.5/index.html)\n2. [Torchvision Models ~ by PyTorch](https://pytorch.org/docs/stable/torchvision/models.html)\n3. [PANDA / submit test ~ Yasufumi Nakama](https://www.kaggle.com/yasufuminakama/panda-submit-test)\n4. [Super-duper fast pytorch tpu kernel... ~ by Abhishek](https://www.kaggle.com/abhishek/super-duper-fast-pytorch-tpu-kernel)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Contents\n\n* [<font size=4>Preparing the ground</font>](#1)\n    * [Set up PyTorch-XLA](#1.1)\n    * [Import libraries](#1.2)\n    * [Set hyperparameters and paths](#1.3)\n    * [Load .csv data](#1.4)\n    * [Convert Gleason scores to list format](#1.5)\n    * [Display few images](#1.6)\n\n    \n* [<font size=4>Modeling</font>](#2)\n    * [Build PyTorch dataset](#2.1)\n    * [Build ResNet model](#2.2)\n    * [Visualize ResNet architecture](#2.3)\n    * [Split train.csv into 8 folds](#2.4)\n    * [Define cross entropy and accuracy](#2.5)\n    * [Define custom PANDA loss for multitask model](#2.6)\n    * [Define helper function for training logs](#2.7)\n    * [Train model on all 8 TPU cores in parallel](#2.8)\n\n\n* [<font size=4>Takeaways</font>](#3)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Preparing the ground <a id=\"1\"></a>","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Set up PyTorch-XLA\n\n* These few lines of code sets up PyTorch XLA for us.\n* We need PyTorch XLA to help us train PyTorch models on TPU.","execution_count":null},{"metadata":{"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\n!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev\n!export XLA_USE_BF16=1\n!pip install -q torchviz","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Import libraries\n\n* We will import several different packages and libraries required for different parts of the project. For example, we import numpy and pandas for data manipulation and torch and torch_xla for modeling.","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport time\nimport numpy as np\nimport pandas as pd\n\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm\nfrom skimage.io import MultiImage\nfrom joblib import Parallel, delayed\n\nfrom sklearn.utils import shuffle\nfrom colorama import Fore, Back, Style\nfrom keras.utils import to_categorical as cat\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nfrom torch import LongTensor as LongTensor\nfrom torch import FloatTensor as FloatTensor\n\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 torchviz import make_dot\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models import resnet34, densenet121, mobilenet_v2\nfrom albumentations import RandomRotate90, Flip, Compose, Normalize, RandomResizedCrop","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.random.seed(42)\ntorch.manual_seed(42)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Set hyperparamerters and paths \n\n* Here, we define the required hyperparameters such as the training batch size, learning rate, training/validation split percentage, etc.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"FOLDS = 8\nEPOCHS = 4\n\nRRC = 1.0\nFLIP = 1.0\nNORM = 1.0\nROTATE = 1.0\nLR = (1e-4, 1e-3)\nMODEL_SAVE_PATH = \"resnet_model\"\n\nWIDTH = 512\nHEIGHT = 512\nBATCH_SIZE = 128\nVAL_BATCH_SIZE = 128\nDATA_PATH = '../input/prostate-cancer-grade-assessment/'\nRESIZED_PATH = '../input/panda-resized-train-data-512x512/train_images/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TEST_DATA_PATH = DATA_PATH + 'test.csv'\nTRAIN_DATA_PATH = DATA_PATH + 'train.csv'\nTEST_IMG_PATH = DATA_PATH + 'test_images/'\nTRAIN_IMG_PATH = RESIZED_PATH + 'train_images/'\nSAMPLE_SUB_PATH = DATA_PATH + 'sample_submission.csv'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load .csv data\n\n* We now load the training and tessting data required for the project using the read_csv function from the pandas library.","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"test_df = pd.read_csv(TEST_DATA_PATH)\ntrain_df = pd.read_csv(TRAIN_DATA_PATH)\nsample_submission = pd.read_csv(SAMPLE_SUB_PATH)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Convert Gleason scores to list format\n\n* Next, we convert the Gleason into the proper format which our model needs.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"gleason_replace_dict = {0:0, 1:1, 3:2, 4:3, 5:4}\n\ndef process_gleason(gleason):\n    if gleason == 'negative': gs = (1, 1)\n    else: gs = tuple(gleason.split('+'))\n    return [gleason_replace_dict[int(g)] for g in gs]\n\ntrain_df.gleason_score = train_df.gleason_score.apply(process_gleason)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Display few images\n\n* Now we display few sample images from the dataset to get an idea of what they look like.\n* In this example, I hve plotted a total of 25 images of biopsy scans in a 5 x 5 grid using matplotlib.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def display_images(num):\n    sq_num = np.sqrt(num)\n    assert sq_num == int(sq_num)\n\n    sq_num = int(sq_num)\n    fig, ax = plt.subplots(nrows=sq_num, ncols=sq_num, figsize=(20, 20))\n\n    for i in range(int(sq_num)):\n        for j in range(int(sq_num)):\n            idx = i*sq_num + j\n            path = TRAIN_IMG_PATH + train_df.image_id[idx]\n    \n            path += '.png'\n            ax[i, j].imshow(cv2.imread(path))\n            ax[i, j].set_title('Image {}'.format(idx), fontsize=12)\n\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"display_images(25)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Modeling <a id=\"2\"></a>","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Build PyTorch dataset\n\n* Now we define a PyTorch Dataset which will help us feed data to the ResNet model for training and inference.\n* We first read images using the OpenCV and skimage python packages, and then normalize and augment the images using albumentations.\n* In this example, I have used the rotation and random cropping augmentations, but more augmentations can be used.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class PANDADataset(Dataset):\n    def __init__(self, data, img_path, is_val=False, is_train=False):\n\n        self.data = data\n        self.is_val = is_val\n        self.is_train = is_train\n        self.image_path = img_path\n        self.image_id = data.image_id\n        self.aug = self.norm = Normalize(p=NORM)\n\n        if is_train or is_val:\n            self.isup_grade = data.isup_grade\n            self.gleason_score = data.gleason_score\n\n            if is_train:\n                self.flip = Flip(p=FLIP)\n                self.rotate = RandomRotate90(p=ROTATE)\n                self.crop = RandomResizedCrop(p=RRC, width=WIDTH, height=HEIGHT)\n                self.aug = Compose([self.flip, self.rotate, self.crop, self.norm], p=1)\n\n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, idx):\n        path = self.image_path + self.image_id[idx]\n\n        if self.is_train or self.is_val:\n            path += '.png'\n            image = cv2.imread(path)\n        else:\n            path += '.tiff'\n            image = MultiImage(path)[-1]\n            image = cv2.resize(image, (HEIGHT, WIDTH))\n\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        image = self.aug(image=image)['image'].reshape((3, HEIGHT, WIDTH))\n        \n        if self.is_train or self.is_val:\n            isup_grade = cat([self.data.isup_grade[idx]], num_classes=6)\n            gleason_0 = cat([self.data.gleason_score[idx][0]], num_classes=5)\n            gleason_1 = cat([self.data.gleason_score[idx][1]], num_classes=5)\n            target = np.concatenate([isup_grade, gleason_0, gleason_1], axis=1)\n            \n        if self.is_train or self.is_val:\n            return FloatTensor(image), FloatTensor(target)\n        else:\n            return FloatTensor(image)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build ResNet model\n\n* Now we build a model based on ResNet-18 from the torchvision package to predict the ISUP grade and the Gleason score.\n* We first get the ResNet-18 backbone and then add three dense heads on this backbone to classify three quantities (ISUP, G0, G1).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class ResNetDetector(nn.Module):\n    def __init__(self):\n        super(ResNetDetector, self).__init__()\n\n        self.softmax = nn.Softmax(dim=1)\n        self.dense_1 = nn.Linear(512, 6)\n        self.dense_2 = nn.Linear(512, 5)\n        self.dense_3 = nn.Linear(512, 5)\n        self.resnet = resnet34(pretrained=True)\n        self.resnet = nn.Sequential(*list(self.resnet.children())[:-1])\n        \n    def forward(self, img):\n        feat = self.resnet(img).squeeze()\n\n        isup_logit = self.dense_1(feat)\n        gleason_logit_0 = self.dense_2(feat)\n        gleason_logit_1 = self.dense_3(feat)\n        \n        isup_prob = self.softmax(isup_logit)\n        gleason_prob_0 = self.softmax(gleason_logit_0)\n        gleason_prob_1 = self.softmax(gleason_logit_1)\n        return torch.cat([isup_prob, gleason_prob_0, gleason_prob_1], axis=1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualize ResNet architecture\n\n* Below is a visualization of the ResNet-18 architecture using torchviz.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"model = ResNetDetector()\nx = torch.randn(2, 3, 32, 32).requires_grad_(True)\ny = model(x)\nmake_dot(y, params=dict(list(model.named_parameters()) + [('x', x)]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del model, x, y\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Split train.csv into 8 folds\n\n* We now split the data into 8 folds using a simple for loop.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"val_sets, train_sets = [], []\nval_splits = np.int32((np.arange(FOLDS + 1)/FOLDS) * len(train_df))\nval_indices = [[val_splits[i], val_splits[i+1]] for i in range(FOLDS)]\n\nfor fold in tqdm(range(FOLDS)):\n    val_idx = val_indices[fold]\n    if fold == FOLDS - 1: val_idx[1] -= 1\n    val_sets.append(train_df[val_idx[0]:val_idx[1]])\n    train_sets.append(pd.concat([train_df[:val_idx[0]], train_df[val_idx[1]:]]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define cross entropy and accuracy\n\n* Here we implement categorical cross entropy and accuracy functions in PyTorch.\n* CEL is the loss function which is commonly used in classification tasks and helps us finetune ResNet's weights.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def cel(inp, targ):\n    _, labels = targ.max(dim=1)\n    return nn.CrossEntropyLoss()(inp, labels)\n\ndef acc(inp, targ):\n    inp_idx = inp.max(axis=1).indices\n    targ_idx = targ.max(axis=1).indices\n    return (inp_idx == targ_idx).float().sum(axis=0)/len(inp_idx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define custom PANDA loss for multitask model\n\n* Now we define a custom loss function for this problem. We basically calculate the cross entropy for each of the three targets and add them.\n* When the model optimizes this new combined CEL function, it will enable the model to learn multiple tasks at the same time (w/ gradients).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def panda_cel(inp, targ):\n    isup_loss = cel(inp[:, :6], targ[:, :6])\n    gleason_loss_0 = cel(inp[:, 6:11], targ[:, 6:11])\n    gleason_loss_1 = cel(inp[:, 11:16], targ[:, 11:16])\n    return [isup_loss, gleason_loss_0, gleason_loss_1],\\\n           isup_loss + gleason_loss_0 + gleason_loss_1\n\ndef panda_acc(inp, targ):\n    isup_accuracy = acc(inp[:, :6], targ[:, :6])\n    gleason_accuracy_0 = acc(inp[:, 6:11], targ[:, 6:11])\n    gleason_accuracy_1 = acc(inp[:, 11:16], targ[:, 11:16])\n    return np.array([isup_accuracy, gleason_accuracy_0, gleason_accuracy_1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define helper function for training logs\n\n* We now define a simple function to get training logs.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def print_metric(data, fold, start, end, metric, typ):\n    r = Fore.RESET\n    n = [\"ISUP\", \"G-0\", \"G-1\"]\n    time = np.round(end - start, 1)\n    time = \"Time: {} s\".format(time)\n    c = [Fore.CYAN, Fore.YELLOW, Fore.MAGENTA]\n    \n    tick = Fore.GREEN + '\\u2714' + Fore.RESET\n    prefix = \"FOLD {} \".format(fold + 1) + tick + \"  \"\n    \n    string = prefix\n    for idx in range(3):\n        value = np.round(data[idx].detach().cpu().numpy(), 3)\n        t = typ, n[idx], metric, c[idx], np.round(value, 3), Fore.RESET\n        string = string + \"{} {} {}: {}{}{}\".format(*t) + \"  \"\n        \n    print(string + time)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train model on all 8 TPU cores in parallel\n\n* We will now train the model in parallel on all 8 TPU cores.\n* In this example, I use the <code>Parallel</code> and <code>delayed</code> functionalities of joblib to run each fold on one TPU core.\n* All 8 folds in our training process and therefore running at the same time on each of the 8 TPU cores available to us on Kaggle kernels.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(fold):\n    val = val_sets[fold]\n    train = train_sets[fold]\n    device = xm.xla_device(fold + 1)\n    \n    def xla(tensor):\n        return tensor.to(device)\n   \n    val = val.reset_index(drop=True)\n    val_set = PANDADataset(val, TRAIN_IMG_PATH, is_val=True)\n    val_loader = DataLoader(val_set, batch_size=VAL_BATCH_SIZE)\n\n    train = train.reset_index(drop=True)\n    train_set = PANDADataset(train, TRAIN_IMG_PATH, is_train=True)\n    train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)\n\n    network = xla(ResNetDetector())\n    optimizer = Adam([{'params': network.resnet.parameters(), 'lr': LR[0]},\n                      {'params': network.dense_1.parameters(), 'lr': LR[1]},\n                      {'params': network.dense_2.parameters(), 'lr': LR[1]},\n                      {'params': network.dense_3.parameters(), 'lr': LR[1]}])\n\n    start = time.time()\n    for epoch in range(EPOCHS):\n\n        batch = 1\n        for train_batch in train_loader:\n            train_img, train_targs = train_batch\n\n            network = xla(network)\n            train_img = xla(train_img)\n            train_targs = xla(train_targs)\n            \n            network.train()\n            train_preds = network.forward(train_img)\n            train_acc = panda_acc(train_preds, train_targs.squeeze())\n            train_loss, total_loss = panda_cel(train_preds, train_targs.squeeze())\n\n            optimizer.zero_grad()\n            total_loss.backward()\n            xm.optimizer_step(optimizer, barrier=True)\n\n            batch = batch + 1\n           \n    network.eval()\n    val_loss, val_acc = 0, 0\n    for val_batch in val_loader:\n\n        img, targ = val_batch\n        with torch.no_grad():\n            network = xla(network)\n            img, targ = xla(img), xla(targ)\n\n            pred = network.forward(img)\n            val_acc += panda_acc(pred, targ.squeeze(dim=1))*len(pred)\n            val_loss += panda_cel(pred, targ.squeeze(dim=1))[1].item()*len(pred)\n       \n    end = time.time()\n    network = network.cpu()\n    val_acc /= len(val_set)\n    val_loss /= len(val_set)\n    path = MODEL_SAVE_PATH + \"_\" + str(fold + 1) + \".pt\"\n    print_metric(val_acc, fold, start, end, metric=\"acc\", typ=\"Val\")\n    torch.save(network.state_dict(), path); del network; gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":false},"cell_type":"code","source":"Parallel(n_jobs=FOLDS, backend=\"threading\")(delayed(train)(i) for i in range(FOLDS))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Takeaways <a id=\"3\"></a>\n\n1. Using all 8 TPU cores in parallel can dramatically speed up KFold training.\n2. Using more complex models (like ResNet-152, DenseNet-201, Efficient-B7, etc) can improve the model's performance.\n3. Separate scripts should be used for training and inference to take full advantage of the TPU for training and GPU for inference.","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}