{"cells":[{"metadata":{},"cell_type":"markdown","source":"## setup"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --upgrade pip\n!pip install -q monai\n!pip install -q git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport sys\n# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # specify GPUs locally\n\n# libraries\nimport time\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport PIL.Image\nimport matplotlib.pyplot as plt\n%matplotlib inline\nfrom tqdm import tqdm as tqdm\nfrom sklearn.metrics import cohen_kappa_score, confusion_matrix\nfrom sklearn.model_selection import train_test_split, StratifiedKFold\nimport random\n\nimport torch\nfrom torch.utils.data import TensorDataset, DataLoader,Dataset\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.transforms as transforms\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data.sampler import SubsetRandomSampler, RandomSampler, SequentialSampler\nfrom torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR\nfrom warmup_scheduler import GradualWarmupScheduler\nfrom sklearn.metrics import roc_auc_score\nimport albumentations\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport monai\nfrom monai.data import NiftiDataset\nfrom monai.transforms import AddChannel, Compose, RandRotate90, Resize, ScaleIntensity, ToTensor\n\n# from apex import amp # I cannot install apex in Kagggle notebook\n\ndevice = torch.device('cuda')\n\ndef set_seed(seed=0):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)  \n    torch.cuda.manual_seed(seed)  \n    torch.cuda.manual_seed_all(seed)  \n    torch.backends.cudnn.deterministic = True\n    \nset_seed(0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DEBUG = True\n\nkernel_type = 'monai3d_160_3ch_1e-5_20ep_aug'\n\nimage_size = 160\nuse_amp = False\ndata_dir = '../input/rsna-str-pe-detection-jpeg-256/train-jpegs'\nnum_workers = 4\ninit_lr = 1e-5\nout_dim = 9\nfreeze_epo = 0\nwarmup_epo = 1\ncosine_epo = 2 if DEBUG else 19\nn_epochs = freeze_epo + warmup_epo + cosine_epo","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"target_cols = [\n        'negative_exam_for_pe', # exam level\n        'rv_lv_ratio_gte_1', # exam level\n        'rv_lv_ratio_lt_1', # exam level\n        'leftsided_pe', # exam level\n        'chronic_pe', # exam level\n        'rightsided_pe', # exam level\n        'acute_and_chronic_pe', # exam level\n        'central_pe', # exam level\n        'indeterminate' # exam level\n    ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('../input/rsna-str-pulmonary-embolism-detection/train.csv')\n\nfrom sklearn.model_selection import GroupKFold\nnp.random.seed(0)\ngroup_kfold = GroupKFold(n_splits=5)\nprint(group_kfold)\n\ndf['fold'] = -1\nfor i, (_, val_index) in enumerate(group_kfold.split(df, groups=df.StudyInstanceUID)):\n    df.loc[val_index, 'fold'] = i\n\ndf.fold.value_counts()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_study = df.drop_duplicates('StudyInstanceUID')[['StudyInstanceUID','SeriesInstanceUID','fold']+target_cols]\n#if DEBUG:\n#    df_study = df_study.head(600)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3D dataset"},{"metadata":{},"cell_type":"markdown","source":"We use the preprocessed data by Dr. Ian Pan https://www.kaggle.com/vaillant/rsna-str-pe-detection-jpeg-256\n1. read all the jpgs in the same study\n2. stack them in order to make it a 3-dimensional array\n3. resize into 160x160x160x3 cube (with 3 channels)"},{"metadata":{"trusted":true},"cell_type":"code","source":"from glob import glob\nfrom monai.transforms import LoadNifti, Randomizable, apply_transform\nfrom monai.transforms import AddChannel, Compose, RandRotate90, Resize, ScaleIntensity, ToTensor, RandAffine\nfrom monai.utils import get_seed\n\nclass RSNADataset3D(torch.utils.data.Dataset, Randomizable):\n    def __init__(self, csv, mode, transform=None):\n\n        self.csv = csv.reset_index()\n        self.mode = mode\n        self.transform = transform\n\n    def __len__(self):\n        return self.csv.shape[0]\n    \n    def randomize(self) -> None:\n        MAX_SEED = np.iinfo(np.uint32).max + 1\n        self._seed = self.R.randint(MAX_SEED, dtype=\"uint32\")    \n\n    def __getitem__(self, index):\n        self.randomize()\n        row = self.csv.iloc[index]\n        jpg_lst = sorted(glob(os.path.join(data_dir, row.StudyInstanceUID, row.SeriesInstanceUID, '*.jpg')))\n        img_lst = [cv2.imread(jpg)[:,:,::-1] for jpg in jpg_lst] \n        img = np.stack([image.astype(np.float32) for image in img_lst], axis=2).transpose(3,0,1,2)\n\n        if self.transform is not None:\n            if isinstance(self.transform, Randomizable):\n                self.transform.set_random_state(seed=self._seed)\n            img = apply_transform(self.transform, img)\n\n        if self.mode == 'test':\n            return img\n        else:\n            return img, torch.tensor(row[target_cols]).float()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There are only 7000 data in train set, so we use heavy augmentations."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_transforms = Compose([ScaleIntensity(), \n                            Resize((image_size, image_size, image_size)), \n                            RandAffine( \n                                      prob=0.5,\n                                      translate_range=(5, 5, 5),\n                                      rotate_range=(np.pi * 4, np.pi * 4, np.pi * 4),\n                                      scale_range=(0.15, 0.15, 0.15),\n                                      padding_mode='border'),\n                            ToTensor()])\nval_transforms = Compose([ScaleIntensity(), Resize((image_size, image_size, image_size)), ToTensor()])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## loss function"},{"metadata":{"trusted":true},"cell_type":"code","source":"bce = nn.BCEWithLogitsLoss()\ndef criterion(logits, target): \n    loss = bce(logits.view(-1), target.view(-1))\n    return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(model, loader, optimizer):\n\n    model.train()\n    train_loss = []\n    bar = tqdm(loader)\n    for (data, target) in bar:\n        data, target = data.to(device), target.to(device)\n        optimizer.zero_grad()\n        logits = model(data)       \n        loss = criterion(logits, target)\n\n        if not use_amp:\n            loss.backward()\n        else:\n            with amp.scale_loss(loss, optimizer) as scaled_loss:\n                scaled_loss.backward()\n\n        optimizer.step()\n\n        loss_np = loss.detach().cpu().numpy()\n        train_loss.append(loss_np)\n        smooth_loss = sum(train_loss[-100:]) / min(len(train_loss), 100)\n        bar.set_description('loss: %.5f, smth: %.5f' % (loss_np, smooth_loss))\n    return train_loss\n\n\ndef val_epoch(model, loader, is_ext=None, n_test=1, get_output=False):\n\n    model.eval()\n    val_loss = []\n    LOGITS = []\n    TARGETS = []\n\n    with torch.no_grad():\n        for (data, target) in tqdm(loader):\n            data, target = data.to(device), target.to(device)\n            logits = model(data)\n            LOGITS.append(logits.detach().cpu())\n            TARGETS.append(target.detach().cpu())\n\n    val_loss = criterion(torch.cat(LOGITS), torch.cat(TARGETS)).numpy()\n    PROBS = torch.sigmoid(torch.cat(LOGITS)).numpy().squeeze()    \n    LOGITS = torch.cat(LOGITS).numpy()\n    TARGETS = torch.cat(TARGETS).numpy()\n    \n    if get_output:\n        return LOGITS, PROBS, TARGETS\n    else:\n        acc = (PROBS.round() == TARGETS).mean() * 100.\n        auc = roc_auc_score(TARGETS, LOGITS)\n        return float(val_loss), acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GradualWarmupSchedulerV2(GradualWarmupScheduler):\n    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):\n        super(GradualWarmupSchedulerV2, self).__init__(optimizer, multiplier, total_epoch, after_scheduler)\n    def get_lr(self):\n        if self.last_epoch > self.total_epoch:\n            if self.after_scheduler:\n                if not self.finished:\n                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]\n                    self.finished = True\n                return self.after_scheduler.get_lr()\n            return [base_lr * self.multiplier for base_lr in self.base_lrs]\n        if self.multiplier == 1.0:\n            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]\n        else:\n            return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(fold):\n    df_train = df_study[(df_study['fold'] != fold)]\n    df_valid = df_study[(df_study['fold'] == fold)]\n\n    dataset_train = RSNADataset3D(df_train, 'train', transform=train_transforms)\n    dataset_valid = RSNADataset3D(df_valid, 'val', transform=val_transforms)\n    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=4, sampler=RandomSampler(dataset_train), num_workers=num_workers)\n    valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=4, num_workers=num_workers)\n\n    model = monai.networks.nets.densenet.densenet121(spatial_dims=3, in_channels=3, out_channels=out_dim).to(device)\n    checkpoint = torch.load('../input/rsna-training-weights/ModelTraining.pth')\n\n    val_loss_best = 1000\n    model_file = f'{kernel_type}_best_fold{fold}.pth'\n\n    optimizer = optim.Adam(model.parameters(), lr=init_lr)\n    if use_amp:\n        model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\")\n        model.load_state_dict(checkpoint['model_state_dict'])\n        optimizer.load_state_dict(checkpoint['checkpoint_state_dict'])\n#     if len(os.environ['CUDA_VISIBLE_DEVICES'].split(',')) > 1:\n#         model = nn.DataParallel(model)         \n        \n    scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, cosine_epo)\n    scheduler_cosine.load_state_dict(scheduler_cosine.state_dict())\n    scheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n    scheduler_warmup.load_state_dict(scheduler_warmup.state_dict())\n\n    print(len(dataset_train), len(dataset_valid))\n\n    for epoch in range(1, n_epochs+1):\n        print(time.ctime(), 'Epoch:', epoch)\n        scheduler_warmup.step(epoch-1)\n\n        train_loss = train_epoch(model, train_loader, optimizer)\n        val_loss, acc, auc = val_epoch(model, valid_loader)\n    \n        content = time.ctime() + ' ' + f'Fold {fold}, Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, train loss: {np.mean(train_loss):.5f}, valid loss: {(val_loss):.5f}, acc: {(acc):.4f}, auc: {(auc):.6f}'\n        print(content)\n        with open(f'log_{kernel_type}.txt', 'a') as appender:\n            appender.write(content + '\\n')             \n            \n        if val_loss < val_loss_best:\n            print('val_loss_best ({:.6f} --> {:.6f}).  Saving model ...'.format(val_loss_best, val_loss))\n            torch.save(model.state_dict(), model_file)\n            val_loss_best = val_loss\n        \n        torch.save({\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'scheduler_cosine_state_dict': scheduler_cosine.state_dict(),\n            'scheduler_warmup_state_dict': scheduler_warmup.state_dict()\n        },'ModelTraining.pth')\n\n    torch.save(model.state_dict(), f'{kernel_type}_model_fold{fold}.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"run(fold=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}