create pl model

model.py

@@ -0,0 +1,111 @@
+import pytorch_lightning as pl
+import torch
+
+from datasets import load_metric
+from torch import nn
+from transformers import SegformerForSemanticSegmentation
+from typing import Dict
+
+
+class SidewalkSegmentationModel(pl.LightningModule):
+    def __init__(
+        self, 
+        num_labels: int, 
+        id2label: Dict[int, str],
+        model_flavor: int = 0,
+        learning_rate: float = 6e-5,
+    ):
+        super().__init__()
+        self.id2label = id2label
+        self.label2id = {v: k for k, v in id2label.items()}
+        self.learning_rate = learning_rate
+        self.metrics = {
+            "train": load_metric("mean_iou"),
+            "val": load_metric("mean_iou"),
+        }
+
+        self.model = SegformerForSemanticSegmentation.from_pretrained(
+            f"nvidia/mit-b{model_flavor}", num_labels=num_labels, id2label=self.id2label, label2id=self.label2id,
+        )
+        self.save_hyperparameters()
+
+    
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    
+    def training_step(self, batch, batch_idx):
+        pixel_values = batch["pixel_values"]
+        labels = batch["labels"]
+
+        outputs = self(pixel_values=pixel_values, labels=labels)
+        loss, logits = outputs.loss, outputs.logits
+
+        self.add_batch_to_metric("train", logits, labels)
+        self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
+        return {"loss": loss}
+
+    
+    def validation_step(self, batch, batch_idx):
+        pixel_values = batch["pixel_values"]
+        labels = batch["labels"]
+
+        outputs = self(pixel_values=pixel_values, labels=labels)
+        loss, logits = outputs.loss, outputs.logits
+
+        self.add_batch_to_metric("val", logits, labels)
+        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
+        return {"val_loss": loss}
+
+    
+    def training_epoch_end(self, training_step_outputs):
+        """
+        Log the training metrics.
+        """
+        metrics = self.metrics["train"].compute(num_labels=len(self.id2label), ignore_index=255, reduce_labels=False)
+        self.log("train_mean_iou", metrics["mean_iou"], prog_bar=True, on_step=False, on_epoch=True)
+        self.log("train_mean_acc", metrics["mean_accuracy"], prog_bar=True, on_step=False, on_epoch=True)
+
+    
+    def validation_epoch_end(self, validation_step_outputs):
+        """
+        Log the validation metrics.
+        """
+        metrics = self.metrics["val"].compute(num_labels=len(self.id2label), ignore_index=255, reduce_labels=False)
+        self.log("val_mean_iou", metrics["mean_iou"], prog_bar=True, on_step=False, on_epoch=True)
+        self.log("val_mean_acc", metrics["mean_accuracy"], prog_bar=True, on_step=False, on_epoch=True)
+
+    
+    def add_batch_to_metric(self, stage: str, logits: torch.Tensor, labels: torch.Tensor):
+        """
+        Add the current batch to the metric.
+
+        Parameters
+        ----------
+        stage : str
+            Stage of the training. Either "train" or "val".
+        logits : torch.Tensor
+            Predicted logits.
+        labels : torch.Tensor
+            Ground truth labels.
+        """
+        with torch.no_grad():
+            upsampled_logits = nn.functional.interpolate(
+                logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
+            )
+            predicted = upsampled_logits.argmax(dim=1)
+            self.metrics[stage].add_batch(
+                predictions=predicted.detach().cpu().numpy(), references=labels.detach().cpu().numpy()
+            )
+
+    
+    def configure_optimizers(self) -> torch.optim.AdamW:
+        """
+        Configure the optimizer.
+
+        Returns
+        -------
+        torch.optim.AdamW
+            Optimizer for the model
+        """
+        return torch.optim.AdamW(self.parameters(), lr=self.learning_rate)