{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"vscode":{"interpreter":{"hash":"a3ea94c989d1c4842e855ce4d11352d386862b4385dba776343ae856b881a440"}},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":39272,"databundleVersionId":4629629,"sourceType":"competition"}],"dockerImageVersionId":30702,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# PyTorch Multi-GPU trainer using 🤗 accelerate + Mixed Precision + W&B Logging\nThis training notebook uses HuggingFace accelerate with mixed precision to train a VIT model on 2x T4 GPUs. The current performance is quite good but I am working on optimizing it.\n\nI am also using Probabilistic F1 score with a `beta=0.5` (in hopes that it will penalize false positive). The entire training pipeline is working well and you could fork the notebook and play the model and other hyperparameters.\n\n**Feel free to fork and change the models and do some preprocessing, but if you do please leave an upvote :)**","metadata":{}},{"cell_type":"markdown","source":"
\n\n
","metadata":{}},{"cell_type":"markdown","source":"## Installation and Imports","metadata":{}},{"cell_type":"code","source":"%%capture\n! pip install timm\n! pip install einops\n! pip install git+https://github.com/huggingface/accelerate","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-04-23T23:40:35.677044Z","iopub.execute_input":"2024-04-23T23:40:35.678178Z","iopub.status.idle":"2024-04-23T23:40:36.912813Z","shell.execute_reply.started":"2024-04-23T23:40:35.67812Z","shell.execute_reply":"2024-04-23T23:40:36.911521Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport glob\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nimport wandb\n\nimport timm\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange\nfrom torch.utils.data import DataLoader, Dataset\nimport sklearn\nfrom sklearn.metrics import f1_score\nfrom sklearn.model_selection import KFold\n\nfrom accelerate.tracking import GeneralTracker\nfrom accelerate import Accelerator, notebook_launcher","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.status.busy":"2024-04-23T23:40:36.915637Z","iopub.execute_input":"2024-04-23T23:40:36.916386Z","iopub.status.idle":"2024-04-23T23:40:36.97011Z","shell.execute_reply.started":"2024-04-23T23:40:36.916335Z","shell.execute_reply":"2024-04-23T23:40:36.967725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utility functions","metadata":{}},{"cell_type":"code","source":"def probabilistic_f1(labels, predictions, beta=0.5):\n y_true_count = 0\n ctp = 0\n cfp = 0\n\n for idx in range(len(labels)):\n prediction = min(max(predictions[idx], 0), 1)\n if (labels[idx]):\n y_true_count += 1\n ctp += prediction\n #cfp += 1 - prediction\n else:\n cfp += prediction\n\n beta_squared = beta * beta\n c_precision = ctp / (ctp + cfp)\n c_recall = ctp / y_true_count\n if (c_precision > 0 and c_recall > 0):\n result = (1 + beta_squared) * (c_precision * c_recall) / (beta_squared * c_precision + c_recall)\n return result\n else:\n return 0\n \ndef wandb_log(**kwargs):\n for k, v in kwargs.items():\n wandb.log({k: v})","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.97163Z","iopub.status.idle":"2024-04-23T23:40:36.972453Z","shell.execute_reply.started":"2024-04-23T23:40:36.972159Z","shell.execute_reply":"2024-04-23T23:40:36.972184Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def probabilistic_f1(truth, probability, threshold=0.5):\n\n if threshold is None:\n predict = [probability]\n else:\n predict = [\n (probability > t).astype(np.float32) for t in threshold\n ]\n\n f1score = []\n for p in predict:\n tp = ((p >= 0.5) & (truth >= 0.5)).sum()\n fp = ((p >= 0.5) & (truth < 0.5)).sum()\n fn = ((p < 0.5) & (truth >= 0.5)).sum()\n\n recall = tp / (tp + fn + 1e-3)\n precision = tp / (tp + fp + 1e-3)\n f1 = 2 * recall * precision / (recall + precision + 1e-3)\n f1score.append(f1)\n f1score = np.array(f1score)\n return f1score","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.973974Z","iopub.status.idle":"2024-04-23T23:40:36.974973Z","shell.execute_reply.started":"2024-04-23T23:40:36.974525Z","shell.execute_reply":"2024-04-23T23:40:36.974548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config and W&B","metadata":{}},{"cell_type":"code","source":"Config = {\n 'TRAIN_BS': 32,\n 'VALID_BS': 32,\n 'MODEL_NAME': 'vit_base_patch16_224',\n 'NUM_WORKERS': 8,\n 'PARENT_PATH': '/data/yijia/code2/code/data/rsna/train_images_processed_512/',\n 'FILE_PATH': '/data/yijia/code2/code/data/rsna/train.csv',\n 'LOSS': 'BCEWithLogitsLoss',\n 'EVAL_METRIC': 'F1',\n 'NB_EPOCHS': 3,\n 'SPLITS': 5,\n 'T_0': 20,\n 'η_min': 1e-4,\n 'fc_dropout': 0.2,\n 'betas': (0.9, 0.999),\n 'N_LABELS': 1,\n 'LR': 2e-4,\n 'competition': 'rsna_mammography',\n '_wandb_kernel': 'tanaym',\n 'logging_dir' : '/data/yijia/code2/code/data/rsna/',\n}","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.976667Z","iopub.status.idle":"2024-04-23T23:40:36.977503Z","shell.execute_reply.started":"2024-04-23T23:40:36.977193Z","shell.execute_reply":"2024-04-23T23:40:36.977241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### About W&B:\n
\"Weights

\n

WandB is a developer tool for companies turn deep learning research projects into deployed software by helping teams track their models, visualize model performance and easily automate training and improving models.\nWe will use their tools to log hyperparameters and output metrics from your runs, then visualize and compare results and quickly share findings with your colleagues.

","metadata":{}},{"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":"markdown","source":"I am making a Custom W&B tracker for easy experiment tracking during training","metadata":{}},{"cell_type":"code","source":"# Start W&B logging\n# W&B Login\n\n\nwandb.login(key='7ac1917aff711d9d9f5d28f4f54627b219144f09')","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.978947Z","iopub.status.idle":"2024-04-23T23:40:36.979703Z","shell.execute_reply.started":"2024-04-23T23:40:36.979451Z","shell.execute_reply":"2024-04-23T23:40:36.979474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset\n\nI am only loading images for now but I am working on extending that to add the meta features available in the dataset","metadata":{}},{"cell_type":"code","source":"class RSNAData(Dataset):\n def __init__(self, df, img_folder, augments=None, is_test=False):\n self.df = df\n self.is_test = is_test\n self.augments = augments\n self.img_folder = img_folder\n \n def __getitem__(self, idx):\n img_path = os.path.join(self.img_folder, self.df['img_name'][idx])\n img = cv2.imread(img_path)\n img = cv2.resize(img, (224, 224))\n if self.augments:\n img = self.augments(image=img)['image']\n img = torch.tensor(img, dtype=torch.float)\n # Rearrange the image dimensions so that channels are first in format\n # This is because VIT Model requires Channels (c) to come first\n img = rearrange(img, 'h w c -> c h w')\n \n if not self.is_test:\n target = self.df['cancer'][idx]\n target = torch.tensor(target, dtype=torch.float)\n return (img, target)\n return (img)\n \n def __len__(self):\n return len(self.df)","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.981133Z","iopub.status.idle":"2024-04-23T23:40:36.981929Z","shell.execute_reply.started":"2024-04-23T23:40:36.981669Z","shell.execute_reply":"2024-04-23T23:40:36.981692Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model\nJust using a simple VIT model, you can fork and extend it to your preference","metadata":{}},{"cell_type":"code","source":"class VITModel(nn.Module):\n def __init__(self, config, pretrained=True):\n super(VITModel, self).__init__()\n self.backbone = timm.create_model(config['MODEL_NAME'], pretrained=pretrained)\n self.backbone.head = nn.Linear(self.backbone.head.in_features, config['N_LABELS'])\n def forward(self, x):\n return self.backbone(x)","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.983308Z","iopub.status.idle":"2024-04-23T23:40:36.984048Z","shell.execute_reply.started":"2024-04-23T23:40:36.983783Z","shell.execute_reply":"2024-04-23T23:40:36.983805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below function uses accelerator scope so that we only download our model once and use it from cache later","metadata":{}},{"cell_type":"code","source":"def init_model(accelerator, config, pretrained=True):\n with accelerator.main_process_first():\n model = VITModel(config=config, pretrained=pretrained)\n return model","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.985505Z","iopub.status.idle":"2024-04-23T23:40:36.986307Z","shell.execute_reply.started":"2024-04-23T23:40:36.986012Z","shell.execute_reply":"2024-04-23T23:40:36.986035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Fit function\n\nThis function houses the entire training and validation code","metadata":{}},{"cell_type":"code","source":"def fit(model, fold, epochs, train_loader, valid_loader, optimizer, train_loss_fn, valid_loss_fn, accelerator):\n for epx in range(epochs):\n accelerator.print(f\"{'='*20} Epoch: {epx+1} {'='*20}\\n\")\n # Training part of the model\n model.train()\n avg_loss = 0\n for idx, (images, targets) in enumerate(train_loader):\n outputs = model(images).view(-1)\n\n loss = train_loss_fn(outputs, targets)\n\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad(set_to_none=True)\n\n avg_loss += loss.item()\n if idx % 100 == 0:\n accelerator.print(f\"batch: {idx}, train_loss: {loss.item():.4f}\")\n\n avg_loss = avg_loss / len(train_loader)\n accelerator.log({'train_loss': avg_loss})\n accelerator.print(f\"\\nEpoch: {epx+1} / {epochs} | Training Loss: {avg_loss:.4f}\\n\")\n\n # Validation part of the model\n model.eval()\n avg_loss = 0\n all_outputs, all_targets = [], []\n with torch.no_grad():\n for idx, (images, targets) in enumerate(valid_loader):\n outputs = model(images).view(-1)\n loss = valid_loss_fn(outputs, targets)\n\n if idx % 10 == 0:\n accelerator.print(f\"batch: {idx}, valid_loss: {loss.item():.4f}\")\n avg_loss += loss.item()\n\n outputs, targets = accelerator.gather_for_metrics((\n outputs, targets\n ))\n all_outputs.extend(torch.sigmoid(outputs).cpu().detach().tolist())\n all_targets.extend(targets.cpu().detach().tolist())\n \n prob_f1_score = probabilistic_f1(all_targets, all_outputs, threshold=0.5)\n avg_loss = avg_loss / len(valid_loader)\n accelerator.log({'val_loss': avg_loss, 'val_prob_f1': prob_f1_score})\n accelerator.print(f\"\\nEpoch: {epx+1} / {epochs} | Validation Loss: {avg_loss:.4f}\")\n accelerator.print(f\"\\nF1 Score for epoch: {epx+1} : {prob_f1_score:.4f}\\n\")\n \n # Save the model\n accelerator.wait_for_everyone() \n model = accelerator.unwrap_model(model)\n accelerator.save(model, f\"fold_{fold}_model.pth\")\n \n # End wandb run\n wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.988457Z","iopub.status.idle":"2024-04-23T23:40:36.990093Z","shell.execute_reply.started":"2024-04-23T23:40:36.989746Z","shell.execute_reply":"2024-04-23T23:40:36.989775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Run function + Training\n\nThe run function is the wrapping function that will invoke the `fit()` function by passing in the data and all the necessary params","metadata":{}},{"cell_type":"code","source":"# Training cell\ndef run(df, config):\n # We are splitting the data naively for now since StratifiedGroupKFold is giving some hard time in multi-GPU setup.\n split = 0.95\n df = df.sample(frac=1).reset_index(drop=True)\n train_samples = int(len(df) * split)\n train_df = df[:train_samples+1].reset_index(drop=True)\n valid_df = df[train_samples:].reset_index(drop=True)\n \n # Initialize Accelerator with Mixed Precision for training. Also init optimizer and loss functions\n accelerator = Accelerator(mixed_precision='fp16', log_with=\"wandb\")\n accelerator.init_trackers(\"rsna_mammography_pytorch\", config=Config)\n model = init_model(accelerator, Config)\n optimizer = torch.optim.Adam(params=model.parameters(), lr=Config['LR'])\n train_loss_fn, valid_loss_fn = nn.BCEWithLogitsLoss(), nn.BCEWithLogitsLoss()\n \n # Load the data into Datasets and then make DataLoaders out of them for training\n train_dataset = RSNAData(\n df = train_df,\n img_folder = Config['PARENT_PATH']\n )\n valid_dataset = RSNAData(\n df = valid_df,\n img_folder = Config['PARENT_PATH']\n )\n train_loader = DataLoader(\n train_dataset,\n batch_size=Config['TRAIN_BS'],\n shuffle=True\n )\n valid_loader = DataLoader(\n valid_dataset,\n batch_size=Config['VALID_BS'],\n shuffle=False\n )\n \n # Send all these things to the prepare function so they can be prepped for Multi-GPU training\n model, optimizer, train_loader, valid_loader = accelerator.prepare(\n model, optimizer, train_loader, valid_loader\n )\n \n # Print out the data sizes we are training on\n accelerator.print(f\"Training on {len(train_df)} samples, Validating on {len(valid_df)} samples\")\n \n # Train the model now\n fit(\n model=model,\n fold=\"single\",\n epochs=Config['NB_EPOCHS'],\n train_loader=train_loader,\n valid_loader=valid_loader,\n optimizer=optimizer,\n train_loss_fn=train_loss_fn,\n valid_loss_fn=valid_loss_fn,\n accelerator=accelerator\n )","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.992077Z","iopub.status.idle":"2024-04-23T23:40:36.992977Z","shell.execute_reply.started":"2024-04-23T23:40:36.992679Z","shell.execute_reply":"2024-04-23T23:40:36.992704Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load the data and pass it onto the training function\ndf = pd.read_csv(\"/data/yijia/code2/code/data/rsna/train.csv\")\ndf['img_name'] = df['patient_id'].astype(str) + \"/\" + df['image_id'].astype(str) + \".png\"\ndf = df.sample(frac=1).reset_index(drop=True)\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.994756Z","iopub.status.idle":"2024-04-23T23:40:36.995605Z","shell.execute_reply.started":"2024-04-23T23:40:36.995343Z","shell.execute_reply":"2024-04-23T23:40:36.995367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Launch training using HuggingFace accelerate on 2x T4 GPUs\nnotebook_launcher(run, args=(df, Config), num_processes=1)","metadata":{"execution":{"iopub.status.busy":"2024-04-23T23:40:36.997079Z","iopub.status.idle":"2024-04-23T23:40:36.997873Z","shell.execute_reply.started":"2024-04-23T23:40:36.997605Z","shell.execute_reply":"2024-04-23T23:40:36.997629Z"},"trusted":true},"execution_count":null,"outputs":[]}]}