{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install /kaggle/input/efficientnet-pytorch/efficientnet_pytorch-0.7.0-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !pip install torch-lr-finder","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!mkdir /root/.cache/torch /root/.cache/torch/hub /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/efficient-b0-b7-from-angtk/*.pth /root/.cache/torch/hub/checkpoints/","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import rasterio\nfrom rasterio.windows import Window","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pathlib, sys, os, random, time\nimport numba, cv2, gc\nfrom efficientnet_pytorch import EfficientNet\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom tqdm.notebook import tqdm\n\nimport albumentations as A","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data as D\n\nimport torchvision\nfrom torchvision import transforms as T","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def set_seeds(seed=17):\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.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n\nset_seeds();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/hubmap-kidney-segmentation'\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# used for converting the decoded image to rle mask\ndef rle_encode(im):\n    '''\n    im: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = im.flatten(order = 'F')\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\ndef rle_decode(mask_rle, shape=(256, 256)):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape, order='F')\n\n@numba.njit()\ndef rle_numba(pixels):\n    size = len(pixels)\n    points = []\n    if pixels[0] == 1: points.append(0)\n    flag = True\n    for i in range(1, size):\n        if pixels[i] != pixels[i-1]:\n            if flag:\n                points.append(i+1)\n                flag = False\n            else:\n                points.append(i+1 - points[-1])\n                flag = True\n    if pixels[-1] == 1: points.append(size-points[-1]+1)    \n    return points\n\ndef rle_numba_encode(image):\n    pixels = image.flatten(order = 'F')\n    points = rle_numba(pixels)\n    return ' '.join(str(x) for x in points)\n\ndef make_grid(shape, window=256, min_overlap=32):\n    \"\"\"\n        Return Array of size (N,4), where N - number of tiles,\n        2nd axis represente slices: x1,x2,y1,y2 \n    \"\"\"\n    x, y = shape\n    nx = x // (window - min_overlap) + 1\n    x1 = np.linspace(0, x, num=nx, endpoint=False, dtype=np.int64)\n    x1[-1] = x - window\n    x2 = (x1 + window).clip(0, x)\n    ny = y // (window - min_overlap) + 1\n    y1 = np.linspace(0, y, num=ny, endpoint=False, dtype=np.int64)\n    y1[-1] = y - window\n    y2 = (y1 + window).clip(0, y)\n    slices = np.zeros((nx,ny, 4), dtype=np.int64)\n    \n    for i in range(nx):\n        for j in range(ny):\n            slices[i,j] = x1[i], x2[i], y1[j], y2[j]    \n    return slices.reshape(nx*ny,4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"identity = rasterio.Affine(1, 0, 0, 0, 1, 0)\n\nclass HubDataset(D.Dataset):\n\n    def __init__(self, root_dir, transform,\n                 window=256, overlap=32, threshold = 100):\n        self.path = pathlib.Path(root_dir)\n        self.overlap = overlap\n        self.window = window\n        self.transform = transform\n        self.csv = pd.read_csv((self.path / 'train.csv').as_posix(),\n                               index_col=[0])\n        self.threshold = threshold\n        \n        self.x, self.y = [], []\n        self.build_slices()\n        self.len = len(self.x)\n        self.as_tensor = T.Compose([\n            T.ToTensor(),\n            T.Normalize([0.6134, 0.4129, 0.6603],\n                        [0.1211, 0.1630, 0.0948]),\n        ])\n        \n    \n    def build_slices(self):\n        self.masks = []\n        self.files = []\n        self.slices = []\n        for i, filename in enumerate(self.csv.index.values):\n            filepath = (self.path /'train'/(filename+'.tiff')).as_posix()\n            self.files.append(filepath)\n            \n            print('Transform', filename)\n            with rasterio.open(filepath, transform = identity) as dataset:\n                self.masks.append(rle_decode(self.csv.loc[filename, 'encoding'], dataset.shape))\n                slices = make_grid(dataset.shape, window=self.window, min_overlap=self.overlap)\n                \n                for slc in tqdm(slices):\n                    x1,x2,y1,y2 = slc\n                    if self.masks[-1][x1:x2,y1:y2].sum() > self.threshold or np.random.randint(100) > 120:\n                        self.slices.append([i,x1,x2,y1,y2])\n                        \n                        image = dataset.read([1,2,3],\n                            window=Window.from_slices((x1,x2),(y1,y2)))\n                        \n#                         if image.std().mean() < 10:\n#                             continue\n                        \n                        # print(image.std().mean(), self.masks[-1][x1:x2,y1:y2].sum())\n                        image = np.moveaxis(image, 0, -1)\n                        self.x.append(image)\n                        self.y.append(self.masks[-1][x1:x2,y1:y2])\n    \n    # get data operation\n    def __getitem__(self, index):\n        image, mask = self.x[index], self.y[index]\n        augments = self.transform(image=image, mask=mask)\n        return self.as_tensor(augments['image']), augments['mask'][None]\n    \n    def __len__(self):\n        \"\"\"\n        Total number of samples in the dataset\n        \"\"\"\n        return self.len\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WINDOW=1024\nMIN_OVERLAP=32\nNEW_SIZE=512\n\ntrfm = A.Compose([\n            A.RandomCrop(NEW_SIZE,NEW_SIZE),\n            A.HorizontalFlip(),\n            A.VerticalFlip(),\n            A.RandomRotate90(),\n            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.05, rotate_limit=15, p=0.5, \n                             border_mode=cv2.BORDER_REFLECT),\n            A.OneOf([\n                A.GaussianBlur(),\n                A.GaussNoise()\n            ], p=0.2),\n    \n            A.OneOf([\n                A.HueSaturationValue(5,5,5),\n                A.CLAHE(clip_limit=2),\n                A.RandomBrightnessContrast(),\n                A.ColorJitter(brightness=0.07, contrast=0.07,\n                   saturation=0.1, hue=0.1, always_apply=False, p=0.1),\n            ], p=0.3),\n        ])\n\nds = HubDataset(DATA_PATH, window=WINDOW, overlap=MIN_OVERLAP, transform=trfm)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # tensor(0.6134) tensor(0.4129) tensor(0.6603)\n# # tensor(0.1211) tensor(0.1630) tensor(0.0948)\n\n# r_m = 0.0\n# g_m = 0.0\n# b_m = 0.0\n# r_s = 0.0\n# g_s = 0.0\n# b_s = 0.0\n# for i,(image, mask) in tqdm(enumerate(ds), total=len(ds)):\n#     r_m+=torch.mean(image[0,:,:],)\n#     g_m+=torch.mean(image[1,:,:],)\n#     b_m+=torch.mean(image[2,:,:],)\n#     r_s+=torch.std(image[0,:,:],)\n#     g_s+=torch.std(image[1,:,:],)\n#     b_s+=torch.std(image[2,:,:],)\n# print(r_m/len(ds),g_m/len(ds),b_m/len(ds))\n# print(r_s/len(ds),g_s/len(ds),b_s/len(ds))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(100,120):\n    image, mask = ds[i]\n    plt.figure(figsize=(16,8))\n    plt.subplot(121)\n    plt.imshow(mask[0], cmap='gray')\n    plt.subplot(122)\n    plt.imshow(image[0]);\n\n# _ = rle_numba_encode(mask[0]) # compile function with numba","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_size = int(0.75 * len(ds))\nvalid_size = len(ds) - train_size\ntrain_ds, valid_ds = torch.utils.data.random_split(ds, [train_size, valid_size])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torchvision.transforms as T\n\n\n# Utility Functions for the model\ndef double_conv(in_,out_,drop): # Double convolution layer for decoder \n\tconv = nn.Sequential(\n\t\tnn.Conv2d(in_,out_,kernel_size=3,padding=(1,1)),\n\t\tnn.ReLU(inplace=True),\n\t\tnn.Conv2d(out_,out_,kernel_size=3,padding=(1,1)),\n\t\tnn.ReLU(inplace=True),\n        nn.Dropout(drop)\n\t\t)\n\treturn conv\n\ndef crop(tensor,target_tensor): # Crop tensor to target tensor size\n\ttarget_shape = target_tensor.shape[2]\n\ttensor_shape = tensor.size()[2]\n\tdelta = (target_shape-tensor_shape)//2\n\treturn tensor[:,:,delta:tensor_shape-delta,delta:tensor_shape-delta]\n\n\n# Hook functions to get values of intermediate layers for cross connection\nhook_values = []\ndef hook(_, input, output):\n\tglobal hook_values\n\thook_values.append(output) # stores values of each layers in hook_values\n\nindices = []\nshapes = []\ndef init_hook(model,device):\n\tglobal shapes, indices, hook_values\n\tshapes = []\n\tindices = []\n\thook_values = []\n\n\tfor i in range(len(model._blocks)):\n\t\tmodel._blocks[i].register_forward_hook(hook) #register hooks\n\t\n\timage = torch.rand([1,3,576,576])\n\timage = image.to(device)\n\tout = model(image) # generate hook values to get shapes\n\t\n\tshape = [i.shape for i in hook_values] # get shape of all layers\n\t\n\tfor i in range(len(shape)-1):\n\t\tif shape[i][2]!=shape[i+1][2]: # get indices of layers only where output dimension change\n\t\t\tindices.append(i)\n\tindices.append(len(shape)-1) # get last layer index\n\t\n\tshapes = [shape[i] for i in indices] # get shapes of required layers\n\tshapes = shapes[::-1]  \n\nencoder_out = []\ndef epoch_hook(model, image):\n\tglobal encoder_out, indices, hook_values\n\thook_values = []\n\n\tout = model(image) # generate layer outputs with current image\n\tencoder_out = [hook_values[i] for i in indices] # get layer outputs for selected indices\n\n\nclass EffUNet(nn.Module):\n\n\tdef __init__(self,model='b0',out_channels=2,dropout=0.1,freeze_backbone=True,pretrained=True,device='cuda'):\n\t\tsuper(EffUNet,self).__init__()\n\t\tglobal layers, shapes\n\n\t\tif model not in set(['b0','b1','b2','b3','b4','b5','b6','b7']):\n\t\t\traise Exception(f'{model} unavailable.')\n\t\tif pretrained:\n\t\t\tself.encoder = EfficientNet.from_pretrained(f'efficientnet-{model}')\n\t\telse:\n\t\t\tself.encoder = EfficientNet.from_name(f'efficientnet-{model}')\n\n\t\t# Disable non required layers by replacing them with identity to save time and memory\n\t\tself.encoder._conv_head=torch.nn.Identity()\n\t\tself.encoder._bn1=torch.nn.Identity()\n\t\tself.encoder._avg_pooling=torch.nn.Identity()\n\t\tself.encoder._dropout=torch.nn.Identity()\n\t\tself.encoder._fc=torch.nn.Identity()\n\t\tself.encoder._swish=torch.nn.Identity()\n\n\t\tif isinstance(device, str):\n\t\t\tself.device = torch.device(device)\n\t\telse:\n\t\t\tself.device = device\n\t\tself.encoder.to(self.device)\n\t\tself.encoder._conv_stem.stride=1 # can't replace this layer with identity, so modify\n\t\tself.encoder._conv_stem.kernel_size=(1,1) # such that it doesn't affect the output shape\n\n\t\t# freeze encoder\n\t\tif freeze_backbone:\n\t\t\tfor param in self.encoder.parameters():\n\t\t\t\tparam.requires_grad = False\n\n\t\t# register hooks & get shapes\n\t\tinit_hook(self.encoder,self.device)\n\n\t\t# Building decoder\n\t\tself.decoder = torch.nn.modules.container.ModuleList()\n\t\tfor i in range(len(shapes)-1):\n\t\t\tself.decoder.append(torch.nn.modules.container.ModuleList())\n\t\t\tself.decoder[i].append(nn.ConvTranspose2d(shapes[i][1],shapes[i][1]-shapes[i+1][1],kernel_size=2,stride=2).to(self.device))\n\t\t\tself.decoder[i].append(double_conv(shapes[i][1],shapes[i+1][1],dropout).to(self.device))\n\n\t\t#output layer\n\t\tself.out = nn.Conv2d(shapes[-1][1],out_channels,kernel_size=1).to(self.device)\n\n\tdef forward(self, image):\n\t\tglobal layers\n\n\t\t# Encoder\n\t\tepoch_hook(self.encoder, image) # required outputs accumulate in \"encoder_out\"\n\n\t\t#Decoder\n\t\tx = encoder_out.pop()\n\t\tfor i in range(len(self.decoder)):\n\t\t\tx = self.decoder[i][0](x) # conv transpose\n\t\t\tprev = encoder_out.pop()\n\t\t\tprev = crop(prev,x) # croping for cross connection\n\t\t\tprev = torch.cat([x,prev],axis=1) # concatenating \n\t\t\tx = self.decoder[i][1](prev) # double conv\n\t\t\n\t\t#out\n\t\tx = self.out(x)\n\t\treturn x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define training and validation data loaders\nBATCH_SIZE = 2\nloader = D.DataLoader(\n    train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=8)\n\nvloader = D.DataLoader(\n    valid_ds, batch_size=BATCH_SIZE, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef validation(model, loader, loss_fn):\n    losses = []\n    model.eval()\n    for image, target in loader:\n        image, target = image.to(DEVICE), target.float().to(DEVICE)\n        output = model(image)\n        loss = loss_fn(output, target)\n        losses.append(loss.item())\n        \n    return np.array(losses).mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# del model\n# gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = EffUNet('b5',out_channels=1,dropout=0.15,freeze_backbone=False,pretrained=True,device=DEVICE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model.load_state_dict(torch.load('../input/effunet-hubmap-weights/model_best_vloss144.pth'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SoftDiceLoss(nn.Module):\n    def __init__(self, smooth=1., dims=(-2,-1)):\n\n        super(SoftDiceLoss, self).__init__()\n        self.smooth = smooth\n        self.dims = dims\n    \n    def forward(self, x, y):\n\n        tp = (x * y).sum(self.dims)\n        fp = (x * (1 - y)).sum(self.dims)\n        fn = ((1 - x) * y).sum(self.dims)\n        \n        dc = (2 * tp + self.smooth) / (2 * tp + fp + fn + self.smooth)\n        dc = dc.mean()\n\n        return 1 - dc\n    \nbce_fn = nn.BCEWithLogitsLoss()\ndice_fn = SoftDiceLoss()\n\ndef loss_fn(y_pred, y_true):\n    dice = dice_fn(y_pred.sigmoid(), y_true)\n    y_true = y_true.type_as(y_pred)\n    bce = bce_fn(y_pred, y_true)\n#     return 0.3*bce + 0.7*dice\n    return 0.5*bce + 0.5*dice","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Ran this Once and found the learning rate around 3e-4\n# from torch_lr_finder import LRFinder\n\n# optimizer = torch.optim.Adam(model.parameters(),lr=1e-6, weight_decay=1e-3)\n# lr_finder = LRFinder(model, optimizer, loss_fn, device=DEVICE)\n# lr_finder.range_test(loader, end_lr=1e-3, num_iter=100)\n# lr_finder.plot() # to inspect the loss-learning rate graph\n# lr_finder.reset() # to reset the model and optimizer to their initial state","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Table for results\nheader = r'''\n        Train | Valid\nEpoch |  Loss |  Loss | Time (s)\n'''\n#          Epoch         metrics            time\nraw_line = '{:6d}' + '\\u2502{:7.3f}'*2 + '\\u2502{:6.2f}'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learning_rate = 3e-4\noptimizer = torch.optim.AdamW(model.parameters(),lr=learning_rate, weight_decay=1e-3)\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=3,\n                                                       verbose=True,min_lr=3e-7)\nearly_stop_count = 0\naccumulation_steps = 32\n\nprint(header)\n\nbest_loss = np.inf\n# best_loss = 0.144\nEPOCHES = 50\nfor epoch in range(1, EPOCHES+1):\n    losses = []\n    start_time = time.time()\n    model.train()\n    for i,(image, target) in enumerate(loader):\n        \n        image, target = image.to(DEVICE), target.float().to(DEVICE)\n        output = model(image)\n        loss = loss_fn(output, target)\n        losses.append(loss.item())\n        loss = loss / accumulation_steps\n        loss.backward()\n        if (i+1) % accumulation_steps == 0:             # Wait for several backward steps\n            optimizer.step()                            # Now do an optimizer step\n            optimizer.zero_grad() \n    vloss = validation(model, vloader, loss_fn)\n    print(raw_line.format(epoch, np.array(losses).mean(), vloss, (time.time()-start_time)))\n    losses = []\n    \n    if vloss < best_loss:\n        best_loss = vloss\n        torch.save(model.state_dict(), 'model_best.pth')\n        print('Loss decreased, model saved.')\n        early_stop_count = 0\n    else:\n        early_stop_count += 1\n    if early_stop_count == 10:\n        print('Early Stopping!')\n        break\n    scheduler.step(vloss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.load_state_dict(torch.load('model_best.pth'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef evaluate(model, loader, loss_fn):\n    dices = []\n    model.eval()\n    for image, target in loader:\n        image, target = image.to(DEVICE), target.float().to(DEVICE)\n        output = model(image)\n        dice = dice_fn(output.sigmoid(), target)\n        dices.append(dice.item())\n        \n    print(f'Mean Dice Score: {1 - np.array(dices).mean():.4f}')\n\nevaluate(model, vloader, loss_fn)","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}