{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This is a training notebook that I use for the [HubMAP competition](https://www.kaggle.com/c/hubmap-kidney-segmentation). You can this notebook to experiment on many other segmentation problem. For convinience, I using Catalyst framework and also segmentation_model_pytorch library and pytorch_toolbelt which is a great library for doing quick development. If you like my notebook, please give an upvote, thank you.","metadata":{}},{"cell_type":"markdown","source":"## Install and import packages","metadata":{}},{"cell_type":"code","source":"!pip install -q catalyst==20.12\n!pip install -q pytorch-toolbelt==0.4.2\n!pip install -q torch-optimizer==0.1.0\n!pip install -q segmentation-models-pytorch==0.1.3\n!pip install -q ttach==0.0.3\n!pip install -q albumentations==0.5.2\n!pip install -q timm","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import segmentation_models_pytorch as smp\nimport torch\nimport torch.nn as nn\nfrom torch.optim.optimizer import Optimizer\nimport torch.nn.functional as F\nimport math\nimport warnings\n\nfrom catalyst.contrib.nn import OneCycleLRWithWarmup\nfrom torch.optim.lr_scheduler import (\n    ExponentialLR,\n    CyclicLR,\n    MultiStepLR,\n    CosineAnnealingLR,\n    CosineAnnealingWarmRestarts,\n    ReduceLROnPlateau\n)\nfrom pytorch_toolbelt.losses import *\nfrom pytorch_toolbelt.utils import image_to_tensor\nfrom pytorch_toolbelt.utils.random import set_manual_seed\nfrom torch.nn import KLDivLoss\nfrom catalyst import utils\nfrom catalyst.contrib.nn import OneCycleLRWithWarmup\nfrom catalyst import dl\nfrom catalyst.contrib.utils.cv import image as cata_image\nfrom catalyst.contrib.callbacks.wandb_logger import WandbLogger\nfrom catalyst.dl import (\n    SupervisedRunner, \n    CriterionCallback, \n    EarlyStoppingCallback, \n    SchedulerCallback, \n    MetricAggregationCallback, \n    IouCallback, \n    DiceCallback\n)\nfrom torch.optim.lr_scheduler import CyclicLR, ReduceLROnPlateau\nfrom torch.utils.data import Dataset, DataLoader\n\nimport albumentations as A\n\nimport os\nimport cv2\nimport json\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import precision_recall_curve, auc, average_precision_score\nfrom pathlib import Path\nfrom tqdm .auto import tqdm\nimport plotly.express as px\nfrom sklearn.model_selection import train_test_split\nfrom pathlib import Path\nfrom collections import OrderedDict\nfrom typing import List,  Optional, Dict\n\n%matplotlib inline","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define data paths","metadata":{}},{"cell_type":"code","source":"MAIN_PATH = Path('../input/hubmap-256x256')\nIMG_PATHS = MAIN_PATH / 'train'\nMASK_PATHS = MAIN_PATH / 'masks'\nLABEL = '../input/hubmap-kidney-segmentation/train.csv'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initial seed","metadata":{}},{"cell_type":"code","source":"SEED = 1999\nset_manual_seed(SEED)   \nutils.set_global_seed(SEED)\nutils.prepare_cudnn(deterministic=False, benchmark=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define config class ","metadata":{}},{"cell_type":"code","source":"class BaseConfig:\n    __basedir__ = MAIN_PATH\n    train_img_path = IMG_PATHS\n    train_mask_path = MASK_PATHS\n    \n    #Data config\n    augmentation = 'medium' #options: normal, easy, medium, advanced\n    scale_size = 256\n    \n    #Train config\n    num_epochs = 10\n    batch_size = 32\n    val_batch_size = 32\n    learning_rate = 1e-5\n    learning_rate_decode = 1e-3\n    weight_decay = 2.5e-5\n    is_fp16 = True\n    \n    #Model config\n    model_name = None\n    model_params = None\n    \n    #Metric config\n    metric = \"dice\"\n    mode = \"max\"\n    \n    #Optimize config\n    criterion = {\"bce\": 0.8, 'log_dice': 0.2}\n    pos_weights = [200]\n    optimizer = \"madgrad\"\n    scheduler = \"simple\"\n\n    resume_path = None #Resume training\n    \n    @classmethod\n    def get_all_attributes(cls):\n        d = {}\n        attributes = dict(cls.__dict__)\n\n        for k, v in attributes.items():\n            if not k.startswith('__') and k != 'get_all_attributes':\n                d[k] = v\n                \n        return d","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-05-27T22:34:51.751984Z","iopub.execute_input":"2021-05-27T22:34:51.752340Z","iopub.status.idle":"2021-05-27T22:34:51.817667Z","shell.execute_reply.started":"2021-05-27T22:34:51.752308Z","shell.execute_reply":"2021-05-27T22:34:51.816310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_config = BaseConfig.get_all_attributes()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"class HUBMAPSegmentation(Dataset):\n    def __init__(self, images: List[Path], masks: List[Path] = None, transform=None, preprocessing_fn=None):\n        self.images = images\n        self.masks = masks\n        self.transform = transform\n        self.preprocessing_fn = preprocessing_fn\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index: int) -> dict:\n        image_path = self.images[index]\n\n        result = OrderedDict()\n        image = cata_image.imread(image_path)\n        result['image'] = image\n        \n        if self.masks is not None:\n            mask = Image.open(self.masks[index]).convert('L')\n            mask = mask.point(lambda x: 255 if x > 0 else 0, '1')\n            mask = np.asarray(mask).astype(np.uint8)\n            result['mask'] = mask\n\n        if self.transform is not None:\n            transformed = self.transform(**result)\n            image = transformed['image']\n            if self.masks is not None:\n                mask = transformed['mask']\n                mask = image_to_tensor(mask, dummy_channels_dim=True).float()\n                result['mask'] = mask\n                \n        if self.preprocessing_fn is not None:\n            image = self.preprocessing_fn(image = image)['image']\n            \n        image = image_to_tensor(image).float()           \n        result['image'] = image\n        result['filename'] = image_path.name\n\n        return result","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get preprocessing function","metadata":{}},{"cell_type":"code","source":"def get_preprocessing_fn():\n    mean = [0.65459856,0.48386562,0.69428385],\n    std = [0.15167958,0.23584107,0.13146145]\n\n    def preprocessing(x, mean=mean, std=std, **kwargs):\n        x = x / 255.0\n        if mean is not None:\n            mean = np.array(mean)\n            x = x - mean\n\n        if std is not None:\n            std = np.array(std)\n            x = x / std\n        return x\n\n    return preprocessing, mean, std\n\npreprocessing_fn, mean, std = get_preprocessing_fn()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get data augmentation","metadata":{}},{"cell_type":"code","source":"class BaseTransform(object):\n    def __init__(self, image_size: int = 1024, preprocessing_fn=None):\n        self.image_size = image_size\n        self.preprocessing_fn = preprocessing_fn\n\n    def pre_transform(self):\n        raise NotImplementedError()\n\n    def hard_transform(self):\n        raise NotImplementedError()\n\n    def resize_transforms(self):\n        raise NotImplementedError()\n\n    def _get_compose(self, transform):\n        result = A.Compose([\n            item for sublist in transform for item in sublist\n        ])\n        return result\n\n    def train_transform(self):\n        return self._get_compose([\n            self.resize_transforms(),\n            self.hard_transform()\n        ])\n\n    def validation_transform(self):\n        return self._get_compose([\n            self.pre_transform()\n        ])\n\n    def test_transform(self):\n        return self.validation_transform()\n\n    def get_preprocessing(self):\n        return A.Compose([\n            A.Lambda(image=self.preprocessing_fn)\n        ])\n    \nclass NormalTransform(BaseTransform):\n    def __init__(self, *args, **kwargs):\n        super(NormalTransform, self).__init__(*args, **kwargs)\n\n    def hard_transform(self):\n        return [\n            A.VerticalFlip(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.RandomRotate90(p=0.7),\n        ]\n    \n    def resize_transforms(self):\n        return [\n            A.LongestMaxSize(self.image_size),\n            A.PadIfNeeded(min_height=self.image_size, min_width=self.image_size,\n                          border_mode=cv2.BORDER_CONSTANT, value=0)\n        ]\n    \n    def pre_transform(self):\n        return self.resize_transforms()\n            \nclass MediumTransform(NormalTransform):\n    def __init__(self, *args, **kwargs):\n        super(MediumTransform, self).__init__(*args, **kwargs)\n\n    def hard_transform(self):\n        return [\n            A.VerticalFlip(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.RandomRotate90(p=0.7),\n            A.OneOf([\n                A.ElasticTransform(alpha=120, sigma=120 * 0.05,\n                                   alpha_affine=120 * 0.03, p=0.5),\n                A.GridDistortion(p=0.5),\n                A.OpticalDistortion(distort_limit=2, shift_limit=0.5, p=0.5)\n            ], p=0.5),\n            A.CLAHE(p=0.5),\n            A.RandomBrightnessContrast(p=0.5),\n            A.RandomGamma(p=0.5)\n        ]\n    \n\ndef get_transform(name):\n    if name == 'normal':\n        return NormalTransform\n    if name == 'easy':\n        return EasyTransform\n    if name == 'medium':\n        return MediumTransform\n    if name == 'advanced':\n        return AdvancedTransform","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Showing some examples from dataset","metadata":{}},{"cell_type":"code","source":"#example of train images with masks\nnormal = get_transform('medium')(train_config['scale_size'], preprocessing_fn)\ntransform = normal.train_transform()\npreprocessing = normal.get_preprocessing()\nds = HUBMAPSegmentation(\n        sorted(train_config['train_img_path'].glob('*.*')), \n        sorted(train_config['train_mask_path'].glob('*.*')), \n        transform,    \n        preprocessing\n)\n\ndl = DataLoader(ds,batch_size=64,shuffle=False,num_workers=2)\nbatch_dict = next(iter(dl))\nimgs = batch_dict['image']\nmasks = batch_dict['mask']\n\nplt.figure(figsize=(16,16))\nfor i,(img,mask) in enumerate(zip(imgs,masks)):\n    img = ((img.permute(1,2,0)*np.array(std) + np.array(mean))*255.0).numpy().astype(np.uint8)\n    plt.subplot(8,8,i+1)\n    plt.imshow(img,vmin=0,vmax=255)\n    plt.imshow(mask.squeeze().numpy(), alpha=0.2)\n    plt.axis('off')\n    plt.subplots_adjust(wspace=None, hspace=None)\n    \ndel ds,dl,imgs,masks","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get data loader","metadata":{}},{"cell_type":"code","source":"def get_loader(\n    images: List[Path],\n    random_state: int,\n    valid_size: float = 0.2,\n    batch_size: int = 4,\n    val_batch_size: int = 8,\n    num_workers: int = 4,\n    train_transforms_fn=None,\n    valid_transforms_fn=None,\n    preprocessing_fn=None,\n    masks: List[Path] = None,\n):    \n    indices = np.arange(len(images))\n\n    train_indices, valid_indices = train_test_split(\n        indices, test_size=valid_size, random_state=random_state, shuffle=True)\n\n    np_images = np.array(images)\n    train_images = np_images[train_indices].tolist()\n    val_images = np_images[valid_indices].tolist()\n    np_masks = np.array(masks)\n    train_masks = np_masks[train_indices].tolist()\n    val_masks = np_masks[valid_indices].tolist()\n\n    train_dataset = HUBMAPSegmentation(\n        sorted(train_images),\n        masks=sorted(train_masks),\n        transform=train_transforms_fn,\n        preprocessing_fn=preprocessing_fn,\n    )\n\n    valid_dataset = HUBMAPSegmentation(\n        sorted(val_images),\n        masks=sorted(val_masks),\n        transform=valid_transforms_fn,\n        preprocessing_fn=preprocessing_fn,\n    )\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        shuffle=True,\n        pin_memory=True,\n        drop_last=True\n    )\n\n    valid_loader = DataLoader(\n        valid_dataset,\n        batch_size=val_batch_size,\n        num_workers=num_workers,\n        shuffle=True,\n        pin_memory=True,\n        drop_last=True\n    )\n\n    loaders = OrderedDict()\n    loaders['train'] = train_loader\n    loaders['valid'] = valid_loader\n\n    return loaders","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get model ","metadata":{}},{"cell_type":"code","source":"def get_model(params, model_name):\n    \n    # Model return logit values\n    model = getattr(smp, model_name)(\n        **params\n    )\n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import TYPE_CHECKING, Any, Callable, Optional\n\nimport torch\nimport torch.optim\n\nif TYPE_CHECKING:\n    from torch.optim.optimizer import _params_t\nelse:\n    _params_t = Any\n\nclass MADGRAD(torch.optim.Optimizer):\n    \"\"\"\n    MADGRAD_: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic \n    Optimization.\n\n    .. _MADGRAD: https://arxiv.org/abs/2101.11075\n\n    MADGRAD is a general purpose optimizer that can be used in place of SGD or\n    Adam may converge faster and generalize better. Currently GPU-only.\n    Typically, the same learning rate schedule that is used for SGD or Adam may\n    be used. The overall learning rate is not comparable to either method and\n    should be determined by a hyper-parameter sweep.\n\n    MADGRAD requires less weight decay than other methods, often as little as\n    zero. Momentum values used for SGD or Adam's beta1 should work here also.\n\n    On sparse problems both weight_decay and momentum should be set to 0.\n\n    Arguments:\n        params (iterable): \n            Iterable of parameters to optimize or dicts defining parameter groups.\n        lr (float): \n            Learning rate (default: 1e-2).\n        momentum (float): \n            Momentum value in  the range [0,1) (default: 0.9).\n        weight_decay (float): \n            Weight decay, i.e. a L2 penalty (default: 0).\n        eps (float): \n            Term added to the denominator outside of the root operation to improve numerical stability. (default: 1e-6).\n    \"\"\"\n\n    def __init__(\n        self, params: _params_t, lr: float = 1e-2, momentum: float = 0.9, weight_decay: float = 0, eps: float = 1e-6,\n    ):\n        if momentum < 0 or momentum >= 1:\n            raise ValueError(f\"Momentum {momentum} must be in the range [0,1]\")\n        if lr <= 0:\n            raise ValueError(f\"Learning rate {lr} must be positive\")\n        if weight_decay < 0:\n            raise ValueError(f\"Weight decay {weight_decay} must be non-negative\")\n        if eps < 0:\n            raise ValueError(f\"Eps must be non-negative\")\n\n        defaults = dict(lr=lr, eps=eps, momentum=momentum, weight_decay=weight_decay)\n        super().__init__(params, defaults)\n\n    @property\n    def supports_memory_efficient_fp16(self) -> bool:\n        return False\n\n    @property\n    def supports_flat_params(self) -> bool:\n        return True\n\n    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:\n        \"\"\"Performs a single optimization step.\n\n        Arguments:\n            closure (callable, optional): A closure that reevaluates the model\n                and returns the loss.\n        \"\"\"\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        # step counter must be stored in state to ensure correct behavior under\n        # optimizer sharding\n        if 'k' not in self.state:\n            self.state['k'] = torch.tensor([0], dtype=torch.long)\n        k = self.state['k'].item()\n\n        for group in self.param_groups:\n            eps = group[\"eps\"]\n            lr = group[\"lr\"] + eps\n            decay = group[\"weight_decay\"]\n            momentum = group[\"momentum\"]\n\n            ck = 1 - momentum\n            lamb = lr * math.pow(k + 1, 0.5)\n\n            for p in group[\"params\"]:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data\n                state = self.state[p]\n\n                if \"grad_sum_sq\" not in state:\n                    state[\"grad_sum_sq\"] = torch.zeros_like(p.data).detach()\n                    state[\"s\"] = torch.zeros_like(p.data).detach()\n                    if momentum != 0:\n                        state[\"x0\"] = torch.clone(p.data).detach()\n\n                if momentum != 0.0 and grad.is_sparse:\n                    raise RuntimeError(\"momentum != 0 is not compatible with sparse gradients\")\n\n                grad_sum_sq = state[\"grad_sum_sq\"]\n                s = state[\"s\"]\n\n                # Apply weight decay\n                if decay != 0:\n                    if grad.is_sparse:\n                        raise RuntimeError(\"weight_decay option is not compatible with sparse gradients\")\n\n                    grad.add_(p.data, alpha=decay)\n\n                if grad.is_sparse:\n                    grad = grad.coalesce()\n                    grad_val = grad._values()\n\n                    p_masked = p.sparse_mask(grad)\n                    grad_sum_sq_masked = grad_sum_sq.sparse_mask(grad)\n                    s_masked = s.sparse_mask(grad)\n\n                    # Compute x_0 from other known quantities\n                    rms_masked_vals = grad_sum_sq_masked._values().pow(1 / 3).add_(eps)\n                    x0_masked_vals = p_masked._values().addcdiv(s_masked._values(), rms_masked_vals, value=1)\n\n                    # Dense + sparse op\n                    grad_sq = grad * grad\n                    grad_sum_sq.add_(grad_sq, alpha=lamb)\n                    grad_sum_sq_masked.add_(grad_sq, alpha=lamb)\n\n                    rms_masked_vals = grad_sum_sq_masked._values().pow_(1 / 3).add_(eps)\n\n                    s.add_(grad, alpha=lamb)\n                    s_masked._values().add_(grad_val, alpha=lamb)\n\n                    # update masked copy of p\n                    p_kp1_masked_vals = x0_masked_vals.addcdiv(s_masked._values(), rms_masked_vals, value=-1)\n                    # Copy updated masked p to dense p using an add operation\n                    p_masked._values().add_(p_kp1_masked_vals, alpha=-1)\n                    p.data.add_(p_masked, alpha=-1)\n                else:\n                    if momentum == 0:\n                        # Compute x_0 from other known quantities\n                        rms = grad_sum_sq.pow(1 / 3).add_(eps)\n                        x0 = p.data.addcdiv(s, rms, value=1)\n                    else:\n                        x0 = state[\"x0\"]\n\n                    # Accumulate second moments\n                    grad_sum_sq.addcmul_(grad, grad, value=lamb)\n                    rms = grad_sum_sq.pow(1 / 3).add_(eps)\n\n                    # Update s\n                    s.data.add_(grad, alpha=lamb)\n\n                    # Step\n                    if momentum == 0:\n                        p.data.copy_(x0.addcdiv(s, rms, value=-1))\n                    else:\n                        z = x0.addcdiv(s, rms, value=-1)\n\n                        # p is a moving average of z\n                        p.data.mul_(1 - ck).add_(z, alpha=ck)\n\n\n        self.state['k'] += 1\n        return loss","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Many optimizers to try","metadata":{}},{"cell_type":"code","source":"def get_optimizer(\n    optimizer_name: str, parameters, learning_rate: float, weight_decay=1e-5, eps=1e-5, **kwargs\n) -> Optimizer:\n    from torch.optim import SGD, Adam, RMSprop, AdamW\n    from torch_optimizer import RAdam, Lamb, DiffGrad, NovoGrad, Ranger, Lookahead\n    \n    lookahead = False\n    if len(optimizer_name.split('_')) > 1:\n        optimizer_name = optimizer_name.split('_')[0]\n        lookahead=True\n\n    if optimizer_name.lower() == \"sgd\":\n        base_optim =  SGD(parameters, learning_rate, momentum=0.9, nesterov=True, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"adam\":\n        # As Jeremy suggests\n        base_optim = Adam(parameters, learning_rate, weight_decay=weight_decay, eps=eps, **kwargs)\n\n    if optimizer_name.lower() == \"rms\":\n        base_optim = RMSprop(parameters, learning_rate, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"adamw\":\n        base_optim =  AdamW(parameters, learning_rate, weight_decay=weight_decay, eps=eps, **kwargs)\n\n    if optimizer_name.lower() == \"radam\":\n        # As Jeremy suggests\n        base_optim =  RAdam(parameters, learning_rate, weight_decay=weight_decay, eps=eps, **kwargs)\n\n    # Optimizers from torch-optimizer\n    if optimizer_name.lower() == \"ranger\":\n        base_optim = Ranger(parameters, learning_rate, eps=eps, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"lamb\":\n        base_optim =  Lamb(parameters, learning_rate, eps=eps, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"diffgrad\":\n        base_optim = DiffGrad(parameters, learning_rate, eps=eps, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"novograd\":\n        base_optim = NovoGrad(parameters, learning_rate, eps=eps, weight_decay=weight_decay, **kwargs)\n\n    if optimizer_name.lower() == \"madgrad\":\n        base_optim = MADGRAD(parameters, learning_rate, eps=eps, weight_decay=weight_decay, **kwargs)\n    else:\n        raise ValueError(\"Unsupported optimizer name \" + optimizer_name)\n    \n    if lookahead:\n        return Lookahead(base_optim)\n    return base_optim","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Many losses to try ","metadata":{}},{"cell_type":"code","source":"class WeightedBCEWithLogits(nn.Module):\n    def __init__(self, pos_weights, ignore_index: Optional[int] = -100, reduction=\"mean\"):\n        super().__init__()\n        self.ignore_index = ignore_index\n        self.reduction = reduction\n        self.pos_weights = pos_weights\n\n    def forward(self, label_input: torch.Tensor, target: torch.Tensor):\n\n        if self.ignore_index is not None:\n            not_ignored_mask = (target != self.ignore_index).float()\n\n        loss = nn.BCEWithLogitsLoss(reduce=None, pos_weight=self.pos_weights)(label_input, target)\n\n        if self.ignore_index is not None:\n            loss = loss * not_ignored_mask.float()\n\n        if self.reduction == \"mean\":\n            loss = loss.mean()\n\n        if self.reduction == \"sum\":\n            loss = loss.sum()\n\n        return loss\n\nclass KLDivLossWithLogits(KLDivLoss):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, input, target):\n\n        # Resize target to size of input\n        target = F.interpolate(target, size=input.size()[2:], mode=\"bilinear\", align_corners=False)\n\n        input = torch.cat([input, 1 - input], dim=1)\n        log_p = F.logsigmoid(input)\n\n        target = torch.cat([target, 1 - target], dim=1)\n\n        loss = F.kl_div(log_p, target, reduction=\"mean\")\n        return loss\n\nclass SymmetricLovasz(nn.Module):\n    def __init__(self):\n        super().__init__()\n    def forward(self, outputs, targets):\n        return 0.5*(lovasz_hinge(outputs, targets) + lovasz_hinge(-outputs, 1.0 - targets))\n\ndef get_loss(loss_name: str, ignore_index=None):\n    if loss_name.lower() == \"kl\":\n        return KLDivLossWithLogits()\n\n    if loss_name.lower() == \"bce\":\n        return SoftBCEWithLogitsLoss(ignore_index=ignore_index)\n\n    if loss_name.lower() == 'wbce':\n        return WeightedBCEWithLogits(ignore_index=ignore_index)\n\n    if loss_name.lower() == \"ce\":\n        return nn.CrossEntropyLoss()\n\n    if loss_name.lower() == \"soft_bce\":\n        return SoftBCEWithLogitsLoss(smooth_factor=0.1, ignore_index=ignore_index)\n\n    if loss_name.lower() == \"focal\":\n        return BinaryFocalLoss(alpha=None, gamma=1.5, ignore_index=ignore_index)\n\n    if loss_name.lower() == \"jaccard\":\n        assert ignore_index is None\n        return JaccardLoss(mode=\"binary\")\n\n    if loss_name.lower() == \"log_jaccard\":\n        assert ignore_index is None\n        return JaccardLoss(mode=\"binary\", log_loss=True)\n\n    if loss_name.lower() == \"dice\":\n        assert ignore_index is None\n        return DiceLoss(mode=\"binary\", log_loss=False)\n\n    if loss_name.lower() == \"log_dice\":\n        assert ignore_index is None\n        return DiceLoss(mode=\"binary\", log_loss=True)\n\n    raise KeyError(loss_name)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Many lr scheduler to try","metadata":{}},{"cell_type":"code","source":"class CosineAnnealingWarmRestartsWithDecay(CosineAnnealingWarmRestarts):\n    def __init__(self, optimizer, T_0, T_mult=1, eta_min=0, last_epoch=-1, gamma=0.9):\n        super().__init__(optimizer, T_0, T_mult, eta_min, last_epoch)\n        self.gamma = gamma\n\n    def get_lr(self):\n        if not self._get_lr_called_within_step:\n            warnings.warn(\n                \"To get the last learning rate computed by the scheduler, \" \"please use `get_last_lr()`.\",\n                DeprecationWarning,\n            )\n\n        return [\n            self.eta_min\n            + (base_lr * self.gamma ** self.last_epoch - self.eta_min)\n            * (1 + math.cos(math.pi * self.T_cur / self.T_i))\n            / 2\n            for base_lr in self.base_lrs\n        ]\n\ndef get_scheduler(scheduler_name: str, optimizer, lr, num_epochs, batches_in_epoch=None, mode=None):\n    if scheduler_name is None or scheduler_name.lower() == \"none\":\n        return None\n\n    if scheduler_name.lower() == \"reduce\":\n        return ReduceLROnPlateau(optimizer, mode=mode, patience=5)\n\n    if scheduler_name.lower() == \"cos\":\n        return CosineAnnealingLR(optimizer, num_epochs, eta_min=1e-6)\n\n    if scheduler_name.lower() == \"cos2\":\n        return CosineAnnealingLR(optimizer, num_epochs, eta_min=float(lr * 0.5))\n\n    if scheduler_name.lower() == \"cosr\":\n        return CosineAnnealingWarmRestarts(optimizer, T_0=max(2, num_epochs // 4), eta_min=1e-6)\n\n    if scheduler_name.lower() == \"cosrd\":\n        return CosineAnnealingWarmRestartsWithDecay(optimizer, T_0=max(2, num_epochs // 6), gamma=0.96, eta_min=1e-6)\n\n    if scheduler_name.lower() in {\"1cycle\", \"one_cycle\"}:\n        return OneCycleLRWithWarmup(\n            optimizer,\n            lr_range=(lr, 1e-6),\n            num_steps=batches_in_epoch * num_epochs,\n            warmup_fraction=0.05,\n            decay_fraction=0.1,\n        )\n\n    if scheduler_name.lower() == \"exp\":\n        return ExponentialLR(optimizer, gamma=0.95)\n\n    if scheduler_name.lower() == \"clr\":\n        return CyclicLR(\n            optimizer,\n            base_lr=1e-6,\n            max_lr=lr,\n            step_size_up=batches_in_epoch // 4,\n            # mode='exp_range',\n            cycle_momentum=True,\n            gamma=0.99,\n        )\n\n    if scheduler_name.lower() == \"multistep\":\n        return MultiStepLR(\n            optimizer, milestones=[int(num_epochs * 0.5), int(num_epochs * 0.7), int(num_epochs * 0.9)], gamma=0.3\n        )\n\n    if scheduler_name.lower() == \"simple\":\n        return MultiStepLR(optimizer, milestones=[int(num_epochs * 0.4), int(num_epochs * 0.7)], gamma=0.1)\n\n    raise KeyError(scheduler_name)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare everything we need for the experiment","metadata":{}},{"cell_type":"code","source":"def prepare_everything(exp_name):\n    print(\"===> Get model\")\n    model = get_model(\n        train_config['model_params'],\n        train_config['model_name']\n    )\n\n    print(\"===> Get transformation\")\n    #Define transform (augemntation)\n    Transform = get_transform(train_config['augmentation'])\n    transforms = Transform(\n        train_config['scale_size'],\n        preprocessing_fn=preprocessing_fn\n    )\n\n    train_transform = transforms.train_transform()\n    val_transform = transforms.validation_transform()\n    preprocessing = transforms.get_preprocessing()\n    \n    print(\"===> Get data loader\")\n    loader = get_loader(\n        images = sorted(train_config['train_img_path'].glob('*.*')),\n        masks = sorted(train_config['train_mask_path'].glob('*.*')),\n        random_state = SEED,\n        valid_size = 0.2,\n        batch_size = train_config['batch_size'],\n        val_batch_size = train_config['val_batch_size'],\n        num_workers = 2,\n        train_transforms_fn=train_transform,\n        valid_transforms_fn=val_transform,\n        preprocessing_fn=preprocessing\n    )\n\n    print(\"===> Get optimizer\")\n    param_group = []\n    if hasattr(model, 'encoder'):\n        encoder_params = filter(lambda p: p.requires_grad, model.encoder.parameters())\n        param_group += [{'params': encoder_params, 'lr': train_config['learning_rate']}]        \n    if hasattr(model, 'decoder'):\n        decoder_params = filter(lambda p: p.requires_grad, model.decoder.parameters())\n        param_group += [{'params': decoder_params}]\n    if hasattr(model, 'segmentation_head'):\n        head_params = filter(lambda p: p.requires_grad, model.segmentation_head.parameters())\n        param_group += [{'params': head_params}]        \n    if len(param_group) == 0:\n        param_group = list(model.parameters())\n        \n    total = int(sum(p.numel() for p in model.parameters()))\n    trainable = int(sum(p.numel() for p in model.parameters() if p.requires_grad))\n    count_parameters = {\"total\": total, \"trainable\": trainable}\n\n    print(\n        f'[INFO] total and trainable parameters in the model {count_parameters}'\n    )\n    \n    #Set optimizer\n    optimizer = get_optimizer(\n        train_config['optimizer'], param_group, train_config['learning_rate_decode'], train_config['weight_decay'])\n    \n    print(\"===> Get shceduler\")\n    scheduler = get_scheduler(\n        train_config['scheduler'], optimizer, train_config['learning_rate'], train_config['num_epochs'],\n        batches_in_epoch=len(loader['train']), mode=train_config['mode']\n    )\n        \n    print(\"===> Get loss\")\n    criterion = {}\n    for loss_name in train_config['criterion']:\n        if loss_name == 'wbce':\n            pos_weights = torch.tensor(train_config['pos_weights'], device=utils.get_device())\n            loss_fn = WeightedBCEWithLogits(pos_weights=pos_weights)\n        else:\n            loss_fn = get_loss(loss_name)\n        criterion[loss_name] = loss_fn\n        \n    print(\"===> Get callbacks\")\n    #Define callbacks\n    callbacks = []\n    losses = []\n    for loss_name, loss_weight in train_config['criterion'].items():\n        criterion_callback = CriterionCallback(\n            input_key=\"mask\",\n            output_key=\"logits\",\n            criterion_key=loss_name,\n            prefix=\"loss_\"+loss_name,\n            multiplier=float(loss_weight)\n        )\n\n        callbacks.append(criterion_callback)\n        losses.append(criterion_callback.prefix)\n\n    callbacks += [MetricAggregationCallback(\n        prefix=\"loss\",\n        mode=\"sum\",\n        metrics=losses\n    )]\n\n    if isinstance(scheduler, (CyclicLR, OneCycleLRWithWarmup)):\n        callbacks += [SchedulerCallback(mode=\"batch\")]\n    elif isinstance(scheduler, (ReduceLROnPlateau)):\n        callbacks += [SchedulerCallback(reduced_metric=train_config['metric'])]\n        \n    early_stopping = EarlyStoppingCallback(\n        patience=10, metric=train_config['metric'], minimize=False)\n\n    iou_scores = IouCallback(\n        input_key=\"mask\",\n        activation=\"Sigmoid\",\n        threshold=0.5\n    )\n\n    dice_scores = DiceCallback(\n        input_key=\"mask\",\n        activation=\"Sigmoid\",\n        threshold=0.5\n    )\n    \n    prefix = exp_name\n    log_dir = os.path.join(\"./\", prefix)\n    if not os.path.exists(log_dir):\n        os.makedirs(log_dir)\n    \n    \n    callbacks += [early_stopping,\n                  iou_scores, dice_scores]\n    \n    print(\"===> Saving config setting\")\n    #Save config as JSON format\n    with open(os.path.join(log_dir, 'config.json'), 'w') as f:\n        save_config = train_config.copy()\n        save_config['train_img_path'] = str(save_config['train_img_path'])\n        save_config['train_mask_path'] = str(save_config['train_mask_path'])\n        json.dump(save_config, f)\n        \n    print(\"===> Done\")\n    \n    return {\n        'model': model,\n        'loader': loader,\n        'optimizer': optimizer,\n        'criterion': criterion,\n        'callbacks': callbacks,\n        'scheduler': scheduler,\n        'log_dir': log_dir\n    }","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model 1","metadata":{}},{"cell_type":"code","source":"train_config['model_name'] =  'UnetPlusPlus'\ntrain_config['model_params'] = {\n  'classes': 1,\n  'decoder_attention_type': 'scse',\n  'decoder_use_batchnorm': True,\n  'encoder_depth': 5,\n  'encoder_name': 'efficientnet-b2',\n  'encoder_weights': 'imagenet',\n  'in_channels': 3\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"runner = SupervisedRunner(\n    device=utils.get_device(), input_key=\"image\", input_target_key=\"mask\")\n\nif train_config['is_fp16']:\n    fp16_params = dict(amp=True)  # params for FP16\nelse:\n    fp16_params = None\n\neverything = prepare_everything(f\"experiment_{train_config['model_name'].lower()}\")\nprint(\"Let's go!!!!\")\nrunner.train(\n    model=everything['model'],\n    criterion=everything['criterion'],\n    optimizer=everything['optimizer'],\n    callbacks=everything['callbacks'],\n    logdir=everything['log_dir'],\n    loaders=everything['loader'],\n    num_epochs=train_config['num_epochs'],\n    scheduler=everything['scheduler'],\n    main_metric=train_config['metric'],\n    minimize_metric=False,\n    timeit=True,\n    fp16=fp16_params,\n    resume=train_config['resume_path'],\n    verbose=False,\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model 2","metadata":{}},{"cell_type":"code","source":"train_config['model_name'] = 'DeepLabV3Plus'\ntrain_config['model_params'] = {\n      'classes': 1,\n      'decoder_atrous_rates': [6, 12, 18],\n      'encoder_depth': 5,\n      'encoder_name': 'se_resnext50_32x4d',\n      'encoder_weights': 'imagenet',\n      'in_channels': 3\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"runner = SupervisedRunner(\n    device=utils.get_device(), input_key=\"image\", input_target_key=\"mask\")\n\nif train_config['is_fp16']:\n    fp16_params = dict(amp=True)  # params for FP16\nelse:\n    fp16_params = None\n\neverything = prepare_everything(f\"experiment_{train_config['model_name'].lower()}\")\nprint(\"Let's go!!!!\")\nrunner.train(\n    model=everything['model'],\n    criterion=everything['criterion'],\n    optimizer=everything['optimizer'],\n    callbacks=everything['callbacks'],\n    logdir=everything['log_dir'],\n    loaders=everything['loader'],\n    num_epochs=train_config['num_epochs'],\n    scheduler=everything['scheduler'],\n    main_metric=train_config['metric'],\n    minimize_metric=False,\n    timeit=True,\n    fp16=fp16_params,\n    resume=train_config['resume_path'],\n    verbose=True,\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## THE END","metadata":{}}]}