{"cells":[{"metadata":{},"cell_type":"markdown","source":"<center><img src=\"https://i.imgur.com/20rO1id.jpg\" width=\"500px\"></center>"},{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nHello everyone! Weclome to the \"Prostate cANcer graDe Assessment (PANDA) Challenge\" competition on Kaggle! In this competition, contestants are challenged to build machine learning models 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\nIn this kernel, I 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. This is the magic of multitask learning!**\n\n<font size=3>I will make use of all 8 cores on the TPU v3-8 to train an 8-fold model in less than 1 hour :D</font>\n\n\n<center><img src=\"https://i.imgur.com/piOxK6F.png\" width=\"750px\"></center>"},{"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)"},{"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)"},{"metadata":{},"cell_type":"markdown","source":"# Preparing the ground <a id=\"1\"></a>"},{"metadata":{},"cell_type":"markdown","source":"## Set up PyTorch-XLA <a id=\"1.1\"></a> <font color=\"forestgreen\" size=4>(inspired by Abhishek's kernel :D)</font>"},{"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 <a id=\"1.2\"></a> <font color=\"forestgreen\" size=4>(for data loading, processing, and modeling on TPU)</font>"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":false},"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 torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torchvision.models import resnet18, 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 <a id=\"1.3\"></a> <font color=\"forestgreen\" size=4>(adjust these to improve CV and LB :D)</font>"},{"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 <a id=\"1.4\"></a> <font color=\"forestgreen\" size=4>(to access image IDs for training and validation)</font>"},{"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 <a id=\"1.5\"></a> <font color=\"forestgreen\" size=4>(to add them as targets for ResNet)</font>"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['gleason_score'] = train_df['gleason_score'].apply(lambda x: \"0+0\" if x==\"negative\" else x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['gleason_score'] ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":""},{"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    print(str(gleason[1]))\n    if gleason == 'negative': gs = (1, 1)\n    else: gs = tuple(str(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 <a id=\"1.6\"></a> <a id=\"1.6\"></a> <font color=\"forestgreen\" size=4>(from <i>train_images</i> directory)</font>"},{"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>"},{"metadata":{},"cell_type":"markdown","source":"## Build PyTorch dataset <a id=\"2.1\"></a> <font color=\"forestgreen\" size=4>(with image transforms and multi-target)</font>"},{"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 <a id=\"2.2\"></a> <font color=\"forestgreen\" size=4>(with 3 dense heads for 3 targets)</font>"},{"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 = resnet18(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<a id=\"2.3\"></a> <font color=\"forestgreen\" size=4>(with pytorchviz)</font>"},{"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":{"_kg_hide-input":true,"_kg_hide-output":true,"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 <a id=\"2.4\"></a> <font color=\"forestgreen\" size=4>(for cross-validation)</font>"},{"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 <a id=\"2.5\"></a> <font color=\"forestgreen\" size=4>(for backpropagation)</font>"},{"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 <a id=\"2.6\"></a> <font color=\"forestgreen\" size=4>(combining CEL for all three targets)</font>"},{"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 [isup_accuracy, gleason_accuracy_0, gleason_accuracy_1]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define helper function for training logs <a id=\"2.7\"></a> <font color=\"forestgreen\" size=4>(to check training status)</font>"},{"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].item(), 3)\n        t = typ, n[idx], metric, c[idx], value, 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 <a id=\"2.8\"></a> <font color=\"forestgreen\" size=4>(one fold per core)</font>"},{"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    scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.5,\n                                  patience=2, verbose=True, eps=1e-6)\n\n    start = time.time()\n    for epoch in range(EPOCHS):\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        for val_batch in val_loader:\n            img, targ = val_batch\n            val_preds, val_targs = [], []\n\n            with torch.no_grad():\n                img = xla(img)\n                network = xla(network)\n                pred = network.forward(img)\n                val_preds.append(pred); val_targs.append(targ)\n        \n        val_preds = torch.cat(val_preds, axis=0)\n        val_targs = torch.cat(val_targs, axis=0)\n        \n        val_targs = xla(val_targs)\n        val_acc = panda_acc(val_preds, val_targs.squeeze())\n        val_loss, _ = panda_cel(val_preds, val_targs.squeeze())\n       \n        scheduler.step(val_loss[0])\n       \n    end = time.time()\n    print_metric(val_acc, fold, start, end, metric=\"acc\", typ=\"Val\")\n    torch.save(network.state_dict(), MODEL_SAVE_PATH + \"_\" + str(fold + 1) + \".pt\")","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":{"trusted":true},"cell_type":"code","source":"plt.plot(train_losses, label='Training loss')\nplt.plot(test_losses, label='Validation loss')\nplt.legend(frameon=False)\nplt.show()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nmodel = ResNetDetector()\nmodel.load_state_dict(torch.load((MODEL_SAVE_PATH + \"_\" + str(fold + 1) + \".pt\")))\nmodel.eval()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef predict_image(image):\n    image_tensor = test_transforms(image).float()\n    image_tensor = image_tensor.unsqueeze_(0)\n    input = Variable(image_tensor)\n    input = input.to(device)\n    output = model(input)\n    index = output.data.cpu().numpy().argmax()\n    return index","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val = val_sets[fold]\ntrain = train_sets[fold]\ndevice = xm.xla_device(fold + 1)\n\ndef xla(tensor):\n    return tensor.to(device)\n\nval = val.reset_index(drop=True)\nval_set = PANDADataset(val, TRAIN_IMG_PATH, is_val=True)\nval_loader = DataLoader(val_set, batch_size=VAL_BATCH_SIZE)\n\n\n\nnetwork = xla(model)\n\n\nstart=time.time()\nfor val_batch in val_loader:\n    img, targ = val_batch\n    val_preds, val_targs = [], []\n\n    with torch.no_grad():\n        img = xla(img)\n        network = xla(network)\n        pred = network.forward(img)\n        val_preds.append(pred); val_targs.append(targ)\n\nval_preds = torch.cat(val_preds, axis=0)\nval_targs = torch.cat(val_targs, axis=0)\n\nval_targs = xla(val_targs)\nval_acc = panda_acc(val_preds, val_targs.squeeze())\nval_loss, _ = panda_cel(val_preds, val_targs.squeeze())\n\n\nend = time.time()\n# print_metric(val_acc, fold, start, end, metric=\"acc\", typ=\"Val\")\n# torch.save(network.state_dict(), MODEL_SAVE_PATH + \"_\" + str(fold + 1) + \".pt\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.dense_2\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.dense_3","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 kernels should be used for training and inference to take full advantage of the TPU for training and GPU for inference."},{"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}