Update model.py

model.py

@@ -6,9 +6,9 @@ class ResNet18Classifier(nn.Module):
     def __init__(
         self,
         num_classes: int,
-        dropout: float = 0.5,
+        dropout: float = 0.4,
         fc_dim: int = 256,
-        freeze_backbone: bool = True,
+        fine_tune_mode: str = "layer4",
     ):
         super().__init__()
 
@@ -17,9 +17,24 @@ class ResNet18Classifier(nn.Module):
 
         in_features = self.backbone.fc.in_features
 
-        if freeze_backbone:
+        # Freeze everything first
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+        # Fine-tuning strategy
+        if fine_tune_mode == "frozen":
+            pass
+
+        elif fine_tune_mode == "layer4":
+            for param in self.backbone.layer4.parameters():
+                param.requires_grad = True
+
+        elif fine_tune_mode == "full":
             for param in self.backbone.parameters():
-                param.requires_grad = False
+                param.requires_grad = True
+
+        else:
+            raise ValueError(f"Unsupported fine_tune_mode: {fine_tune_mode}")
 
         self.backbone.fc = nn.Sequential(
             nn.Dropout(dropout),
@@ -29,6 +44,7 @@ class ResNet18Classifier(nn.Module):
             nn.Linear(fc_dim, num_classes),
         )
 
+        # Always train classifier head
         for param in self.backbone.fc.parameters():
             param.requires_grad = True