train_hls_burns_scars.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import segmentation_models_pytorch as smp\n",
    "import matplotlib.pyplot as plt\n",
    "from torchmetrics import MetricCollection\n",
    "from torchmetrics.classification import Accuracy, FBetaScore, JaccardIndex\n",
    "from torchmetrics.wrappers import ClasswiseWrapper\n",
    "\n",
    "from torchgeo.trainers import SemanticSegmentationTask\n",
    "from torchgeo.datasets.utils import RGBBandsMissingError, unbind_samples\n",
    "\n",
    "from prithvi_pytorch import PrithviEncoderDecoder\n",
    "from prithvi_pytorch.datasets import HLSBurnScarsDataModule\n",
    "\n",
    "CKPT_PATH = \"weights/Prithvi_100M.pt\"\n",
    "CFG_PATH = \"weights/Prithvi_100M_config.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FocalJaccardLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ce_loss = smp.losses.FocalLoss(mode=\"multiclass\", normalized=True)\n",
    "        self.jaccard_loss = smp.losses.JaccardLoss(mode=\"multiclass\")\n",
    "\n",
    "    def forward(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:\n",
    "        return self.ce_loss(preds, targets) + self.jaccard_loss(preds, targets)\n",
    "\n",
    "\n",
    "class PrithviSegmentationTask(SemanticSegmentationTask):\n",
    "    def configure_losses(self):\n",
    "        self.criterion = FocalJaccardLoss()\n",
    "\n",
    "    def configure_metrics(self):\n",
    "        num_classes = self.hparams[\"num_classes\"]\n",
    "        ignore_index = self.hparams[\"ignore_index\"]\n",
    "\n",
    "        self.train_metrics = MetricCollection(\n",
    "            {\n",
    "                \"OverallAccuracy\": Accuracy(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    average=\"micro\",\n",
    "                    multidim_average=\"global\",\n",
    "                ),\n",
    "                \"OverallF1Score\": FBetaScore(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    beta=1.0,\n",
    "                    average=\"micro\",\n",
    "                    multidim_average=\"global\",\n",
    "                ),\n",
    "                \"OverallIoU\": JaccardIndex(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    ignore_index=ignore_index,\n",
    "                    average=\"micro\",\n",
    "                ),\n",
    "                \"AverageAccuracy\": Accuracy(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    average=\"macro\",\n",
    "                    multidim_average=\"global\",\n",
    "                ),\n",
    "                \"AverageF1Score\": FBetaScore(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    beta=1.0,\n",
    "                    average=\"macro\",\n",
    "                    multidim_average=\"global\",\n",
    "                ),\n",
    "                \"AverageIoU\": JaccardIndex(\n",
    "                    task=\"multiclass\",\n",
    "                    num_classes=num_classes,\n",
    "                    ignore_index=ignore_index,\n",
    "                    average=\"macro\",\n",
    "                ),\n",
    "                \"Accuracy\": ClasswiseWrapper(\n",
    "                    Accuracy(\n",
    "                        task=\"multiclass\",\n",
    "                        num_classes=num_classes,\n",
    "                        average=\"none\",\n",
    "                        multidim_average=\"global\",\n",
    "                    ),\n",
    "                ),\n",
    "                \"F1Score\": ClasswiseWrapper(\n",
    "                    FBetaScore(\n",
    "                        task=\"multiclass\",\n",
    "                        num_classes=num_classes,\n",
    "                        beta=1.0,\n",
    "                        average=\"none\",\n",
    "                        multidim_average=\"global\",\n",
    "                    ),\n",
    "                ),\n",
    "                \"IoU\": ClasswiseWrapper(\n",
    "                    JaccardIndex(\n",
    "                        task=\"multiclass\", num_classes=num_classes, average=\"none\"\n",
    "                    ),\n",
    "                ),\n",
    "            },\n",
    "            prefix=\"train_\",\n",
    "        )\n",
    "        self.val_metrics = self.train_metrics.clone(prefix=\"val_\")\n",
    "        self.test_metrics = self.train_metrics.clone(prefix=\"test_\")\n",
    "\n",
    "    def configure_models(self):\n",
    "        self.model = PrithviEncoderDecoder(\n",
    "            num_classes=self.hparams[\"num_classes\"],\n",
    "            cfg_path=CFG_PATH,\n",
    "            ckpt_path=CKPT_PATH,\n",
    "            in_chans=self.hparams[\"in_channels\"],\n",
    "            img_size=224,\n",
    "            freeze_encoder=False,\n",
    "            num_neck_filters=32,\n",
    "        )\n",
    "\n",
    "    def training_step(self, batch, batch_idx, dataloader_idx=0):\n",
    "        x = batch[\"image\"]\n",
    "        y = batch[\"mask\"]\n",
    "        y_hat = self(x)\n",
    "        y_hat_hard = y_hat.argmax(dim=1)\n",
    "        loss = self.criterion(y_hat, y)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        self.train_metrics(y_hat_hard, y)\n",
    "        self.log_dict({f\"{k}\": v for k, v in self.train_metrics.compute().items()})\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx, dataloader_idx=0):\n",
    "        x = batch[\"image\"]\n",
    "        y = batch[\"mask\"]\n",
    "        y_hat = self(x)\n",
    "        y_hat_hard = y_hat.argmax(dim=1)\n",
    "        loss = self.criterion(y_hat, y)\n",
    "        self.log(\"val_loss\", loss)\n",
    "        self.val_metrics(y_hat_hard, y)\n",
    "        self.log_dict({f\"{k}\": v for k, v in self.val_metrics.compute().items()})\n",
    "\n",
    "        # Plot samples\n",
    "        if (\n",
    "            batch_idx < 10\n",
    "            and hasattr(self.trainer, \"datamodule\")\n",
    "            and hasattr(self.trainer.datamodule, \"plot\")\n",
    "            and self.logger\n",
    "            and hasattr(self.logger, \"experiment\")\n",
    "            and hasattr(self.logger.experiment, \"add_figure\")\n",
    "        ):\n",
    "            datamodule = self.trainer.datamodule\n",
    "            batch[\"prediction\"] = y_hat_hard\n",
    "            for key in [\"image\", \"mask\", \"prediction\"]:\n",
    "                batch[key] = batch[key].cpu()\n",
    "            sample = unbind_samples(batch)[0]\n",
    "\n",
    "            fig = None\n",
    "            try:\n",
    "                fig = datamodule.plot(sample)\n",
    "            except RGBBandsMissingError:\n",
    "                pass\n",
    "\n",
    "            if fig:\n",
    "                summary_writer = self.logger.experiment\n",
    "                summary_writer.add_figure(\n",
    "                    f\"image/{batch_idx}\", fig, global_step=self.global_step\n",
    "                )\n",
    "                plt.close()\n",
    "\n",
    "    def test_step(self, batch, batch_idx, dataloader_idx=0):\n",
    "        x = batch[\"image\"]\n",
    "        y = batch[\"mask\"]\n",
    "        y_hat = self(x)\n",
    "        y_hat_hard = y_hat.argmax(dim=1)\n",
    "        loss = self.criterion(y_hat, y)\n",
    "        self.log(\"test_loss\", loss)\n",
    "        self.test_metrics(y_hat_hard, y)\n",
    "        self.log_dict({f\"{k}\": v for k, v in self.test_metrics.compute().items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "module = PrithviSegmentationTask(in_channels=6, num_classes=2, lr=1e-4, patience=10)\n",
    "datamodule = HLSBurnScarsDataModule(\n",
    "    root=\"data/hls_burn_scars\",\n",
    "    batch_size=8,\n",
    "    num_workers=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit Automatic Mixed Precision (AMP)\n",
      "GPU available: True (cuda), 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"
     ]
    }
   ],
   "source": [
    "trainer = lightning.Trainer(\n",
    "    accelerator=\"gpu\", devices=[0], logger=True, max_epochs=20, precision=\"16-mixed\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing logger folder: /workspace/storage/github/prithvi-pytorch/lightning_logs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "\n",
      "  | Name          | Type                  | Params\n",
      "--------------------------------------------------------\n",
      "0 | criterion     | FocalJaccardLoss      | 0     \n",
      "1 | train_metrics | MetricCollection      | 0     \n",
      "2 | val_metrics   | MetricCollection      | 0     \n",
      "3 | test_metrics  | MetricCollection      | 0     \n",
      "4 | model         | PrithviEncoderDecoder | 113 M \n",
      "--------------------------------------------------------\n",
      "112 M     Trainable params\n",
      "252 K     Non-trainable params\n",
      "113 M     Total params\n",
      "452.012   Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc3dfbbe4ea342a092bff86e4447df33",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ffe72461e2744829b8b7004c40e52391",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c3babe2b4c674bd0a28f291de2108868",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9dabcb1b8194d67971884acc38454d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0293e92881a40269ecdbe6933ae408e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64ef1b659aa644d19cedb6dd30400b57",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82f86fb49f3c43f1b3dc144a35baeb7e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aaaf7fc407f1489d839883292ccd1e7b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9aee320000240748384c6d2726dde56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77355ce9a532417b82476327ef7f23fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1905687eab9b455e9a9a42e10c4b6bff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1c2aa1d4bda4942be0e8a98fc109c7a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "819a125abe584172af3122c1a0f976e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a761ab6f30cd4dde9c9a8d5e6fed4137",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.fit(model=module, datamodule=datamodule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.test(datamodule=datamodule)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torchenv",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}