torchgeo_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d0f533d-16af-4b24-91f0-31b338c91fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demo using the torchgeo package to initialize a SatlasPretrain model and finetune\n",
    "# on the UCMerced dataset.\n",
    "#\n",
    "# SETUP - this demo requires a DIFFERENT conda environment than the SatlasPretrain demo\n",
    "# conda create --name torchgeodemo\n",
    "# conda activate torchgeodemo\n",
    "# NOTE: Satlas weights will be a part of the 0.6.0 release and the current version is 0.5.1, so install from git for now.\n",
    "# pip install git+https://github.com/microsoft/torchgeo "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "291fd9c9-66cd-4643-a387-b6a8a7ec7754",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/piperw/opt/anaconda3/envs/torchgeotest/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import tempfile\n",
    "from typing import Optional\n",
    "from lightning.pytorch import Trainer\n",
    "\n",
    "from torchgeo.models import Swin_V2_B_Weights, swin_v2_b\n",
    "from torchgeo.datamodules import UCMercedDataModule\n",
    "from torchgeo.trainers import ClassificationTask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3ed5b26a-e461-4258-936a-d26c79d1e4f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experiment Arguments\n",
    "batch_size = 8\n",
    "num_workers = 2\n",
    "max_epochs = 10\n",
    "fast_dev_run = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4bb46a02-8fa4-4e29-8f0d-1c0ce3de10f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Torchgeo lightning datamodule to initialize dataset\n",
    "root = os.path.join(tempfile.gettempdir(), \"ucm\")\n",
    "datamodule = UCMercedDataModule(\n",
    "    root=root, batch_size=batch_size, num_workers=num_workers, download=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aac6a7a3-684a-4f8a-aece-6d4eaa55772e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom ClassificationTask to load in the SatlasPretrain model\n",
    "class SatlasClassificationTask(ClassificationTask):\n",
    "    def configure_models(self):\n",
    "        weights = Swin_V2_B_Weights.SENTINEL2_RGB_SI_SATLAS\n",
    "        self.model = swin_v2_b(weights)\n",
    "\n",
    "        # Replace first layer's input channels with the task's number input channels.\n",
    "        first_layer = self.model.features[0][0]\n",
    "        self.model.features[0][0] = torch.nn.Conv2d(3,\n",
    "                                    first_layer.out_channels,\n",
    "                                    kernel_size=first_layer.kernel_size,\n",
    "                                    stride=first_layer.stride,\n",
    "                                    padding=first_layer.padding,\n",
    "                                    bias=(first_layer.bias is not None))\n",
    "\n",
    "        # Replace last layer's output features with the number classes.\n",
    "        self.model.head = torch.nn.Linear(in_features=1024, out_features=self.hparams[\"num_classes\"], bias=True)\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        # Accessing metrics logged during the current validation epoch\n",
    "        val_loss = self.trainer.callback_metrics.get('val_loss', 'N/A')\n",
    "        val_acc = self.trainer.callback_metrics.get('val_OverallAccuracy', 'N/A')\n",
    "        print(f\"Epoch {self.current_epoch} Validation - Loss: {val_loss}, Accuracy: {val_acc}\")\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        # Accessing metrics logged during the current validation epoch\n",
    "        val_loss = self.trainer.callback_metrics.get('val_loss', 'N/A')\n",
    "        val_acc = self.trainer.callback_metrics.get('val_OverallAccuracy', 'N/A')\n",
    "        print(f\"Epoch {self.current_epoch} Validation - Loss: {val_loss}, Accuracy: {val_acc}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b4c09027-6c47-42f6-8d06-e3fbeb9d8cab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Classifcation Task\n",
    "task = SatlasClassificationTask(num_classes=21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f88e3e92-d30e-48d0-9c74-20efc91f4459",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "# Initialize the training code.\n",
    "accelerator = \"gpu\" if torch.cuda.is_available() else \"cpu\"\n",
    "default_root_dir = os.path.join(tempfile.gettempdir(), \"experiments\")\n",
    "\n",
    "trainer = Trainer(\n",
    "    accelerator=accelerator,\n",
    "    default_root_dir=default_root_dir,\n",
    "    fast_dev_run=fast_dev_run,\n",
    "    log_every_n_steps=1,\n",
    "    min_epochs=1,\n",
    "    max_epochs=max_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aa6a6684-cf14-4fa2-b801-f3d35b94519d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name          | Type             | Params\n",
      "---------------------------------------------------\n",
      "0 | criterion     | CrossEntropyLoss | 0     \n",
      "1 | train_metrics | MetricCollection | 0     \n",
      "2 | val_metrics   | MetricCollection | 0     \n",
      "3 | test_metrics  | MetricCollection | 0     \n",
      "4 | model         | SwinTransformer  | 86.9 M\n",
      "---------------------------------------------------\n",
      "86.9 M    Trainable params\n",
      "0         Non-trainable params\n",
      "86.9 M    Total params\n",
      "347.709   Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sanity Checking DataLoader 0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  0.60it/s]Epoch 0 Validation - Loss: 3.1346988677978516, Accuracy: 0.0\n",
      "                                                                                                                                                                                                               \r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/piperw/opt/anaconda3/envs/torchgeotest/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:436: Consider setting `persistent_workers=True` in 'train_dataloader' to speed up the dataloader worker initialization.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  78%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                                | 124/158 [10:06<02:46,  0.20it/s, v_num=1]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/piperw/opt/anaconda3/envs/torchgeotest/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py:54: Detected KeyboardInterrupt, attempting graceful shutdown...\n"
     ]
    }
   ],
   "source": [
    "# Train\n",
    "trainer.fit(model=task, datamodule=datamodule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3034caca-1f6b-434d-9e01-6cadf98fb1ac",
   "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}