{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"
\n \n
","metadata":{}},{"cell_type":"markdown","source":"

Goal

\n\nThe scope of this notebook is to train people on how to train the baseline of a CoAt for this dataset.
\nCoAT currently has the highest submission, so give this notebook a try.
\nIf you really want a medal, you have to tweek hyperparameters and work hard.
\nThe [Inference notebook](https://www.kaggle.com/code/alincijov/inference-hubmap-coat) for this notebook.
\nIf you appreciate the work, don't forget to upvote !","metadata":{}},{"cell_type":"markdown","source":"

Installation

","metadata":{}},{"cell_type":"code","source":"!pip install -qq torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html\n!pip install -qq git+https://github.com/qubvel/segmentation_models.pytorch\n!pip install -qq timm==0.4.12\n!pip install -qq einops","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-09-01T23:41:06.546685Z","iopub.execute_input":"2022-09-01T23:41:06.547325Z","iopub.status.idle":"2022-09-01T23:45:02.361198Z","shell.execute_reply.started":"2022-09-01T23:41:06.547262Z","shell.execute_reply":"2022-09-01T23:45:02.359807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Libraries

","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nimport sys\nimport glob\n\nimport torch\nimport torch.nn as nn\nimport albumentations as A\n\nimport cv2\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport tqdm\nimport segmentation_models_pytorch as smp\nfrom sklearn.model_selection import StratifiedKFold\n\nimport tifffile as tiff\n\ntorch.backends.cudnn.benchmark = True","metadata":{"execution":{"iopub.status.busy":"2022-09-02T00:02:47.718959Z","iopub.execute_input":"2022-09-02T00:02:47.719564Z","iopub.status.idle":"2022-09-02T00:02:47.729464Z","shell.execute_reply.started":"2022-09-02T00:02:47.719514Z","shell.execute_reply":"2022-09-02T00:02:47.728224Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Configuration

","metadata":{}},{"cell_type":"code","source":"fold = 0\nnfolds = 5\nimsize = 384\ntrain_csv = '../input/hubmap-organ-segmentation/train.csv'\nBATCH_SIZE = 8\nDEVICE = ('cuda' if torch.cuda.is_available() else 'cpu')\nEPOCHS = 5\nNUM_WORKERS = 4\nSEED = 24\nTRAIN_PATH = '../input/mmsegmentation256x256/train/'\nMASK_PATH = '../input/mmsegmentation256x256/masks/'","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:55:04.251372Z","iopub.execute_input":"2022-09-01T23:55:04.251822Z","iopub.status.idle":"2022-09-01T23:55:04.266970Z","shell.execute_reply.started":"2022-09-01T23:55:04.251775Z","shell.execute_reply":"2022-09-01T23:55:04.265880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(seed=12):\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n np.random.seed(seed)\n torch.backends.cudnn.deterministic = True\nset_seed(12)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-09-01T23:45:06.490243Z","iopub.execute_input":"2022-09-01T23:45:06.490629Z","iopub.status.idle":"2022-09-01T23:45:06.499519Z","shell.execute_reply.started":"2022-09-01T23:45:06.490577Z","shell.execute_reply":"2022-09-01T23:45:06.498628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Dataset

","metadata":{}},{"cell_type":"code","source":"class HuBMAPDataset(torch.utils.data.Dataset):\n def __init__(self, fold=fold, train=True, tfms=None):\n self.train = train\n ids = pd.read_csv(train_csv).id.values\n labels = pd.read_csv(train_csv).organ.values\n kf = StratifiedKFold(n_splits=nfolds,random_state=SEED,shuffle=True)\n ids = (ids[list(kf.split(ids,labels))[fold][0 if train else 1]]).tolist()\n self.fnames = [fname for fname in os.listdir(TRAIN_PATH) if int(fname.split('_')[0]) in ids]\n self.image_size = imsize\n self.tfms = tfms\n \n def img2tensor(self, img,dtype:np.dtype=np.float32):\n if img.ndim==2 : img = np.expand_dims(img,2)\n img = np.transpose(img,(2,0,1)) # C , H , W\n return torch.from_numpy(img.astype(dtype, copy=False))\n \n def __len__(self):\n return len(self.fnames)\n \n def resize(self, img, interp):\n return cv2.resize(\n img, (self.image_size, self.image_size), interpolation=interp)\n \n def __getitem__(self, idx):\n fname = self.fnames[idx]\n img = cv2.cvtColor(cv2.imread(TRAIN_PATH + fname), cv2.COLOR_BGR2RGB)\n mask = cv2.imread((MASK_PATH + fname),cv2.IMREAD_GRAYSCALE)\n if self.tfms is not None:\n augmented = self.tfms(image=img,mask=mask)\n img,mask = augmented['image'],augmented['mask']\n return self.img2tensor(self.resize(img , cv2.INTER_NEAREST)) , self.img2tensor(self.resize(mask , cv2.INTER_NEAREST))","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:45:06.501074Z","iopub.execute_input":"2022-09-01T23:45:06.501544Z","iopub.status.idle":"2022-09-01T23:45:06.515861Z","shell.execute_reply.started":"2022-09-01T23:45:06.501506Z","shell.execute_reply":"2022-09-01T23:45:06.515074Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Transformer

","metadata":{}},{"cell_type":"code","source":"def transformer(p=1.0):\n return A.Compose([\n A.HorizontalFlip(),\n A.VerticalFlip(),\n A.RandomRotate90(),\n A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.9, \n border_mode=cv2.BORDER_REFLECT),\n A.OneOf([\n A.OpticalDistortion(p=0.3),\n A.GridDistortion(p=.1),\n A.PiecewiseAffine(p=0.3),\n ], p=0.3),\n A.OneOf([\n A.HueSaturationValue(10,15,10),\n A.CLAHE(clip_limit=2),\n A.RandomBrightnessContrast(), \n ], p=0.3),\n ], p=p)","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:53:41.142393Z","iopub.execute_input":"2022-09-01T23:53:41.142777Z","iopub.status.idle":"2022-09-01T23:53:41.150909Z","shell.execute_reply.started":"2022-09-01T23:53:41.142747Z","shell.execute_reply":"2022-09-01T23:53:41.149466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Viewing Data

","metadata":{}},{"cell_type":"code","source":"ds = HuBMAPDataset(tfms=transformer())\ndl = torch.utils.data.DataLoader(ds,batch_size=64,shuffle=False,num_workers=NUM_WORKERS)\nit = iter(dl)\nimgs,masks = next(it)","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:53:41.605654Z","iopub.execute_input":"2022-09-01T23:53:41.606050Z","iopub.status.idle":"2022-09-01T23:53:53.409722Z","shell.execute_reply.started":"2022-09-01T23:53:41.606018Z","shell.execute_reply":"2022-09-01T23:53:53.407777Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(16,16))\nfor i,(img,mask) in enumerate(zip(imgs,masks)):\n img = ((img.permute(1,2,0))).numpy().astype(np.uint8) # H , W , C\n plt.subplot(8,8,i+1)\n plt.imshow(img,vmin=0,vmax=255)\n plt.imshow(mask.squeeze().numpy(), alpha=0.2)\n plt.axis('off')\n plt.subplots_adjust(wspace=None, hspace=None)\n \ndel ds,dl,imgs,masks","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:53:53.421680Z","iopub.execute_input":"2022-09-01T23:53:53.430480Z","iopub.status.idle":"2022-09-01T23:54:08.295140Z","shell.execute_reply.started":"2022-09-01T23:53:53.430418Z","shell.execute_reply":"2022-09-01T23:54:08.293880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Importing CoAT

