for_developers/add_custom_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to add custom model\n",
    "\n",
    "This is a simple Jupyter notebook example to show how the developer can add a new model for the OTX task (Multi class classification at this time) and execute model training.\n",
    "First let me start with importing everything we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from typing import TYPE_CHECKING, Any\n",
    "\n",
    "import torch\n",
    "from otx.core.data.entity.classification import (\n",
    "    MulticlassClsBatchDataEntity,\n",
    "    MulticlassClsBatchPredEntity,\n",
    ")\n",
    "from otx.core.model.entity.classification import OTXMulticlassClsModel\n",
    "from torch import nn\n",
    "from torchvision.models.resnet import ResNet50_Weights, resnet50\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from otx.core.data.entity.base import OTXBatchLossEntity\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have everything we need. Before we start, please keep in mind that this is not our end image. The training from Python API design is not determined yet and this is very first place.\n",
    "\n",
    "The first thing is that we need to develop the actual PyTorch Model which should be created in `OTXModel._create_model()` function.\n",
    "As you know, `OTXModel` is required to produce the task losses in the training.\n",
    "On the other hand, it should produce the model predictions from the image in the evaluation.\n",
    "Therefore, this `nn.Module` should be able to compute the task losses.\n",
    "This is important thing you have to notice.\n",
    "Let's see the code now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNet50WithLossComputation(nn.Module):\n",
    "    def __init__(self, num_classes: int) -> None:\n",
    "        super().__init__()\n",
    "        self.num_classes = num_classes\n",
    "        net = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)\n",
    "        net.fc = nn.Linear(\n",
    "            in_features=net.fc.in_features, out_features=self.num_classes\n",
    "        )\n",
    "        self.net = net\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, images: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n",
    "        logits = self.net(images)\n",
    "\n",
    "        if self.training:\n",
    "            return self.criterion(logits, labels)\n",
    "\n",
    "        return self.softmax(logits)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next thing is that we need to develop the one derived from `OTXModel`.\n",
    "However, in this example, we want to add the multi class classification model.\n",
    "We should implement the class derived from `OTXClassificationModel`.\n",
    "For another OTX task, such as `OTXTaskType.DETECTION`, we might be able to make a custom model by deriving from `OTXDetectionModel`.\n",
    "\n",
    "Since every `OTXModel` is an abstract class, it is designed to require a developer to implement three abstract functions:\n",
    "\n",
    "1) `_create_model()`\n",
    "2) `_customize_inputs()`\n",
    "3) `_customize_outputs()`\n",
    "\n",
    "You can see that the following example is exactly implementing those three functions.\n",
    "Let's see together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OTXResNet50(OTXMulticlassClsModel):\n",
    "    def __init__(self, num_classes: int) -> None:\n",
    "        super().__init__(num_classes=num_classes)\n",
    "        self.register_buffer(\n",
    "            \"mean\",\n",
    "            torch.FloatTensor([123.675, 116.28, 103.53]).view(-1, 1, 1),\n",
    "            False,\n",
    "        )\n",
    "        self.register_buffer(\n",
    "            \"std\",\n",
    "            torch.FloatTensor([58.395, 57.12, 57.375]).view(-1, 1, 1),\n",
    "            False,\n",
    "        )\n",
    "\n",
    "    def _create_model(self) -> nn.Module:\n",
    "        # ResNet50_Weights.IMAGENET1K_V2 is a really powerful pretrained model equipped with the modern training scheme:\n",
    "        # ImageNet-1K acc@1: 80.858, acc@5\": 95.434.\n",
    "        return ResNet50WithLossComputation(num_classes=self.num_classes)\n",
    "\n",
    "    def _customize_inputs(self, inputs: MulticlassClsBatchDataEntity) -> dict[str, Any]:\n",
    "        images = inputs.images.to(dtype=torch.float32)\n",
    "        images = (images - self.mean) / self.std\n",
    "        return {\n",
    "            \"images\": images,\n",
    "            \"labels\": torch.cat(inputs.labels, dim=0),\n",
    "        }\n",
    "\n",
    "    def _customize_outputs(\n",
    "        self, outputs: Any, inputs: MulticlassClsBatchDataEntity\n",
    "    ) -> MulticlassClsBatchPredEntity | OTXBatchLossEntity:\n",
    "        if self.training:\n",
    "            return {\"loss\": outputs}\n",
    "\n",
    "        # To list, batch-wise\n",
    "        scores = torch.unbind(outputs, 0)\n",
    "\n",
    "        return MulticlassClsBatchPredEntity(\n",
    "            batch_size=inputs.batch_size,\n",
    "            images=inputs.images,\n",
    "            imgs_info=inputs.imgs_info,\n",
    "            scores=scores,\n",
    "            labels=inputs.labels,\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have our own custom model which can be used for the OTX training process.\n",
    "However, there are many things and configurations for the model training.\n",
    "**Setting these things from scratch is so sick.**\n",
    "Therefore, we will **borrow the configurations from the similar model for the multi class classification task OTX provided: `classification/otx_efficientnet_b0`.**\n",
    "We just override our custom model on top of that.\n",
    "\n",
    "It means that we can compose **1) Data transform from MMPretrain and 2) Custom model from TorchVision**.\n",
    "As you know, our design let them independent each other.\n",
    "Therefore, any composition, not just this example, such as 1) Data transform from TorchVision and 2) Custom model from Detectron, is possible.\n",
    "\n",
    "Please see the following how we do that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:The corresponding keys in config are not used.: ['verbose', 'data_root', 'task', 'seed', 'callback_monitor', 'resume', 'disable_infer_num_classes']\n",
      "WARNING:root:Set Default Optimizer: {'class_path': 'torch.optim.SGD', 'init_args': {'lr': 0.0049, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0.0001, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False}}\n",
      "WARNING:root:Set Default Scheduler: {'class_path': 'lightning.pytorch.cli.ReduceLROnPlateau', 'init_args': {'monitor': 'train/loss', 'mode': 'min', 'factor': 0.5, 'patience': 1, 'threshold': 0.0001, 'threshold_mode': 'rel', 'cooldown': 0, 'min_lr': 0, 'eps': 1e-08, 'verbose': False}}\n",
      "INFO:pytorch_lightning.utilities.rank_zero:GPU available: True (cuda), used: True\n",
      "INFO:pytorch_lightning.utilities.rank_zero:TPU available: False, using: 0 TPU cores\n",
      "INFO:pytorch_lightning.utilities.rank_zero:IPU available: False, using: 0 IPUs\n",
      "INFO:pytorch_lightning.utilities.rank_zero:HPU available: False, using: 0 HPUs\n",
      "INFO:pytorch_lightning.utilities.rank_zero:You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "\n",
      "  | Name        | Type               | Params\n",
      "---------------------------------------------------\n",
      "0 | model       | OTXResNet50        | 23.5 M\n",
      "1 | val_metric  | MulticlassAccuracy | 0     \n",
      "2 | test_metric | MulticlassAccuracy | 0     \n",
      "---------------------------------------------------\n",
      "23.5 M    Trainable params\n",
      "0         Non-trainable params\n",
      "23.5 M    Total params\n",
      "94.049    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                           \r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/harimkan/workspace/repo/otx-regression/venv/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py:293: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: 100%|██████████| 1/1 [00:00<00:00,  8.83it/s, v_num=9, train/loss=0.563, val/accuracy=1.000]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:pytorch_lightning.utilities.rank_zero:`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: 100%|██████████| 1/1 [00:00<00:00,  2.92it/s, v_num=9, train/loss=0.563, val/accuracy=1.000]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'train/loss': tensor(0.5628), 'val/accuracy': tensor(1.)}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from otx.engine import Engine\n",
    "\n",
    "data_dir = \"../tests/assets/classification_dataset\"\n",
    "otx_model = OTXResNet50(num_classes=2)\n",
    "\n",
    "engine = Engine(\n",
    "    data_root=data_dir,\n",
    "    model=otx_model,\n",
    "    device=\"gpu\",\n",
    "    work_dir=\"otx-workspace\",\n",
    ")\n",
    "\n",
    "engine.train(max_epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Saying again. This is not the end image of the OTX training API. We will continue to strive to improve it so that users can use it conveniently. And, I believe that it is not difficult since we already have a solid core design and it is just an entrypoint.**"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "otx-v2",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}