{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Summary\nThis is an adaptation of a great notebook [Single Fold Training of Resnet200d [LB0.965]](https://www.kaggle.com/underwearfitting/single-fold-training-of-resnet200d-lb0-965) for usage with TPU. I had to decrease image size to 128 to prevent memory overflow."},{"metadata":{},"cell_type":"markdown","source":"## Configuration"},{"metadata":{"trusted":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version 1.7 > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"#fold_id = 0\n\nimage_size = 128 #512\nseed = 42\nwarmup_epo = 1\ninit_lr = 1e-4/3\nbatch_size = 64\nvalid_batch_size = 32\nn_epochs = 30\nwarmup_factor = 10\nnum_workers = 4\n\ndebug = True # change this to run on full data and multiple epochs\nearly_stop = 5\n\nkernel_type = 'resnet200d'\ndata_dir = '../input/ranzcr-clip-catheter-line-classification/train'\nmodel_dir = f'weights/'\n! mkdir $model_dir","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Imports"},{"metadata":{"trusted":true},"cell_type":"code","source":"! pip install -q pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport sys\nsys.path.append('../input/pytorch-image-models/pytorch-image-models-master')\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\nimport time\nimport cv2\nimport PIL.Image\nimport random\nfrom sklearn.metrics import accuracy_score\nimport torch\nfrom torch.utils.data import 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.lr_scheduler import CosineAnnealingLR \nfrom warmup_scheduler import GradualWarmupScheduler\nimport albumentations\nfrom albumentations import *\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\nimport gc\nfrom sklearn.metrics import roc_auc_score\nimport seaborn as sns\nfrom pylab import rcParams\nimport timm\n\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.utils.serialization as xser\n\nfrom warnings import filterwarnings\nfilterwarnings(\"ignore\")\n\n #torch.device('cuda')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RANZCRResNet200D(nn.Module):\n    def __init__(self, model_name='resnet200d', out_dim=11, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=False)\n        if pretrained:\n            pretrained_path = '../input/resnet200d-pretrained-weight/resnet200d_ra2-bdba9bf9.pth'\n            self.model.load_state_dict(torch.load(pretrained_path))\n        n_features = self.model.fc.in_features\n        self.model.global_pool = nn.Identity()\n        self.model.fc = nn.Identity()\n        self.pooling = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Linear(n_features, out_dim)\n\n    def forward(self, x):\n        bs = x.size(0)\n        features = self.model(x)\n        pooled_features = self.pooling(features).view(bs, -1)\n        output = self.fc(pooled_features)\n        return output","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Utils"},{"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]\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    \nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Transforms"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n   albumentations.RandomResizedCrop(image_size, image_size, scale=(0.9, 1), p=1), \n   albumentations.HorizontalFlip(p=0.5),\n   albumentations.ShiftScaleRotate(p=0.5),\n   albumentations.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, p=0.7),\n   albumentations.RandomBrightnessContrast(brightness_limit=(-0.2,0.2), contrast_limit=(-0.2, 0.2), p=0.7),\n   albumentations.CLAHE(clip_limit=(1,4), p=0.5),\n   albumentations.OneOf([\n       albumentations.OpticalDistortion(distort_limit=1.0),\n       albumentations.GridDistortion(num_steps=5, distort_limit=1.),\n       albumentations.ElasticTransform(alpha=3),\n   ], p=0.2),\n   albumentations.OneOf([\n       albumentations.GaussNoise(var_limit=[10, 50]),\n       albumentations.GaussianBlur(),\n       albumentations.MotionBlur(),\n       albumentations.MedianBlur(),\n   ], p=0.2),\n  albumentations.Resize(image_size, image_size),\n  albumentations.OneOf([\n      JpegCompression(),\n      Downscale(scale_min=0.1, scale_max=0.15),\n  ], p=0.2),\n  IAAPiecewiseAffine(p=0.2),\n  IAASharpen(p=0.2),\n  albumentations.Cutout(max_h_size=int(image_size * 0.1), max_w_size=int(image_size * 0.1), num_holes=5, p=0.5),\n  albumentations.Normalize(),\n])\n\ntransforms_valid = albumentations.Compose([\n    albumentations.Resize(image_size, image_size),\n    albumentations.Normalize()\n])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RANZERDataset(Dataset):\n    def __init__(self, df, mode, transform=None):\n        \n        self.df = df.reset_index(drop=True)\n        self.mode = mode\n        self.transform = transform\n        self.labels = df[target_cols].values\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        row = self.df.loc[index]\n        img = cv2.imread(row.file_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        if self.transform is not None:\n            res = self.transform(image=img)\n            img = res['image']\n                \n        img = img.astype(np.float32)\n        img = img.transpose(2,0,1)\n        label = torch.tensor(self.labels[index]).float()\n        if self.mode == 'test':\n            return torch.tensor(img).float()\n        else:\n            return torch.tensor(img).float(), label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = pd.read_csv('../input/how-to-properly-split-folds/train_folds.csv')\ndf_train['file_path'] = df_train.StudyInstanceUID.apply(lambda x: os.path.join(data_dir, f'{x}.jpg'))\nif debug:\n    df_train = df_train.sample(frac=0.1)\ntarget_cols = df_train.iloc[:, 1:12].columns.tolist()\ndataset = RANZERDataset(df_train, 'train', transform=transforms_train)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"def macro_multilabel_auc(label, pred):\n    aucs = []\n    for i in range(len(target_cols)):\n        if len(np.unique(label[:, i])) != 1:\n            aucs.append(roc_auc_score(label[:, i], pred[:, i]))\n    return np.mean(aucs)\n\n\ndef train_func(train_loader, model, criterion, optimizer):\n    model.train()\n    bar = tqdm(train_loader, desc=\"Training\", disable=not xm.is_master_ordinal())\n    losses = []\n    for batch_idx, (images, targets) in enumerate(bar):  \n        \n        logits = model(images)\n        loss = criterion(logits, targets)\n        loss.backward()\n        xm.optimizer_step(optimizer)\n        loss_reduced = xm.mesh_reduce('loss_reduce',loss, lambda x: sum(x) / len(x))\n        bar.set_postfix(loss=loss_reduced.item())\n        losses.append(loss_reduced.item())\n\n    loss_train = np.mean(losses)\n\n    return loss_train\n\n\ndef valid_func(valid_loader, model, criterion, optimizer):\n    model.eval()\n    bar = tqdm(valid_loader, desc=\"Validation\", disable=not xm.is_master_ordinal())\n\n    PROB = []\n    TARGETS = []\n    losses = []\n    PREDS = []\n    \n    with torch.no_grad():\n        for batch_idx, (images, targets) in enumerate(bar):\n\n            #images, targets = images.to(device), targets.to(device)\n            logits = model(images)\n            PREDS += [logits.sigmoid()]\n            TARGETS += [targets.detach().cpu()]\n            loss = criterion(logits, targets)\n            loss_reduced= xm.mesh_reduce('loss_reduce',loss, lambda x: sum(x) / len(x)) \n            bar.set_postfix(loss=loss_reduced.item())\n            losses.append(loss_reduced.item())\n\n    loss_valid = np.mean(losses)\n    \n    PREDS = torch.cat(PREDS).cpu().numpy()\n    TARGETS = torch.cat(TARGETS).cpu().numpy()\n    roc_auc = macro_multilabel_auc(TARGETS, PREDS)\n    roc_auc_reduced = xm.mesh_reduce('roc_auc_reduce', roc_auc, lambda x: sum(x) / len(x))\n    xm.master_print(f'val_roc_auc = {roc_auc_reduced}')\n    \n    return loss_valid, roc_auc_reduced","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(rank, fold_id):\n    MX = xmp.MpModelWrapper(RANZCRResNet200D(out_dim=len(target_cols), pretrained=True))\n    device = xm.xla_device()\n    model = MX.to(device)\n\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.Adam(model.parameters(), lr=init_lr/warmup_factor**xm.xrt_world_size())\n    scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs, eta_min=1e-7)\n    scheduler_warmup = GradualWarmupSchedulerV2(optimizer, multiplier=10, total_epoch=warmup_epo, after_scheduler=scheduler_cosine)\n\n    df_train_this = df_train[df_train['fold'] != fold_id]\n    df_valid_this = df_train[df_train['fold'] == fold_id]\n\n    dataset_train = RANZERDataset(df_train_this, 'train', transform=transforms_train)\n    dataset_valid = RANZERDataset(df_valid_this, 'valid', transform=transforms_valid)\n\n    # special sampler needed for distributed/multi-core (divides dataset among the replicas/cores/devices)\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n            dataset_train,\n            num_replicas=xm.xrt_world_size(), #divide dataset among this many replicas\n            rank=xm.get_ordinal(), #which replica/device/core\n            shuffle=True)\n\n    # same as train but with valid data\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n            dataset_valid,\n            num_replicas=xm.xrt_world_size(),\n            rank=xm.get_ordinal(),\n            shuffle=False)\n\n    train_loader = torch.utils.data.DataLoader(dataset_train, \n                                               batch_size=batch_size,\n                                               sampler=train_sampler, \n                                               num_workers=num_workers)\n    valid_loader = torch.utils.data.DataLoader(dataset_valid, \n                                               batch_size=valid_batch_size,\n                                               sampler=valid_sampler, \n                                               num_workers=num_workers)\n\n    train_loader = pl.MpDeviceLoader(train_loader, device) # puts the train data onto the current TPU core\n    valid_loader = pl.MpDeviceLoader(valid_loader, device) # puts the valid data onto the current TPU core\n\n    log = {}\n    roc_auc_max = 0.\n    loss_min = 99999\n    not_improving = 0\n    \n    xm.master_print(f'Training fold {fold_id} for {n_epochs} epochs')\n    for epoch in range(1, n_epochs+1):\n        scheduler_warmup.step(epoch-1)\n        loss_train = train_func(train_loader, model, criterion, optimizer)\n        loss_valid, roc_auc = valid_func(valid_loader, model, criterion, optimizer)\n\n        log['loss_train'] = log.get('loss_train', []) + [loss_train]\n        log['loss_valid'] = log.get('loss_valid', []) + [loss_valid]\n        log['lr'] = log.get('lr', []) + [optimizer.param_groups[0][\"lr\"]]\n        log['roc_auc'] = log.get('roc_auc', []) + [roc_auc]\n\n        content = time.ctime() + ' ' + f'Fold {fold_id}, Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, loss_train: {loss_train:.5f}, loss_valid: {loss_valid:.5f}, roc_auc: {roc_auc:.6f}.'\n        xm.master_print(content)\n        not_improving += 1\n\n        if roc_auc > roc_auc_max:\n            xm.master_print(f'roc_auc_max ({roc_auc_max:.6f} --> {roc_auc:.6f}). Saving model ...')\n            xm.rendezvous('save_model')\n            xm.save(model.state_dict(), f'{model_dir}{kernel_type}_fold{fold_id}_best_AUC.pth')\n            roc_auc_max = roc_auc\n            not_improving = 0\n\n        if loss_valid < loss_min:\n            loss_min = loss_valid\n            xm.rendezvous('save_model')\n            xm.save(model.state_dict(), f'{model_dir}{kernel_type}_fold{fold_id}_best_loss.pth')\n\n        if not_improving == early_stop:\n            xm.master_print('Early Stopping...')\n            break\n\n        if debug: ## run only 1 epoch if in debug mode\n            break\n    \n    xm.rendezvous('save_model')\n    xm.save(model.state_dict(), f'{model_dir}{kernel_type}_fold{fold_id}_final.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(5):\n    start_time = time.time()\n    xmp.spawn(run, args=(i,), nprocs=8, start_method='fork')\n    print('time taken: ', time.time()-start_time)\n    print('==============================================================================')","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}