model_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import itertools\n",
    "from functools import partial\n",
    "from concurrent.futures import ProcessPoolExecutor\n",
    "from typing import List\n",
    "from typing_extensions import override\n",
    "\n",
    "import numpy as np \n",
    "import pandas as pd \n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import lightning as pl\n",
    "from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor, EarlyStopping, RichProgressBar\n",
    "from lightning.pytorch.loggers import WandbLogger\n",
    "\n",
    "import subprocess\n",
    "import timm\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import monai as mn\n",
    "from transforms.Transform4ClassifierBase import Transform4ClassifierBase\n",
    "from models.ClassifierBase import Classifier\n",
    "\n",
    "SEED = 5566\n",
    "pl.seed_everything(SEED)\n",
    "torch.set_float32_matmul_precision('medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_dict_part(df_part):\n",
    "    \"Important! Modify this function\"\n",
    "    \n",
    "    BASE_PATH =  #edit\n",
    "    IMG_PATH_COLUMN_NAME = # edit\n",
    "    \n",
    "    data_dict = list()\n",
    "    for i in tqdm(range(len(df_part)), desc=\"Processing part\"):\n",
    "        row = df_part.iloc[i]\n",
    "\n",
    "        data_dict.append({\n",
    "            'img':f'{BASE_PATH}/'+row[f\"{IMG_PATH_COLUMN_NAME}\"],\n",
    "            #\"label\": row[f\"{LABEL_COLUMN_NAME}\"],\n",
    "            \"paths\": f'{BASE_PATH}/'+row[f\"{IMG_PATH_COLUMN_NAME}\"],\n",
    "            ** {l: np.array([row[f'{l}']]) for l in LABEL_COLUMN_NAMES}\n",
    "        })\n",
    "    \n",
    "    return data_dict\n",
    "\n",
    "def get_data_dict(df, num_cores=2):\n",
    "    parts = np.array_split(df, num_cores)\n",
    "    func = partial(get_data_dict_part)\n",
    "    \n",
    "    with ProcessPoolExecutor(num_cores) as executor:\n",
    "        data_dicts = executor.map(func, parts)\n",
    "    \n",
    "    return list(itertools.chain(*data_dicts))\n",
    "\n",
    "def split_data(df, group_column, n_splits):\n",
    "    frac = 1/n_splits\n",
    "    val_idx = df[group_column].drop_duplicates().sample(frac=frac)\n",
    "    df_temp = df.set_index(group_column)\n",
    "    df_val = df_temp.loc[val_idx,:].reset_index()\n",
    "    df_train = df_temp.drop(index=val_idx).reset_index()\n",
    "    return df_train, df_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTANT BEFORE PROCEEDING --> DO YOU WANT TO DELETE CACHE???\n",
    "DELETE_CACHE = False\n",
    "\n",
    "INPUT = './Train.csv' # #edit \n",
    "\n",
    "TIMM_MODEL = \"hf-hub:timm/convnext_base.fb_in22k_ft_in1k\"\n",
    "\n",
    "LABEL_COLUMN_NAMES = #edit ex: ['Pneumothorax','Cardiomegaly']\n",
    "\n",
    "PROJECT =  #edit \n",
    "TEST_NAME =  #edit \n",
    "MONAI_CACHE_DIR = f'./cache/{TEST_NAME}' #edit \n",
    "IMG_SIZE = 256 #edit \n",
    "BATCH_SIZE = 16 #edit \n",
    "PRECISION = 'bf16-mixed' \n",
    "LEARNING_RATE = 1e-5 #edit \n",
    "EPOCHS = 300 #edit \n",
    "WEIGHT_PATH = f'./weights/{TEST_NAME}' \n",
    "os.makedirs(WEIGHT_PATH, exist_ok=True)\n",
    "\n",
    "ENTITY =  #edit, wandb id \n",
    "\n",
    "os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '1' #edit \n",
    "os.environ['WANDB_API_KEY']= '' #edit\n",
    "os.environ['WANDB_SILENT']='true'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if DELETE_CACHE:\n",
    "    if os.path.exists(MONAI_CACHE_DIR):\n",
    "        subprocess.call(['rm', '-rf', f'{MONAI_CACHE_DIR}'])\n",
    "        print(f\"MONAI's {MONAI_CACHE_DIR} cache directory removed successfully!\")\n",
    "    else:\n",
    "        print(f\"MONAI's {MONAI_CACHE_DIR} cache directory does not exist!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read input file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(INPUT)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split train and val data\n",
    "\n",
    "PATIENT_ID_COLUMN = #edit ex:'empi_anon'\n",
    "\n",
    "train_df, val_df = split_data(df, n_splits=10, group_column=PATIENT_ID_COLUMN)\n",
    "\n",
    "val_df.to_csv(f\"val_data_{TEST_NAME}.csv\", index=False)\n",
    "\n",
    "print(len(train_df), len(val_df))\n",
    "\n",
    "train_dict = get_data_dict(train_df)\n",
    "val_dict = get_data_dict(val_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define transforms\n",
    "\n",
    "train_transforms = Transform4ClassifierBase(IMG_SIZE, LABEL_COLUMN_NAMES).train\n",
    "val_transforms = Transform4ClassifierBase(IMG_SIZE, LABEL_COLUMN_NAMES).val\n",
    "\n",
    "# define datasets\n",
    "\n",
    "train_ds = mn.data.PersistentDataset(data=train_dict, transform=train_transforms, cache_dir=f\"{MONAI_CACHE_DIR}/train\")\n",
    "val_ds = mn.data.PersistentDataset(data=val_dict, transform=val_transforms, cache_dir=f\"{MONAI_CACHE_DIR}/val\")\n",
    "\n",
    "# define data loader\n",
    "\n",
    "train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, persistent_workers=True, num_workers=2, drop_last=True)\n",
    "val_dl = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, drop_last=False, persistent_workers=True)\n",
    "\n",
    "# instantiate the model\n",
    "model = Classifier(TIMM_MODEL=TIMM_MODEL, num_classes=len(LABEL_COLUMN_NAMES), LEARNING_RATE=LEARNING_RATE, BATCH_SIZE=BATCH_SIZE, use_ema=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SPOT CHECK\n",
    "test_ds=mn.data.Dataset(data=train_dict, transform=train_transforms)\n",
    "\n",
    "for _ in range(3):\n",
    "    random_i = np.random.randint(0, len(test_ds))\n",
    "    for data_ in test_ds[random_i:random_i+1]:\n",
    "        \n",
    "        print(f\"{data_['paths']}\")\n",
    "        plt.imshow(np.flipud(np.rot90(np.squeeze(np.array(data_['img'])))), cmap='gray')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_monitor = LearningRateMonitor(logging_interval=\"step\")\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(dirpath=f\"{WEIGHT_PATH}\",\n",
    "                                    filename=f'{TEST_NAME}_{{epoch}}_{{valid_loss:0.4F}}',\n",
    "                                    monitor=\"valid_loss\",\n",
    "                                    mode=\"min\",\n",
    "                                    save_last=False,\n",
    "                                    save_top_k=1)\n",
    "\n",
    "early_stop_callback = EarlyStopping(monitor='valid_loss',\n",
    "                                    min_delta=0.00001,\n",
    "                                    patience=5,\n",
    "                                    verbose=False,\n",
    "                                    mode='min')\n",
    "\n",
    "wandb_logger = WandbLogger(save_dir=f\"{WEIGHT_PATH}\",\n",
    "                           name=f'{TEST_NAME}',\n",
    "                           project=PROJECT,\n",
    "                           entity=ENTITY,\n",
    "                           offline=False,\n",
    "                           log_model=False,\n",
    "                           config={\"Creator\": \"HITI\"})\n",
    "\n",
    "# csv_logger = CSVLogger(\"logs\", name=\"demo\", flush_logs_every_n_steps=10)\n",
    "\n",
    "progress_bar = RichProgressBar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate trainer\n",
    "\n",
    "trainer = pl.Trainer(gradient_clip_val=1.0,\n",
    "                    callbacks=[progress_bar, lr_monitor, checkpoint_callback, early_stop_callback],\n",
    "                    logger= wandb_logger,\n",
    "                    precision = PRECISION,\n",
    "                    accelerator = \"gpu\",\n",
    "                    devices=1,\n",
    "                    log_every_n_steps=1,\n",
    "                    default_root_dir= WEIGHT_PATH,\n",
    "                    max_epochs=EPOCHS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# train the model\n",
    "\n",
    "trainer.fit(model, train_dataloaders=train_dl, val_dataloaders=val_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dlpipelines",
   "language": "python",
   "name": "dlpipelines"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}