mlp layer + new grid

src/eval_wilds.py

@@ -127,7 +127,14 @@ def main(args):
         crop_size=crop_size,
         model_name=model_name,
     )
-    model = ViTClassifier(encoder, num_classes, embed_dim).to(device)
+    model = ViTClassifier(
+        encoder,
+        num_classes,
+        embed_dim,
+        probe_type=meta_args.get("probe_type", "linear"),
+        mlp_hidden_dim=meta_args.get("mlp_hidden_dim"),
+        dropout=meta_args.get("dropout", 0.0),
+    ).to(device)
 
     checkpoint_path = _resolve_checkpoint_path(meta_args)
     _load_model_state(model, checkpoint_path, device)

src/models/head.py

@@ -3,13 +3,37 @@ import torch.nn as nn
 
 
 class ViTClassifier(nn.Module):
-    def __init__(self, encoder, num_classes, embed_dim):
+    def __init__(
+        self,
+        encoder,
+        num_classes,
+        embed_dim,
+        probe_type="linear",
+        mlp_hidden_dim=None,
+        dropout=0.0,
+    ):
         super().__init__()
         self.encoder = encoder
-        self.head = nn.Linear(embed_dim, num_classes)
 
-        nn.init.trunc_normal_(self.head.weight, std=0.01)
-        nn.init.zeros_(self.head.bias)
+        probe_type = str(probe_type).lower()
+        if probe_type == "linear":
+            self.head = nn.Linear(embed_dim, num_classes)
+        elif probe_type == "mlp":
+            if mlp_hidden_dim is None:
+                raise ValueError("mlp_hidden_dim must be set for probe_type='mlp'")
+            self.head = nn.Sequential(
+                nn.Linear(embed_dim, mlp_hidden_dim),
+                nn.GELU(),
+                nn.Dropout(dropout),
+                nn.Linear(mlp_hidden_dim, num_classes),
+            )
+        else:
+            raise ValueError(f"Unknown probe_type: {probe_type}")
+
+        for module in self.head.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.trunc_normal_(module.weight, std=0.01)
+                nn.init.zeros_(module.bias)
 
     def forward(self, x):
         # ViT -> (B, N, D)

src/train_supervised.py

@@ -207,7 +207,14 @@ def main(args, resume_preempt=False):
     )
 
     embed_dim = m_args["embed_dim"]
-    model = ViTClassifier(encoder, m_args["num_classes"], embed_dim).to(device)
+    model = ViTClassifier(
+        encoder,
+        m_args["num_classes"],
+        embed_dim,
+        probe_type=m_args.get("probe_type", "linear"),
+        mlp_hidden_dim=m_args.get("mlp_hidden_dim"),
+        dropout=m_args.get("dropout", 0.0),
+    ).to(device)
 
     if o_args["freeze_weights"]:
         logger.info("Freezing encoder weights (Linear Probing mode)")