{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Fine-tuning a SegFormer model with Pytorch Lightning\n", "\n", "This notebook fine-tunes a SegFormer model for semantic segmentation using PyTorch Lightning.\n", "\n", "The original notebook can be found [here](https://colab.research.google.com/drive/1250K828ixr-sG2xzLVYYN5AVa3cTaDEF).\n", "\n", "![](../images/lightning-balloons-semseg.png)\n", "\n", "\n", "\n", "In this tutorial we fine-tune a pre-trained SegFormer model for semantic segmentation on a custom dataset.\n", "We integrate with 3LC by creating Tables up front, initializing a Run, and collecting per-sample\n", "predicted masks from raw Lightning hooks.\n", "\n", "This notebook demonstrates:\n", "\n", "+ Training a SegFormer model on a custom dataset with PyTorch Lightning.\n", "+ Registering train/val/test sets into 3LC Tables.\n", "+ Collecting per-sample semantic segmentation predicted masks via `tlc.collect_metrics`\n", " called from `on_train_epoch_end` / `on_train_end`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Project Setup" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "parameters" ] }, "outputs": [], "source": [ "PROJECT_NAME = \"3LC Tutorials - Image Segmentation\"\n", "RUN_NAME = \"Train Balloon SegFormer\"\n", "DESCRIPTION = \"Train a SegFormer model using PyTorch Lightning\"\n", "TRAIN_DATASET_NAME = \"balloons-train\"\n", "VAL_DATASET_NAME = \"balloons-val\"\n", "TEST_DATASET_NAME = \"balloons-test\"\n", "SEGFORMER_MODEL_ID = \"nvidia/mit-b0\"\n", "DATA_PATH = \"../../data\"\n", "EPOCHS = 10\n", "BATCH_SIZE = 8\n", "NUM_WORKERS = 0\n", "DEVICE = None\n", "LEARNING_RATE = 2e-05\n", "INSTALL_DEPENDENCIES = True" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Install dependencies" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if INSTALL_DEPENDENCIES:\n", " %pip install -q 3lc[huggingface]\n", " %pip install -q pytorch-lightning\n", " %pip install -q matplotlib" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Imports" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "import numpy as np\n", "import pytorch_lightning as pl\n", "import tlc\n", "import torch\n", "from evaluate import load\n", "from matplotlib import pyplot as plt\n", "from PIL import Image\n", "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n", "from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\n", "from tlc.integration.torch.samplers import create_sampler\n", "from torch import nn\n", "from torch.utils.data import DataLoader\n", "from torchvision.datasets import VisionDataset\n", "from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if DEVICE is None:\n", " if torch.cuda.is_available():\n", " device = \"cuda:0\"\n", " elif torch.backends.mps.is_available():\n", " # Disable MPS due to tensor view issues with SegFormer\n", " device = \"cpu\"\n", " else:\n", " device = \"cpu\"\n", "else:\n", " device = DEVICE\n", "\n", "device = torch.device(device)\n", "print(f\"Using device: {device}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup Datasets and Training helpers\n", "\n", "We will create a Table with the images and their associated masks.\n", "\n", "Moreover, we will also define helpers to pre-process this dataset into a suitable form for training and collecting metrics.\n", "\n", "To finish, we define a Pytorch LightningModule to define the steps for training, validation and test.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class TLCSemanticSegmentationDataset(VisionDataset):\n", " \"\"\"Image (semantic) segmentation dataset.\"\"\"\n", "\n", " def __init__(self, root_dir):\n", " super().__init__(root_dir)\n", " self.root_dir = root_dir\n", " image_file_names = [f for f in os.listdir(self.root_dir) if \".jpg\" in f]\n", " mask_file_names = [f for f in os.listdir(self.root_dir) if \".png\" in f]\n", " self.images = sorted(image_file_names)\n", " self.masks = sorted(mask_file_names)\n", "\n", " def __len__(self):\n", " return len(self.images)\n", "\n", " def __getitem__(self, idx):\n", " image = Image.open(os.path.join(self.root_dir, self.images[idx]))\n", " segmentation_map = Image.open(os.path.join(self.root_dir, self.masks[idx]))\n", " return image, segmentation_map, segmentation_map.width, segmentation_map.height" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_id2label(root_dir):\n", " classes_csv_file = os.path.join(root_dir, \"_classes.csv\")\n", " with open(classes_csv_file) as fid:\n", " data = [line.split(\",\") for idx, line in enumerate(fid) if idx != 0]\n", " # transformers >=5 requires int keys for `id2label`.\n", " return {int(float(x[0])): x[1].strip() for x in data}\n", "\n", "\n", "image_processor = SegformerImageProcessor.from_pretrained(SEGFORMER_MODEL_ID)\n", "image_processor.do_reduce_labels = False\n", "image_processor.size = 128\n", "dataset_location = tlc.Url(DATA_PATH + \"/balloons-mask-segmentation\").to_absolute()\n", "id2label = get_id2label(f\"{dataset_location}/train/\") # Assuming the same classes for train, val, and test\n", "\n", "model = SegformerForSemanticSegmentation.from_pretrained(\n", " SEGFORMER_MODEL_ID,\n", " num_labels=len(id2label.keys()),\n", " id2label=id2label,\n", " label2id={v: k for k, v in id2label.items()},\n", " ignore_mismatched_sizes=True,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "schema = {\n", " \"image\": tlc.schemas.ImageSchema(),\n", " \"segmentation_map\": tlc.schemas.SemanticSegmentationSchema(classes=id2label),\n", " \"mask_width\": tlc.schemas.Int32Schema(),\n", " \"mask_height\": tlc.schemas.Int32Schema(),\n", " # `TableWriter` doesn't add a weight column by default; we add one explicitly\n", " # so the 3LC sampler has weights to draw from.\n", " \"weight\": tlc.schemas.SampleWeightSchema(),\n", "}\n", "\n", "\n", "def mc_preprocess_fn(batch, predictor_output):\n", " \"\"\"Transform a batch of inputs and model outputs to a format expected by the metrics collector.\"\"\"\n", "\n", " original_mask_size = batch[\"mask_size\"].tolist()\n", " outputs = predictor_output.forward\n", "\n", " predicted_masks = image_processor.post_process_semantic_segmentation(\n", " outputs=outputs,\n", " target_sizes=original_mask_size,\n", " )\n", " return batch, predicted_masks\n", "\n", "\n", "segmentation_metrics_collector = tlc.metrics.SegmentationMetricsCollector(\n", " label_map=id2label,\n", " preprocess_fn=mc_preprocess_fn,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "################## 3LC ##################\n", "# Create the 3LC Tables up front, OUTSIDE the LightningModule. This sidesteps DDP\n", "# coordination inside `train_dataloader()`: each rank simply opens the same Table from disk.\n", "\n", "\n", "def transforms(sample):\n", " encoded_inputs = image_processor(sample[\"image\"], sample[\"segmentation_map\"], return_tensors=\"pt\")\n", " for k in encoded_inputs:\n", " encoded_inputs[k] = encoded_inputs[k].squeeze() # remove batch dimension\n", "\n", " # Add the original mask size so we can resize the predicted mask back later.\n", " encoded_inputs[\"mask_size\"] = torch.tensor([sample[\"mask_width\"], sample[\"mask_height\"]])\n", " return dict(encoded_inputs)\n", "\n", "\n", "def write_balloons_table(root_dir, dataset_name):\n", " \"\"\"Stream a balloons split into a 3LC Table via TableWriter.\"\"\"\n", " dataset = TLCSemanticSegmentationDataset(root_dir)\n", " writer = tlc.TableWriter(\n", " project_name=PROJECT_NAME,\n", " dataset_name=dataset_name,\n", " schema=schema,\n", " if_exists=\"overwrite\",\n", " )\n", " images, masks, widths, heights = [], [], [], []\n", " for image, segmentation_map, mask_width, mask_height in dataset:\n", " images.append(image)\n", " masks.append(segmentation_map)\n", " widths.append(mask_width)\n", " heights.append(mask_height)\n", " writer.add_batch(\n", " {\n", " \"image\": images,\n", " \"segmentation_map\": masks,\n", " \"mask_width\": widths,\n", " \"mask_height\": heights,\n", " \"weight\": [1.0] * len(images),\n", " }\n", " )\n", " return writer.finalize()\n", "\n", "\n", "train_table = write_balloons_table(f\"{dataset_location}/train/\", TRAIN_DATASET_NAME)\n", "val_table = write_balloons_table(f\"{dataset_location}/valid/\", VAL_DATASET_NAME)\n", "test_table = write_balloons_table(f\"{dataset_location}/test/\", TEST_DATASET_NAME)\n", "#########################################" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class SegformerFinetuner(pl.LightningModule):\n", " def __init__(\n", " self,\n", " model,\n", " id2label,\n", " train_table,\n", " val_table,\n", " test_table,\n", " metrics_interval=100,\n", " learning_rate=2e-05,\n", " ):\n", " super().__init__()\n", " self.save_hyperparameters(ignore=[\"model\", \"id2label\", \"train_table\", \"val_table\", \"test_table\"])\n", " self.train_table = train_table\n", " self.val_table = val_table\n", " self.test_table = test_table\n", " self.metrics_interval = metrics_interval\n", " self.learning_rate = learning_rate\n", "\n", " self.id2label = id2label\n", " self.num_classes = len(id2label.keys())\n", "\n", " self.model = model\n", "\n", " self.train_mean_iou = load(\"mean_iou\")\n", " self.val_mean_iou = load(\"mean_iou\")\n", " self.test_mean_iou = load(\"mean_iou\")\n", "\n", " self.training_step_outputs = []\n", " self.validation_step_outputs = []\n", " self.test_step_outputs = []\n", "\n", " self.tlc_run: tlc.Run | None = None\n", "\n", " def forward(self, images, masks=None):\n", " outputs = self.model(images, masks)\n", " return outputs\n", "\n", " def training_step(self, batch, batch_idx):\n", " images, masks = batch[\"pixel_values\"], batch[\"labels\"]\n", " outputs = self(images, masks)\n", " loss, logits = outputs[0], outputs[1]\n", " upsampled_logits = nn.functional.interpolate(\n", " logits, size=masks.shape[-2:], mode=\"bilinear\", align_corners=False\n", " )\n", " predicted = upsampled_logits.argmax(dim=1)\n", " self.train_mean_iou.add_batch(\n", " predictions=predicted.detach().cpu().numpy(),\n", " references=masks.detach().cpu().numpy(),\n", " )\n", " if batch_idx % self.metrics_interval == 0:\n", " metrics = self.train_mean_iou.compute(\n", " num_labels=self.num_classes,\n", " ignore_index=255,\n", " reduce_labels=False,\n", " )\n", " metrics = {\n", " \"loss\": loss,\n", " \"mean_iou\": metrics[\"mean_iou\"],\n", " \"mean_accuracy\": metrics[\"mean_accuracy\"],\n", " }\n", " for k, v in metrics.items():\n", " self.log(k, v, prog_bar=True)\n", "\n", " tlc.log(\n", " {\n", " **{k: v.item() for k, v in metrics.items()},\n", " \"step\": self.global_step,\n", " }\n", " )\n", " else:\n", " metrics = {\"loss\": loss}\n", "\n", " self.training_step_outputs.append(metrics)\n", " return metrics\n", "\n", " def validation_step(self, batch):\n", " images, masks = batch[\"pixel_values\"], batch[\"labels\"]\n", " outputs = self(images, masks)\n", " loss, logits = outputs[0], outputs[1]\n", " upsampled_logits = nn.functional.interpolate(\n", " logits, size=masks.shape[-2:], mode=\"bilinear\", align_corners=False\n", " )\n", " predicted = upsampled_logits.argmax(dim=1)\n", " self.val_mean_iou.add_batch(\n", " predictions=predicted.detach().cpu().numpy(),\n", " references=masks.detach().cpu().numpy(),\n", " )\n", " self.validation_step_outputs.append(loss)\n", " return {\"val_loss\": loss}\n", "\n", " def on_validation_epoch_end(self):\n", " metrics = self.val_mean_iou.compute(\n", " num_labels=self.num_classes,\n", " ignore_index=255,\n", " reduce_labels=False,\n", " )\n", "\n", " avg_val_loss = torch.stack(self.validation_step_outputs).mean()\n", " val_mean_iou = metrics[\"mean_iou\"]\n", " val_mean_accuracy = metrics[\"mean_accuracy\"]\n", "\n", " metrics = {\n", " \"val_loss\": avg_val_loss,\n", " \"val_mean_iou\": val_mean_iou,\n", " \"val_mean_accuracy\": val_mean_accuracy,\n", " }\n", " for k, v in metrics.items():\n", " self.log(k, v, prog_bar=True)\n", "\n", " self.validation_step_outputs.clear()\n", "\n", " if not self.trainer.sanity_checking:\n", " tlc.log(\n", " {\n", " **{k: v.item() for k, v in metrics.items()},\n", " \"step\": self.global_step,\n", " }\n", " )\n", "\n", " return metrics\n", "\n", " def test_step(self, batch):\n", " images, masks = batch[\"pixel_values\"], batch[\"labels\"]\n", " outputs = self(images, masks)\n", " loss, logits = outputs[0], outputs[1]\n", " upsampled_logits = nn.functional.interpolate(\n", " logits, size=masks.shape[-2:], mode=\"bilinear\", align_corners=False\n", " )\n", " predicted = upsampled_logits.argmax(dim=1)\n", " self.test_mean_iou.add_batch(\n", " predictions=predicted.detach().cpu().numpy(),\n", " references=masks.detach().cpu().numpy(),\n", " )\n", " self.test_step_outputs.append(loss)\n", " return {\"test_loss\": loss}\n", "\n", " def on_test_epoch_end(self):\n", " metrics = self.test_mean_iou.compute(\n", " num_labels=self.num_classes,\n", " ignore_index=255,\n", " reduce_labels=False,\n", " )\n", "\n", " avg_test_loss = torch.stack(self.test_step_outputs).mean()\n", " test_mean_iou = metrics[\"mean_iou\"]\n", " test_mean_accuracy = metrics[\"mean_accuracy\"]\n", " metrics = {\n", " \"test_loss\": avg_test_loss,\n", " \"test_mean_iou\": test_mean_iou,\n", " \"test_mean_accuracy\": test_mean_accuracy,\n", " }\n", " for k, v in metrics.items():\n", " self.log(k, v)\n", " self.test_step_outputs.clear()\n", " return metrics\n", "\n", " def configure_optimizers(self):\n", " return torch.optim.Adam(\n", " [p for p in self.parameters() if p.requires_grad],\n", " lr=self.learning_rate,\n", " eps=1e-08,\n", " )\n", "\n", " def train_dataloader(self):\n", " return DataLoader(\n", " self.train_table.with_transform(transforms),\n", " sampler=create_sampler(self.train_table, weighted=True, exclude_zero_weights=True),\n", " batch_size=BATCH_SIZE,\n", " num_workers=NUM_WORKERS,\n", " )\n", "\n", " def val_dataloader(self):\n", " return DataLoader(\n", " self.val_table.with_transform(transforms),\n", " batch_size=BATCH_SIZE,\n", " num_workers=NUM_WORKERS,\n", " )\n", "\n", " def test_dataloader(self):\n", " return DataLoader(\n", " self.test_table.with_transform(transforms),\n", " batch_size=BATCH_SIZE,\n", " num_workers=NUM_WORKERS,\n", " )\n", "\n", " ################## 3LC ##################\n", " def on_train_start(self):\n", " super().on_train_start()\n", " self.tlc_run = tlc.init(\n", " project_name=PROJECT_NAME,\n", " run_name=RUN_NAME,\n", " description=DESCRIPTION,\n", " parameters=dict(self.hparams_initial),\n", " if_exists=\"overwrite\",\n", " )\n", " self.tlc_run.set_status_running()\n", "\n", " def on_train_end(self):\n", " super().on_train_end()\n", " self._collect_3lc_metrics()\n", " if self.tlc_run is not None:\n", " self.tlc_run.set_status_completed()\n", "\n", " def _collect_3lc_metrics(self):\n", " predictor = tlc.metrics.Predictor(self)\n", " for split, table in [(\"train\", self.train_table), (\"val\", self.val_table)]:\n", " tlc.collect_metrics(\n", " table=table.with_transform(transforms),\n", " metrics_collectors=[segmentation_metrics_collector],\n", " predictor=predictor,\n", " split=split,\n", " constants={\"epoch\": self.current_epoch},\n", " exclude_zero_weights=True,\n", " )\n", "\n", " #########################################\n", "\n", "\n", "segformer_finetuner = SegformerFinetuner(\n", " model=model,\n", " id2label=id2label,\n", " train_table=train_table,\n", " val_table=val_table,\n", " test_table=test_table,\n", " learning_rate=LEARNING_RATE,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Training the model" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "early_stop_callback = EarlyStopping(\n", " monitor=\"val_loss\",\n", " patience=5,\n", " verbose=True,\n", ")\n", "\n", "checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor=\"val_loss\", save_last=True)\n", "\n", "trainer = pl.Trainer(\n", " accelerator=\"cpu\", # CPU avoids MPS tensor view issues with SegFormer\n", " callbacks=[early_stop_callback, checkpoint_callback],\n", " max_epochs=EPOCHS,\n", " val_check_interval=1.0, # validate once per training epoch\n", " log_every_n_steps=7,\n", ")\n", "\n", "trainer.fit(segformer_finetuner)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Checking results " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "res = trainer.test(ckpt_path=\"last\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "\n", "mapped_test_dataloader = segformer_finetuner.test_dataloader()\n", "\n", "batch = next(iter(mapped_test_dataloader))\n", "images, masks = batch[\"pixel_values\"].to(device), batch[\"labels\"].to(device)\n", "segformer_finetuner.eval().to(device)\n", "outputs = segformer_finetuner.model(images, masks, return_dict=True)\n", "batch_prediction = image_processor.post_process_semantic_segmentation(outputs, batch[\"mask_size\"].tolist())\n", "\n", "n_rows = len(images)\n", "n_cols = 3\n", "fig_width = n_cols * 5\n", "fig_height = n_rows * 5\n", "fig, ax = plt.subplots(n_rows, n_cols, figsize=(fig_width, fig_height))\n", "fig.suptitle(\"Test Batch Predictions\", fontsize=16)\n", "plt.tight_layout(pad=3.0, h_pad=-1.0, w_pad=1.0, rect=[0, 0, 1, 1])\n", "\n", "for i in range(n_rows):\n", " for j in range(3):\n", " ax[i, j].axis(\"off\")\n", "\n", " ax[i, 0].imshow(masks[i, :, :].cpu().numpy(), cmap=\"gray\")\n", " ax[i, 0].set_title(f\"Ground Truth (id={i})\", fontsize=14)\n", "\n", " ax[i, 1].imshow(batch_prediction[i].cpu().numpy(), cmap=\"gray\")\n", " ax[i, 1].set_title(\"Predicted mask (latest model)\", fontsize=14)\n", "\n", " im = tlc.active_run().metrics_tables[-1][i][\"predicted_mask\"]\n", " ax[i, 2].imshow(np.array(im), cmap=\"gray\")\n", " ax[i, 2].set_title(\"Predicted mask (3LC metrics)\", fontsize=14)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [], "toc_visible": true }, "kernelspec": { "display_name": ".venv", "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.13.3" } }, "nbformat": 4, "nbformat_minor": 0 }