remove mlp + add batch norm + last 4 avg pool layer as in the paper of IJEPA

configs/eval_wilds_vith14.yaml

@@ -6,9 +6,8 @@ meta:
   patch_size: 14
   crop_size: 224
   use_bfloat16: true
-  probe_type: linear
-  mlp_hidden_dim:
-  dropout: 0.0
+  representation_type: last_avgpool
+  head_type: linear
   checkpoint_folder: experiment_logs/supervised/
   read_checkpoint: linear_probe-best.pth.tar
 

configs/grids/lp_grid.yaml

@@ -10,9 +10,8 @@ grid:
   optimization.momentum: [0.9, 0.99, 0.5]
   optimization.lr_schedule: [cosine, step]
   data.batch_size: [32, 64, 128, 256]
-  meta.probe_type: [linear]
-  # meta.mlp_hidden_dim: [512, 1024, 2048]
-  # meta.dropout: [0.0, 0.1, 0.3]
+  meta.representation_type: [last_avgpool]
+  meta.head_type: [linear]
 
 launch:
   folder: submitit_logs/

configs/supervised_wilds_vitb16-448_ep600-lp.yaml

@@ -6,9 +6,8 @@ meta:
   read_checkpoint: jepa-ep150.pth.tar
   use_bfloat16: true
   num_classes: 182
-  probe_type: linear
-  mlp_hidden_dim: 1024
-  dropout: 0.3
+  representation_type: last_avgpool
+  head_type: linear
 
 data:
   batch_size: 128

configs/supervised_wilds_vith14_ep300-lp.yaml

@@ -6,9 +6,8 @@ meta:
   read_checkpoint: jepa-ep300.pth.tar
   use_bfloat16: true
   num_classes: 182
-  probe_type: mlp
-  mlp_hidden_dim: 1024
-  dropout: 0.3
+  representation_type: last_avgpool
+  head_type: linear
 
 data:
   batch_size: 256