{"cells":[{"metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"!pip install albumentations==0.4.6\n!pip install git+https://github.com/qubvel/segmentation_models.pytorch","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 \n\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.model_selection import KFold\n\nimport cv2\nimport tifffile\n\nimport os\nimport time\nfrom IPython.display import clear_output\nimport torch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Config:\n    root_path = '../input/hubmap-kidney-segmentation'\n    seed = 55\n    custom_colors = ['#35FCFF', '#FF355A', '#96C503', '#C5035B', '#28B463', '#35FFAF', '#8000FF', '#F400FF']\n    images_path = '../input/hubmap-256x256/train'\n    masks_path = '../input/hubmap-256x256/masks'\n    img_size = 256\n    pretrained_model_path = None\n    train_logs_path = None\n    test_csv_path = ''\n    \ndef seed_everything(seed: int):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    np.random.seed(seed)\n       \nconfig = Config()\nseed_everything(config.seed) ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(config.root_path, 'train.csv'))\ntrain_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":false},"cell_type":"code","source":"def get_mask(rle_mask, shape=(1600, 256)):\n\n    mask = np.zeros((shape[0]*shape[1]), dtype=np.uint8)\n    \n    rle_mask = rle_mask.split()\n    positions = map(int, rle_mask[::2])\n    lengths = map(int, rle_mask[1::2])\n    for pos, le in zip(positions, lengths):\n        mask[pos-1:pos+le-1] = 1\n    mask = mask.reshape((shape[0], shape[1]))\n    \n    return mask.T\n\nimg = tifffile.imread(os.path.join(config.root_path, \"train/aaa6a05cc.tiff\"))\nmask = get_mask(train_df.loc[train_df['id'] == 'aaa6a05cc', 'encoding'].values[0],\n                (img.shape[1], img.shape[0]))\nprint(\"img shape ->\", img.shape)\n\nplt.figure(figsize=(15, 15))\nplt.imshow(img)\nplt.imshow(np.ma.masked_where(mask == False, mask), alpha=0.7, cmap='winter')\nplt.axis(\"off\");","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"meta_df = pd.read_csv(os.path.join(config.root_path,'HuBMAP-20-dataset_information.csv'))\nmeta_df.sample(5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"image_dimensions = meta_df.drop_duplicates(subset=['image_file']).groupby(['width_pixels', 'height_pixels']).count()['image_file'].to_dict()\nprint(\"Unique Image dimensions:\")\nimage_dimensions","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"race_dist = meta_df['race'].value_counts().to_dict()\n\nfig, ax = plt.subplots(figsize=(10, 6))\nax.bar(np.arange(len(race_dist)), race_dist.values(), color=config.custom_colors)\nax.set_xticks(np.arange(len(race_dist)))\nax.set_xticklabels(race_dist, fontsize=12)\nax.set_ylabel(\"count\", fontsize=15)\nax.set_xlabel(\"race\", fontsize=18)\nax.set_title(\"Race Distribution\", fontsize=15);\n\nfor percentage, p in zip(race_dist.values(), ax.patches):\n    percentage = f'{percentage} people'\n    x = p.get_x() + p.get_width() / 2 - 0.15\n    y = p.get_y() + p.get_height()\n    ax.annotate(str(percentage), (x, y), fontsize=12, fontweight='bold')","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(figsize=(10, 6))\nsns.countplot(meta_df['sex'], palette=config.custom_colors, ax=ax)\n\nax.set_xlabel(ax.get_xlabel(), fontsize=15)\nax.set_ylabel(ax.get_ylabel(), fontsize=15)\n\nax.set_xticklabels(['Female', 'Male'], fontsize=12, rotation=0)\nax.set_xlabel('sex', fontsize=18)\nax.set_ylabel(ax.get_ylabel(), fontsize=15)\nax.set_title(\"Sex Distribution\", fontsize=15);\n\nfor count, p in zip([6, 7], ax.patches):\n    x = p.get_x() + p.get_width() / 2 - 0.1\n    y = p.get_y() + p.get_height()\n    ax.annotate(str(count)+ \" people\", (x, y), fontsize=12, fontweight='bold')","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(figsize=(10, 6))\nsns.countplot(meta_df['age'], palette=config.custom_colors, ax=ax)\n\nax.set_xlabel(ax.get_xlabel(), fontsize=15)\nax.set_ylabel(ax.get_ylabel(), fontsize=15)\nax.set_title(\"Age Distribution\", fontsize=15);","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"sns.catplot(x='race',y='age', hue='sex', data=meta_df, kind=\"bar\", height=6, aspect=2, palette=\"cool\", capsize=.05)\nplt.ylabel('Age', fontsize=15)\nplt.xlabel('race vs sex', fontsize=15);\nplt.title(\"Distribution by Gender, Race and Age.\", fontsize=20);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"sns.catplot(x='sex',y='age', hue='weight_kilograms', data=meta_df, kind=\"bar\", height=6, aspect=2, palette=\"cool\", capsize=.05)\nplt.ylabel('Age', fontsize=15)\nplt.xlabel('sex vs weight_kilograms', fontsize=15);\nplt.title(\"Distribution by Gender, weight kilograms and Age.\", fontsize=20);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"create dataframe with additional metadata"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_all_json_file(root: str) -> list:\n    \"\"\"Extraction all unique ids from file names.\"\"\"\n    ids = []\n    for dirname, _, filenames in os.walk(root):\n        for filename in filenames:\n            path = os.path.join(dirname, filename)\n            if path.endswith(\".json\"):\n                ids.append(path) \n    ids = list(set(filter(None, ids)))\n    print(f\"Extracted {len(ids)} json files.\")\n    return ids\n\n\ntest_json_files = get_all_json_file('../input/hubmap-kidney-segmentation/test')\ntrain_json_files = get_all_json_file('../input/hubmap-kidney-segmentation/train')\n\nall_json_files = train_json_files + test_json_files\n\ntrain_ids = [x.split(\"/\")[-1].replace('.', '-').split(\"-\")[0] for x in train_json_files]\ntest_ids = [x.split(\"/\")[4].split(\"-\")[0] for x in test_json_files]\nall_ids = train_ids + test_ids","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"df = pd.read_json(all_json_files[0])\ndf['ID'] = all_ids[0]\n\nfor i, js in enumerate(all_json_files[1:], 1):\n    df_ = pd.read_json(js)\n    df_['ID'] = all_ids[i]\n    df = pd.concat([df, df_], ignore_index=True, sort=False)\n    \ndf = df.drop('geometry', 1).assign(**df.geometry.dropna().apply(pd.Series))\ndf = df.drop('properties', 1).assign(**df.properties.dropna().apply(pd.Series))\ndf = df.drop('classification', 1).assign(**df.classification.dropna().apply(pd.Series))\n\ndf['is_train'] = 0\ndf.loc[df['ID'].isin(train_ids), 'is_train'] = 1\n\ndf.to_csv(\"meta_data2.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"meta_df2 = pd.read_csv('meta_data2.csv')\nmeta_df2.sample(5)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"percentages = [c / meta_df2.shape[0] * 100 for c in meta_df2['name'].value_counts()]\n\nfig, ax = plt.subplots(figsize=(10, 6))\nsns.countplot(meta_df2['name'], order=meta_df2['name'].value_counts().index,\n              ax=ax, palette=config.custom_colors)\n\nax.set_xlabel(ax.get_xlabel(), fontsize=15)\nax.set_ylabel(ax.get_ylabel(), fontsize=15)\n\nax.set_xticklabels(ax.get_xticklabels(), fontsize=12, rotation=0)\nax.set_xlabel('name', fontsize=18)\nax.set_ylabel(ax.get_ylabel(), fontsize=15)\nax.set_title(\"Distribution by classification in metadata\", fontsize=15);\n\nfor percentage, count, p in zip(percentages,\n                               meta_df2['name'].value_counts(sort=True).values,\n                               ax.patches):\n    percentage = f'{np.round(percentage, 2)}%'\n    x = p.get_x() + p.get_width() / 2 - 0.4\n    y = p.get_y() + p.get_height()\n    ax.annotate(str(percentage)+\" / \"+str(count), (x, y), fontsize=12, fontweight='bold')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training Process"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nimport albumentations as A\n\nfrom segmentation_models_pytorch.unet import Unet","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset and Dataloader\n\nDataset - https://www.kaggle.com/iafoss/256x256-images"},{"metadata":{"trusted":true},"cell_type":"code","source":"class HuBMAPDataset(Dataset):\n    def __init__(self, ids, phase):\n        self.ids = ids\n        self.augmentations = get_augmentations(phase)\n        \n    def __getitem__(self, idx):\n        name = self.ids[idx]\n        img = cv2.imread(f\"{config.images_path}/{name}\").astype(\"float32\")\n        img /= 255.\n        mask = cv2.imread(f\"{config.masks_path}/{name}\")[:,:,0:1]\n\n        augmented = self.augmentations(image=img, mask=mask)\n        img = augmented['image']\n        mask = augmented['mask']\n        img = img.transpose(2,0,1).astype('float32')\n        mask = mask.transpose(2,0,1).astype('float32')\n        return img, mask\n\n    def __len__(self):\n        return len(self.ids)\n\n    \ndef get_augmentations(phase: str='train'):\n    if phase == 'train':\n        list_transform = [\n            A.HorizontalFlip(),\n            A.OneOf([\n                A.RandomContrast(),\n                A.RandomGamma(),\n                A.RandomBrightness(),\n                ], p=0.3),\n            A.OneOf([\n                A.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),\n                A.GridDistortion(),\n                A.OpticalDistortion(distort_limit=2, shift_limit=0.5),\n                ], p=0.3),\n            A.ShiftScaleRotate(p=0.2),\n            A.Resize(config.img_size,config.img_size,always_apply=True),\n        ]\n    else:\n        list_transform = [\n            A.Resize(config.img_size,config.img_size)\n        ]\n    return A.Compose(list_transform)\n\n\ndef get_dataloader(\n    phase: str ='train',\n    batch_size: int = 8,\n    num_workers: int = 6,\n    val_size: float = 0.2,\n    fold: int = 0, \n):\n    '''Returns: dataloader for the model training'''\n    \n    ids = os.listdir(config.images_path)\n    train_data = pd.DataFrame(ids, columns=['ids'])\n    skf = KFold(\n        n_splits=7, random_state=config.seed, shuffle=True\n    )\n\n    for i, (train_index, val_index) in enumerate(\n            skf.split(train_data, train_data)\n            ):\n            train_data.loc[val_index, \"fold\"] = i\n    \n    train_df = train_data.loc[train_data['fold'] != fold].reset_index(drop=True)\n    val_df = train_data.loc[train_data['fold'] == fold].reset_index(drop=True)\n\n    ids = train_df['ids'].tolist() if phase == \"train\" else val_df['ids'].tolist()\n    image_dataset = HuBMAPDataset(ids, phase)\n    dataloader = DataLoader(\n        image_dataset,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        pin_memory=True,\n        shuffle=True,   \n    )\n\n    return dataloader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_dataloader = get_dataloader(phase='train')\nimgs, masks = next(iter(val_dataloader))\nprint(\"images ->\", imgs.shape, imgs.dtype)\nprint(\"masks ->\", masks.shape, masks.dtype)\n\n#fig, ax = plt.subplots(1,2, figsize=(15, 7))\n#ax[0].imshow(imgs[0].permute(1,2,0))\n#ax[1].imshow(masks[0][0,:,:])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Exploring several augmentations "},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def plot_with_augmentation(image, mask, augment):\n    \"\"\"\n    Wrapper for `visualize` function.\n    \"\"\"\n    augmented = augment(image=image, mask=mask)\n    aug_image = augmented['image']\n    aug_mask = augmented['mask']\n    \n    f, ax = plt.subplots(2, 2, figsize=(15, 15))\n    fontsize = 20\n\n    ax[0, 0].imshow(image)\n    ax[0, 0].set_title('Original image', fontsize=fontsize)\n    ax[0, 1].imshow(mask)\n    ax[0, 1].set_title('Original mask', fontsize=fontsize)\n    \n    ax[1, 0].imshow(aug_image)\n    ax[1, 0].set_title('Transformed image', fontsize=fontsize)\n    ax[1, 1].imshow(aug_mask)\n    ax[1, 1].set_title('Transformed mask', fontsize=fontsize)\n    \n    plt.tight_layout()\n    plt.show()\n\nimg = cv2.imread(os.path.join(config.images_path, '0486052bb_150.png'))\nmask = cv2.imread(os.path.join(config.masks_path, '0486052bb_150.png')).astype(\"float32\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_with_augmentation(img, mask,  A.HorizontalFlip(p=1))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_with_augmentation(img, mask,  A.ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_with_augmentation(img, mask,  A.OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_with_augmentation(img, mask,  A.RandomBrightness(p=1))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Metric and Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"def dice_coef_metric(probabilities: torch.Tensor,\n                     truth: torch.Tensor,\n                     treshold: float = 0.5,\n                     eps: float = 1e-9) -> np.ndarray:\n    \"\"\"\n    Calculate Dice score for data batch.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: truth values.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        Returns: dice score aka f1.\n    \"\"\"\n    scores = []\n    num = probabilities.shape[0]\n    predictions = (probabilities >= treshold).float()\n    assert(predictions.shape == truth.shape)\n    for i in range(num):\n        prediction = predictions[i]\n        truth_ = truth[i]\n        intersection = 2.0 * (truth_ * prediction).sum()\n        union = truth_.sum() + prediction.sum()\n        if truth_.sum() == 0 and prediction.sum() == 0:\n            scores.append(1.0)\n        else:\n            scores.append((intersection + eps) / union)\n    return np.mean(scores)\n\n\ndef jaccard_coef_metric(probabilities: torch.Tensor,\n               truth: torch.Tensor,\n               treshold: float = 0.5,\n               eps: float = 1e-9) -> np.ndarray:\n    \"\"\"\n    Calculate Jaccard index for data batch.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: truth values.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        Returns: jaccard score aka iou.\"\n    \"\"\"\n    scores = []\n    num = probabilities.shape[0]\n    predictions = (probabilities >= treshold).float()\n    assert(predictions.shape == truth.shape)\n\n    for i in range(num):\n        prediction = predictions[i]\n        truth_ = truth[i]\n        intersection = (prediction * truth_).sum()\n        union = (prediction.sum() + truth_.sum()) - intersection + eps\n        if truth_.sum() == 0 and prediction.sum() == 0:\n            scores.append(1.0)\n        else:\n            scores.append((intersection + eps) / union)\n    return np.mean(scores)\n\n\nclass Meter:\n    '''factory for storing and updating iou and dice scores.'''\n    def __init__(self, treshold: float = 0.5):\n        self.threshold: float = treshold\n        self.dice_scores: list = []\n        self.iou_scores: list = []\n    \n    def update(self, logits: torch.Tensor, targets: torch.Tensor):\n        \"\"\"\n        Takes: logits from output model and targets,\n        calculates dice and iou scores, and stores them in lists.\n        \"\"\"\n        probs = torch.sigmoid(logits)\n        dice = dice_coef_metric(probs, targets, self.threshold)\n        iou = jaccard_coef_metric(probs, targets, self.threshold)\n        \n        self.dice_scores.append(dice)\n        self.iou_scores.append(iou)\n    \n    def get_metrics(self) -> np.ndarray:\n        \"\"\"\n        Returns: the average of the accumulated dice and iou scores.\n        \"\"\"\n        dice = np.mean(self.dice_scores)\n        iou = np.mean(self.iou_scores)\n        return dice, iou\n    \n\nclass DiceLoss(nn.Module):\n    \"\"\"Calculate dice loss.\"\"\"\n    def __init__(self, eps: float = 1e-9):\n        super(DiceLoss, self).__init__()\n        self.eps = eps\n        \n    def forward(self,\n                logits: torch.Tensor,\n                targets: torch.Tensor) -> torch.Tensor:\n        \n        num = targets.size(0)\n        probability = torch.sigmoid(logits)\n        probability = probability.view(num, -1)\n        targets = targets.view(num, -1)\n        assert(probability.shape == targets.shape)\n        \n        intersection = 2.0 * (probability * targets).sum()\n        union = probability.sum() + targets.sum()\n        dice_score = (intersection + self.eps) / union\n        #print(\"intersection\", intersection, union, dice_score)\n        return 1.0 - dice_score\n        \n        \nclass BCEDiceLoss(nn.Module):\n    \"\"\"Compute objective loss: BCE loss + DICE loss.\"\"\"\n    def __init__(self):\n        super(BCEDiceLoss, self).__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.dice = DiceLoss()\n        \n    def forward(self, \n                logits: torch.Tensor,\n                targets: torch.Tensor) -> torch.Tensor:\n        assert(logits.shape == targets.shape)\n        dice_loss = self.dice(logits, targets)\n        bce_loss = self.bce(logits, targets)\n        \n        return bce_loss + dice_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Unet('efficientnet-b2', encoder_weights=\"imagenet\", classes=1, activation=None)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Trainer:\n    \"\"\"\n    Factory for training proccess.\n    Args:\n        display_plot: if True - plot train history after each epoch.\n        net: neural network for mask prediction.\n        criterion: factory for calculating objective loss.\n        optimizer: optimizer for weights updating.\n        phases: list with train and validation phases.\n        dataloaders: dict with data loaders for train and val phases.\n        imgs_dir: path to folder with images.\n        masks_dir: path to folder with imasks.\n        path_to_csv: path to csv file.\n        meter: factory for storing and updating metrics.\n        batch_size: data batch size for one step weights updating.\n        num_epochs: num weights updation for all data.\n        accumulation_steps: the number of steps after which the optimization step can be taken\n                    (https://www.kaggle.com/c/understanding_cloud_organization/discussion/105614).\n        lr: learning rate for optimizer.\n        scheduler: scheduler for control learning rate.\n        losses: dict for storing lists with losses for each phase.\n        jaccard_scores: dict for storing lists with jaccard scores for each phase.\n        dice_scores: dict for storing lists with dice scores for each phase.\n    \"\"\"\n    def __init__(self,\n                 net: nn.Module,\n                 criterion: nn.Module,\n                 lr: float,\n                 accumulation_steps: int,\n                 batch_size: int,\n                 num_epochs: int,\n                 display_plot: bool = True,\n                 fold: int = 0,\n                ):\n\n        \"\"\"Initialization.\"\"\"\n        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n        print(\"device:\", self.device)\n        self.display_plot = display_plot\n        self.net = net\n        self.net = self.net.to(self.device)\n        self.criterion = criterion\n        self.optimizer = Adam(self.net.parameters(), lr=lr)\n        self.scheduler = ReduceLROnPlateau(self.optimizer, mode=\"min\",\n                                           patience=3, verbose=True)\n        self.accumulation_steps = accumulation_steps // batch_size\n        self.phases = [\"train\", \"val\"]\n        self.num_epochs = num_epochs\n\n        self.dataloaders = {\n            phase: get_dataloader(\n                phase = phase,\n                batch_size = 8,\n                num_workers = 4\n            )\n            for phase in self.phases\n        }\n        self.best_loss = float(\"inf\")\n        self.losses = {phase: [] for phase in self.phases}\n        self.dice_scores = {phase: [] for phase in self.phases}\n        self.jaccard_scores = {phase: [] for phase in self.phases}\n         \n    def _compute_loss_and_outputs(self,\n                                  images: torch.Tensor,\n                                  targets: torch.Tensor):\n        images = images.to(self.device)\n        targets = targets.to(self.device)\n        logits = self.net(images)\n        loss = self.criterion(logits, targets)\n        return loss, logits\n        \n    def _do_epoch(self, epoch: int, phase: str):\n        print(f\"{phase} epoch: {epoch} | time: {time.strftime('%H:%M:%S')}\")\n\n        self.net.train() if phase == \"train\" else self.net.eval()\n        meter = Meter()\n        dataloader = self.dataloaders[phase]\n        total_batches = len(dataloader)\n        running_loss = 0.0\n        self.optimizer.zero_grad()\n        for itr, (images, targets) in enumerate(dataloader):\n            loss, logits = self._compute_loss_and_outputs(images, targets)\n            loss = loss / self.accumulation_steps\n            if phase == \"train\":\n                loss.backward()\n                if (itr + 1) % self.accumulation_steps == 0:\n                    self.optimizer.step()\n                    self.optimizer.zero_grad()\n            running_loss += loss.item()\n            meter.update(logits.detach().cpu(),\n                         targets.detach().cpu()\n                        )\n            \n        epoch_loss = (running_loss * self.accumulation_steps) / total_batches\n        epoch_dice, epoch_iou = meter.get_metrics()\n        \n        self.losses[phase].append(epoch_loss)\n        self.dice_scores[phase].append(epoch_dice)\n        self.jaccard_scores[phase].append(epoch_iou)\n\n        return epoch_loss\n        \n    def run(self):\n        for epoch in range(self.num_epochs):\n            self._do_epoch(epoch, \"train\")\n            with torch.no_grad():\n                val_loss = self._do_epoch(epoch, \"val\")\n                self.scheduler.step(val_loss)\n            if self.display_plot:\n                self._plot_train_history()\n                \n            if val_loss < self.best_loss:\n                print(f\"\\n{'#'*20}\\nSaved new checkpoint\\n{'#'*20}\\n\")\n                self.best_loss = val_loss\n                torch.save(self.net.state_dict(), \"best_model.pth\")\n            print()\n        self._save_train_history()\n            \n    def _plot_train_history(self):\n        data = [self.losses, self.dice_scores, self.jaccard_scores]\n        colors = ['deepskyblue', \"crimson\"]\n        labels = [\n            f\"\"\"\n            train loss {self.losses['train'][-1]}\n            val loss {self.losses['val'][-1]}\n            \"\"\",\n            \n            f\"\"\"\n            train dice score {self.dice_scores['train'][-1]}\n            val dice score {self.dice_scores['val'][-1]} \n            \"\"\", \n                  \n            f\"\"\"\n            train jaccard score {self.jaccard_scores['train'][-1]}\n            val jaccard score {self.jaccard_scores['val'][-1]}\n            \"\"\",\n        ]\n        \n        clear_output(True)\n        with plt.style.context(\"seaborn-dark-palette\"):\n            fig, axes = plt.subplots(3, 1, figsize=(8, 10))\n            for i, ax in enumerate(axes):\n                ax.plot(data[i]['val'], c=colors[0], label=\"val\")\n                ax.plot(data[i]['train'], c=colors[-1], label=\"train\")\n                ax.set_title(labels[i])\n                ax.legend(loc=\"upper right\")\n                \n            plt.tight_layout()\n            plt.show()\n            \n    def load_predtrain_model(self,\n                             state_path: str):\n        self.net.load_state_dict(torch.load(state_path))\n        print(\"Predtrain model loaded\")\n        \n    def _save_train_history(self):\n        \"\"\"writing model weights and training logs to files.\"\"\"\n        torch.save(self.net.state_dict(),\n                   f\"last_epoch_model.pth\")\n\n        logs_ = [self.losses, self.dice_scores, self.jaccard_scores]\n        log_names_ = [\"_loss\", \"_dice\", \"_jaccard\"]\n        logs = [logs_[i][key] for i in list(range(len(logs_)))\n                         for key in logs_[i]]\n        log_names = [key+log_names_[i] \n                     for i in list(range(len(logs_))) \n                     for key in logs_[i]\n                    ]\n        pd.DataFrame(\n            dict(zip(log_names, logs))\n        ).to_csv(\"train_log.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"trainer = Trainer(net=model,\n                  criterion=BCEDiceLoss(),\n                  lr=5e-4,\n                  accumulation_steps=32,\n                  batch_size=8,\n                  num_epochs=30,\n                  fold=0,\n                 )\n\nif config.pretrained_model_path is not None:\n    trainer.load_predtrain_model(config.pretrained_model_path)\n    \n    # if need - load the logs.      \n    train_logs = pd.read_csv(config.train_logs_path)\n    trainer.losses[\"train\"] =  train_logs.loc[:, \"train_loss\"].to_list()\n    trainer.losses[\"val\"] =  train_logs.loc[:, \"val_loss\"].to_list()\n    trainer.dice_scores[\"train\"] = train_logs.loc[:, \"train_dice\"].to_list()\n    trainer.dice_scores[\"val\"] = train_logs.loc[:, \"val_dice\"].to_list()\n    trainer.jaccard_scores[\"train\"] = train_logs.loc[:, \"train_jaccard\"].to_list()\n    trainer.jaccard_scores[\"val\"] = train_logs.loc[:, \"val_jaccard\"].to_list()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrainer.run()","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}