{"cells":[{"metadata":{},"cell_type":"markdown","source":"In previous notebooks, we discussed, \n- [Explorative data analysis](https://www.kaggle.com/zij212/human-protein-atlas-eda)\n- [Train validation split for imbalance set](https://www.kaggle.com/zij212/human-protein-atlas-train-val-split)\n- [Get statistics of image set](https://www.kaggle.com/zij212/human-protein-atlas-data-stats)\n\nNow that we have split out the training and validation data, and got the stats of the training set, we are ready to training a model.\n\nIn this notebook, we will,\n\n- [Preprocess the data (using data augmentation, and scaling)](#Data-Preprocessing)\n- [Build a model](#Modeling)\n    - [Use Pretrained Resnet34 for transfer learning](#Transfer-Learning-using-Resnet34)\n    - [Define the Focal loss funtion](#Focal-Loss)\n    - [Define evaluation metric](#F-score)\n- [Training the model](#Model-Training)\n    - [Train weights of the first Conv2D layer and the output layer (freeze all other layers)](#Train-Output-and-1st-Conv2d-Layer)\n    - [Train all weights of the model (unfreeze all layers)](#Train-all-Model-Weights)\n    - [plot out the train/validation score, loss, and learning rate](#Plot-Learning-Process)"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as T\nfrom PIL import Image\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models as models\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nfrom typing import List, Dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_CSV = '/kaggle/input/human-protein-atlas-train-val-split/train_df.csv'\nVAL_CSV = '/kaggle/input/human-protein-atlas-train-val-split/test_df.csv'\nTEST_CSV = '/kaggle/input/human-protein-atlas-image-classification/sample_submission.csv'\n\nTRAIN_DIR = '/kaggle/input/human-protein-atlas-image-classification/train'\nTEST_DIR = '/kaggle/input/human-protein-atlas-image-classification/test'\n\nSTATS_DIR = '/kaggle/input/human-protein-atlas-data-stats/stats.pt'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(TRAIN_CSV)\nval_df = pd.read_csv(VAL_CSV)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"text_labels = {\n0:  \"Nucleoplasm\", \n1:  \"Nuclear membrane\",   \n2:  \"Nucleoli\",   \n3:  \"Nucleoli fibrillar center\" ,  \n4:  \"Nuclear speckles\",\n5:  \"Nuclear bodies\",\n6:  \"Endoplasmic reticulum\",   \n7:  \"Golgi apparatus\",\n8:  \"Peroxisomes\",\n9:  \"Endosomes\",\n10:  \"Lysosomes\",\n11:  \"Intermediate filaments\",   \n12:  \"Actin filaments\",\n13:  \"Focal adhesion sites\",   \n14:  \"Microtubules\",\n15:  \"Microtubule ends\",   \n16:  \"Cytokinetic bridge\",   \n17:  \"Mitotic spindle\",\n18:  \"Microtubule organizing center\",  \n19:  \"Centrosome\",\n20:  \"Lipid droplets\",   \n21:  \"Plasma membrane\",   \n22:  \"Cell junctions\", \n23:  \"Mitochondria\",\n24:  \"Aggresome\",\n25:  \"Cytosol\",\n26:  \"Cytoplasmic bodies\",   \n27:  \"Rods & rings\" \n}\n\nNUM_LABELS = len(text_labels)\nprint(f\"There are {NUM_LABELS} labels\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Preprocessing"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"FILTERS = ['red', 'green', 'blue', 'yellow']\n\ndef load_image(image_id, ddir):\n    \"\"\"\n    return: 4-channel PIL Image\n    \"\"\"\n    return Image.merge('RGBA', [Image.open(f\"{TRAIN_DIR}/{image_id}_{f}.png\") for f in FILTERS])\n\n\ndef encode(image_labels: str):\n    \"\"\"\n    image_labels: label(s) of an image, e.g. \"25 0\"\n    return: tensor of size (28)\n    \"\"\"\n    target = torch.zeros(NUM_LABELS)\n    for label in image_labels.split():\n        target[int(label)] = 1\n    return target","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ProteinLocalizationDataset(Dataset):\n    def __init__(self, df, ddir, transform=None):\n        self.df = df\n        self.ddir = ddir\n        self.transform = transform\n    \n    \n    def __len__(self):\n        return len(self.df)\n    \n    \n    def __getitem__(self, idx):\n        image_id, image_labels = self.df.loc[idx]\n        image = load_image(image_id, self.ddir)\n        if self.transform:\n            image = self.transform(image)\n        return image, encode(image_labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"stats = torch.load(STATS_DIR)\nstats","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_tfms = T.Compose([\n    T.RandomRotation(10),\n    T.RandomHorizontalFlip(),\n    T.ToTensor(), \n    T.Normalize(*stats, inplace=True)\n])\n\ntest_tfms = T.Compose([\n    T.ToTensor(), \n    T.Normalize(*stats, inplace=True)\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = ProteinLocalizationDataset(\n    train_df, TRAIN_DIR, transform=train_tfms)\nval_ds = ProteinLocalizationDataset(\n    val_df, TRAIN_DIR, transform=test_tfms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dl = DataLoader(\n    train_ds, batch_size, shuffle=True, num_workers=2, pin_memory=True)\nval_dl = DataLoader(\n    val_ds, batch_size*2, shuffle=False, num_workers=2, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Modeling"},{"metadata":{},"cell_type":"markdown","source":"## Transfer Learning using Resnet34"},{"metadata":{"trusted":true},"cell_type":"code","source":"class MultiLocalizationClassification(nn.Module):\n    def training_step(self, batch):\n        imgs, targets = batch\n        out = self(imgs)\n        loss = CRITERION(out, targets)\n        return loss\n  \n\n    def validation_step(self, batch):\n        imgs, targets = batch\n        out = self(imgs)\n        loss = CRITERION(out, targets)\n        score = f_score(out, targets)\n        return {'val_loss': loss.detach(), 'val_score': score.detach()}\n    \n    \n    @staticmethod\n    def validation_epoch_end(outputs: List):\n        val_losses = [o['val_loss'] for o in outputs]\n        val_scores = [o['val_score'] for o in outputs]\n        \n        val_loss = torch.mean(torch.stack(val_losses))\n        val_score = torch.mean(torch.stack(val_scores))\n        return {'val_loss': val_loss.item(), 'val_score': val_score.item()}\n\n    \n    @staticmethod\n    def epoch_end(epoch_num: int, result: Dict):\n        train_loss, val_loss, val_score = result['train_loss'], result['val_loss'], result['val_score']\n        print(f\"Epoch {epoch_num}, train_loss: {train_loss}, val_loss: {val_loss}, val_score:{val_score}\")\n\n        \n        \nclass Resnet34(MultiLocalizationClassification):\n    def __init__(self):\n        super().__init__()\n        self.network = models.resnet34(pretrained=True)\n        # weight for RGB is from Resnet34, weight for Y is set to mean(weight of RGB)\n        weight = self.network.conv1.weight.clone()\n        self.network.conv1 = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n        with torch.no_grad():\n            self.network.conv1.weight[:,:3] = weight\n            self.network.conv1.weight[:, 3] = torch.mean(weight, dim=1)\n        # update out_features to NUM_LABELS\n        in_features = self.network.fc.in_features\n        self.network.fc = nn.Linear(in_features, NUM_LABELS)\n        \n            \n    def forward(self, xb):\n        return torch.sigmoid(self.network(xb))\n    \n    \n    def freeze(self):\n        for param in self.network.parameters():\n            param.requires_grad = False\n        for param in self.network.fc.parameters():\n            param.requires_grad = True\n        for param in self.network.conv1.parameters():\n            param.requires_grad = True\n    \n    \n    def unfreeze(self):\n        for param in self.network.parameters():\n            param.requires_grad = True","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Focal Loss\n\n[Focal loss](https://arxiv.org/pdf/1708.02002.pdf) is a variation of cross entropy loss.\n\nLet's see what's the difference between the binary cross entropy loss and the focal loss\n\n$$\n\\begin{align}\n\\text{BCELoss}&=-[y \\cdot ln(p) + (1-y) \\cdot ln(1-p)] \\\\\n\\text{FocalLoss} &= -[(1-p)^\\gamma \\cdot y \\cdot ln(p) + p^\\gamma \\cdot (1-y) \\cdot ln(1-p)]\n\\end{align}\n$$\n\nSetting $\\gamma \\gt 0$, reduce the relative weight of well classified samples, and focus more on the harder samples.\n\nIn the implementation below, we will add a small number $\\epsilon$ to $p$ before taking the log to avoid nan."},{"metadata":{"trusted":true},"cell_type":"code","source":"class FocalLoss(nn.Module):\n    def __init__(self, gamma, eps=1e-7):\n        super().__init__()\n        self.gamma = gamma\n        self.eps = eps  \n        \n    def forward(self, preds, targets):\n        preds = preds.clamp(self.eps, 1 - self.eps)\n        loss = (1 - preds) ** self.gamma * targets * torch.log(preds)  \\\n               + preds ** self.gamma * (1 - targets) * torch.log(1 - preds) \n    \n        return -torch.mean(loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# CRITERION = F.binary_cross_entropy\nCRITERION = FocalLoss(gamma=1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## F-score"},{"metadata":{"trusted":true},"cell_type":"code","source":"EPSILON = 1e-6\n\ndef f_score(pred, target, threshold=0.5, beta=1):\n    target = target > threshold\n    pred = pred > threshold\n    \n    TP = (pred & target).sum(1, dtype=float)\n    FP = (pred & ~target).sum(1, dtype=float)\n    FN = (~pred & target).sum(1, dtype=float)\n    \n    precision = TP / (TP + FP + EPSILON)\n    recall = TP / (TP + FN + EPSILON)\n    f_scores = (1 + beta ** 2) * precision * recall / (\n        beta ** 2 * precision + recall + EPSILON)\n    \n    return f_scores.mean()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Before moving on to train the model, we need to move the data and the model to gpu if available."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_default_device():\n    if torch.cuda.is_available():\n        return torch.device('cuda')\n    return torch.device('cpu')\n\ndef to_device(data, device):\n    if isinstance(data, (list, tuple)):\n        return [to_device(d, device) for d in data]\n    return data.to(device, non_blocking=True)\n\nclass DeviceDataLoader():\n    def __init__(self, dl, device):\n        self.dl = dl\n        self.device = device\n    \n    def __len__(self):\n        return len(self.dl)\n    \n    def __iter__(self):\n        for batch in self.dl:\n            yield to_device(batch, self.device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = get_default_device()\n\nmodel = to_device(Resnet34(), device)\n\ntrain_dl = DeviceDataLoader(train_dl, device)\nval_dl = DeviceDataLoader(val_dl, device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef evaluate(model, val_loader):\n    model.eval()\n    outputs = [model.validation_step(batch) for batch in val_loader]\n    return model.validation_epoch_end(outputs)\n\n\ndef get_lr(optimizer):\n    for param_group in optimizer.param_groups:\n        return param_group['lr']\n\n    \ndef fit_one_cycle(epochs, model, train_dl, val_dl, max_lr,weight_decay=0, \n                  grad_clip=None, opt_func=torch.optim.SGD):\n    \n    torch.cuda.empty_cache()\n    history = []\n    \n    optimizer = opt_func(model.parameters(), max_lr, weight_decay=weight_decay)\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n        optimizer, max_lr, epochs=epochs, steps_per_epoch=len(train_dl))\n    \n    for epoch in range(epochs):\n        model.train()\n        train_losses = []\n        lrs = []\n        \n        for batch in train_dl:\n            loss = model.training_step(batch)\n            train_losses.append(loss.detach())\n            loss.backward()\n            \n            if grad_clip:\n                nn.utils.clip_grad_value_(model.parameters(), grad_clip)\n            \n            optimizer.step()\n            optimizer.zero_grad()\n            lrs.append(get_lr(optimizer))\n            \n            scheduler.step()\n        \n        result = evaluate(model, val_dl)\n        result['train_loss'] = torch.stack(train_losses).mean().item()\n        result['lrs'] = lrs\n        model.epoch_end(epoch, result)\n        history.append(result)    \n    return history   \n            ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nhistory = [evaluate(model, val_dl)]\nhistory","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Output and 1st Conv2d Layer"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.freeze()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 10\nmax_lr = 0.01\ngrad_clip = 0.1\nweight_decay = 1e-4\nopt_func = torch.optim.Adam","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nhistory += fit_one_cycle(epochs, model, train_dl, val_dl, max_lr, \n                         grad_clip=grad_clip, \n                         weight_decay=weight_decay, \n                         opt_func=opt_func)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train all Model Weights"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.unfreeze()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 10\nmax_lr = 0.01\ngrad_clip = 0.1\nweight_decay = 1e-4\nopt_func = torch.optim.Adam","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nhistory += fit_one_cycle(epochs, model, train_dl, val_dl, max_lr, \n                         grad_clip=grad_clip, \n                         weight_decay=weight_decay, \n                         opt_func=opt_func)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 10\nmax_lr = 0.005\ngrad_clip = 0.1\nweight_decay = 1e-4\nopt_func = torch.optim.Adam","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nhistory += fit_one_cycle(epochs, model, train_dl, val_dl, max_lr, \n                         grad_clip=grad_clip, \n                         weight_decay=weight_decay, \n                         opt_func=opt_func)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plot Learning Process"},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_lrs(history):\n    scores = [x.get('lrs') for x in history]\n    plt.plot(scores, '-x')\n    plt.xlabel('epoch')\n    plt.ylabel('learning rate')\n    plt.title('Learning rate vs. No. of epochs');\n    \ndef plot_scores(history):\n    scores = [x.get('val_score') for x in history]\n    plt.plot(scores, '-x')\n    plt.xlabel('epoch')\n    plt.ylabel('score')\n    plt.title('F1 score vs. No. of epochs');\n\n    \ndef plot_losses(history):\n    train_losses = [x.get('train_loss') for x in history]\n    val_losses = [x.get('val_loss') for x in history]\n    plt.plot(train_losses, '-bx')\n    plt.plot(val_losses, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('loss')\n    plt.legend(['Training', 'Validation'])\n    plt.title('Loss vs. No. of epochs');\n    \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_losses(history)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_scores(history)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_lrs(history)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"weights_fname = 'protein-resnet.pth'\ntorch.save(model.state_dict(), weights_fname)","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}