{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch Modelling GPU "},{"metadata":{},"cell_type":"markdown","source":"Models are from **https://github.com/qubvel/segmentation_models.pytorch** \nSome of the lossses are from **https://github.com/JunMa11/SegLoss**\n\nThe base codes are from here -> https://www.kaggle.com/vineeth1999/hubmap-pytorch-efficientunet-offline\nPlease upvote this notebook as well"},{"metadata":{},"cell_type":"markdown","source":"## Things that should be updated\n\n>- parameter settings\n>- loss functions\n>- other optimizers\n>- pth file save(find best pth in a fold)\n>- more models(change encoder)\n>- working without the Internet\n>- inference\n>- etc"},{"metadata":{},"cell_type":"markdown","source":"## Imports"},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.model_selection import GroupKFold\nimport torch\nfrom torch import nn\nimport torchvision\nimport cv2\nimport os\nimport numpy as np\nimport pandas as pd\nfrom torchvision import transforms\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau, CosineAnnealingWarmRestarts\nfrom scipy.ndimage.interpolation import zoom\nimport albumentations as A\nfrom torch.nn import functional as F\nfrom albumentations.pytorch import ToTensorV2\n\nimport matplotlib.pyplot as plt\nimport sys\nimport time\nimport random","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CFG:\n    data = 512 #256\n    debug=False\n    apex=False\n    print_freq=100\n    num_workers=4\n    img_size=224 # appropriate input size for encoder \n    scheduler='CosineAnnealingWarmRestarts' # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts']\n    epoch=5 # Change epochs\n    criterion= 'Lovasz' #'DiceBCELoss' # ['DiceLoss', 'Hausdorff', 'Lovasz']\n    base_model='FPN' # ['Unet']\n    encoder = 'se_resnet50' # ['efficientnet-b5'] or other encoders from smp\n    lr=1e-4\n    min_lr=1e-6\n    batch_size=4\n    weight_decay=1e-6\n    gradient_accumulation_steps=1\n    seed=2021\n    n_fold=5\n    trn_fold=[0, 1, 2, 3, 4]\n    train=True\n    inference=False\n    optimizer = 'Adam'\n    T_0=10\n    # N=5 \n    # M=9\n    T_max=10\n    #factor=0.2\n    #patience=4\n    #eps=1e-6\n    smoothing=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_torch(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    \nseed_torch(seed=CFG.seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"RandAugment - I tried to use it but got some errors when used with albumentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"'''\nsys.path.append('../input/randaug')\nfrom RandAugment import RandAugment\n'''","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"base_transform = A.Compose([\n        A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n        A.HorizontalFlip(),\n        A.VerticalFlip(),\n        A.RandomRotate90(),\n        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=20, p=0.9, \n                         border_mode=cv2.BORDER_REFLECT),\n        A.OneOf([\n            A.OpticalDistortion(p=0.4),\n            A.GridDistortion(p=.1),\n            A.IAAPiecewiseAffine(p=0.4),\n        ], p=0.3),\n        A.OneOf([\n            A.HueSaturationValue(10,15,10),\n            A.CLAHE(clip_limit=3),\n            A.RandomBrightnessContrast(),            \n        ], p=0.5),\n        ToTensorV2()\n    ], p=1.0)\n\n'''\nrand_transform = A.Compose([\n            RandAugment(CFG.N, CFG.M),\n            A.Transpose(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=0, p=0.85),\n            A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n            A.Normalize(),\n            ToTensorV2()\n        ])\n'''\n\nstrong_transform = A.Compose([\n            A.Transpose(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03, p=0.5),\n            A.OneOf([\n                    A.RandomGamma(),\n                    A.GaussNoise()           \n                ], p=0.5),\n            A.OneOf([\n                    A.OpticalDistortion(p=0.4),\n                    A.GridDistortion(p=0.2),\n                    A.IAAPiecewiseAffine(p=0.4),\n                ], p=0.5),\n            A.OneOf([\n                    A.HueSaturationValue(10,15,10),\n                    A.CLAHE(clip_limit=4),\n                    A.RandomBrightnessContrast(),            \n                ], p=0.5),\n    \n            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=0, p=0.85),\n            A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n            ToTensorV2()\n        ])\n\nweak_transform = A.Compose([\n        A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n        A.HorizontalFlip(),\n        A.VerticalFlip(),\n        A.RandomRotate90(),\n        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.9, \n                         border_mode=cv2.BORDER_REFLECT),\n        ToTensorV2()\n    ], p=1.0)\n\nval_transform = A.Compose([\n        A.Resize(CFG.img_size, CFG.img_size, p=1.0),\n        ToTensorV2()\n    ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mean = np.array([0.65459856,0.48386562,0.69428385])\nstd = np.array([0.15167958,0.23584107,0.13146145])\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self,df, train='train', augment='weak', transform=True):\n        ids = df.id.values\n        #kf = KFold(n_splits=nfolds,random_state=SEED,shuffle=True)\n        #ids = set(ids[list(kf.split(ids))[fold][0 if train else 1]])\n        if CFG.data==512:\n            self.fnames = [fname for fname in os.listdir('../input/hubmap-512x512/train') if fname.split('_')[0] in ids]\n        elif CFG.data==256:\n            self.fnames = [fname for fname in os.listdir('../input/hubmap-256x256/train') if fname.split('_')[0] in ids]\n        self.train = train\n        self.augment=augment\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.fnames)\n    \n    def __getitem__(self, idx):\n        fname = self.fnames[idx]\n        if CFG.data==512:\n            img = cv2.cvtColor(cv2.imread(os.path.join('../input/hubmap-512x512/train',fname)), cv2.COLOR_BGR2RGB)\n            mask = cv2.imread(os.path.join('../input/hubmap-512x512/masks',fname),cv2.IMREAD_GRAYSCALE)\n        elif CFG.data==256:\n            img = cv2.cvtColor(cv2.imread(os.path.join('../input/hubmap-256x256/train',fname)), cv2.COLOR_BGR2RGB)\n            mask = cv2.imread(os.path.join('../input/hubmap-256x256/masks',fname),cv2.IMREAD_GRAYSCALE)\n        \n        if self.train == 'train':\n            if self.transform == True:\n                if self.augment == 'base':\n                    augmented = base_transform(image=img,mask=mask)\n                    img,mask = augmented['image'],augmented['mask']\n                elif self.augment == 'weak':\n                    augmented = weak_transform(image=img,mask=mask)\n                    img,mask = augmented['image'],augmented['mask']\n                elif self.augment == 'strong':\n                    augmented = strong_transform(image=img,mask=mask)\n                    img,mask = augmented['image'],augmented['mask']\n                    \n        elif self.train == 'val':\n            transformed = val_transform(image=img,mask=mask)\n            img,mask = transformed['image'],transformed['mask']\n            \n        img = img.type('torch.FloatTensor')\n        img = img/255\n        mask = mask.type('torch.FloatTensor')\n\n        return img, mask\n          ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('../input/hubmap-kidney-segmentation/train.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Visualize augmented data(weak) and mask"},{"metadata":{"trusted":true},"cell_type":"code","source":"# weak augmentation\ntrain_dataset = HuBMAPDataset(train_df, train='train', augment='weak', transform=True)\nplt.figure(figsize=(15,10))\nfor i in range(6):\n    image, mask = train_dataset[i]\n    plt.subplot(3,4,2*i+1)\n    plt.imshow(np.transpose((image), (1,2,0)))\n    plt.subplot(3,4,2*i+2)\n    plt.imshow(mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Visualize augmented data and mask"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = HuBMAPDataset(train_df, train='train', augment='base', transform=True)\nplt.figure(figsize=(15,10))\nfor i in range(6):\n    image, mask = train_dataset[i]\n    plt.subplot(3,4,2*i+1)\n    plt.imshow(np.transpose((image), (1,2,0)))\n    plt.subplot(3,4,2*i+2)\n    plt.imshow(mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Visualize augmented data(strong) and mask"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = HuBMAPDataset(train_df, train='train', augment='strong', transform=True)\nplt.figure(figsize=(15,10))\nfor i in range(6):\n    image, mask = train_dataset[i]\n    plt.subplot(3,4,2*i+1)\n    plt.imshow(np.transpose((image), (1,2,0)))\n    plt.subplot(3,4,2*i+2)\n    plt.imshow(mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## DataLoader"},{"metadata":{"trusted":true},"cell_type":"code","source":"if CFG.data==512:\n    directory_list = os.listdir('../input/hubmap-512x512/train')\nelif CFG.data==256:\n    directory_list = os.listdir('../input/hubmap-256x256/train')\ndirectory_list = [fnames.split('_')[0] for fnames in directory_list]\ndir_df = pd.DataFrame(directory_list, columns=['id'])\ndir_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_train_valid_dataloader(df, fold):\n    train_ids = df[~df.Folds.isin(fold)]\n    val_ids = df[df.Folds.isin(fold)]\n    \n    train_ds = HuBMAPDataset(train_ids, train='train', weak=False, transform=True)\n    val_ds = HuBMAPDataset(val_ids, train='train', weak=False, transform=True)\n    train_loader = DataLoader(train_ds, batch_size=CFG.batch_size, pin_memory=True, shuffle=True, num_workers=CFG.num_workers)\n    val_loader = DataLoader(val_ds, batch_size=CFG.batch_size, pin_memory=True, shuffle=False, num_workers=CFG.num_workers)\n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ../input/segmentationmodelspytorch/segmentation_models/timm-0.1.20-py3-none-any.whl\n!pip install ../input/segmentationmodelspytorch/segmentation_models/segmentation_models_pytorch-0.1.2-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sys.path.append('../input/segmentationmodelspytorch')\nimport segmentation_models_pytorch as smp\nfrom segmentation_models_pytorch import Unet, FPN","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if CFG.base_model =='Unet':\n    base_model = smp.Unet(CFG.encoder, encoder_weights='imagenet', classes=1)\nif CFG.base_model =='FPN':\n    base_model = smp.FPN(CFG.encoder, encoder_weights='imagenet', classes=1)\nprint(base_model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class HuBMAP(nn.Module):\n    def __init__(self):\n        super(HuBMAP, self).__init__()\n        self.cnn_model = base_model\n        \n        #self.cnn_model.decoder.blocks.append(self.cnn_model.decoder.blocks[-1])\n        #self.cnn_model.decoder.blocks[-2] = self.cnn_model.decoder.blocks[-3]\n    \n    def forward(self, imgs):\n        img_segs = self.cnn_model(imgs)\n        return img_segs","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"sys.path.append('../input/segloss/Segmentationloss/SegLoss-master')\nfrom losses_pytorch.hausdorff import HausdorffDTLoss\nfrom losses_pytorch.lovasz_loss import LovaszSoftmax","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=CFG.smoothing):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        \n        return dice\n    \n    \n    \nclass DiceBCELoss(nn.Module):\n    # Formula Given above.\n    def __init__(self, weight=None, size_average=True):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=CFG.smoothing):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).mean()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(inputs.mean() + targets.mean() + smooth)  \n        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')\n        Dice_BCE = BCE + dice_loss\n        \n        return Dice_BCE.mean()\n    \n    \nclass Hausdorff_loss(nn.Module):\n    def __init__(self):\n        super(Hausdorff_loss, self).__init__()\n        \n    def forward(self, inputs, targets):\n        return HausdorffDTLoss()(inputs, targets)\n    \n    \nclass Lovasz_loss(nn.Module):\n    def __init__(self):\n        super(Lovasz_loss, self).__init__()\n        \n    def forward(self, inputs, targets):\n        return LovaszSoftmax()(inputs, targets)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if CFG.criterion == 'DiceBCELoss':\n    criterion = DiceBCELoss()\nelif CFG.criterion == 'DiceLoss':\n    criterion = DiceLoss()\nelif CFG.criterion == 'Hausdorff':\n    criterion = Hausdorff_loss()\nelif CFG.criterion == 'Lovasz':\n    criterion = Lovasz_loss()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def HuBMAPLoss(images, targets, model, device, loss_func=criterion):\n    model.to(device)\n    images = images.to(device)\n    targets = targets.to(device)\n    outputs = model(images)\n    loss_func = loss_func\n    loss = loss_func(outputs, targets)\n    return loss, outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(epoch, model, device, optimizer, scheduler, trainloader):\n    model.train()\n    t = time.time()\n    total_loss = 0\n    for step, (images, targets) in enumerate(trainloader):\n        loss, outputs = HuBMAPLoss(images, targets, model, device)\n        loss.backward()\n        if ((step+1)%4==0 or (step+1)==len(trainloader)):\n            optimizer.step()\n            scheduler.step()\n            optimizer.zero_grad()\n        loss = loss.detach().item()\n        total_loss += loss\n        if ((step+1)%10==0 or (step+1)==len(trainloader)):\n            print(\n                    f'epoch {epoch} train step {step+1}/{len(trainloader)}, ' + \\\n                    f'loss: {total_loss/len(trainloader):.4f}, ' + \\\n                    f'time: {(time.time() - t):.4f}', end= '\\r' if (step + 1) != len(trainloader) else '\\n'\n                )\n\n            \n        \ndef valid_one_epoch(epoch, model, device, optimizer, scheduler, validloader):\n    model.eval()\n    t = time.time()\n    total_loss = 0\n    for step, (images, targets) in enumerate(validloader):\n        loss, outputs = HuBMAPLoss(images, targets, model, device)\n        loss = loss.detach().item()\n        total_loss += loss\n        if ((step+1)%4==0 or (step+1)==len(validloader)):\n            scheduler.step(total_loss/len(validloader))\n        if ((step+1)%10==0 or (step+1)==len(validloader)):\n            print(\n                    f'epoch {epoch} trainz step {step+1}/{len(validloader)}, ' + \\\n                    f'loss: {total_loss/len(validloader):.4f}, ' + \\\n                    f'time: {(time.time() - t):.4f}', end= '\\r' if (step + 1) != len(validloader) else '\\n'\n                )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Creating Folds Column"},{"metadata":{"trusted":true},"cell_type":"code","source":"FOLDS = 5\ngkf = GroupKFold(FOLDS)\ndir_df['Folds'] = 0\nfor fold, (tr_idx, val_idx) in enumerate(gkf.split(dir_df, groups=dir_df[dir_df.columns[0]].values)):\n    dir_df.loc[val_idx, 'Folds'] = fold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dir_df","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The Real Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_train_valid_dataloader(df, fold):\n    train_ids = df[~df.Folds.isin(fold)]\n    val_ids = df[df.Folds.isin(fold)]\n    \n    train_ds = HuBMAPDataset(train_ids, train='train', augment='base', transform=True)\n    val_ds = HuBMAPDataset(val_ids, train='val', augment='base', transform=True)\n    train_loader = DataLoader(train_ds, batch_size=4, pin_memory=True, shuffle=True, num_workers=4)\n    val_loader = DataLoader(val_ds, batch_size=4, pin_memory=True, shuffle=False, num_workers=4)\n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nmodel = HuBMAP().to(device)\n#optimizer\noptimizer = Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay, amsgrad=False)\n\n# scheduler setting\nif CFG.scheduler == 'CosineAnnealingWarmRestarts':\n    scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)\nelif CFG.scheduler == 'ReduceLROnPlateau':\n    scheduler = ReduceLROnPlateauReduceLROnPlateau(optimizer, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)\nelif CFG.scheduler == 'CosineAnnealingLR':\n    scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold, (tr_idx, val_idx) in enumerate(gkf.split(dir_df, groups=dir_df[dir_df.columns[0]].values)):\n    if fold>1:\n        break\n    trainloader, validloader = prepare_train_valid_dataloader(dir_df, [fold])\n\n    num_epochs = CFG.epoch\n    #num_epochs = 2\n    for epoch in range(num_epochs):\n        train_one_epoch(epoch, model, device, optimizer, scheduler, trainloader)\n        with torch.no_grad():\n            valid_one_epoch(epoch, model, device, optimizer, scheduler, validloader)\n    torch.save(model.state_dict(),f'FOLD-{fold}-model.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}