","metadata":{}},{"cell_type":"code","source":"sys.path.append('../input/hubmap-coat/')\n\nfrom coat import *\nfrom daformer import *\nfrom helper import *","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.296726Z","iopub.execute_input":"2022-09-01T23:54:08.298103Z","iopub.status.idle":"2022-09-01T23:54:08.332736Z","shell.execute_reply.started":"2022-09-01T23:54:08.298063Z","shell.execute_reply":"2022-09-01T23:54:08.331764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Modifying CoAT

","metadata":{}},{"cell_type":"code","source":"class Net(nn.Module):\n \n def __init__(self,\n encoder=coat_lite_medium,\n decoder=daformer_conv3x3,\n encoder_cfg={},\n decoder_cfg={},\n ):\n \n super(Net, self).__init__()\n decoder_dim = decoder_cfg.get('decoder_dim', 320)\n\n self.encoder = encoder\n\n self.rgb = RGB()\n\n encoder_dim = self.encoder.embed_dims\n # [64, 128, 320, 512]\n\n self.decoder = decoder(\n encoder_dim=encoder_dim,\n decoder_dim=decoder_dim,\n )\n self.logit = nn.Sequential(\n nn.Conv2d(decoder_dim, 1, kernel_size=1),\n nn.Upsample(scale_factor = 4, mode='bilinear', align_corners=False),\n )\n\n def forward(self, batch):\n\n x = self.rgb(batch)\n\n B, C, H, W = x.shape\n encoder = self.encoder(x)\n\n last, decoder = self.decoder(encoder)\n logit = self.logit(last)\n\n output = {}\n probability_from_logit = torch.sigmoid(logit)\n output['probability'] = probability_from_logit\n\n return output","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.335285Z","iopub.execute_input":"2022-09-01T23:54:08.335910Z","iopub.status.idle":"2022-09-01T23:54:08.345045Z","shell.execute_reply.started":"2022-09-01T23:54:08.335870Z","shell.execute_reply":"2022-09-01T23:54:08.344179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_model():\n encoder = coat_lite_medium()\n checkpoint = '../input/hubmap-coat-medium/coat_lite_medium_384x384_f9129688.pth'\n checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage)\n state_dict = checkpoint['model']\n encoder.load_state_dict(state_dict,strict=False)\n \n net = Net(encoder=encoder).cuda()\n \n return net","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.346401Z","iopub.execute_input":"2022-09-01T23:54:08.346875Z","iopub.status.idle":"2022-09-01T23:54:08.356254Z","shell.execute_reply.started":"2022-09-01T23:54:08.346819Z","shell.execute_reply":"2022-09-01T23:54:08.355348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Loss functions

