{"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

Brain Tumor Video Classification using PyTorch + W&B Tracking ✨

\n
\n\n

\nI'm approaching this problem as a video classification problem, using all images in the FLAIR folder. Thanks to Ayush Thakur's Notebook doing the same but in Tensorflow.
\nI have also used Weights and Biases tracking to keep track of the training process and the experiments I am conducting.\n

\n\n
\n You can upvote this kernel, if you found it useful!\n
","metadata":{}},{"cell_type":"markdown","source":"
\n

📔 Imports and Installation

\n
","metadata":{}},{"cell_type":"code","source":"%%sh\npip install -q timm\npip install -q einops\npip install -q rich\npip install -q wandb --upgrade;","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport sys\nimport re\nimport gc\nimport platform\nimport random\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n\nimport einops\n\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import StratifiedKFold, train_test_split\n\nimport timm\nimport glob\nimport cv2\n\nfrom rich import print as _pprint\nfrom rich.progress import track\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport wandb\n\nimport warnings\nwarnings.simplefilter('ignore')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

⛽ Utility Functions

\n
","metadata":{}},{"cell_type":"code","source":"def get_patient_id(patient_id):\n if patient_id < 10:\n return '0000'+str(patient_id)\n elif patient_id >= 10 and patient_id < 100:\n return '000'+str(patient_id)\n elif patient_id >= 100 and patient_id < 1000:\n return '00'+str(patient_id)\n else:\n return '0'+str(patient_id)\n\ndef get_path(row):\n patient_id = get_patient_id(row.BraTS21ID)\n return f'../input/rsna-miccai-png/train/{patient_id}/FLAIR/'\n\ndef wandb_log(**kwargs):\n \"\"\"\n Logs a key-value pair to W&B\n \"\"\"\n for k, v in kwargs.items():\n wandb.log({k: v})\n \ndef cprint(string):\n \"\"\"\n Utility function for beautiful colored printing.\n \"\"\"\n _pprint(f\"[black]{string}[/black]\")","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

🚀 Config Dictionary and W&B Integration

\n
","metadata":{}},{"cell_type":"code","source":"Config = dict(\n MAX_FRAMES = 12,\n EPOCHS = 5,\n LR = 2e-4,\n IMG_SIZE = (224, 224),\n FEATURE_EXTRACTOR = 'resnext50_32x4d',\n DR_RATE = 0.35,\n NUM_CLASSES = 1,\n RNN_HIDDEN_SIZE = 100,\n RNN_LAYERS = 1,\n TRAIN_BS = 4,\n VALID_BS = 4,\n NUM_WORKERS = 4,\n infra = \"Kaggle\",\n competition = 'rsna_miccai',\n _wandb_kernel = 'tanaym'\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To login to W&B, you can use below snippet.\n\n```python\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n```\nMake sure you have your W&B key stored as `WANDB_API_KEY` under Add-ons -> Secrets\n\nYou can view [this](https://www.kaggle.com/ayuraj/experiment-tracking-with-weights-and-biases) notebook to learn more about W&B tracking.\n\nIf you don't want to login to W&B, the kernel will still work and log everything to W&B in anonymous mode.","metadata":{}},{"cell_type":"code","source":"from kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WBWB\")\n\nwandb.login(key=wb_key)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run = wandb.init(\n project='pytorch',\n config=Config,\n group='vision',\n job_type='train',\n anonymous='allow'\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Augments:\n \"\"\"\n Contains Train, Validation Augments\n \"\"\"\n train_augments = A.Compose([\n ToTensorV2(p=1.0),\n ],p=1.)\n \n valid_augments = A.Compose([\n ToTensorV2(p=1.0),\n ], p=1.)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

💻 Custom Dataset Class

\n
\n\n
\n 📌 In this custom Dataset, I am essentially reading \"MAX_FRAMES\" number of images from a patient's FLAIR folder and making list of those frames and converting it to torch tensor.\n
","metadata":{}},{"cell_type":"code","source":"class RSNADataset(Dataset):\n def __init__(self, df, augments=None, is_test=False):\n self.df = df\n self.augments = augments\n self.is_test = is_test\n \n def __getitem__(self, idx):\n row = self.df.loc[idx]\n paths = self.getPaths(row)\n frames = []\n for path in paths:\n img = cv2.imread(path)\n img = cv2.resize(img, Config['IMG_SIZE'])\n img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n frames.append(img)\n\n frames_tr = np.stack(frames, axis=2)\n \n if self.augments:\n for frame in frames:\n frame = self.augments(image=frame)['image']\n frames_tr.append(frame)\n \n if self.is_test:\n return frames_tr\n else:\n label = torch.tensor(row['MGMT_value']).float()\n return frames_tr, label\n \n def __len__(self):\n return len(self.df)\n \n def getPaths(self, row):\n paths = glob.glob(row['path'] + '*.png')\n sortedPaths = self.sort(paths)\n maxWindowStart = len(sortedPaths) - Config['MAX_FRAMES']\n start = 0 # np.random.randint(1, maxWindowStart)\n paths = sortedPaths[start:Config['MAX_FRAMES']]\n \n return paths\n \n def sort(self, entry):\n # https://stackoverflow.com/a/2669120/7636462\n convert = lambda text: int(text) if text.isdigit() else text \n alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]\n \n return sorted(entry, key = alphanum_key)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

📈 Model Class with ResNext Backbone

\n
","metadata":{}},{"cell_type":"code","source":"class ResNextModel(nn.Module):\n def __init__(self):\n super(ResNextModel, self).__init__()\n self.backbone = timm.create_model(Config['FEATURE_EXTRACTOR'], in_chans=1)\n def forward(self, x):\n return self.backbone(x)\n\nclass Identity(nn.Module):\n def __init__(self):\n super(Identity, self).__init__()\n def forward(self, x):\n return x\n\nclass RSNAModel(nn.Module):\n def __init__(self, pretrained=True):\n super(RSNAModel, self).__init__()\n self.backbone = ResNextModel()\n num_features = self.backbone.backbone.fc.in_features\n \n self.backbone.backbone.fc = Identity()#将backbone模型的最后一个全连接层替换为一个恒等映射模块，从而移除了最后一层全连接层。\n self.dropout= nn.Dropout(Config['DR_RATE'])\n self.rnn = nn.LSTM(num_features, Config['RNN_HIDDEN_SIZE'], Config['RNN_LAYERS'])\n self.fc1 = nn.Linear(Config['RNN_HIDDEN_SIZE'], Config['NUM_CLASSES'])\n \n def forward(self, x):\n b_z, fr, h, w = x.shape\n ii = 0\n in_pass = x[:, ii].unsqueeze_(1)\n y = self.backbone((in_pass))\n output, (hn, cn) = self.rnn(y.unsqueeze(1))\n for ii in range(1, fr):\n y = self.backbone((x[:, ii].unsqueeze_(1)))\n out, (hn, cn) = self.rnn(y.unsqueeze(1), (hn, cn))\n out = self.dropout(out[:, -1])\n out = self.fc1(out)\n return out","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

🏴‍☠️ Training and Validation Functions

\n
","metadata":{}},{"cell_type":"code","source":"def train_one_epoch(model, train_dataloader, optimizer, loss_fn, epoch, device, log_wandb=True, verbose=False):\n \"\"\"\n Trains model for one epoch\n \"\"\"\n model.train()\n running_loss = 0\n prog_bar = tqdm(enumerate(train_dataloader), total=len(train_dataloader))\n for batch, (frames, targets) in prog_bar:\n optimizer.zero_grad()\n \n frames = frames.to(device, torch.float)\n targets = targets.to(device, torch.float)\n \n # Re arrange the frames in the format our model wants to recieve\n frames = einops.rearrange(frames, 'b h w f -> b f h w')\n\n preds = model(frames).view(-1)\n loss = loss_fn(preds, targets)\n \n loss.backward()\n optimizer.step()\n \n loss_item = loss.item()\n running_loss += loss_item\n \n prog_bar.set_description(f\"loss: {loss_item:.4f}\")\n \n if log_wandb == True:\n wandb_log(\n batch_train_loss=loss_item\n )\n \n if verbose == True and batch % 20 == 0:\n print(f\"Batch: {batch}, Loss: {loss_item}\")\n \n avg_loss = running_loss / len(train_dataloader)\n \n return avg_loss\n\n@torch.no_grad()\ndef valid_one_epoch(model, valid_dataloader, loss_fn, epoch, device, log_wandb=True, verbose=False):\n \"\"\"\n Validates the model for one epoch\n \"\"\"\n model.eval()\n running_loss = 0\n prog_bar = tqdm(enumerate(valid_dataloader), total=len(valid_dataloader))\n for batch, (frames, targets) in prog_bar:\n frames = frames.to(device, torch.float)\n targets = targets.to(device, torch.float)\n \n # Re arrange the frames in the format our model wants to recieve\n frames = einops.rearrange(frames, 'b h w f -> b f h w')\n preds = model(frames).view(-1)\n loss = loss_fn(preds, targets)\n \n loss_item = loss.item()\n running_loss += loss_item\n \n prog_bar.set_description(f\"val_loss: {loss_item:.4f}\")\n \n if log_wandb == True:\n wandb_log(\n batch_val_loss=loss_item\n )\n \n if verbose == True and batch % 10 == 0:\n print(f\"Batch: {batch}, Loss: {loss_item}\")\n \n avg_val_loss = running_loss / len(valid_dataloader)\n \n return avg_val_loss","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n

🏗 Training and Validating the Model

\n
","metadata":{}},{"cell_type":"code","source":"log_wandb = True\nif torch.cuda.is_available():\n print(\"Using GPU: {}\\n\".format(torch.cuda.get_device_name()))\n device = torch.device('cuda')\nelse:\n print(\"\\nGPU not found. Using CPU: {}\\n\".format(platform.processor()))\n device = torch.device('cpu')\n\n\n# Load training csv file\ndf = pd.read_csv('../input/rsna-miccai-brain-tumor-radiogenomic-classification/train_labels.csv')\ndf['path'] = df.apply(lambda row: get_path(row), axis=1)\n\n# Removing two patient ids from the dataframe since there are not FLAIR directories for these ids. \ndf = df.loc[df.BraTS21ID!=109]\ndf = df.loc[df.BraTS21ID!=709]\ndf = df.reset_index(drop=True)\n\ntrain_df, valid_df = train_test_split(df, test_size=0.1, stratify=df.MGMT_value.values)\ntrain_df = train_df.reset_index(drop=True)\nvalid_df = valid_df.reset_index(drop=True)\n\nprint(f'Size of Training Set: {len(train_df)}, Validation Set: {len(valid_df)}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nmodel = RSNAModel()\nmodel = model.to(device)\noptimizer = torch.optim.AdamW(model.parameters(), lr=Config['LR'])\n\ntrain_loss_fn = nn.BCEWithLogitsLoss()\nvalid_loss_fn = nn.BCEWithLogitsLoss()\n\nprint(f\"\\nUsing Backbone: {Config['FEATURE_EXTRACTOR']}\")\nprint(train_df)\ntrain_data = RSNADataset(train_df)\nvalid_data = RSNADataset(valid_df)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#对每个dataset中的数据进行检查是否尺寸已经归一化\nfor i in range(5):\n for d,m in enumerate(train_data):\n data = m[0]\n label = m[1]\n print('data:',data.shape)\n print('label',label)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\ntrain_loader = DataLoader(\n train_data,\n batch_size=Config['TRAIN_BS'], \n shuffle=True,\n num_workers=Config['NUM_WORKERS']\n)\n\nvalid_loader = DataLoader(\n valid_data, \n batch_size=Config['VALID_BS'], \n shuffle=False,\n num_workers=Config['NUM_WORKERS']\n)\n\ncurrent_loss = 1000\nfor epoch in range(Config['EPOCHS']):\n print(f\"\\n{'--'*8} EPOCH: {epoch+1} {'--'*8}\\n\")\n\n train_loss = train_one_epoch(model, train_loader, optimizer, train_loss_fn, epoch=epoch, device=device, log_wandb=log_wandb)\n\n valid_loss = valid_one_epoch(model, valid_loader, valid_loss_fn, epoch=epoch, device=device, log_wandb=log_wandb)\n\n print(f\"val_loss: {valid_loss:.4f}\")\n\n if log_wandb == True:\n wandb_log(\n train_loss=train_loss,\n valid_loss=valid_loss\n )\n\n if valid_loss < current_loss:\n current_loss = valid_loss\n torch.save(model.state_dict(), f\"model_{Config['FEATURE_EXTRACTOR']}.pt\")","metadata":{"execution":{"iopub.status.busy":"2023-06-09T07:08:09.889807Z","iopub.execute_input":"2023-06-09T07:08:09.890206Z","iopub.status.idle":"2023-06-09T07:25:18.080425Z","shell.execute_reply.started":"2023-06-09T07:08:09.890167Z","shell.execute_reply":"2023-06-09T07:25:18.074882Z"}}},{"cell_type":"markdown","source":"

View the complete dashboard here ✨

\n\n![Results](https://i.imgur.com/gcn59xp.gif)","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}