{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"DEBUG = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install torchsummary","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport sys\nsys.path = [\n    '../input/efficientnet-pytorch/EfficientNet-PyTorch/EfficientNet-PyTorch-master',\n] + sys.path","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import time\nimport skimage.io\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport PIL.Image\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SubsetRandomSampler, RandomSampler, SequentialSampler\nfrom warmup_scheduler import GradualWarmupScheduler\nfrom efficientnet_pytorch import model as enet\nimport albumentations\nfrom sklearn.model_selection import StratifiedKFold\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import cohen_kappa_score\nfrom tqdm import tqdm_notebook as tqdm\nfrom torchsummary import summary","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Config","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"data_dir = '../input/prostate-cancer-grade-assessment'\ndf_train = pd.read_csv(os.path.join(data_dir, 'train.csv'))\nimage_folder = os.path.join(data_dir, 'train_images')\n\nkernel_type = 'enet_b0'\n\nenet_type = 'efficientnet-b0'\nfold = 0\ntile_size = 256\nimage_size = 256\nn_tiles = 36\nbatch_size = 2\nnum_workers = 4\nout_dim = 5\ninit_lr = 3e-4\nwarmup_factor = 10\n\nwarmup_epo = 1\nn_epochs = 1 if DEBUG else 3\ndf_train = df_train.sample(100).reset_index(drop=True) if DEBUG else df_train\n\ndevice = torch.device('cuda')\n\nprint(image_folder)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Folds","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"skf = StratifiedKFold(5, shuffle=True, random_state=42)\ndf_train['fold'] = -1\nfor i, (train_idx, valid_idx) in enumerate(skf.split(df_train, df_train['isup_grade'])):\n    df_train.loc[valid_idx, 'fold'] = i\ndf_train.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"pretrained_model = {\n    'efficientnet-b0': '../input/efficientnet-pytorch/efficientnet-b0-08094119.pth'\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class enetv2(nn.Module):\n    def __init__(self, backbone, out_dim):\n        super(enetv2, self).__init__()\n        self.enet = enet.EfficientNet.from_name(backbone)\n        self.enet.load_state_dict(torch.load(pretrained_model[backbone]))\n\n        self.myfc = nn.Linear(self.enet._fc.in_features, out_dim)\n        self.enet._fc = nn.Identity()\n\n    def extract(self, x):\n        return self.enet(x)\n\n    def forward(self, x):\n        x = self.extract(x)\n        x = self.myfc(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_tiles(img, mode=0):\n        result = []\n        h, w, c = img.shape\n        pad_h = (tile_size - h % tile_size) % tile_size\n        pad_w = (tile_size - w % tile_size) % tile_size\n\n        img2 = np.pad(img,[[pad_h // 2, pad_h - pad_h // 2], [pad_w // 2,pad_w - pad_w//2], [0,0]], constant_values=255)\n        img3 = img2.reshape(img2.shape[0] // tile_size, tile_size, img2.shape[1] // tile_size, tile_size, 3)\n\n        img3 = img3.transpose(0,2,1,3,4).reshape(-1, tile_size, tile_size,3)\n        if len(img3) < n_tiles:\n            img3 = np.pad(img3,[[0,n_tiles-len(img3)],[0,0],[0,0],[0,0]], constant_values=255)\n        idxs = np.argsort(img3.reshape(img3.shape[0],-1).sum(-1))[:n_tiles]\n        img3 = img3[idxs]\n        for i in range(len(img3)):\n            result.append({'img':img3[i], 'idx':i})\n        return result\n\n\nclass PANDADataset(Dataset):\n    def __init__(self,df,image_size,n_tiles=n_tiles,tile_mode=0,rand=False,transform=None):\n\n        self.df = df.reset_index(drop=True)\n        self.image_size = image_size\n        self.n_tiles = n_tiles\n        self.tile_mode = tile_mode\n        self.rand = rand\n        self.transform = transform\n\n    def __len__(self):\n        return self.df.shape[0]\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        img_id = row.image_id\n        \n        tiff_file = os.path.join(image_folder, f'{img_id}.tiff')\n        image = skimage.io.MultiImage(tiff_file)[1]\n        tiles = get_tiles(image, self.tile_mode)\n\n        if self.rand:\n            idxes = np.random.choice(list(range(self.n_tiles)), self.n_tiles, replace=False)\n        else:\n            idxes = list(range(self.n_tiles))\n\n        n_row_tiles = int(np.sqrt(self.n_tiles))\n        images = np.zeros((image_size * n_row_tiles, image_size * n_row_tiles, 3))\n        for h in range(n_row_tiles):\n            for w in range(n_row_tiles):\n                i = h * n_row_tiles + w\n    \n                if len(tiles) > idxes[i]:\n                    this_img = tiles[idxes[i]]['img']\n                else:\n                    this_img = np.ones((self.image_size, self.image_size, 3)).astype(np.uint8) * 255\n                this_img = 255 - this_img\n                if self.transform is not None:\n                    this_img = self.transform(image=this_img)['image']\n                h1 = h * image_size\n                w1 = w * image_size\n                images[h1:h1+image_size, w1:w1+image_size] = this_img\n\n        if self.transform is not None:\n            images = self.transform(image=images)['image']\n        images = images.astype(np.float32)\n        images /= 255\n        images = images.transpose(2, 0, 1)\n\n        label = np.zeros(5).astype(np.float32)\n        label[:row.isup_grade] = 1.\n        return torch.tensor(images), torch.tensor(label)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    albumentations.Transpose(p=0.5),\n    albumentations.VerticalFlip(p=0.5),\n    albumentations.HorizontalFlip(p=0.5),\n])\ntransforms_val = albumentations.Compose([])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_show = PANDADataset(df_train, image_size, n_tiles, 0, transform=transforms_train)\nfrom pylab import rcParams\nrcParams['figure.figsize'] = 20,10\nfor i in range(2):\n    f, axarr = plt.subplots(1,5)\n    for p in range(5):\n        idx = np.random.randint(0, len(dataset_show))\n        img, label = dataset_show[idx]\n        axarr[p].imshow(1. - img.transpose(0, 1).transpose(1,2).squeeze())\n        axarr[p].set_title(str(sum(label)))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"criterion = nn.BCEWithLogitsLoss()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train & Val","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(loader, optimizer):\n\n    model.train()\n    train_loss = []\n    bar = tqdm(loader)\n    for (data, target) in bar:\n        \n        data, target = data.to(device), target.to(device)\n        loss_func = criterion\n        optimizer.zero_grad()\n        logits = model(data)\n        loss = loss_func(logits, target)\n        loss.backward()\n        optimizer.step()\n\n        loss_np = loss.detach().cpu().numpy()\n        train_loss.append(loss_np)\n        smooth_loss = sum(train_loss[-100:]) / min(len(train_loss), 100)\n        bar.set_description('loss: %.5f, smth: %.5f' % (loss_np, smooth_loss))\n    return train_loss\n\n\ndef val_epoch(loader, get_output=False):\n\n    model.eval()\n    val_loss = []\n    LOGITS = []\n    PREDS = []\n    TARGETS = []\n\n    with torch.no_grad():\n        for (data, target) in tqdm(loader):\n            data, target = data.to(device), target.to(device)\n            logits = model(data)\n\n            loss = criterion(logits, target)\n\n            pred = logits.sigmoid().sum(1).detach().round()\n            LOGITS.append(logits)\n            PREDS.append(pred)\n            TARGETS.append(target.sum(1))\n\n            val_loss.append(loss.detach().cpu().numpy())\n        val_loss = np.mean(val_loss)\n\n    LOGITS = torch.cat(LOGITS).cpu().numpy()\n    PREDS = torch.cat(PREDS).cpu().numpy()\n    TARGETS = torch.cat(TARGETS).cpu().numpy()\n    acc = (PREDS == TARGETS).mean() * 100.\n    \n    qwk = cohen_kappa_score(PREDS, TARGETS, weights='quadratic')\n    qwk_k = cohen_kappa_score(PREDS[df_valid['data_provider'] == 'karolinska'], df_valid[df_valid['data_provider'] == 'karolinska'].isup_grade.values, weights='quadratic')\n    qwk_r = cohen_kappa_score(PREDS[df_valid['data_provider'] == 'radboud'], df_valid[df_valid['data_provider'] == 'radboud'].isup_grade.values, weights='quadratic')\n    print('qwk', qwk, 'qwk_k', qwk_k, 'qwk_r', qwk_r)\n\n    if get_output:\n        return LOGITS\n    else:\n        return val_loss, acc, qwk\n\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Dataloader & Model & Optimizer","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"train_idx = np.where((df_train['fold'] != fold))[0]\nvalid_idx = np.where((df_train['fold'] == fold))[0]\n\ndf_this  = df_train.loc[train_idx]\ndf_valid = df_train.loc[valid_idx]\n\ndataset_train = PANDADataset(df_this , image_size, n_tiles, transform=transforms_train)\ndataset_valid = PANDADataset(df_valid, image_size, n_tiles, transform=transforms_val)\n\ntrain_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, sampler=RandomSampler(dataset_train), num_workers=num_workers)\nvalid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=batch_size, sampler=SequentialSampler(dataset_valid), num_workers=num_workers)\n\n# checkpoint = torch.load('../input/panda-checkpoint-demo/model_checkpoint_epoch2.pth')\n#pretrained_model = checkpoint\n#enet_type = 'model_state_dict'\nmodel = enetv2(enet_type, out_dim=out_dim)\nmodel = model.to(device)\n\noptimizer = optim.Adam(model.parameters(), lr=init_lr/warmup_factor)\nscheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs-warmup_epo)\nscheduler = GradualWarmupScheduler(optimizer, multiplier=warmup_factor, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n\nprint(len(dataset_train), len(dataset_valid))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"next(iter(train_loader))[0].shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"summary(model, (3, 1536, 1536))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Run Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"for param in model.enet.parameters():\n    param.requires_grad = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"summary(model, (3, 1536, 1536))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\n    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n    prev_epoch = checkpoint['epoch']\n    tloss = checkpoint['loss']\nexcept:\n    prev_epoch = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"qwk_max = 0.\nbest_file = f'{kernel_type}_best_fold{fold}.pth'\nfor epoch in range(1, n_epochs+1):\n    print(time.ctime(), 'Epoch:', epoch)\n    scheduler.step(epoch-1)\n\n    train_loss = train_epoch(train_loader, optimizer)\n    val_loss, acc, qwk = val_epoch(valid_loader)\n\n    content = time.ctime() + ' ' + f'Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, train loss: {np.mean(train_loss):.5f}, val loss: {np.mean(val_loss):.5f}, acc: {(acc):.5f}, qwk: {(qwk):.5f}'\n    print(content)\n    with open(f'log_{kernel_type}.txt', 'a') as appender:\n        appender.write(content + '\\n')\n    \n    curr_epoch = epoch+prev_epoch+1\n    if qwk > qwk_max:\n        print('score2 ({:.6f} --> {:.6f}).  Saving model ...'.format(qwk_max, qwk))\n        torch.save({\n            'epoch': curr_epoch,\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'loss': tloss\n            }, f'{kernel_type}_final_fold{fold}_epoch_{curr_epoch}.pth')\n        qwk_max = qwk\n        \n#torch.save(model.state_dict(), os.path.join(f'{kernel_type}_final_fold{fold}.pth'))","execution_count":null,"outputs":[]},{"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}