2-full-demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_data_dir = \"/media/curttigges/project-files/datasets/ms-coco/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import math\n",
    "import timm\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "import wandb\n",
    "\n",
    "from torch.optim.lr_scheduler import OneCycleLR\n",
    "from torchmetrics.functional import accuracy, precision\n",
    "import torchmetrics.functional as tf\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning.callbacks import TQDMProgressBar\n",
    "from pytorch_lightning.callbacks import StochasticWeightAveraging\n",
    "from positional_encodings.torch_encodings import PositionalEncodingPermute2D, Summer\n",
    "\n",
    "from q2l_labeller.data.coco_data_module import COCODataModule\n",
    "from q2l_labeller.pl_modules.query2label_train_module import Query2LabelTrainModule\n",
    "from q2l_labeller.data.cutmix import CutMixCriterion\n",
    "from q2l_labeller.loss_modules.simple_asymmetric_loss import AsymmetricLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimmBackbone(nn.Module):\n",
    "    def __init__(self, model_name):\n",
    "        super().__init__()\n",
    "\n",
    "        # Creating the model in this way produces unpooled, unclassified features\n",
    "        self.model = timm.create_model(\n",
    "            model_name, \n",
    "            pretrained=True,\n",
    "            num_classes=0,\n",
    "            global_pool='')        \n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        out = self.model(x)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Query2Label(nn.Module):\n",
    "    \"\"\"Modified Query2Label model\n",
    "\n",
    "    Unlike the model described in the paper (which uses a modified DETR \n",
    "    transformer), this version uses a standard, unmodified Pytorch Transformer. \n",
    "    Learnable label embeddings are passed to the decoder module as the target \n",
    "    sequence (and ultimately is passed as the Query to MHA).\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self, model, conv_out, num_classes, hidden_dim=256, nheads=8, \n",
    "        encoder_layers=6, decoder_layers=6, use_pos_encoding=False):\n",
    "        \"\"\"Initializes model\n",
    "\n",
    "        Args:\n",
    "            model (str): Timm model descriptor for backbone.\n",
    "            conv_out (int): Backbone output channels.\n",
    "            num_classes (int): Number of possible label classes\n",
    "            hidden_dim (int, optional): Hidden channels from linear projection of\n",
    "            backbone output. Defaults to 256.\n",
    "            nheads (int, optional): Number of MHA heads. Defaults to 8.\n",
    "            encoder_layers (int, optional): Number of encoders. Defaults to 6.\n",
    "            decoder_layers (int, optional): Number of decoders. Defaults to 6.\n",
    "            use_pos_encoding (bool, optional): Flag for use of position encoding. \n",
    "            Defaults to False.\n",
    "        \"\"\"        \n",
    "        \n",
    "        super().__init__()\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.use_pos_encoding = use_pos_encoding\n",
    "\n",
    "        self.backbone = TimmBackbone(model)\n",
    "        self.conv = nn.Conv2d(conv_out, hidden_dim, 1)\n",
    "        self.transformer = nn.Transformer(\n",
    "            hidden_dim, nheads, encoder_layers, decoder_layers)\n",
    "\n",
    "        if self.use_pos_encoding:\n",
    "            # returns the encoding object\n",
    "            self.pos_encoder = PositionalEncodingPermute2D(hidden_dim)\n",
    "\n",
    "            # returns the summing object\n",
    "            self.encoding_adder = Summer(self.pos_encoder)\n",
    "\n",
    "        # prediction head\n",
    "        self.classifier = nn.Linear(num_classes * hidden_dim, num_classes)\n",
    "\n",
    "        # learnable label embedding\n",
    "        self.label_emb = nn.Parameter(torch.rand(1, num_classes, hidden_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Passes batch through network\n",
    "\n",
    "        Args:\n",
    "            x (Tensor): Batch of images\n",
    "\n",
    "        Returns:\n",
    "            Tensor: Output of classification head\n",
    "        \"\"\"        \n",
    "        # produces output of shape [N x C x H x W]\n",
    "        out = self.backbone(x)\n",
    "        \n",
    "        # reduce number of feature planes for the transformer\n",
    "        h = self.conv(out)\n",
    "        B, C, H, W = h.shape\n",
    "\n",
    "        # add position encodings\n",
    "        if self.use_pos_encoding:\n",
    "            \n",
    "            # input with encoding added\n",
    "            h = self.encoding_adder(h*0.1)\n",
    "\n",
    "        # convert h from [N x C x H x W] to [H*W x N x C] (N=batch size)\n",
    "        # this corresponds to the [SIZE x BATCH_SIZE x EMBED_DIM] dimensions \n",
    "        # that the transformer expects\n",
    "        h = h.flatten(2).permute(2, 0, 1)\n",
    "        \n",
    "        # image feature vector \"h\" is sent in after transformation above; we \n",
    "        # also convert label_emb from [1 x TARGET x (hidden)EMBED_SIZE] to \n",
    "        # [TARGET x BATCH_SIZE x (hidden)EMBED_SIZE]\n",
    "        label_emb = self.label_emb.repeat(B, 1, 1)\n",
    "        label_emb = label_emb.transpose(0, 1)\n",
    "        h = self.transformer(h, label_emb).transpose(0, 1)\n",
    "        \n",
    "        # output from transformer was of dim [TARGET x BATCH_SIZE x EMBED_SIZE];\n",
    "        # however, we transposed it to [BATCH_SIZE x TARGET x EMBED_SIZE] above.\n",
    "        # below we reshape to [BATCH_SIZE x TARGET*EMBED_SIZE].\n",
    "        #\n",
    "        # next, we project transformer outputs to class labels\n",
    "        h = torch.reshape(h,(B, self.num_classes * self.hidden_dim))\n",
    "\n",
    "        return self.classifier(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Query2LabelTrainModule(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        data,\n",
    "        backbone_desc,\n",
    "        conv_out_dim,\n",
    "        hidden_dim,\n",
    "        num_encoders,\n",
    "        num_decoders,\n",
    "        num_heads,\n",
    "        batch_size,\n",
    "        image_dim,\n",
    "        learning_rate,\n",
    "        momentum,\n",
    "        weight_decay,\n",
    "        n_classes,\n",
    "        thresh=0.5,\n",
    "        use_cutmix=False,\n",
    "        use_pos_encoding=False,\n",
    "        loss=\"BCE\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        # Key parameters\n",
    "        self.save_hyperparameters(ignore=[\"model\", \"data\"])\n",
    "        self.data = data\n",
    "        self.model = Query2Label(\n",
    "            model=backbone_desc,\n",
    "            conv_out=conv_out_dim,\n",
    "            num_classes=n_classes,\n",
    "            hidden_dim=hidden_dim,\n",
    "            nheads=num_heads,\n",
    "            encoder_layers=num_encoders,\n",
    "            decoder_layers=num_decoders,\n",
    "            use_pos_encoding=use_pos_encoding,\n",
    "        )\n",
    "        if loss == \"BCE\":\n",
    "            self.base_criterion = nn.BCEWithLogitsLoss()\n",
    "        elif loss == \"ASL\":\n",
    "            self.base_criterion = AsymmetricLoss(gamma_neg=1, gamma_pos=0)\n",
    "\n",
    "        self.criterion = CutMixCriterion(self.base_criterion)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        return x\n",
    "\n",
    "    def evaluate(self, batch, stage=None):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.base_criterion(y_hat, y.type(torch.float))\n",
    "\n",
    "        rmap = tf.retrieval_average_precision(y_hat, y.type(torch.int))\n",
    "\n",
    "        category_prec = precision(\n",
    "            y_hat,\n",
    "            y.type(torch.int),\n",
    "            average=\"macro\",\n",
    "            num_classes=self.hparams.n_classes,\n",
    "            threshold=self.hparams.thresh,\n",
    "            multiclass=False,\n",
    "        )\n",
    "        category_recall = tf.recall(\n",
    "            y_hat,\n",
    "            y.type(torch.int),\n",
    "            average=\"macro\",\n",
    "            num_classes=self.hparams.n_classes,\n",
    "            threshold=self.hparams.thresh,\n",
    "            multiclass=False,\n",
    "        )\n",
    "        category_f1 = tf.f1_score(\n",
    "            y_hat,\n",
    "            y.type(torch.int),\n",
    "            average=\"macro\",\n",
    "            num_classes=self.hparams.n_classes,\n",
    "            threshold=self.hparams.thresh,\n",
    "            multiclass=False,\n",
    "        )\n",
    "\n",
    "        overall_prec = precision(\n",
    "            y_hat, y.type(torch.int), threshold=self.hparams.thresh, multiclass=False\n",
    "        )\n",
    "        overall_recall = tf.recall(\n",
    "            y_hat, y.type(torch.int), threshold=self.hparams.thresh, multiclass=False\n",
    "        )\n",
    "        overall_f1 = tf.f1_score(\n",
    "            y_hat, y.type(torch.int), threshold=self.hparams.thresh, multiclass=False\n",
    "        )\n",
    "\n",
    "        if stage:\n",
    "            self.log(f\"{stage}_loss\", loss, prog_bar=True)\n",
    "            self.log(f\"{stage}_rmap\", rmap, prog_bar=True, on_step=False, on_epoch=True)\n",
    "\n",
    "            self.log(f\"{stage}_cat_prec\", category_prec, prog_bar=True)\n",
    "            self.log(f\"{stage}_cat_recall\", category_recall, prog_bar=True)\n",
    "            self.log(f\"{stage}_cat_f1\", category_f1, prog_bar=True)\n",
    "\n",
    "            self.log(f\"{stage}_ovr_prec\", overall_prec, prog_bar=True)\n",
    "            self.log(f\"{stage}_ovr_recall\", overall_recall, prog_bar=True)\n",
    "            self.log(f\"{stage}_ovr_f1\", overall_f1, prog_bar=True)\n",
    "\n",
    "            # log prediction examples to wandb\n",
    "            \"\"\"\n",
    "            pred = self.model(x)\n",
    "            pred_keys = pred[0].sigmoid().tolist()\n",
    "            pred_keys = [0 if p < self.hparams.thresh else 1 for p in pred_keys]\n",
    "\n",
    "\n",
    "            mapper = cc.COCOCategorizer()\n",
    "            pred_lbl = mapper.get_labels(pred_keys)\n",
    "            \n",
    "            try:\n",
    "                self.logger.experiment.log({\"val_pred_examples\": [wandb.Image(x[0], caption=pred_lbl)]})\n",
    "            except AttributeError:\n",
    "                pass\n",
    "            \"\"\"\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        if self.hparams.use_cutmix:\n",
    "            x, y = batch\n",
    "            y_hat = self(x)\n",
    "            # y1, y2, lam = y\n",
    "            loss = self.criterion(y_hat, y)\n",
    "\n",
    "        else:\n",
    "            x, y = batch\n",
    "            y_hat = self(x)\n",
    "            loss = self.base_criterion(y_hat, y.type(torch.float))\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"val\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            self.parameters(),\n",
    "            lr=self.hparams.learning_rate,\n",
    "            betas=(0.9, 0.999),\n",
    "            weight_decay=self.hparams.weight_decay,\n",
    "        )\n",
    "\n",
    "        lr_scheduler_dict = {\n",
    "            \"scheduler\": OneCycleLR(\n",
    "                optimizer,\n",
    "                self.hparams.learning_rate,\n",
    "                epochs=self.trainer.max_epochs,\n",
    "                steps_per_epoch=len(self.data.train_dataloader()),\n",
    "                anneal_strategy=\"cos\",\n",
    "            ),\n",
    "            \"interval\": \"step\",\n",
    "        }\n",
    "        return {\"optimizer\": optimizer, \"lr_scheduler\": lr_scheduler_dict}\n",
    "        # return optimizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "42"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pl.seed_everything(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_dict = {\n",
    "    \"backbone_desc\":\"resnet18\",\n",
    "    \"conv_out_dim\":512,\n",
    "    \"hidden_dim\":256,\n",
    "    \"num_encoders\":1,\n",
    "    \"num_decoders\":2,\n",
    "    \"num_heads\":8,\n",
    "    \"batch_size\":64,\n",
    "    \"image_dim\":224,\n",
    "    \"learning_rate\":0.0001, \n",
    "    \"momentum\":0.9,\n",
    "    \"weight_decay\":0.01, \n",
    "    \"n_classes\":80,\n",
    "    \"thresh\":0.5,\n",
    "    \"use_cutmix\":True,\n",
    "    \"use_pos_encoding\":False,\n",
    "    \"loss\":\"ASL\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "coco = COCODataModule(\n",
    "    img_data_dir,\n",
    "    img_size=param_dict[\"image_dim\"],\n",
    "    batch_size=param_dict[\"batch_size\"],\n",
    "    num_workers=24,\n",
    "    use_cutmix=param_dict[\"use_cutmix\"],\n",
    "    cutmix_alpha=1.0)\n",
    "param_dict[\"data\"] = coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl_model = Query2LabelTrainModule(**param_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m (\u001b[33mascendant\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.13.1 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.21"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>training/logs/wandb/run-20220812_175544-1bcn1xv5</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/ascendant/coco-labeling/runs/1bcn1xv5\" target=\"_blank\">super-waterfall-116</a></strong> to <a href=\"https://wandb.ai/ascendant/coco-labeling\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: logging graph, to disable use `wandb.watch(log_graph=False)`\n",
      "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",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=7.97s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=4.74s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name           | Type           | Params\n",
      "--------------------------------------------------\n",
      "0 | model          | Query2Label    | 17.4 M\n",
      "1 | base_criterion | AsymmetricLoss | 0     \n",
      "--------------------------------------------------\n",
      "17.4 M    Trainable params\n",
      "0         Non-trainable params\n",
      "17.4 M    Total params\n",
      "34.881    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:   0%|          | 0/1927 [00:00<?, ?it/s]                          "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/curttigges/miniconda3/envs/cv-env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:724: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...\n",
      "  rank_zero_warn(\"Detected KeyboardInterrupt, attempting graceful shutdown...\")\n"
     ]
    }
   ],
   "source": [
    "# Comment out if not using wandb\n",
    "wandb_logger = WandbLogger(\n",
    "    project=\"coco-labeling\", \n",
    "    save_dir=\"training/logs/\",\n",
    "    log_model=True)\n",
    "wandb_logger.watch(pl_model, log=\"all\")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=24,\n",
    "    precision=16,\n",
    "    accelerator='gpu', \n",
    "    devices=1,\n",
    "    logger=wandb_logger, # Comment out if not using wandb\n",
    "    default_root_dir=\"training/checkpoints/\",\n",
    "    callbacks=[TQDMProgressBar(refresh_rate=10)])\n",
    "trainer.fit(pl_model, param_dict[\"data\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.10 ('cv-env')",
   "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.7.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "86c88f9a3588ea9343d84fb206e74c1b312a3b0d43eff9010fec7c5800ba29d3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}