effnetb7-fold3_focal_jarvis_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.007246,
     "end_time": "2022-08-16T19:02:52.302567",
     "exception": false,
     "start_time": "2022-08-16T19:02:52.295321",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thu Sep 22 15:56:46 2022       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 470.42.01    Driver Version: 470.42.01    CUDA Version: 11.4     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Quadro RTX 5000     Off  | 00000000:1B:00.0 Off |                  Off |\n",
      "| 33%   30C    P5    13W / 230W |      0MiB / 16125MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "papermill": {
     "duration": 5.945852,
     "end_time": "2022-08-16T19:02:58.254572",
     "exception": false,
     "start_time": "2022-08-16T19:02:52.308720",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import random\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.cuda.amp as amp\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import RandomSampler \n",
    "from torch.utils.data import SequentialSampler\n",
    "import torch.nn.functional as F\n",
    "# from torchmetrics.functional import dice_score\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "import tifffile\n",
    "from fastai.vision.all import *\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "import torch\n",
    "from torch.optim.optimizer import Optimizer\n",
    "import itertools as it\n",
    "\n",
    "is_amp = True\n",
    "import logging\n",
    "\n",
    "from sklearn.model_selection import KFold\n",
    "\n",
    "\n",
    "from itertools import repeat\n",
    "import collections.abc\n",
    "import math\n",
    "import torch\n",
    "from torch import Tensor\n",
    "from torch.optim.optimizer import Optimizer, required\n",
    "from collections import defaultdict\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import sys\n",
    "sys.path.append('/home/timm-pytorch-image-models/pytorch-image-models-master')\n",
    "import timm\n",
    "\n",
    "sys.path.append('/home/efficientnet-pytorch/EfficientNet-PyTorch/EfficientNet-PyTorch-master')\n",
    "\n",
    "from efficientnet_pytorch import EfficientNet\n",
    "from efficientnet_pytorch.utils import url_map, url_map_advprop, get_model_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005742,
     "end_time": "2022-08-16T19:02:58.266495",
     "exception": false,
     "start_time": "2022-08-16T19:02:58.260753",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "papermill": {
     "duration": 2.014493,
     "end_time": "2022-08-16T19:03:00.286680",
     "exception": false,
     "start_time": "2022-08-16T19:02:58.272187",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘/home/lakshita/somusan/hubmap_kaggle/nbs/result’: No such file or directory\n",
      "mkdir: cannot create directory ‘/home/lakshita/somusan/hubmap_kaggle/nbs/checkpoint’: No such file or directory\n"
     ]
    }
   ],
   "source": [
    "!mkdir /home/lakshita/somusan/hubmap_kaggle/nbs/result\n",
    "!mkdir /home/lakshita/somusan/hubmap_kaggle/nbs/checkpoint\n",
    "\n",
    "# root_dir = '/home/lakshita/somusan/hubmap_kaggle/nbs'\n",
    "# #pretrain_dir = '/kaggle/input/swin-tiny-small-22k-pretrained/'\n",
    "\n",
    "# TRAIN = '/home/lakshita/somusan/hubmap_kaggle/nbs/aug_lung_data/hubmap-22-aug-pixel-size'\n",
    "# LABELS = '/home/lakshita/somusan/hubmap_kaggle/nbs/aug_lung_data/only_lung.csv' #only lungs\n",
    "\n",
    "root_dir = '/home/'\n",
    "#pretrain_dir = '/kaggle/input/swin-tiny-small-22k-pretrained/'\n",
    "\n",
    "TRAIN = 'hubmap-22-aug-pixel-size/fixed_blur_imgs'\n",
    "LABELS = '/home/random-files/train.csv' #only lungs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005637,
     "end_time": "2022-08-16T19:03:00.299522",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.293885",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "papermill": {
     "duration": 0.053647,
     "end_time": "2022-08-16T19:03:00.359131",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.305484",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def image_to_tensor(image, mode='bgr'): #image mode\n",
    "    if mode=='bgr':\n",
    "        image = image[:,:,::-1]\n",
    "    x = image\n",
    "    x = x.transpose(2,0,1)\n",
    "    x = np.ascontiguousarray(x)\n",
    "    x = torch.tensor(x, dtype = torch.float)\n",
    "    return x\n",
    "\n",
    "def mask_to_tensor(mask):\n",
    "    x = mask\n",
    "    #x = x.transpose(2, 0, 1)\n",
    "    x = torch.tensor(x, dtype = torch.float)\n",
    "    return x\n",
    "\n",
    "\n",
    "class RGB(nn.Module):\n",
    "    IMAGE_RGB_MEAN = [0.485, 0.456, 0.406] #[0.5, 0.5, 0.5]\n",
    "    IMAGE_RGB_STD  = [0.229, 0.224, 0.225] #[0.5, 0.5, 0.5]\n",
    "\n",
    "    def __init__(self,):\n",
    "        super(RGB, self).__init__()\n",
    "        self.register_buffer('mean', torch.zeros(1,3,1,1))\n",
    "        self.register_buffer('std', torch.ones(1,3,1,1))\n",
    "        self.mean.data = torch.FloatTensor(self.IMAGE_RGB_MEAN).view(self.mean.shape)\n",
    "        self.std.data = torch.FloatTensor(self.IMAGE_RGB_STD).view(self.std.shape)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = (x-self.mean)/self.std\n",
    "        return x\n",
    "\n",
    "# def message(mode='print'):\n",
    "#     asterisk = ' '\n",
    "#     if mode==('print'):\n",
    "#         loss = batch_loss\n",
    "#     if mode==('log'):\n",
    "#         loss = train_loss\n",
    "#         if (iteration % iter_save == 0): asterisk = '*'\n",
    "    \n",
    "#     print(\"=---------------------DEBUG\")\n",
    "#     print(loss)\n",
    "#     print()\n",
    "#     print(loss.shape)\n",
    "#     print(\"=---------------------DEBUG\")\n",
    "    \n",
    "#     text = \\\n",
    "#         ('%0.2e   %08d%s %6.2f | '%(rate, iteration, asterisk, epoch,)).replace('e-0','e-').replace('e+0','e+') + \\\n",
    "#         '%4.3f  %4.3f   | '%(*valid_loss,) + \\\n",
    "#         '%4.3f  %4.3f | '%(*dict(loss)) + \\\n",
    "#         '%s' % ((time.time() - start_timer))\n",
    "#     return text\n",
    "\n",
    "def message(mode='print'):\n",
    "    asterisk = ' '\n",
    "    if mode==('print'):\n",
    "        loss = batch_loss\n",
    "    if mode==('log'):\n",
    "        loss = train_loss\n",
    "        if (iteration % iter_save == 0): asterisk = '*'\n",
    "\n",
    "    # text = \\\n",
    "    #     ('%0.2e   %08d%s %6.2f | '%(rate, iteration, asterisk, epoch,)).replace('e-0','e-').replace('e+0','e+') + \\\n",
    "    #     '%4.3f  %4.3f  %4.4f  %4.3f   | '%(*valid_loss,) + \\\n",
    "    #     )\n",
    "    # print(\"=---------------------DEBUG\")\n",
    "    # print(loss)\n",
    "    # print(loss.shape)\n",
    "    # print(\"=---------------------DEBUG\")\n",
    "    text = f'{rate}  {iteration}, {epoch}  | {round(valid_loss[0], 4)} {round(valid_loss[1], 4)}  \\\n",
    "            |  {round((loss[0]),2)} {round((loss[1]),2)} | {round((time.time() - start_timer), 3)}'\n",
    "    # print(('%0.2e   %08d%s %6.2f | '%(rate, iteration, asterisk, epoch,)).replace('e-0','e-').replace('e+0','e+'))\n",
    "    # print('VALID LOSS: %4.3f  %4.3f | '%(*valid_loss,))\n",
    "    # print()\n",
    "    # print('TRAINING LOSS: %4.3f  %4.3f   | '%(*loss,))\n",
    "    # print()\n",
    "    # print('%s' % ((time.time() - start_timer)))\n",
    "    # print('\\n')\n",
    "    return text \n",
    "\n",
    "def rle_decode(mask_rle, shape):\n",
    "    '''\n",
    "    mask_rle: run-length as string formated (start length)\n",
    "    shape: (height,width) of array to return \n",
    "    Returns numpy array, 1 - mask, 0 - background\n",
    "\n",
    "    '''\n",
    "    s = mask_rle.split()\n",
    "    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n",
    "    starts -= 1\n",
    "    ends = starts + lengths\n",
    "    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n",
    "    for lo, hi in zip(starts, ends):\n",
    "        img[lo:hi] = 1\n",
    "    return img.reshape(shape).T\n",
    "\n",
    "\n",
    "# ref.: https://www.kaggle.com/stainsby/fast-tested-rle\n",
    "def rle_encode(img):\n",
    "    '''\n",
    "    img: numpy array, 1 - mask, 0 - background\n",
    "    Returns run length as string formated\n",
    "    '''\n",
    "    pixels = img.flatten()\n",
    "    pixels = np.concatenate([[0], pixels, [0]])\n",
    "    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n",
    "    runs[1::2] -= runs[::2]\n",
    "    return ' '.join(str(x) for x in runs)\n",
    "\n",
    "\n",
    "def get_learning_rate(optimizer):\n",
    "    return optimizer.param_groups[0]['lr']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005519,
     "end_time": "2022-08-16T19:03:00.370356",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.364837",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Augments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "papermill": {
     "duration": 0.016754,
     "end_time": "2022-08-16T19:03:00.392910",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.376156",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def valid_augment5(image, mask, organ):\n",
    "    #image, mask  = do_crop(image, mask, image_size, xy=(None,None))\n",
    "    return image, mask\n",
    "from numpy.random import choice\n",
    "def train_augment5b(image, mask, organ):\n",
    "    # image, mask = do_random_flip(image, mask)\n",
    "    # image, mask = do_random_rot90(image, mask)\n",
    "    \n",
    "\n",
    "    for fn in np.random.choice([\n",
    "        lambda image, mask: (image, mask),\n",
    "        #lambda image, mask: do_random_noise(image, mask, mag=0.1),\n",
    "        lambda image, mask: do_random_contast(image, mask, mag=0.40),\n",
    "        lambda image, mask: do_random_hsv(image, mask, mag=[0.40, 0.40, 0])\n",
    "    ], 2): image, mask = fn(image, mask)\n",
    "\n",
    "    for fn in np.random.choice([\n",
    "        lambda image, mask: (image, mask),\n",
    "        # lambda image, mask: do_random_rotate_scale(image, mask, angle=45, scale=[0.50, 2.0]),\n",
    "    ], 1): image, mask = fn(image, mask)\n",
    "\n",
    "    return image, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "papermill": {
     "duration": 0.024978,
     "end_time": "2022-08-16T19:03:00.423769",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.398791",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def do_random_flip(image, mask):\n",
    "    if np.random.rand()>0.5:\n",
    "        image = cv2.flip(image,0)\n",
    "        mask = cv2.flip(mask,0)\n",
    "    if np.random.rand()>0.5:\n",
    "        image = cv2.flip(image,1)\n",
    "        mask = cv2.flip(mask,1)\n",
    "    if np.random.rand()>0.5:\n",
    "        image = image.transpose(1,0,2)\n",
    "        mask = mask.transpose(1,0)\n",
    "    \n",
    "    image = np.ascontiguousarray(image)\n",
    "    mask = np.ascontiguousarray(mask)\n",
    "    return image, mask\n",
    "\n",
    "def do_random_rot90(image, mask):\n",
    "    r = np.random.choice([\n",
    "        0,\n",
    "        cv2.ROTATE_90_CLOCKWISE,\n",
    "        cv2.ROTATE_90_COUNTERCLOCKWISE,\n",
    "        cv2.ROTATE_180,\n",
    "    ])\n",
    "    if r==0:\n",
    "        return image, mask\n",
    "    else:\n",
    "        image = cv2.rotate(image, r)\n",
    "        mask = cv2.rotate(mask, r)\n",
    "        return image, mask\n",
    "    \n",
    "def do_random_contast(image, mask, mag=0.3): #this thing kills the image and sets all pixels to 1\n",
    "    alpha = 1 + random.uniform(-1,1)*mag\n",
    "    image = image * alpha\n",
    "    image = np.clip(image,0,1)\n",
    "    return image, mask\n",
    "\n",
    "def do_random_hsv(image, mask, mag=[0.15,0.25,0.25]):\n",
    "    image = (image*255).astype(np.uint8)\n",
    "    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n",
    "\n",
    "    h = hsv[:, :, 0].astype(np.float32)  # hue\n",
    "    s = hsv[:, :, 1].astype(np.float32)  # saturation\n",
    "    v = hsv[:, :, 2].astype(np.float32)  # value\n",
    "    h = (h*(1 + random.uniform(-1,1)*mag[0]))%180\n",
    "    s =  s*(1 + random.uniform(-1,1)*mag[1])\n",
    "    v =  v*(1 + random.uniform(-1,1)*mag[2])\n",
    "\n",
    "    hsv[:, :, 0] = np.clip(h,0,180).astype(np.uint8)\n",
    "    hsv[:, :, 1] = np.clip(s,0,255).astype(np.uint8)\n",
    "    hsv[:, :, 2] = np.clip(v,0,255).astype(np.uint8)\n",
    "    image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)\n",
    "    image = image.astype(np.float32)/255\n",
    "    return image, mask\n",
    "\n",
    "def do_random_noise(image, mask, mag=0.1): #also seems to kill the image and set to 1\n",
    "    height, width = image.shape[:2]\n",
    "    noise = np.random.uniform(-1,1, (height, width,1))*mag\n",
    "    image = image + noise\n",
    "    image = np.clip(image,0,1)\n",
    "    return image, mask\n",
    "\n",
    "def do_random_rotate_scale(image, mask, angle=30, scale=[0.8,1.2] ):\n",
    "    angle = np.random.uniform(-angle, angle)\n",
    "    scale = np.random.uniform(*scale) if scale is not None else 1\n",
    "    \n",
    "    height, width = image.shape[:2]\n",
    "    center = (height // 2, width // 2)\n",
    "    \n",
    "    transform = cv2.getRotationMatrix2D(center, angle, scale)\n",
    "    image = cv2.warpAffine( image, transform, (width, height), flags=cv2.INTER_LINEAR,\n",
    "                            borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0))\n",
    "    mask  = cv2.warpAffine( mask, transform, (width, height), flags=cv2.INTER_LINEAR,\n",
    "                            borderMode=cv2.BORDER_CONSTANT, borderValue=0)\n",
    "    return image, mask\n",
    "\n",
    "\n",
    "# image = cv2.imread(\"/home/lakshita/somusan/hubmap_kaggle/hubmap_data/hubmap-organ-segmentation/train_images/686.tiff\")\n",
    "# mask = cv2.imread(\"/home/lakshita/somusan/hubmap_kaggle/hubmap_data/mask_png/train_binary_masks/686.png\",0)\n",
    "# aug_img, aug_mask = do_random_rotate_scale(image, mask)\n",
    "\n",
    "\n",
    "# plt.figure(figsize=(16,18))\n",
    "# plt.subplot(1,2,1)\n",
    "# plt.imshow(aug_img)\n",
    "\n",
    "# plt.subplot(1,2,2)\n",
    "# plt.imshow(aug_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "papermill": {
     "duration": 0.012714,
     "end_time": "2022-08-16T19:03:00.442141",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.429427",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# dummy_mask = np.zeros( (2023, 2023, 1) )\n",
    "# dummy_mask[ dummy_mask.shape[0] // 2 : dummy_mask.shape[0] // 2 + 100, \n",
    "#            dummy_mask.shape[0] // 2 : dummy_mask.shape[0] // 2 + 100] = 1\n",
    "# # dummy_mask[150 : 170, 220 : 240] = 1\n",
    "# # dummy_mask[300 : 320, 150 : 170] = 1\n",
    "# dummy_mask[1000 : 1100, 450 : 550] = 1\n",
    "# dummy_mask[1300 : 1400, 500 : 600] = 1\n",
    "# dummy_mask[1100 : 1200, 700 : 800] = 1\n",
    "# dummy_mask[800 : 900, 1000 : 1100] = 1\n",
    "# dummy_mask[1350 : 1450, 1250 : 1350] = 1\n",
    "\n",
    "# plt.imshow(dummy_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005493,
     "end_time": "2022-08-16T19:03:00.454279",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.448786",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "papermill": {
     "duration": 0.02051,
     "end_time": "2022-08-16T19:03:00.480870",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.460360",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_size = 512 #1024 #512 #512\n",
    "\n",
    "class HubmapDataset(Dataset):\n",
    "    def __init__(self, df, augment=None):\n",
    "\n",
    "        self.df = df\n",
    "        self.augment = augment\n",
    "        self.length = len(self.df)\n",
    "        self.organ_to_label = {'kidney' : 0,\n",
    "                               'prostate' : 1,\n",
    "                               'largeintestine' : 2,\n",
    "                               'spleen' : 3,\n",
    "                               'lung' : 4}\n",
    "\n",
    "    def __str__(self):\n",
    "        string = ''\n",
    "        string += '\\tlen = %d\\n' % len(self)\n",
    "\n",
    "        d = self.df.organ.value_counts().to_dict()\n",
    "        for k in ['kidney', 'prostate', 'largeintestine', 'spleen', 'lung']:\n",
    "            string +=  '%24s %3d (%0.3f) \\n'%(k, d.get(k,0), d.get(k,0)/len(self.df))\n",
    "        return string\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        d = self.df.iloc[index]\n",
    "        img_height = self.df.loc[index, 'img_height']\n",
    "        img_width = self.df.loc[index, 'img_width']\n",
    "        organ = self.organ_to_label[d.organ]\n",
    "\n",
    "        image = cv2.cvtColor(cv2.imread(os.path.join(TRAIN, f'{d.id}.png')), cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        rle_mask = self.df.loc[index, 'rle']\n",
    "        mask = rle_decode(rle_mask, (img_height, img_width))\n",
    "        #mask = cv2.cvtColor(mask, cv2.IMREAD_GRAYSCALE)\n",
    "        #mask = cv2.imread(os.path.join(MASKS,fname),cv2.IMREAD_GRAYSCALE)\n",
    "        # mask = np.expand_dims(mask, axis = 2)\n",
    "        # print(mask.shape)\n",
    "        \n",
    "        image = image.astype(np.float32)/255\n",
    "        #mask  = mask.astype(np.float32)/255\n",
    "        mask = mask.astype(np.float32)\n",
    "\n",
    "        s = d.pixel_size/0.4 * (image_size/3000)\n",
    "        image = cv2.resize(image,dsize=(image_size,image_size),interpolation=cv2.INTER_LINEAR)\n",
    "        mask  = cv2.resize(mask, dsize=(image_size,image_size),interpolation=cv2.INTER_LINEAR)\n",
    "\n",
    "        if self.augment is not None:\n",
    "            image, mask = self.augment(image, mask, organ)\n",
    "\n",
    "        mask = np.expand_dims(mask, axis = 0)\n",
    "\n",
    "        r ={}\n",
    "        r['index']= index\n",
    "        r['organ'] = torch.tensor([organ], dtype=torch.long)\n",
    "        r['image'] = image_to_tensor(image)\n",
    "        r['mask' ] = mask_to_tensor(mask)\n",
    "        return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "papermill": {
     "duration": 0.01215,
     "end_time": "2022-08-16T19:03:00.498755",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.486605",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1024, 1024])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"/home/random-files/train.csv\")\n",
    "ds = HubmapDataset(df)\n",
    "\n",
    "ds[10]['mask'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.00545,
     "end_time": "2022-08-16T19:03:00.509765",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.504315",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "papermill": {
     "duration": 0.049669,
     "end_time": "2022-08-16T19:03:00.565172",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.515503",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class FPN(nn.Module):\n",
    "    def __init__(self, input_channels:list, output_channels:list):\n",
    "        super().__init__()\n",
    "        self.convs = nn.ModuleList(\n",
    "            [nn.Sequential(nn.Conv2d(in_ch, out_ch*2, kernel_size=3, padding=1),\n",
    "             nn.ReLU(inplace=True), nn.BatchNorm2d(out_ch*2),\n",
    "             nn.Conv2d(out_ch*2, out_ch, kernel_size=3, padding=1))\n",
    "            for in_ch, out_ch in zip(input_channels, output_channels)])\n",
    "        \n",
    "    def forward(self, xs:list, last_layer):\n",
    "        hcs = [F.interpolate(c(x),scale_factor=2**(len(self.convs)-i),mode='bilinear') \n",
    "               for i,(c,x) in enumerate(zip(self.convs, xs))]\n",
    "        hcs.append(last_layer)\n",
    "        return torch.cat(hcs, dim=1)\n",
    "\n",
    "class UnetBlock(nn.Module):\n",
    "    def __init__(self, up_in_c:int, x_in_c:int, nf:int=None, blur:bool=False,\n",
    "                 self_attention:bool=False, **kwargs):\n",
    "        super().__init__()\n",
    "        self.shuf = PixelShuffle_ICNR(up_in_c, up_in_c//2, blur=blur, **kwargs)\n",
    "        self.bn = nn.BatchNorm2d(x_in_c)\n",
    "        ni = up_in_c//2 + x_in_c\n",
    "        nf = nf if nf is not None else max(up_in_c//2,32)\n",
    "        self.conv1 = ConvLayer(ni, nf, norm_type=None, **kwargs)\n",
    "        self.conv2 = ConvLayer(nf, nf, norm_type=None,\n",
    "            xtra=SelfAttention(nf) if self_attention else None, **kwargs)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "\n",
    "    def forward(self, up_in:Tensor, left_in:Tensor) -> Tensor:\n",
    "        s = left_in\n",
    "        up_out = self.shuf(up_in)\n",
    "        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))\n",
    "        return self.conv2(self.conv1(cat_x))\n",
    "        \n",
    "class _ASPPModule(nn.Module):\n",
    "    def __init__(self, inplanes, planes, kernel_size, padding, dilation, groups=1):\n",
    "        super().__init__()\n",
    "        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,\n",
    "                stride=1, padding=padding, dilation=dilation, bias=False, groups=groups)\n",
    "        self.bn = nn.BatchNorm2d(planes)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "        self._init_weight()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.atrous_conv(x)\n",
    "        x = self.bn(x)\n",
    "\n",
    "        return self.relu(x)\n",
    "\n",
    "    def _init_weight(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                torch.nn.init.kaiming_normal_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "class ASPP(nn.Module):\n",
    "    def __init__(self, inplanes=512, mid_c=256, dilations=[6, 12, 18, 24], out_c=None):\n",
    "        super().__init__()\n",
    "        self.aspps = [_ASPPModule(inplanes, mid_c, 1, padding=0, dilation=1)] + \\\n",
    "            [_ASPPModule(inplanes, mid_c, 3, padding=d, dilation=d,groups=4) for d in dilations]\n",
    "        self.aspps = nn.ModuleList(self.aspps)\n",
    "        self.global_pool = nn.Sequential(nn.AdaptiveMaxPool2d((1, 1)),\n",
    "                        nn.Conv2d(inplanes, mid_c, 1, stride=1, bias=False),\n",
    "                        nn.BatchNorm2d(mid_c), nn.ReLU())\n",
    "        out_c = out_c if out_c is not None else mid_c\n",
    "        self.out_conv = nn.Sequential(nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False),\n",
    "                                    nn.BatchNorm2d(out_c), nn.ReLU(inplace=True))\n",
    "        self.conv1 = nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False)\n",
    "        self._init_weight()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x0 = self.global_pool(x)\n",
    "        xs = [aspp(x) for aspp in self.aspps]\n",
    "        x0 = F.interpolate(x0, size=xs[0].size()[2:], mode='bilinear', align_corners=True)\n",
    "        x = torch.cat([x0] + xs, dim=1)\n",
    "        return self.out_conv(x)\n",
    "    \n",
    "    def _init_weight(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                torch.nn.init.kaiming_normal_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "\n",
    "class config:\n",
    "    pretrained_root = 'efficientnet-pytorch/'\n",
    "    efficient_net_encoders = {\n",
    "        \"efficientnet-b0\": {\n",
    "            \"out_channels\": (3, 32, 24, 40, 112, 320),\n",
    "            \"stage_idxs\": (3, 5, 9, 16),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b0-08094119.pth\"\n",
    "        },\n",
    "        \"efficientnet-b1\": {\n",
    "            \"out_channels\": (3, 32, 24, 40, 112, 320),\n",
    "            \"stage_idxs\": (5, 8, 16, 23),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b1-dbc7070a.pth\"\n",
    "        },\n",
    "        \"efficientnet-b2\": {\n",
    "            \"out_channels\": (3, 32, 24, 48, 120, 352),\n",
    "            \"stage_idxs\": (5, 8, 16, 23),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b2-27687264.pth\"\n",
    "        },\n",
    "        \"efficientnet-b3\": {\n",
    "            \"out_channels\": (3, 40, 32, 48, 136, 384),\n",
    "            \"stage_idxs\": (5, 8, 18, 26),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b3-c8376fa2.pth\"\n",
    "        },\n",
    "        \"efficientnet-b4\": {\n",
    "            \"out_channels\": (3, 48, 32, 56, 160, 448),\n",
    "            \"stage_idxs\": (6, 10, 22, 32),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b4-e116e8b3.pth\"\n",
    "        },\n",
    "        \"efficientnet-b5\": {\n",
    "            \"out_channels\": (3, 48, 40, 64, 176, 512),\n",
    "            \"stage_idxs\": (8, 13, 27, 39),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b5-586e6cc6.pth\"\n",
    "        },\n",
    "        \"efficientnet-b6\": {\n",
    "            \"out_channels\": (3, 56, 40, 72, 200, 576),\n",
    "            \"stage_idxs\": (9, 15, 31, 45),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b6-c76e70fd.pth\"\n",
    "        },\n",
    "        \"efficientnet-b7\": {\n",
    "            \"out_channels\": (3, 64, 48, 80, 224, 640),\n",
    "            \"stage_idxs\": (11, 18, 38, 55),\n",
    "            \"weight_path\": pretrained_root + \"efficientnet-b7-dcc49843.pth\"\n",
    "        }\n",
    "    }\n",
    "    model = 'efficientnet-b7'\n",
    "    \n",
    "class EfficientNetEncoder(EfficientNet):\n",
    "    def __init__(self, stage_idxs, out_channels, model_name, depth=5):\n",
    "\n",
    "        blocks_args, global_params = get_model_params(model_name, override_params=None)\n",
    "        super().__init__(blocks_args, global_params)\n",
    "        \n",
    "        cfg = config.efficient_net_encoders[model_name]\n",
    "\n",
    "        self._stage_idxs = stage_idxs\n",
    "        self._out_channels = out_channels\n",
    "        self._depth = depth\n",
    "        self._in_channels = 3\n",
    "\n",
    "        del self._fc\n",
    "        self.load_state_dict(torch.load(cfg['weight_path']))\n",
    "\n",
    "    def get_stages(self):\n",
    "        return [\n",
    "            nn.Identity(),\n",
    "            nn.Sequential(self._conv_stem, self._bn0, self._swish),\n",
    "            self._blocks[:self._stage_idxs[0]],\n",
    "            self._blocks[self._stage_idxs[0]:self._stage_idxs[1]],\n",
    "            self._blocks[self._stage_idxs[1]:self._stage_idxs[2]],\n",
    "            self._blocks[self._stage_idxs[2]:],\n",
    "        ]\n",
    "\n",
    "    def forward(self, x):\n",
    "        stages = self.get_stages()\n",
    "\n",
    "        block_number = 0.\n",
    "        drop_connect_rate = self._global_params.drop_connect_rate\n",
    "\n",
    "        features = []\n",
    "        for i in range(self._depth + 1):\n",
    "\n",
    "            # Identity and Sequential stages\n",
    "            if i < 2:\n",
    "                x = stages[i](x)\n",
    "\n",
    "            # Block stages need drop_connect rate\n",
    "            else:\n",
    "                for module in stages[i]:\n",
    "                    drop_connect = drop_connect_rate * block_number / len(self._blocks)\n",
    "                    block_number += 1.\n",
    "                    x = module(x, drop_connect)\n",
    "\n",
    "            features.append(x)\n",
    "\n",
    "        return features\n",
    "\n",
    "    def load_state_dict(self, state_dict, **kwargs):\n",
    "        state_dict.pop(\"_fc.bias\")\n",
    "        state_dict.pop(\"_fc.weight\")\n",
    "        super().load_state_dict(state_dict, **kwargs)  \n",
    "        \n",
    "\n",
    "class EffUnet(nn.Module):\n",
    "    def __init__(self, model_name, stride=1):\n",
    "        super().__init__()\n",
    "        self.focal_loss = FocalLoss()\n",
    "        cfg = config.efficient_net_encoders[model_name]\n",
    "        stage_idxs = cfg['stage_idxs']\n",
    "        out_channels = cfg['out_channels']\n",
    "        \n",
    "        self.encoder = EfficientNetEncoder(stage_idxs, out_channels, model_name)\n",
    "\n",
    "        #aspp with customized dilatations\n",
    "        self.aspp = ASPP(out_channels[-1], 256, out_c=384, \n",
    "                         dilations=[stride*1, stride*2, stride*3, stride*4])\n",
    "        self.drop_aspp = nn.Dropout2d(0.5)\n",
    "        #decoder\n",
    "        self.dec4 = UnetBlock(384, out_channels[-2], 256)\n",
    "        self.dec3 = UnetBlock(256, out_channels[-3], 128)\n",
    "        self.dec2 = UnetBlock(128, out_channels[-4], 64)\n",
    "        self.dec1 = UnetBlock(64, out_channels[-5], 32)\n",
    "        self.fpn = FPN([384, 256, 128, 64], [16]*4)\n",
    "        self.drop = nn.Dropout2d(0.1)\n",
    "        self.final_conv = ConvLayer(32+16*4, 1, ks=1, norm_type=None, act_cls=None)\n",
    "        \n",
    "        self.rgb = RGB()\n",
    "        \n",
    "    def forward(self, batch):\n",
    "        x = batch['image']\n",
    "        B, C, H, W = x.shape\n",
    "        x = self.rgb(x)\n",
    "        enc0, enc1, enc2, enc3, enc4 = self.encoder(x)[-5:]\n",
    "        enc5 = self.aspp(enc4)\n",
    "        dec3 = self.dec4(self.drop_aspp(enc5), enc3)\n",
    "        dec2 = self.dec3(dec3,enc2)\n",
    "        dec1 = self.dec2(dec2,enc1)\n",
    "        dec0 = self.dec1(dec1,enc0)\n",
    "        x = self.fpn([enc5, dec3, dec2, dec1], dec0)\n",
    "        x = self.final_conv(self.drop(x))\n",
    "        x = F.interpolate(x, size = 1024, mode='bilinear', align_corners=False)\n",
    "\n",
    "\n",
    "        # print(\"s shape\",x.shape)\n",
    "        output = {}\n",
    "        # if 'loss' in self.output_type:\n",
    "        output['bce_loss'] = F.binary_cross_entropy_with_logits(x,batch['mask'])\n",
    "        # for i in range(4):\n",
    "        #     output['aux%d_loss'%i] = criterion_aux_loss(self.aux[i](decoder[i]),batch['mask'])\n",
    "        output['focal_loss'] = self.focal_loss(x,batch['mask'])\n",
    "        # print(output)\n",
    "        # if 'inference' in self.output_type:\n",
    "        #     output['probability'] = torch.sigmoid(x)\n",
    "        return x, output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "papermill": {
     "duration": 0.013343,
     "end_time": "2022-08-16T19:03:00.584292",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.570949",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# FOCAL LOSS ----------------\n",
    "ALPHA = 0.8\n",
    "GAMMA = 2\n",
    "\n",
    "class FocalLoss(nn.Module):\n",
    "    def __init__(self, weight=None, size_average=True):\n",
    "        super(FocalLoss, self).__init__()\n",
    "\n",
    "    def forward(self, inputs, targets, alpha=ALPHA, gamma=GAMMA, smooth=1):\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",
    "        #first compute binary cross-entropy \n",
    "        BCE = F.binary_cross_entropy_with_logits(inputs, targets, reduction='mean')\n",
    "        BCE_EXP = torch.exp(-BCE)\n",
    "        focal_loss = alpha * (1-BCE_EXP)**gamma * BCE\n",
    "                       \n",
    "        return focal_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_check_net():\n",
    "    batch_size = 2\n",
    "    image_size = 1024 #768 #512\n",
    "\n",
    "    #---\n",
    "    batch = {\n",
    "        'image' : torch.from_numpy( np.random.uniform(-1,1,(batch_size,3,image_size,image_size)) ).float(),\n",
    "        'mask'  : torch.from_numpy( np.random.choice(2,(batch_size,1,image_size,image_size)) ).float(),\n",
    "        'organ' : torch.from_numpy( np.random.choice(5,(batch_size)) ).long(),\n",
    "    }\n",
    "    batch = {k:v.cuda() for k,v in batch.items()}\n",
    "\n",
    "    net = EffUnet(config.model).cuda()\n",
    "    # net.load_pretrain()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        with torch.cuda.amp.autocast(enabled=True):\n",
    "            output, loss = net(batch)\n",
    "\n",
    "    print('batch')\n",
    "    for k,v in batch.items():\n",
    "        print('%32s :'%k, v.shape)\n",
    "\n",
    "    print('output')\n",
    "    print(output.shape)\n",
    "    print(loss['bce_loss'])\n",
    "    print(loss['focal_loss'])\n",
    "\n",
    "    \n",
    "    # for k,v in output.items():\n",
    "    #     if 'loss' not in k:\n",
    "    #         print('%32s :'%k, v.shape)\n",
    "    # for k,v in output.items():\n",
    "    #     if 'loss' in k:\n",
    "    #         print('%32s :'%k, v.item())\n",
    "\n",
    "# run_check_net()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.00547,
     "end_time": "2022-08-16T19:03:00.617292",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.611822",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Folds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "papermill": {
     "duration": 0.017484,
     "end_time": "2022-08-16T19:03:00.640381",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.622897",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def make_fold(fold = 3):\n",
    "    df = pd.read_csv(LABELS)\n",
    "    df = df[df['organ'] == 'lung']\n",
    "    print(len(df))\n",
    "\n",
    "    num_fold =48\n",
    "    skf = KFold(n_splits = num_fold, shuffle = True,random_state = 42)\n",
    "\n",
    "    df.loc[:,'fold']=-1\n",
    "    for f,(t_idx, v_idx) in enumerate(skf.split(X=df['id'], y=df['organ'])):\n",
    "        df.iloc[v_idx,-1]=f\n",
    "\n",
    "    #check\n",
    "    if 0:\n",
    "        for f in range(num_fold):\n",
    "            train_df=df[df.fold!=f].reset_index(drop=True)\n",
    "            valid_df=df[df.fold==f].reset_index(drop=True)\n",
    "\n",
    "            print('fold %d'%f)\n",
    "            t = train_df.organ.value_counts().to_dict()\n",
    "            v = valid_df.organ.value_counts().to_dict()\n",
    "            for k in ['kidney', 'prostate', 'largeintestine', 'spleen', 'lung']:\n",
    "                print('%32s %3d (%0.3f)  %3d (%0.3f)'%(k,t.get(k,0),t.get(k,0)/len(train_df),v.get(k,0),v.get(k,0)/len(valid_df)))\n",
    "\n",
    "            print('')\n",
    "            zz=0\n",
    "\n",
    "    train_df=df[df.fold!=fold].reset_index(drop=True)\n",
    "    valid_df=df[df.fold==fold].reset_index(drop=True)\n",
    "    return train_df,valid_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005895,
     "end_time": "2022-08-16T19:03:00.651958",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.646063",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Competition metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "papermill": {
     "duration": 0.013577,
     "end_time": "2022-08-16T19:03:00.671364",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.657787",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_dice_score(probability, mask, smooth = 1):\n",
    "    N = len(probability)\n",
    "    p = probability.reshape(N,-1)\n",
    "    t = mask.reshape(N,-1)\n",
    "\n",
    "    p = p>0.5\n",
    "    t = t>0.5\n",
    "    uion = p.sum(-1) + t.sum(-1)\n",
    "    overlap = (p*t).sum(-1)\n",
    "    dice = 2*overlap/(uion+0.0001)\n",
    "    return dice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005378,
     "end_time": "2022-08-16T19:03:00.682366",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.676988",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "papermill": {
     "duration": 0.01812,
     "end_time": "2022-08-16T19:03:00.706195",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.688075",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def validate(net, valid_loader, debug = False):\n",
    "    \n",
    "    valid_num = 0\n",
    "    valid_probability = []\n",
    "    valid_mask = []\n",
    "    valid_loss = 0\n",
    "    \n",
    "    criterion = nn.BCEWithLogitsLoss()\n",
    "    \n",
    "    net = net.eval()\n",
    "    start_timer = time.time()\n",
    "    \n",
    "    for t, batch in enumerate(valid_loader):\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            with amp.autocast(enabled = is_amp):\n",
    "                \n",
    "                batch_size = len(batch['index'])\n",
    "                batch['image'] = batch['image'].cuda()\n",
    "                batch['mask' ] = batch['mask' ].cuda()\n",
    "                batch['organ'] = batch['organ'].cuda()\n",
    "\n",
    "                output, loss = net(batch)\n",
    "                loss0 = loss['bce_loss'].mean()\n",
    "                # batch['mask'] = batch['mask']\n",
    "                # loss = criterion(output, batch['mask'])\n",
    "\n",
    "        valid_probability.append(output.data.cpu().numpy())\n",
    "        valid_mask.append(batch['mask'].data.cpu().numpy())\n",
    "        valid_num += batch_size\n",
    "        valid_loss += batch_size * loss0.item()\n",
    "        \n",
    "        print('\\r %8d / %d  %s'%(valid_num, len(valid_loader.dataset),(time.time() - start_timer)),end='',flush=True)\n",
    "    \n",
    "    assert(valid_num == len(valid_loader.dataset))\n",
    "\n",
    "    probability = np.concatenate(valid_probability)\n",
    "    mask = np.concatenate(valid_mask)\n",
    "\n",
    "    loss = valid_loss/valid_num\n",
    "\n",
    "    dice = compute_dice_score(probability, mask)\n",
    "    dice = dice.mean()\n",
    "    \n",
    "    return [dice, loss]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "df_lung = df[df['organ']=='lung']\n",
    "train_df_, valid_df_ = train_test_split(df_lung, test_size=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.00551,
     "end_time": "2022-08-16T19:03:00.717363",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.711853",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "papermill": {
     "duration": 8.210001,
     "end_time": "2022-08-16T19:03:08.932978",
     "exception": false,
     "start_time": "2022-08-16T19:03:00.722977",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48\n",
      "numper of training and validation samples 47 1\n"
     ]
    }
   ],
   "source": [
    "CFG = {\n",
    "    'fold_num': 5,\n",
    "    'seed': 719,\n",
    "    'model_arch': 'tf_efficientnet_b4_ns',\n",
    "    'img_size': 512,\n",
    "    'epochs': 10,\n",
    "    'train_bs': 16,\n",
    "    'valid_bs': 32,\n",
    "    'T_0': 10,\n",
    "    'lr': 1e-4,\n",
    "    'min_lr': 1e-6,\n",
    "    'weight_decay':1e-6,\n",
    "    'num_workers': 4,\n",
    "    'accum_iter': 2, # suppoprt to do batch accumulation for backprop with effectively larger batch size\n",
    "    'verbose_step': 1,\n",
    "    'device': 'cuda:0'\n",
    "}\n",
    "\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "\n",
    "fold = 0\n",
    "\n",
    "out_dir = root_dir + '/result/effnetb7/focal-fold-%d' % (fold)\n",
    "initial_checkpoint = None\n",
    "\n",
    "start_lr = 5e-5\n",
    "batch_size = 2\n",
    "\n",
    "## setup  ----------------------------------------\n",
    "for f in ['checkpoint','train','valid','backup'] : os.makedirs(out_dir +'/'+f, exist_ok=True)\n",
    "\n",
    "    \n",
    "log = open(out_dir+'/log.train.txt',mode='a')\n",
    "log.write('\\n--- [START %s] %s\\n\\n' % ('EfficientNet-b7', '-' * 64))\n",
    "log.write('\\n')\n",
    "\n",
    "## dataset ----------------------------------------\n",
    "log.write('** dataset setting **\\n')\n",
    "\n",
    "\n",
    "train_df, valid_df = make_fold(fold) #train_df_, valid_df_ #make_fold(fold)\n",
    "print(\"numper of training and validation samples\",len(train_df), len(valid_df))\n",
    "\n",
    "train_dataset = HubmapDataset(train_df, train_augment5b)\n",
    "valid_dataset = HubmapDataset(valid_df, valid_augment5)\n",
    "\n",
    "# train_loader = torch.utils.data.DataLoader(\n",
    "#                dataset=train_dataset,\n",
    "#                batch_size=batch_size,\n",
    "#                shuffle=False,\n",
    "#                num_workers=8,\n",
    "#                drop_last   = True)\n",
    "\n",
    "train_loader  = DataLoader(\n",
    "    train_dataset,\n",
    "    sampler = RandomSampler(train_dataset),\n",
    "    batch_size  = batch_size,\n",
    "    drop_last   = False,\n",
    "    num_workers = 8,\n",
    "    pin_memory  = False,\n",
    "    worker_init_fn = lambda id: np.random.seed(torch.initial_seed() // 2 ** 32 + id)\n",
    ")\n",
    "\n",
    "valid_loader = DataLoader(\n",
    "    valid_dataset,\n",
    "    sampler = SequentialSampler(valid_dataset),\n",
    "    batch_size  = batch_size,\n",
    "    drop_last   = False,\n",
    "    num_workers = 4,\n",
    "    pin_memory  = False\n",
    ")\n",
    "\n",
    "\n",
    "log.write('fold = %s\\n'%str(fold))\n",
    "log.write('train_dataset : \\n%s\\n'%(train_dataset))\n",
    "log.write('valid_dataset : \\n%s\\n'%(valid_dataset))\n",
    "log.write('\\n')\n",
    "\n",
    "\n",
    "## net ----------------------------------------\n",
    "log.write('** net setting **\\n')\n",
    "\n",
    "scaler = amp.GradScaler(enabled = is_amp)\n",
    "net = EffUnet(config.model).cuda()\n",
    "\n",
    "if initial_checkpoint is not None:\n",
    "    f = torch.load(initial_checkpoint, map_location=lambda storage, loc: storage)\n",
    "    start_iteration = f['iteration']\n",
    "    start_epoch = f['epoch']\n",
    "    state_dict  = f['state_dict']\n",
    "    net.load_state_dict(state_dict,strict=False)  #True\n",
    "else:\n",
    "    start_iteration = 0\n",
    "    start_epoch = 0\n",
    "\n",
    "\n",
    "log.write('\\tinitial_checkpoint = %s\\n' % initial_checkpoint)\n",
    "log.write('\\n')\n",
    "\n",
    "\n",
    "## optimiser ----------------------------------\n",
    "if 0: ##freeze\n",
    "    for p in net.stem.parameters():   p.requires_grad = False\n",
    "    pass\n",
    "\n",
    "def freeze_bn(net):\n",
    "    for m in net.modules():\n",
    "        if isinstance(m, nn.BatchNorm2d):\n",
    "            m.eval()\n",
    "            m.weight.requires_grad = False\n",
    "            m.bias.requires_grad = False\n",
    "            \n",
    "#freeze_bn(net)\n",
    "\n",
    "#-----------------------------------------------\n",
    "\n",
    "# optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr)\n",
    "# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=start_lr,weight_decay=1e-5)\n",
    "optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr)\n",
    "\n",
    "#optimizer = Lookahead(RAdam(filter(lambda p: p.requires_grad, net.parameters()),lr=start_lr), alpha=0.5, k=5)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=CFG['T_0'], T_mult=1, eta_min=CFG['min_lr'], last_epoch=-1)\n",
    "log.write('optimizer\\n  %s\\n'%(optimizer))\n",
    "log.write('\\n')\n",
    "\n",
    "#num_iteration = 1000*len(train_loader)\n",
    "num_iteration = 12000\n",
    "iter_log   = 100 #len(train_loader)*3 #479\n",
    "iter_valid = iter_log\n",
    "iter_save  = iter_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pandas as pd\n",
    "# pd.read_csv(\"/home/lakshita/somusan/hubmap_kaggle/nbs/aug_lung_data/only_lung.csv\").__len__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next(iter(valid_loader))[\"mask\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# items = next(iter(train_loader))\n",
    "# imgs = items[\"image\"].permute((0, 2, 3, 1))\n",
    "# msks = items[\"mask\"].permute((0,2,3,1))\n",
    "# print(imgs.size(), msks.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# items = next(iter(train_loader))\n",
    "# imgs = items[\"image\"].permute((0, 2, 3, 1))\n",
    "# msks = items[\"mask\"].permute((0, 2, 3, 1))\n",
    "# print(imgs.size(), msks.size())\n",
    "\n",
    "# # torch.unique(msks)\n",
    "\n",
    "# import matplotlib.pyplot as plt\n",
    "# def plot_batch(imgs, msks, size=3):\n",
    "#     for idx in range(size):\n",
    "#         plt.figure(figsize=(4*3, 5))\n",
    "\n",
    "#         plt.subplot(1, 3, 1); plt.imshow(imgs[idx])\n",
    "#         plt.title('image', fontsize=15)\n",
    "#         plt.axis('OFF')\n",
    "\n",
    "#         plt.subplot(1, 3, 2); plt.imshow(msks[idx])\n",
    "#         plt.title('mask', fontsize=15)\n",
    "#         plt.axis('OFF')\n",
    "            \n",
    "#         plt.subplot(1, 3, 3); plt.imshow(imgs[idx]); plt.imshow(msks[idx], alpha=0.3)\n",
    "#         plt.title('overlay', fontsize=15)\n",
    "#         plt.axis('OFF')\n",
    "        \n",
    "#         plt.tight_layout()\n",
    "#         plt.show()\n",
    "\n",
    "# plot_batch(imgs, msks, size=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.006146,
     "end_time": "2022-08-16T19:03:08.945358",
     "exception": false,
     "start_time": "2022-08-16T19:03:08.939212",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "papermill": {
     "duration": 2131.600859,
     "end_time": "2022-08-16T19:38:40.552279",
     "exception": false,
     "start_time": "2022-08-16T19:03:08.951420",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5e-05  22, 0.9166666666666663  | 0.0 0.129              |  0.05000000074505806 0.0 | 20.6481482582 | 17.39525355 | 6.369"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [22]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m amp\u001b[38;5;241m.\u001b[39mautocast(enabled \u001b[38;5;241m=\u001b[39m is_amp):\n\u001b[0;32m---> 86\u001b[0m         output, loss_ \u001b[38;5;241m=\u001b[39m \u001b[43mnet\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     87\u001b[0m         loss0  \u001b[38;5;241m=\u001b[39m loss_[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mbce_loss\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mmean()\n\u001b[1;32m     88\u001b[0m         loss1  \u001b[38;5;241m=\u001b[39m loss_[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mfocal_loss\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mmean()\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36mEffUnet.forward\u001b[0;34m(self, batch)\u001b[0m\n\u001b[1;32m    220\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrgb(x)\n\u001b[1;32m    221\u001b[0m enc0, enc1, enc2, enc3, enc4 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mencoder(x)[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m5\u001b[39m:]\n\u001b[0;32m--> 222\u001b[0m enc5 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maspp\u001b[49m\u001b[43m(\u001b[49m\u001b[43menc4\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    223\u001b[0m dec3 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdec4(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdrop_aspp(enc5), enc3)\n\u001b[1;32m    224\u001b[0m dec2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdec3(dec3,enc2)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36mASPP.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[0;32m---> 75\u001b[0m     x0 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mglobal_pool\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     76\u001b[0m     xs \u001b[38;5;241m=\u001b[39m [aspp(x) \u001b[38;5;28;01mfor\u001b[39;00m aspp \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maspps]\n\u001b[1;32m     77\u001b[0m     x0 \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39minterpolate(x0, size\u001b[38;5;241m=\u001b[39mxs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39msize()[\u001b[38;5;241m2\u001b[39m:], mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mbilinear\u001b[39m\u001b[38;5;124m'\u001b[39m, align_corners\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/container.py:141\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    140\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 141\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    142\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/batchnorm.py:168\u001b[0m, in \u001b[0;36m_BatchNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    161\u001b[0m     bn_training \u001b[38;5;241m=\u001b[39m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_mean \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;129;01mand\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_var \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    163\u001b[0m \u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    164\u001b[0m \u001b[38;5;124;03mBuffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be\u001b[39;00m\n\u001b[1;32m    165\u001b[0m \u001b[38;5;124;03mpassed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are\u001b[39;00m\n\u001b[1;32m    166\u001b[0m \u001b[38;5;124;03mused for normalization (i.e. in eval mode when buffers are not None).\u001b[39;00m\n\u001b[1;32m    167\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 168\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    169\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    170\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# If buffers are not to be tracked, ensure that they won't be updated\u001b[39;49;00m\n\u001b[1;32m    171\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_mean\u001b[49m\n\u001b[1;32m    172\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\n\u001b[1;32m    173\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    174\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_var\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    175\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    176\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbn_training\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[43m    \u001b[49m\u001b[43mexponential_average_factor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    179\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    180\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:2409\u001b[0m, in \u001b[0;36mbatch_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[1;32m   2396\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m   2397\u001b[0m         batch_norm,\n\u001b[1;32m   2398\u001b[0m         (\u001b[38;5;28minput\u001b[39m, running_mean, running_var, weight, bias),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   2406\u001b[0m         eps\u001b[38;5;241m=\u001b[39meps,\n\u001b[1;32m   2407\u001b[0m     )\n\u001b[1;32m   2408\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m training:\n\u001b[0;32m-> 2409\u001b[0m     \u001b[43m_verify_batch_size\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2411\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mbatch_norm(\n\u001b[1;32m   2412\u001b[0m     \u001b[38;5;28minput\u001b[39m, weight, bias, running_mean, running_var, training, momentum, eps, torch\u001b[38;5;241m.\u001b[39mbackends\u001b[38;5;241m.\u001b[39mcudnn\u001b[38;5;241m.\u001b[39menabled\n\u001b[1;32m   2413\u001b[0m )\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:2377\u001b[0m, in \u001b[0;36m_verify_batch_size\u001b[0;34m(size)\u001b[0m\n\u001b[1;32m   2375\u001b[0m     size_prods \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m=\u001b[39m size[i \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m2\u001b[39m]\n\u001b[1;32m   2376\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size_prods \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m-> 2377\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExpected more than 1 value per channel when training, got input size \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(size))\n",
      "\u001b[0;31mValueError\u001b[0m: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])"
     ]
    }
   ],
   "source": [
    "log.write('** start training here! **\\n')\n",
    "log.write('   batch_size = %d \\n'%(batch_size))\n",
    "log.write('                     |-------------- VALID---------|---- TRAIN/BATCH ----------------\\n')\n",
    "log.write('rate     iter  epoch | dice   loss   tp     tn     | loss           | time           \\n')\n",
    "log.write('-------------------------------------------------------------------------------------\\n')\n",
    "\n",
    "valid_loss = np.zeros(4,np.float32)\n",
    "train_loss = np.zeros(2,np.float32)\n",
    "batch_loss = np.zeros_like(train_loss)\n",
    "sum_train_loss = np.zeros_like(train_loss)\n",
    "sum_train = 0\n",
    "\n",
    "start_timer = time.time()\n",
    "iteration = start_iteration\n",
    "epoch = start_epoch\n",
    "rate = 0\n",
    "accum_iter = 4\n",
    "schd_batch_update = False\n",
    "\n",
    "rate = 0\n",
    "mix_proba = 0.4\n",
    "mix_alpha = 0.4\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "while iteration < num_iteration:\n",
    "    for t, batch in enumerate(train_loader):\n",
    "\n",
    "        if iteration%iter_save==0:\n",
    "            if iteration != start_iteration:\n",
    "                torch.save({\n",
    "                    'state_dict': net.state_dict(),\n",
    "                    'iteration': iteration,\n",
    "                    'epoch': epoch,\n",
    "                }, out_dir + '/checkpoint/%08d.model.pth' %  (iteration))\n",
    "                pass\n",
    "\n",
    "\n",
    "        if (iteration%iter_valid==0):\n",
    "            valid_loss = validate(net, valid_loader)\n",
    "            pass\n",
    "\n",
    "\n",
    "        if (iteration%iter_log==0) or (iteration%iter_valid==0):\n",
    "            print('\\r', end='', flush=True)\n",
    "            log.write(message(mode='log') + '\\n')\n",
    "\n",
    "\n",
    "        # learning rate schduler ------------\n",
    "        rate = get_learning_rate(optimizer)\n",
    "\n",
    "        # one iteration update  -------------\n",
    "        batch_size = len(batch['index'])\n",
    "        \n",
    "        #print(batch_size, iteration, epoch)\n",
    "        batch['image'] = batch['image'].half().cuda()\n",
    "        batch['mask' ] = batch['mask' ].half().cuda()\n",
    "        batch['organ'] = batch['organ'].cuda()\n",
    "\n",
    "\n",
    "        net.train()\n",
    "        net.output_type = ['loss']\n",
    "        # if 1:\n",
    "        #     with amp.autocast(enabled = is_amp):\n",
    "        #         output = net(batch)\n",
    "        #         batch['mask'] = batch['mask']\n",
    "        #         loss = criterion(output, batch['mask'])\n",
    "        #         loss = loss / accum_iter \n",
    "                \n",
    "            \n",
    "        #     scaler.scale(loss).backward()\n",
    "\n",
    "            \n",
    "            \n",
    "        #     if ((t + 1) % accum_iter == 0) or (t + 1 == len(train_loader)):\n",
    "        #         scaler.unscale_(optimizer)\n",
    "        #         scaler.step(optimizer)\n",
    "        #         scaler.update()\n",
    "        #         optimizer.zero_grad()\n",
    "        #         if scheduler is not None and schd_batch_update:\n",
    "        #             scheduler.step()\n",
    "\n",
    "        net.train()\n",
    "        net.output_type = ['loss']\n",
    "        if 1:\n",
    "            with amp.autocast(enabled = is_amp):\n",
    "                output, loss_ = net(batch)\n",
    "                loss0  = loss_['bce_loss'].mean()\n",
    "                loss1  = loss_['focal_loss'].mean()\n",
    "\n",
    "            loss = loss0+0.2*loss1\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            scaler.scale(loss).backward()\n",
    "\n",
    "            scaler.unscale_(optimizer)\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "           \n",
    "\n",
    "        # print statistics  --------\n",
    "        # batch_loss = loss.item()\n",
    "        batch_loss[:2] = [loss0.item(),loss1.item()]\n",
    "        sum_train_loss += batch_loss\n",
    "        sum_train += 1\n",
    "        if t % 100 == 0:\n",
    "            train_loss = sum_train_loss / (sum_train + 1e-12)\n",
    "            sum_train_loss = 0\n",
    "            sum_train = 0\n",
    "\n",
    "        print('\\r', end='', flush=True)\n",
    "        print(message(mode='print'), end='', flush=True)\n",
    "        epoch += 1 / len(train_loader)\n",
    "        iteration += 1\n",
    "    \n",
    "    torch.cuda.empty_cache()\n",
    "    \n",
    "log.write('\\n')\n",
    "log.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### valid_loss\n",
    "- comp data\n",
    "- ~~mask data~~ \n",
    "- effnet data\n",
    "- "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "21a2a557654fc1676068684031cf9bb9dfda94e124d3623f4e9c9ed764d794ac"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}