few_shot_clip.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fc110348-9d48-4a0d-b3fe-4d2ccbdbf7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!gpustat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7f156307-0a4e-4cdb-a2cf-65770a918466",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "075645d4-1f4a-4f25-86ad-e4265df5ea32",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_models_to_fp32(model): \n",
    "    for p in model.parameters(): \n",
    "        p.data = p.data.float() \n",
    "        if p.grad is not None:\n",
    "            p.grad.data = p.grad.data.float() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "56fcb3bb-303f-4242-ad80-de9a7ac91fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "from few_shot_clip_utils import CLIPClassifier\n",
    "from clip_utils import ImageDataset, FinetuneDataModule\n",
    "\n",
    "\n",
    "\n",
    "dataset_name = \"mimic-cxr\"\n",
    "model_name = \"densenet\" #\"ViT-B/16, densenet, ViT-L/14\n",
    "mode = \"train_mlp_norm\" # [\"freeze\", \"train_norm\", \"train_full\", \"train_mlp_norm\"]\n",
    "\n",
    "# do not change\n",
    "use_augs = True\n",
    "max_epochs = 20\n",
    "precision = 16\n",
    "use_pos_weight = True\n",
    "# some pre-set settings\n",
    "if mode == \"torchxrayvision\":\n",
    "    model_name = \"densenet\"\n",
    "    lr = 0.001\n",
    "    weight_decay = 0.0\n",
    "    batch_size = 64\n",
    "    \n",
    "    rot_aug = 45\n",
    "    shift_aug = 0.15\n",
    "    scale_aug = 0.1\n",
    "    \n",
    "    img_value_scale = (1024, 1024)\n",
    "    input_res = (224, 224)\n",
    "    \n",
    "    lr_scale_when_plateau = 0.1\n",
    "    lr_patience = 5\n",
    "    # check code of torchxrayvisoin for default learning rate and for num epochs until decay\n",
    "    \n",
    "\n",
    "if model_name == \"ViT-L/14\":\n",
    "    batch_size = 8 # max 32 for single GPU CL on VitB16, 16 for ViT-L/14\n",
    "elif model_name == \"ViT-B/16\":\n",
    "    batch_size = 32\n",
    "else:\n",
    "    batch_size = 32\n",
    "\n",
    "if model_name == \"densenet\":\n",
    "    lr = 1e-6\n",
    "    weight_decay = 0.2\n",
    "    mode = \"full\"\n",
    "else:\n",
    "    lr = 1e-6\n",
    "    weight_decay = 0.2 # decay 0.05 in convnext, 0.2 in CLIP training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8bb94d7b-ac86-499b-a95a-6467e1582e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#convert_models_to_fp32(model)\n",
    "num_labels = 14 #data_module.num_labels\n",
    "\n",
    "if model_name == \"densenet_224\":\n",
    "    import torchxrayvision as xrv\n",
    "    model = xrv.models.DenseNet(weights=\"densenet121-res224-mimic_ch\")\n",
    "    labels_to_remove = [\"No finding\", \"Support devices\", \"Pleural other\"]\n",
    "    print(FinetuneDataModule.feature_names)\n",
    "    \n",
    "elif model_name == \"densenet\":\n",
    "    import torchvision\n",
    "    import torchvision.transforms as TF\n",
    "    # get model and re-init out layer\n",
    "    model = torchvision.models.densenet121(pretrained=True) \n",
    "    model.classifier = torch.nn.Linear(1024, num_labels)\n",
    "    model.num_labels = num_labels\n",
    "    # create transform\n",
    "    normalize = TF.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                             std=[0.229, 0.224, 0.225])\n",
    "    densenet_size = 256\n",
    "    transform = TF.Compose([TF.Resize(size=densenet_size, \n",
    "                                               interpolation=TF.InterpolationMode.BILINEAR),\n",
    "                                     TF.CenterCrop(size=(densenet_size, densenet_size)),\n",
    "                                     TF.ToTensor(),\n",
    "                                     normalize])\n",
    "else:\n",
    "    from clip_utils import load_clip\n",
    "    clip_base_model, transform, clip_name = load_clip(model_name, device=\"cpu\")\n",
    "    model = CLIPClassifier(clip_base_model, mode, num_labels)\n",
    "\n",
    "data_module = FinetuneDataModule(model, transform, \n",
    "                                 dataset_name=dataset_name, mode=mode, use_augs=use_augs,\n",
    "                                batch_size=batch_size)\n",
    "model.label_names = data_module.label_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "03c5fd99-9fe6-41d8-8aaa-5e301a96dfd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#path = data_module.train_ds.paths[0]\n",
    "#print(path)\n",
    "#from PIL import Image\n",
    "#pil_img = Image.open(path)\n",
    "#pil_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9eef903e-df9a-4803-9b26-6eb117e1ff62",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "wandb: Currently logged in as: antonius (use `wandb login --relogin` to force relogin)\n",
      "wandb: wandb version 0.12.10 is available!  To upgrade, please run:\n",
      "wandb:  $ pip install wandb --upgrade\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                    Syncing run <strong><a href=\"https://wandb.ai/antonius/early_tests/runs/1bbyneiz\" target=\"_blank\">warm-firefly-19</a></strong> to <a href=\"https://wandb.ai/antonius/early_tests\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">docs</a>).<br/>\n",
       "\n",
       "                "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit native Automatic Mixed Precision (AMP)\n",
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n"
     ]
    }
   ],
   "source": [
    "import pytorch_lightning\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "\n",
    "wandb_logger = pytorch_lightning.loggers.WandbLogger(name=None, save_dir=None, offline=False, id=None, \n",
    "                                      anonymous=None, version=None, project=\"early_tests\", \n",
    "                                      log_model=False, experiment=None, prefix='')\n",
    "wandb_logger.log_hyperparams({\"mode\": mode,\n",
    "                             \"dataset_name\": dataset_name,\n",
    "                             \"use_augs\": use_augs,\n",
    "                             \"use_pos_weight\": use_pos_weight,\n",
    "                             \"batch_size\": batch_size,\n",
    "                             \"model_name\": model_name,\n",
    "                             })\n",
    "\n",
    "\n",
    "trainer = pytorch_lightning.Trainer(val_check_interval=300,\n",
    "                                    precision=precision,\n",
    "                                    logger=wandb_logger,\n",
    "                                    max_epochs=max_epochs,\n",
    "                                    gpus=int(torch.cuda.is_available()),\n",
    "                                    #overfit_batches=1, \n",
    "                                    benchmark=True,\n",
    "                                    num_sanity_val_steps=0,\n",
    "                                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "03d8ec4d-5bd7-4ec9-9860-bb1c7126e364",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pos weight:  tensor([  4.8115,   4.8980,  24.9163,   9.4579,  36.8328,  48.8305,  34.1711,\n",
      "          3.9655,   1.6122,   3.9360, 110.1615,  13.4747,  25.3513,   3.5054],\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import few_shot_clip_utils\n",
    "importlib.reload(few_shot_clip_utils)\n",
    "from few_shot_clip_utils import LitCLIP\n",
    "\n",
    "pos_fraction = data_module.pos_fraction\n",
    "#print(pos_fraction)\n",
    "lit_model = LitCLIP(model, max_epochs, learning_rate=lr, \n",
    "                    steps_per_epoch=data_module.steps_per_epoch, \n",
    "                    weight_decay=weight_decay,\n",
    "                    pos_fraction=pos_fraction,\n",
    "                    use_pos_weight=use_pos_weight)#0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2758d6-29d8-4042-8e23-b36ffd7ebdfc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc07a9bc-a484-40e2-b473-08d00bcf6e9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [7]\n",
      "\n",
      "  | Name      | Type              | Params\n",
      "------------------------------------------------\n",
      "0 | model     | DenseNet          | 7.0 M \n",
      "1 | loss_func | BCEWithLogitsLoss | 0     \n",
      "------------------------------------------------\n",
      "7.0 M     Trainable params\n",
      "0         Non-trainable params\n",
      "7.0 M     Total params\n",
      "13.936    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "790ee05d19744def996621fb8528e066",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#data_module.prepare_data()\n",
    "#data_module.setup(stage=\"fit\")\n",
    "#trainer.validate(model)\n",
    "trainer.fit(lit_model, data_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d7979cf-2ee8-4489-843a-1ba812120a05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: uniformly calc all metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e35dfa89-148f-4aeb-a53b-aecea7fbea67",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}