Adds model factory with several model supports

src/model.py

@@ -2,42 +2,26 @@ import lightning as L
 import torch
 from torch import nn
 from torchmetrics.functional import accuracy, cohen_kappa
-from torchvision import models
+from src.models.factory import ModelFactory
 
 
 class DRModel(L.LightningModule):
     def __init__(
-        self, num_classes: int, learning_rate: float = 2e-4, class_weights=None
+        self,
+        num_classes: int,
+        model_name: str = "densenet121",
+        learning_rate: float = 3e-4,
+        class_weights=None,
+        use_scheduler: bool = True,
     ):
         super().__init__()
         self.save_hyperparameters()
         self.num_classes = num_classes
         self.learning_rate = learning_rate
+        self.use_scheduler = use_scheduler
 
         # Define the model
-        # self.model = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
-        # self.model = models.densenet169(weights=models.DenseNet169_Weights.DEFAULT)
-        # self.model = models.densenet161(weights=models.DenseNet161_Weights.DEFAULT)
-        self.model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)
-        # self.model = models.vit_b_32(weights=models.ViT_B_32_Weights.DEFAULT)
-
-        # freeze the feature extractor
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-        # self.model.head.weight.requires_grad = True
-        # self.model.head.bias.requires_grad = True
-
-        # Change the output layer to have the number of classes
-        # in_features = self.model.classifier.in_features
-        in_features = 768
-        self.model.heads = nn.Sequential(
-            # self.model.classifier = nn.Sequential(
-            nn.Linear(in_features, in_features // 2),
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(in_features // 2, num_classes),
-        )
+        self.model = ModelFactory(name=model_name, num_classes=num_classes)()
 
         # Define the loss function
         self.criterion = nn.CrossEntropyLoss(weight=class_weights)
@@ -70,16 +54,20 @@ class DRModel(L.LightningModule):
         self.log("val_kappa", kappa, on_step=True, on_epoch=True, prog_bar=True)
 
     def configure_optimizers(self):
-        # optimizer = torch.optim.Adam(
-        #     self.parameters(), lr=self.learning_rate, weight_decay=1e-4
-        # )
-
         optimizer = torch.optim.AdamW(
             self.parameters(), lr=self.learning_rate, weight_decay=0.05
         )
-        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
-        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+
+        configuration = {
+            "optimizer": optimizer,
+            "monitor": "val_loss",  # monitor validation loss
+        }
+
+        if self.use_scheduler:
+            # Add lr scheduler
+            # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
+            # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                 optimizer,
                 mode="min",  # or "max" if you're maximizing a metric
                 factor=0.1,  # factor by which the learning rate will be reduced
@@ -87,12 +75,7 @@ class DRModel(L.LightningModule):
                 verbose=True,  # print a message when learning rate is reduced
                 threshold=0.001,  # threshold for measuring the new optimum, to only focus on significant changes
             )
-        return {
-            "optimizer": optimizer,
-            "lr_scheduler": {
-                "scheduler": scheduler,
-                "interval": "epoch",
-                "monitor": "val_loss",
-            },
-        }
-        # return optimizer
+
+            configuration["lr_scheduler"] = scheduler
+
+        return configuration

src/models/factory.py

@@ -0,0 +1,141 @@
+from torchvision import models
+from torch import nn
+
+model_mapping = {
+    "densenet121": (
+        models.densenet121,
+        {"weights": models.DenseNet121_Weights.DEFAULT, "family": "densenet"},
+    ),
+    "densenet161": (
+        models.densenet161,
+        {"weights": models.DenseNet161_Weights.DEFAULT, "family": "densenet"},
+    ),
+    "densenet169": (
+        models.densenet169,
+        {"weights": models.DenseNet169_Weights.DEFAULT, "family": "densenet"},
+    ),
+    "densenet201": (
+        models.densenet201,
+        {"weights": models.DenseNet201_Weights.DEFAULT, "family": "densenet"},
+    ),
+    "resnet50": (
+        models.resnet50,
+        {"weights": models.ResNet50_Weights.IMAGENET1K_V2, "family": "resnet"},
+    ),
+    "resnet101": (
+        models.resnet101,
+        {"weights": models.ResNet101_Weights.IMAGENET1K_V2, "family": "resnet"},
+    ),
+    "resnet152": (
+        models.resnet152,
+        {"weights": models.ResNet152_Weights.IMAGENET1K_V2, "family": "resnet"},
+    ),
+    "vit-b-16": (
+        models.vit_b_16,
+        {"weights": models.ViT_B_16_Weights.DEFAULT, "family": "vit"},
+    ),
+    "vit-b-32": (
+        models.vit_b_32,
+        {"weights": models.ViT_B_32_Weights.DEFAULT, "family": "vit"},
+    ),
+    # Add more models as needed with their respective configurations.
+}
+
+
+class Model(nn.Module):
+    """Moodel definition."""
+
+    def __init__(self, model_name: str, num_classes: int):
+        """
+        Initialize Model instance.
+
+        Args:
+            model_name (str): Name of the model architecture.
+            num_classes (int): Number of output classes.
+        """
+        super(Model, self).__init__()
+
+        model_class, model_config = model_mapping[model_name]
+        self.model = model_class(weights=model_config["weights"])
+
+        # Freeze model parameters
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        in_features = self._get_in_features(model_config["family"])
+
+        if model_config["family"] == "densenet":
+            self.model.classifier = self._create_classifier(in_features, num_classes)
+        elif model_config["family"] == "resnet":
+            self.model.fc = self._create_classifier(in_features, num_classes)
+        elif model_config["family"] == "vit":
+            self.model.heads = self._create_classifier(in_features, num_classes)
+
+    def forward(self, x):
+        """Forward pass through the model."""
+        return self.model(x)
+
+    def _get_in_features(self, family: str) -> int:
+        """Return the number of input features for the classifier."""
+        if family == "densenet":
+            return self.model.classifier.in_features
+        elif family == "resnet":
+            return self.model.fc.in_features
+        elif family == "vit":
+            return self.model.heads.head.in_features
+
+    def _create_classifier(self, in_features: int, num_classes: int) -> nn.Sequential:
+        """Create the classifier module."""
+        return nn.Sequential(
+            nn.Linear(in_features, in_features // 2),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(in_features // 2, num_classes),
+        )
+
+
+class ModelFactory:
+    """
+    Factory for creating different models based on their names.
+
+    Args:
+        name (str): The name of the model factory.
+        num_classes (int): The number of output classes.
+
+    Raises:
+        ValueError: If the specified model factory is not implemented.
+    """
+
+    def __init__(self, name: str, num_classes: int):
+        """
+        Initialize ModelFactory instance.
+
+        Args:
+            name (str): The name of the model.
+            num_classes (int): The number of output classes.
+        """
+        self.name = name
+        self.num_classes = num_classes
+
+    def __call__(self):
+        """
+        Create a model instance based on the provided name.
+
+        Args:
+            model_name (str): Name of the model architecture.
+            num_classes (int): Number of output classes.
+
+        Returns:
+            Model: An instance of the selected model.
+        """
+        if self.name not in model_mapping:
+            valid_options = ", ".join(model_mapping.keys())
+            raise ValueError(
+                f"Invalid model name: '{self.name}'. Available options: {valid_options}"
+            )
+
+        return Model(self.name, self.num_classes)
+
+
+if __name__ == "__main__":
+    model = ModelFactory("resnet50", 5)()