","metadata":{}},{"cell_type":"code","source":"class CustomLoss(nn.Module):\n def __init__(self):\n super(CustomLoss,self).__init__()\n self.diceloss = smp.losses.DiceLoss(mode='binary')\n self.binloss = smp.losses.SoftBCEWithLogitsLoss(reduction = 'mean' , smooth_factor = 0.1)\n\n def forward(self, output, mask):\n dice = self.diceloss(outputs,mask)\n bce = self.binloss(outputs , mask)\n loss = dice * 0.7 + bce * 0.3\n return loss","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.357841Z","iopub.execute_input":"2022-09-01T23:54:08.358252Z","iopub.status.idle":"2022-09-01T23:54:08.367743Z","shell.execute_reply.started":"2022-09-01T23:54:08.358218Z","shell.execute_reply":"2022-09-01T23:54:08.366611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DiceCoef(nn.Module):\n def __init__(self, weight=None, size_average=True):\n super().__init__()\n\n def forward(self, y_pred, y_true, smooth=1.):\n y_true = y_true.view(-1)\n y_pred = y_pred.view(-1)\n \n #Round off y_pred\n y_pred = torch.round((y_pred - y_pred.min()) / (y_pred.max() - y_pred.min()))\n \n intersection = (y_true * y_pred).sum()\n dice = (2.0*intersection + smooth)/(y_true.sum() + y_pred.sum() + smooth)\n \n return dice","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.368908Z","iopub.execute_input":"2022-09-01T23:54:08.369845Z","iopub.status.idle":"2022-09-01T23:54:08.381264Z","shell.execute_reply.started":"2022-09-01T23:54:08.369799Z","shell.execute_reply":"2022-09-01T23:54:08.380369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Additionals

","metadata":{}},{"cell_type":"code","source":"def plot_df(df):\n fig,ax = plt.subplots(1,2,figsize=(15,5))\n ax[0].plot(df['Train_loss'])\n ax[0].plot(df['Val_loss'])\n ax[0].legend()\n ax[0].set_title('Loss')\n ax[1].plot(df['Train_Dice'])\n ax[1].plot(df['Val_Dice'])\n ax[1].legend()\n ax[1].set_title('Dice')","metadata":{"execution":{"iopub.status.busy":"2022-09-01T23:54:08.382551Z","iopub.execute_input":"2022-09-01T23:54:08.383088Z","iopub.status.idle":"2022-09-01T23:54:08.392400Z","shell.execute_reply.started":"2022-09-01T23:54:08.382948Z","shell.execute_reply":"2022-09-01T23:54:08.391430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

