{"cells":[{"metadata":{},"cell_type":"markdown","source":"This is my baseline Unet architecture for the competition segmentation task. I'm currently using [this](https://www.kaggle.com/xhlulu/hubmap-512x512-full-size-tiles) dataset. Please give my notebook an upvote if you find it useful. \n\nI will try to update this notebook continuously. This version has the follwoing features:\n\n- Unet with Attention Gates as skip connections\n- Data Augmentations with Albumentation\n- Resnest as Unet Encoder (The current version support almost all *Res* models though)\n- Gradient Accumulation\n- Mixed Precison"},{"metadata":{},"cell_type":"markdown","source":"### Installing necessary libraries"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"! pip install -q timm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Config\n\nI'll convert this cell to a **class** later. You can change any parameters such as learning rate, batch size, encoder model from here."},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import cv2\nimport albumentations as A\nfrom albumentations.augmentations.transforms import Equalize, Posterize, Downscale\nfrom albumentations import (\n    PadIfNeeded, HorizontalFlip, VerticalFlip, CenterCrop,    \n    RandomCrop, Resize, Crop, Compose, HueSaturationValue,\n    Transpose, RandomRotate90, ElasticTransform, GridDistortion, \n    OpticalDistortion, RandomSizedCrop, Resize, CenterCrop,\n    VerticalFlip, HorizontalFlip, OneOf, CLAHE, Normalize,\n    RandomBrightnessContrast, Cutout, RandomGamma, ShiftScaleRotate ,\n    GaussNoise, Blur, MotionBlur, GaussianBlur, \n)\n\nSEED = 69\nn_epochs = 10\ndevice = 'cuda:0'\ndata_dir = '../input/hubmap-512x512-full-size-tiles'\nloss_thr = 1e6\nimg_path = f'{data_dir}/train'\nlabel_path = f'{data_dir}/masks'\nencoder_model = 'resnest50_fast_1s1x64d'\nmodel_name= 'ResUnest50' # Will come up with a better name later\nmodel_dir = 'model_dir'\nhistory_dir = 'history_dir'\nload_model = False\napply_log = False\nimg_dim = 320\nbatch_size = 24\naccum_step = 2\nlearning_rate = 2.50e-3\nnum_workers = 4\nmixed_precision = True\npatience = 3\ntrain_aug = A.Compose([A.CenterCrop(p=0.3, height=300, width=300),\nA.augmentations.transforms.RandomCrop(280, 280, p=0.3),\nA.augmentations.transforms.Rotate(limit=30, interpolation=1, border_mode=4, value=None, mask_value=None, always_apply=False, p=0.5),\nA.augmentations.transforms.Resize(320, 320, interpolation=1, always_apply=True, p=0.6),\nCutout(num_holes=8, max_h_size=20, max_w_size=20, fill_value=0, always_apply=False, p=0.2),\n# A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, brightness_by_max=True, always_apply=False, p=0.3),\n# A.augmentations.transforms.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=20, always_apply=False, p=0.4),\nOneOf([\n        GaussNoise(var_limit=0.1),\n        Blur(),\n        GaussianBlur(blur_limit=3),\n        # RandomGamma(p=0.7),\n        ], p=0.3),\nA.HorizontalFlip(p=0.3)])\nval_aug = Compose([Normalize(always_apply=True)])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset Class: Borrowed from [here](https://www.kaggle.com/orkatz2/hubmap-res34unet-baseline-train)"},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function, division\nimport numpy as np\nimport os\nimport random\n# os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"-1\"\nimport torch\nimport numpy as np\nimport cv2\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, utils\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self, ids, transforms=None, preprocessing=None):\n        self.ids = ids\n        self.transforms = transforms\n        self.preprocessing = preprocessing\n    def __getitem__(self, idx):\n        name = self.ids[idx]\n        img = cv2.imread(f\"{img_path}/{name}\")\n        img = cv2.resize(img, (img_dim, img_dim))/255.\n        mask = cv2.imread(f\"{label_path}/{name}\")[:,:,0:1]\n        mask = cv2.resize(mask, (img_dim, img_dim), interpolation = cv2.INTER_AREA)\n        if self.transforms:\n            augmented = self.transforms(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n        if self.preprocessing:\n            preprocessed = self.preprocessing(image=img, mask=mask)\n            img = preprocessed['image']\n            mask = preprocessed['mask']\n        return img.reshape(img_dim, img_dim, 3).transpose(2, 0, 1), mask\n\n    def __len__(self):\n        return len(self.ids)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Unet with Resnet\n\nI've created an Unet architecture that supports [Fast-Resnest](https://arxiv.org/abs/2004.08955) as it's encoder. I've found out that, this model is shallow, converges faster and gives better results. I've also replaces the skip connections with [Attention Gates](https://github.com/LeeJunHyun/Image_Segmentation/blob/master/network.py#L108). Attention technique has proven to be greatly useful for image classification and segmentation tasks recently. If you are more interested, please read this paper: [\nAttention U-Net: Learning Where to Look for the Pancreas](https://arxiv.org/abs/1804.03999)."},{"metadata":{},"cell_type":"markdown","source":"### Model Utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\n\nclass DecoderBlock(nn.Module):\n    def __init__(self, in_channels, n_filters, is_deconv=False, scale=True):\n        super().__init__()\n\n        # B, C, H, W -> B, C/4, H, W\n        self.conv1 = nn.Conv2d(in_channels, in_channels // 4, 1)\n        self.norm1 = nn.BatchNorm2d(in_channels // 4)\n        nonlinearity = nn.ReLU\n        self.relu1 = nonlinearity(inplace=True)\n\n        if scale:\n            # B, C/4, H, W -> B, C/4, H, W\n            if is_deconv:\n                self.upscale = nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 3,\n                                                  stride=2, padding=1, output_padding=1)\n            else:\n                self.upscale = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n        else:\n            self.upscale = nn.Conv2d(in_channels // 4, in_channels // 4, 3, padding=1)\n        self.norm2 = nn.BatchNorm2d(in_channels // 4)\n        self.relu2 = nonlinearity(inplace=True)\n\n        # B, C/4, H, W -> B, C, H, W\n        self.conv3 = nn.Conv2d(in_channels // 4, n_filters, 1)\n        self.norm3 = nn.BatchNorm2d(n_filters)\n        self.relu3 = nonlinearity(inplace=True)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.norm1(x)\n        x = self.relu1(x)\n        x = self.upscale(x)\n        x = self.norm2(x)\n        x = self.relu2(x)\n        x = self.conv3(x)\n        x = self.norm3(x)\n        x = self.relu3(x)\n        return x\n\nclass Attention_block(nn.Module):\n    def __init__(self,F_g,F_l,F_int):\n        super(Attention_block,self).__init__()\n        self.W_g = nn.Sequential(\n            nn.Conv2d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(F_int)\n            )\n        \n        self.W_x = nn.Sequential(\n            nn.Conv2d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(F_int)\n        )\n\n        self.psi = nn.Sequential(\n            nn.Conv2d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True),\n            nn.BatchNorm2d(1),\n            nn.Sigmoid()\n        )\n        \n        self.relu = nn.ReLU(inplace=True)\n        \n    def forward(self,g,x):\n        g1 = self.W_g(g)\n        x1 = self.W_x(x)\n        psi = self.relu(g1+x1)\n        psi = self.psi(psi)\n\n        return x*psi\n\nclass conv_block(nn.Module):\n    def __init__(self,ch_in,ch_out):\n        super(conv_block,self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),\n            nn.BatchNorm2d(ch_out),\n            nn.ReLU(inplace=True),\n        )\n    def forward(self,x):\n        x = self.conv(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_model(valid_loss, valid_dice, best_valid_loss, best_valid_dice, best_state, savepath):\n    if valid_loss<best_valid_loss:\n        print(f'Validation loss has decreased from:  {best_valid_loss:.4f} to: {valid_loss:.4f}. Saving checkpoint')\n        torch.save(best_state, savepath+'_loss.pth')\n        best_valid_loss = valid_loss\n    if valid_dice>best_valid_dice:\n        print(f'Validation dice has increased from:  {best_valid_dice:.4f} to: {valid_dice:.4f}. Saving checkpoint')\n        torch.save(best_state, savepath + '_dice.pth')\n        best_valid_dice = valid_dice\n    else:\n        torch.save(best_state, savepath + '_last.pth')\n    return best_valid_loss, best_valid_dice ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\nimport timm\n\nclass Resnest(nn.Module):\n\n    def __init__(self, model_name, out_neurons=600):\n        super().__init__()\n        try:\n            self.backbone = timm.create_model(model_name, pretrained=True)\n        except:\n            self.backbone = torch.hub.load('zhanghang1989/ResNeSt', model_name, pretrained=True)\n            \n        self.in_features = 2048\n        \n    def forward(self, x):\n        x = self.backbone.conv1(x)\n        x = self.backbone.bn1(x)\n        try:\n            x = self.backbone.act1(x)\n        except:\n            x = self.backbone.relu(x)\n        x = self.backbone.maxpool(x)\n\n        layer1 = self.backbone.layer1(x)\n        layer2 = self.backbone.layer2(layer1)\n        layer3 = self.backbone.layer3(layer2)\n\n        return x, layer1, layer2, layer3\n\nclass ResnestDecoder(nn.Module):\n    def __init__(self):\n        super().__init__()\n        nonlinearity = nn.ReLU\n        self.decode1 = DecoderBlock(1024, 512)\n        self.decode2 = DecoderBlock(512, 256)\n        self.decode3 = DecoderBlock(256, 64)\n        self.decode4 = DecoderBlock(64, 32)\n        self.decode5 = DecoderBlock(32, 16)\n        self.conv1 = conv_block(1024, 512)\n        self.conv2 = conv_block(512, 256)\n        self.conv3 = conv_block(128, 64)\n        self.Att1 = Attention_block(512, 512, 256)\n        self.Att2 = Attention_block(256, 256, 64)\n        self.Att3 = Attention_block(64, 64, 32)\n        self.Att4 = Attention_block(64, 64, 32)\n        self.conv4 = nn.Conv2d(64, 64, 3, 2, 1)\n        self.finalconv2 = nn.Conv2d(16, 4, 3, padding=1)\n        self.finalrelu2 = nonlinearity(inplace=True)\n        self.finalconv3 = nn.Conv2d(4, 1, 3, padding=1)\n    \n    def forward(self, x, l1, l2, l3):\n        d1 = self.decode1(l3)\n        l2 = self.Att1(d1, l2)\n        d1 = torch.cat((l2,d1),dim=1)\n        d1 = self.conv1(d1)\n        d2 = self.decode2(d1)\n        l1 = self.Att2(d2, l1)\n        d2 = torch.cat((l1,d2),dim=1)\n        d2 = self.conv2(d2)\n        d3 = self.decode3(d2)\n        d3 = self.conv4(d3)\n        x = self.Att3(d3, x)\n        d3 = torch.cat((x,d3),dim=1)\n        d3 = self.conv3(d3)\n        d4 = self.decode4(d3)\n        d5 = self.decode5(d4)\n        out = self.finalconv2(d5)\n        out = self.finalrelu2(out)\n        out = self.finalconv3(out)\n        return out\n\nclass resUnest(nn.Module):\n    def __init__(self, encoder_model):\n        super().__init__()\n        self.resnest = Resnest(model_name=encoder_model)\n        self.decoder = ResnestDecoder()\n        \n    def forward(self, x):\n        x, l1, l2, l3 = self.resnest(x)\n        out = self.decoder(x, l1, l2, l3)\n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Metric and Loss Function"},{"metadata":{},"cell_type":"markdown","source":"Since competion metric is [Dice Score](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient) and it is differentiable, I decided to use it our loss function."},{"metadata":{"trusted":true},"cell_type":"code","source":"def dice_value(pred, target, threshold = 0.5, smooth = 1e-6):\n    if threshold is not None:\n        pred = (torch.sigmoid(pred) > threshold).float()\n    else:\n        pred = torch.sigmoid(pred)\n    num = pred.size(0)\n    m1 = pred.view(num, -1)  # Flatten\n    m2 = target.view(num, -1)  # Flatten\n    intersection = (m1 * m2).sum(dim=1)\n    dice_score = (2. * intersection + smooth) / (m1.sum(dim=1) + m2.sum(dim=1) + smooth) \n    return dice_score.mean()\n\ndef dice_coeff(pred, target, threshold = None, smooth = 1e-6):\n    pred = torch.sigmoid(pred)\n    if threshold:\n        pred = (pred>threshold).float()\n    num = pred.size(0)\n    m1 = pred.view(num, -1)  # Flatten\n    m2 = target.view(num, -1)  # Flatten\n    intersection = (m1 * m2).sum()\n\n    return (2. * intersection + smooth) / (m1.sum() + m2.sum() + smooth)\n\n\ndef DiceLoss(input, target, apply_log=False):\n    if apply_log:\n        loss = - torch.log(dice_coeff(input, target))\n    else:\n        loss = 1 - dice_coeff(input, target)\n    return loss.mean()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"import logging\nlogging.basicConfig(level=logging.ERROR)\nfrom functools import partial\nfrom collections import Counter\nimport gc\nimport time\nimport pandas as pd\nimport torch\nfrom torch import nn\nfrom torch import optim\n\nm_p = mixed_precision\nif m_p:\n  scaler = torch.cuda.amp.GradScaler() \n\nnp.random.seed(SEED)\n\nfile_ids = os.listdir(img_path)\ntrain_len = int(0.8*len(file_ids))\ntrain_ds = HuBMAPDataset(file_ids[:train_len], train_aug)\ntrain_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n\nval_ds = HuBMAPDataset(file_ids[train_len:])\nvalid_loader = torch.utils.data.DataLoader(\nval_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n\nos.makedirs(model_dir, exist_ok=True)\nos.makedirs(history_dir, exist_ok=True)\n\nresult = pd.DataFrame(columns=['name', 'prediction', 'label', 'difference'])\nif os.path.exists(f'{history_dir}/history_{model_name}_{img_dim}.csv'):\n    history = pd.read_csv(f'{history_dir}/history_{model_name}_{img_dim}.csv')\nelse:\n    history = pd.DataFrame(columns=['train_loss','train_time','val_loss','val_dice', 'val_time'])\n\nmodel = resUnest(encoder_model=encoder_model).to(device)\ncriterion = partial(DiceLoss, apply_log=apply_log)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Data Visualization "},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.utils import make_grid\nfrom matplotlib import pyplot as plt\nfrom matplotlib.pyplot import figure\nimgs, masks = iter(valid_loader).next()\ngrid_imgs = make_grid(imgs, nrow=5)\nfigure(num=None, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')\nplt.imshow(np.transpose(grid_imgs.numpy(), (1, 2, 0)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# grid_imgs = make_grid(masks, nrow=5)\n# print(grid_imgs.size())\n# figure(num=None, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')\n# plt.imshow(np.transpose(grid_imgs.numpy(), (0, 1, 2)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I'm going to train a model for 10 epochs and save models at the end of each epoch, the models that have the lowest loss and highest dice scores."},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_val(epoch, dataloader, optimizer, pretrained=None, train=True, mode='train', record=True):\n    global m_p\n    global result\n    global batch_size\n    global accum_step\n    t1 = time.time()\n    running_loss = 0\n    epoch_samples = 0\n    dice_scores = 0\n    raw_dice_coeff = 0\n    if pretrained:\n        model.load_state_dict(pretrained)\n    if train:\n        model.train()\n        print(\"Initiating train phase ...\")\n    else:\n        model.eval()\n        print(\"Initiating val phase ...\")\n    for idx, (img, labels) in enumerate(dataloader):\n        with torch.set_grad_enabled(train):\n            img = img.to(device, dtype=torch.float32)\n            labels = labels.to(device, dtype=torch.float32)\n            epoch_samples += len(img)\n            optimizer.zero_grad()\n            with torch.cuda.amp.autocast(m_p):\n                if m_p:\n                    img = img.half()\n                else:\n                    img = img.float()\n                outputs = model(img)\n\n                loss = criterion(outputs, labels).sum()\n                running_loss += loss.item()*len(img)\n                loss = loss/accum_step\n      \n                if train:\n                     if m_p:\n                         scaler.scale(loss).backward()\n                         if (idx+1) % accum_step == 0:\n                             scaler.step(optimizer)\n                             scaler.update() \n                             optimizer.zero_grad()\n                      # cyclic_scheduler.step()\n                     else:\n                         loss.backward()\n                         if (idx+1) % accum_step == 0:\n                             optimizer.step()\n                             optimizer.zero_grad()\n\n        elapsed = int(time.time() - t1)\n        eta = int(elapsed / (idx+1) * (len(dataloader)-(idx+1)))\n        dice_val = dice_value(outputs, labels).cpu().numpy()\n        raw_dice_val = dice_value(outputs, labels, None).data.cpu().numpy()\n        dice_scores += dice_val * len(labels) \n        raw_dice_coeff += raw_dice_val * len(labels)\n        # if mode != 'train':\n            # result = analyzer(idx, name, torch.sigmoid(outputs), labels, loss, result)\n        if train:\n            msg = f\"Epoch: {epoch} Progress: [{idx}/{len(dataloader)}] loss: {(running_loss/epoch_samples):.4f} Time: {elapsed}s ETA: {eta} s\"\n        else:\n            msg = f'Epoch {epoch} Progress: [{idx}/{len(dataloader)}] loss: {(running_loss/epoch_samples):.4f} Time: {elapsed}s ETA: {eta} s'\n        print(msg, end= '\\r')\n    history.loc[epoch, f'{mode}_loss'] = running_loss/epoch_samples\n    history.loc[epoch, f'{mode}_time'] = elapsed\n    if mode=='val' or mode=='test':\n        val_dice = dice_scores/epoch_samples\n        raw_val_dice = raw_dice_coeff/epoch_samples\n        lr_reduce_scheduler.step(running_loss)\n        msg = f'{mode} Loss: {running_loss/epoch_samples:.4f} \\n {mode} Dice Score: {val_dice:.4f} \\n {mode} Raw Dice Score:{raw_val_dice:.4f}'\n        print(msg)\n        history.loc[epoch, f'{mode}_loss'] = running_loss/epoch_samples\n        history.loc[epoch, f'{mode}_dice'] = val_dice\n        history.loc[epoch, f'Raw_{mode}_dice'] = raw_val_dice\n        # NaN check\n        if running_loss/epoch_samples > loss_thr or running_loss!=running_loss:\n            print('\\033[91mMixed Precision\\033[0m rendering nan value. Forcing \\033[91mMixed Precision\\033[0m to be False ...')\n            m_p = False\n            batch_size = batch_size//2\n            accum_step = accum_step*2\n            print('Loading last best model ...')\n            tmp = torch.load(os.path.join(model_dir, model_name+'_loss.pth'))\n            model.load_state_dict(tmp['model'])\n            optimizer.load_state_dict(tmp['optim'])\n            lr_reduce_scheduler.load_state_dict(tmp['scheduler'])\n            # cyclic_scheduler.load_state_dict(tmp['cyclic_scheduler'])\n            del tmp\n            \n        if record:\n            history.to_csv(f'{history_dir}/history_{model_name}_{img_dim}.csv', index=False)\n        else:\n            result.to_csv('Result.csv', index=False)\n            # confusion_matrix_generator(result)\n        return running_loss/epoch_samples, raw_val_dice, val_dice\n\n\nplist = [ \n        {'params': model.resnest.parameters(),  'lr': learning_rate/100},\n        {'params': model.decoder.parameters(),  'lr': learning_rate}\n    ]\noptimizer = optim.Adam(plist, lr=learning_rate)\nlr_reduce_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=patience, verbose=True, threshold=1e-4, threshold_mode='rel', cooldown=0, min_lr=1e-7, eps=1e-08)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def main():\n    prev_epoch_num = 0\n    best_valid_loss = np.inf\n    best_valid_dice = 0.0\n    best_raw_val_dice = 0.0\n\n    if load_model:\n        tmp = torch.load(os.path.join(model_dir, model_name+'_loss.pth'))\n        model.load_state_dict(tmp['model'])\n        optimizer.load_state_dict(tmp['optim'])\n        lr_reduce_scheduler.load_state_dict(tmp['scheduler'])\n        # cyclic_scheduler.load_state_dict(tmp['cyclic_scheduler'])\n        if m_p:\n            try:\n                scaler.load_state_dict(tmp['scaler'])\n            except: pass\n        prev_epoch_num = tmp['epoch']\n        best_valid_loss = tmp['best_loss']\n        best_valid_loss, best_raw_val_dice, best_valid_dice = train_val(prev_epoch_num+1, valid_loader, optimizer=optimizer, train=False, mode='val', record=False)\n        del tmp\n        print('Model Loaded!')\n  \n    for epoch in range(prev_epoch_num, n_epochs):\n        torch.cuda.empty_cache()\n        print(gc.collect())\n        train_val(epoch, train_loader, optimizer=optimizer, train=True, mode='train')\n        valid_loss, raw_valid_dice, valid_dice = train_val(epoch, valid_loader, optimizer=optimizer, train=False, mode='val')\n        print(\"#\"*20)\n        print(f\"Epoch {epoch} Report:\")\n        print(f\"Validation Loss: {valid_loss :.4f} Validation dice: {valid_dice :.4f} Raw Validation Dice {raw_valid_dice  :.4f}\")\n        if m_p:\n            best_state = {'model': model.state_dict(), 'optim': optimizer.state_dict(), 'scheduler':lr_reduce_scheduler.state_dict(), \n            # 'cyclic_scheduler':cyclic_scheduler.state_dict(), \n            'scaler': scaler.state_dict(), 'best_loss':valid_loss, 'best_dice':valid_dice, 'epoch':epoch}\n        else:\n            best_state = {'model': model.state_dict(), 'optim': optimizer.state_dict(), 'scheduler':lr_reduce_scheduler.state_dict(), \n            # 'cyclic_scheduler':cyclic_scheduler.state_dict(), \n            'best_loss':valid_loss, 'best_dice':valid_dice, 'epoch':epoch}\n        best_valid_loss, best_valid_dice = save_model(valid_loss, valid_dice, best_valid_loss, best_valid_dice, best_state, os.path.join(model_dir, model_name))\n        print(\"#\"*20)   \nif __name__== '__main__':\n    main()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Future Tasks\nI intend to add the following features in the upcoming verisons of this notebook:\n\n- Data Preprocessing\n- Stratified Train-Validation Split\n- Performance Visualiztion through **[Wandb](wandb.ai)**"}],"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}