Updates models

src/model.py

@@ -1,7 +1,7 @@
 import lightning as L
 import torch
 from torch import nn
-from torchmetrics.functional import accuracy
+from torchmetrics.functional import accuracy, cohen_kappa
 from torchvision import models
 
 
@@ -19,15 +19,15 @@ class DRModel(L.LightningModule):
         self.model = models.densenet169(weights=models.DenseNet169_Weights.DEFAULT)
         # self.model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)
         # freeze the feature extractor
-        for param in self.model.parameters():
-            param.requires_grad = False
+        # for param in self.model.parameters():
+        #     param.requires_grad = False
         # Change the output layer to have the number of classes
         in_features = self.model.classifier.in_features  
         # in_features = 768  
         self.model.classifier = nn.Sequential(
             nn.Linear(in_features, in_features // 2),
             nn.ReLU(),
-            nn.Dropout(0.1),
+            nn.Dropout(0.5),
             nn.Linear(in_features // 2, num_classes),
         )
 
@@ -50,14 +50,22 @@ class DRModel(L.LightningModule):
         loss = self.criterion(logits, y)
         preds = torch.argmax(logits, dim=1)
         acc = accuracy(preds, y, task="multiclass", num_classes=self.num_classes)
+        kappa = cohen_kappa(preds, y, task="multiclass", num_classes=self.num_classes, weights="quadratic")
         self.log("val_loss", loss, prog_bar=True)
         self.log("val_acc", acc, prog_bar=True)
+        self.log("val_kappa", kappa, prog_bar=True)
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(
+        # optimizer = torch.optim.Adam(
+        #     self.parameters(), lr=self.learning_rate, weight_decay=1e-4
+        # )
+        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.1, patience=3, verbose=True)
+
+        optimizer = torch.optim.AdamW(
             self.parameters(), lr=self.learning_rate, weight_decay=1e-4
         )
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
+        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.1, patience=3, verbose=True)
         return {
             "optimizer": optimizer,
             "lr_scheduler": {