Training

","metadata":{}},{"cell_type":"code","source":"print(f\"Running on device : {DEVICE}\" )\nfor fold in range(nfolds):\n \n val_losses = []\n losses = []\n train_scores=[]\n val_scores = []\n best_loss = 999\n best_score = 0\n \n ds_train = HuBMAPDataset(fold=fold, train=True, tfms=transformer())\n ds_val = HuBMAPDataset(fold=fold, train=False)\n \n dataloader_train = torch.utils.data.DataLoader(ds_train,batch_size=BATCH_SIZE, shuffle=True,num_workers=NUM_WORKERS)\n dataloader_val = torch.utils.data.DataLoader(ds_val,batch_size=BATCH_SIZE, shuffle=False,num_workers=NUM_WORKERS)\n \n model = init_model().to(DEVICE)\n \n optimizer = torch.optim.Adam([\n {'params': model.decoder.parameters(), 'lr': 5e-5}, \n {'params': model.encoder.parameters(), 'lr': 8e-5}, \n ])\n \n scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optimizer, pct_start=0.1, div_factor=1e3, \n max_lr=1e-3, epochs=EPOCHS, steps_per_epoch=len(dataloader_train))\n \n loss_func = CustomLoss()\n dice_coe = DiceCoef()\n \n print(f\"######## FOLD: {fold} ##############\")\n \n for epoch in tqdm.notebook.tqdm(range(EPOCHS)):\n \n ### Train ###########################################################################################\n \n model.train()\n train_loss = 0\n score = 0\n \n for data in tqdm.notebook.tqdm(dataloader_train ,total = len(dataloader_train)):\n optimizer.zero_grad()\n img, mask = data\n img = img.to(DEVICE)\n mask = mask.to(DEVICE)\n \n outputs = model(img)['probability'] \n\n loss = loss_func(outputs, mask)\n loss.backward()\n optimizer.step()\n scheduler.step()\n train_loss += loss.item()\n score += dice_coe(outputs,mask).item()\n \n train_loss /= len(dataloader_train)\n score /= len(dataloader_train)\n losses.append(train_loss)\n train_scores.append(score)\n print(f\"FOLD: {fold}, EPOCH: {epoch + 1}, train_loss: {train_loss} , Dice coe : {score} \") #\n \n \n gc.collect()\n torch.cuda.empty_cache()\n \n ### Validation ####################################################################################\n \n model.eval()\n \n with torch.no_grad():\n \n valid_loss = 0\n val_score = 0\n \n for data in dataloader_val:\n \n img, mask = data\n img = img.to(DEVICE)\n mask = mask.to(DEVICE)\n\n outputs = model(img)['probability']\n\n loss = loss_func(outputs, mask)\n valid_loss += loss.item()\n val_score += dice_coe(outputs,mask).item()\n \n valid_loss /= len(dataloader_val)\n val_losses.append(valid_loss)\n \n val_score /= len(dataloader_val)\n val_scores.append(val_score)\n \n print(f\"FOLD: {fold}, EPOCH: {epoch + 1}, valid_loss: {valid_loss} , Val Dice COE : {val_score}\") #\n \n gc.collect()\n torch.cuda.empty_cache()\n \n if val_score > best_score:\n best_score = val_score\n torch.save(model.state_dict(), f\"/kaggle/working/FOLD{fold}_best_score.pth\")\n print(f\"Saved model for best score : FOLD{fold}_best_score.pth\")\n \n if valid_loss < best_loss:\n best_loss = valid_loss\n torch.save(model.state_dict(), f\"/kaggle/working/FOLD{fold}_best_loss.pth\")\n print(f\"Saved model for best loss : FOLD{fold}_best_loss.pth\") \n\n \n column_names = ['Train_loss','Val_loss','Train_Dice','Val_Dice']\n df = pd.DataFrame(np.stack([losses,val_losses,train_scores,val_scores],axis=1),columns=column_names)\n print(f\" ################# FOLD {fold} #####################\")\n plot_df(df)\n plt.show(block=False)\n df.to_csv(f\"logs_fold{fold}.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-09-02T00:02:54.255760Z","iopub.execute_input":"2022-09-02T00:02:54.256140Z"},"trusted":true},"execution_count":null,"outputs":[]}]}