clean up configs

README.md

@@ -13,7 +13,7 @@ Reference: official I-JEPA README https://github.com/facebookresearch/ijepa/blob
 ## Models
 
 - ViT-H, 14x14 patches, 224x224 resolution (trained)
-- ViT-H, 16x16 patches, 448x448 resolution (planned)
+- ViT-H, 16x16 patches, 448x448 resolution (trained)
 - Plan: add a graph comparing models with the WILDS leaderboard https://wilds.stanford.edu/leaderboard/#with-unlabeled-data-1
 
 <!-- Optional: add a WILDS leaderboard comparison graph here -->
@@ -25,7 +25,7 @@ Reference: official I-JEPA README https://github.com/facebookresearch/ijepa/blob
 - `src/train_supervised.py`: supervised training loop
 - `configs/`: training configs
 - `configs/wilds_vith14_ep300.yaml`: SSL config used here
-- `configs/supervised_wilds_vith14_ep300.yaml`: supervised config used here
+- `configs/supervised_vith14_224.yaml`: supervised config used here (see `configs/` for all supervised linear-probe configs)
 - `main_distributed.py`: entrypoint for distributed SSL training
 - `main_distributed_supervised.py`: entrypoint for distributed supervised training
 - `requirements.txt`: dependencies
@@ -50,7 +50,7 @@ python3 main_distributed.py --fname configs/wilds_vith14_ep300.yaml --folder $su
 Supervised fine-tuning:
 
 ```
-python3 main_distributed_supervised.py --fname configs/supervised_wilds_vith14_ep300.yaml --folder $submitit_folder --partition $slurm_partition --nodes $nodes --tasks-per-node $tasks_per_node --time $time
+python3 main_distributed_supervised.py --fname configs/supervised_vith14_224.yaml --folder $submitit_folder --partition $slurm_partition --nodes $nodes --tasks-per-node $tasks_per_node --time $time
 ```
 
 Evaluation on iWildCam test split:

configs/supervised_vitb16_448.yaml

@@ -10,12 +10,11 @@ meta:
   head_type: linear
 
 data:
-  batch_size: 1024
+  batch_size: 128
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
   crop_size: 448
-  crop_scale: [1.0, 1.0]
   use_random_resized_crop: false
   use_horizontal_flip: false
   use_color_distortion: false
@@ -26,25 +25,14 @@ mask:
   patch_size: 16
 
 optimization:
-  optimizer: adamw # 'adamw', 'sgd', or 'lars'
-  freeze_weights: true # true for linear probing, false for full fine-tuning
+  optimizer: adamw
+  freeze_weights: true
   epochs: 300
   lr: 0.01
-  weight_decay: 0
+  weight_decay: 5.0e-4
   lr_schedule: cosine
-  step_milestones: [15, 30, 45]
-  step_gamma: 0.1
-  start_lr: 0.0
+  gradient_accumulation_steps: 1
   final_lr: 0.0
-  warmup: 0
-  momentum: 0.9
-  lars_eta: 0.001
-  lars_eps: 1.0e-8
-  lars_exclude_bias_and_norm: true
-  ipe_scale: 1.0
-  # Sweep suggestions (manual edits):
-  # - lr: 0.01 | 0.05 | 0.001
-  # - weight_decay: 5.0e-4 | 0.0
   early_stopping:
     enabled: true
     patience: 10

configs/supervised_vitg16_224_in22k.yaml

@@ -3,12 +3,12 @@ meta:
   tag: in22k
   embed_dim: 1408
   load_checkpoint: true
-  checkpoint_folder: experiment_logs/imagenet-vith16.448/
+  checkpoint_folder: experiment_logs/imagenet/
   read_checkpoint: IN22K-vit.g.16-600e.pth.tar
   use_bfloat16: true
   num_classes: 182
   representation_type: last_avgpool
-  head_type: linear
+  head_type: bn_linear
   use_gradient_checkpointing: true
 
 data:
@@ -17,7 +17,6 @@ data:
   num_workers: 10
   pin_mem: true
   crop_size: 224
-  crop_scale: [1.0, 1.0]
   use_random_resized_crop: false
   use_horizontal_flip: false
   use_color_distortion: false
@@ -31,20 +30,11 @@ optimization:
   optimizer: adamw
   freeze_weights: true
   epochs: 300
-  lr: 0.01
+  lr: 0.001
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  step_milestones: [15, 30, 45]
-  step_gamma: 0.1
-  start_lr: 0.0
+  gradient_accumulation_steps: 16
   final_lr: 0.0
-  warmup: 0
-  momentum: 0.99
-  lars_eta: 0.001
-  lars_eps: 1.0e-8
-  lars_exclude_bias_and_norm: true
-  ipe_scale: 1.0
-  gradient_accumulation_steps: 1
   early_stopping:
     enabled: true
     patience: 10

configs/supervised_vith14_224.yaml

@@ -0,0 +1,48 @@
+meta:
+  model_name: vit_huge
+  embed_dim: 1280
+  load_checkpoint: true
+  checkpoint_folder: experiment_logs/vith14.224-bs.128-ep.300/
+  read_checkpoint: jepa-ep300.pth.tar
+  use_bfloat16: true
+  num_classes: 182
+  representation_type: last4_avgpool_concat
+  head_type: bn_linear
+
+data:
+  batch_size: 16
+  root_path: ./wilds_data
+  num_workers: 10
+  pin_mem: true
+  crop_size: 224
+  use_random_resized_crop: false
+  use_horizontal_flip: false
+  use_color_distortion: false
+  color_jitter_strength: 0.0
+  use_gaussian_blur: false
+
+mask:
+  patch_size: 14
+
+optimization:
+  optimizer: adamw
+  freeze_weights: true
+  epochs: 300
+  lr: 0.001
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 32
+  final_lr: 0.0
+  early_stopping:
+    enabled: true
+    patience: 10
+    min_delta: 1.0e-4
+    min_epochs: 15
+    restore_best_weights: true
+
+validation:
+  eval_every: 1
+
+logging:
+  write_tag: linear_probe
+  auto_folder: true

configs/supervised_vith14_224_in1k.yaml

@@ -0,0 +1,49 @@
+meta:
+  model_name: vit_huge
+  tag: in1k
+  embed_dim: 1280
+  load_checkpoint: true
+  checkpoint_folder: experiment_logs/imagenet/
+  read_checkpoint: IN1K-vit.h.14-300e.pth.tar
+  use_bfloat16: true
+  num_classes: 182
+  representation_type: last_avgpool
+  head_type: linear
+
+data:
+  batch_size: 16
+  root_path: ./wilds_data
+  num_workers: 10
+  pin_mem: true
+  crop_size: 224
+  use_random_resized_crop: false
+  use_horizontal_flip: false
+  use_color_distortion: false
+  color_jitter_strength: 0.0
+  use_gaussian_blur: false
+
+mask:
+  patch_size: 14
+
+optimization:
+  optimizer: adamw
+  freeze_weights: true
+  epochs: 300
+  lr: 0.01
+  weight_decay: 0
+  lr_schedule: cosine
+  gradient_accumulation_steps: 64
+  final_lr: 0.0
+  early_stopping:
+    enabled: true
+    patience: 10
+    min_delta: 1.0e-4
+    min_epochs: 15
+    restore_best_weights: true
+
+validation:
+  eval_every: 1
+
+logging:
+  write_tag: linear_probe
+  auto_folder: true

configs/supervised_vith14_224_in22k.yaml

@@ -3,20 +3,19 @@ meta:
   tag: in22k
   embed_dim: 1280
   load_checkpoint: true
-  checkpoint_folder: experiment_logs/imagenet-vith16.448/
+  checkpoint_folder: experiment_logs/imagenet/
   read_checkpoint: IN22K-vit.h.14-900e.pth.tar
   use_bfloat16: true
   num_classes: 182
   representation_type: last_avgpool
-  head_type: linear
+  head_type: bn_linear
 
 data:
-  batch_size: 256
+  batch_size: 16
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
   crop_size: 224
-  crop_scale: [1.0, 1.0]
   use_random_resized_crop: false
   use_horizontal_flip: false
   use_color_distortion: false
@@ -27,25 +26,14 @@ mask:
   patch_size: 14
 
 optimization:
-  optimizer: adamw # 'adamw', 'sgd', or 'lars'
-  freeze_weights: true # true for linear probing, false for full fine-tuning
+  optimizer: adamw
+  freeze_weights: true
   epochs: 300
-  lr: 0.01
+  lr: 0.001
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  step_milestones: [15, 30, 45]
-  step_gamma: 0.1
-  start_lr: 0.0
+  gradient_accumulation_steps: 32
   final_lr: 0.0
-  warmup: 0
-  momentum: 0.99
-  lars_eta: 0.001
-  lars_eps: 1.0e-8
-  lars_exclude_bias_and_norm: true
-  ipe_scale: 1.0
-  # Sweep suggestions (manual edits):
-  # - lr: 0.01 | 0.05 | 0.001
-  # - weight_decay: 5.0e-4 | 0.0
   early_stopping:
     enabled: true
     patience: 10

configs/supervised_vith16_448.yaml

@@ -7,15 +7,14 @@ meta:
   use_bfloat16: true
   num_classes: 182
   representation_type: last_avgpool
-  head_type: linear
+  head_type: bn_linear
 
 data:
-  batch_size: 1024
+  batch_size: 16
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
   crop_size: 448
-  crop_scale: [1.0, 1.0]
   use_random_resized_crop: false
   use_horizontal_flip: false
   use_color_distortion: false
@@ -26,28 +25,17 @@ mask:
   patch_size: 16
 
 optimization:
-  optimizer: adamw # 'adamw', 'sgd', or 'lars'
-  freeze_weights: true # true for linear probing, false for full fine-tuning
+  optimizer: adamw
+  freeze_weights: true
   epochs: 300
   lr: 0.01
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  step_milestones: [15, 30, 45]
-  step_gamma: 0.1
-  start_lr: 0.0
+  gradient_accumulation_steps: 64
   final_lr: 0.0
-  warmup: 0
-  momentum: 0.99
-  lars_eta: 0.001
-  lars_eps: 1.0e-8
-  lars_exclude_bias_and_norm: true
-  ipe_scale: 1.0
-  # Sweep suggestions (manual edits):
-  # - lr: 0.01 | 0.05 | 0.001
-  # - weight_decay: 5.0e-4 | 0.0
   early_stopping:
     enabled: true
-    patience: 15
+    patience: 10
     min_delta: 1.0e-4
     min_epochs: 15
     restore_best_weights: true

configs/supervised_vith16_448_in1k.yaml

@@ -0,0 +1,50 @@
+meta:
+  model_name: vit_huge
+  tag: in1k
+  embed_dim: 1280
+  load_checkpoint: true
+  checkpoint_folder: experiment_logs/imagenet/
+  read_checkpoint: IN1K-vit.h.16-448px-300e.pth.tar
+  use_bfloat16: true
+  num_classes: 182
+  representation_type: last4_avgpool_concat
+  head_type: bn_linear
+  use_gradient_checkpointing: true
+
+data:
+  batch_size: 64
+  root_path: ./wilds_data
+  num_workers: 10
+  pin_mem: true
+  crop_size: 448
+  use_random_resized_crop: false
+  use_horizontal_flip: false
+  use_color_distortion: false
+  color_jitter_strength: 0.0
+  use_gaussian_blur: false
+
+mask:
+  patch_size: 16
+
+optimization:
+  optimizer: adamw
+  freeze_weights: true
+  epochs: 300
+  lr: 0.001
+  weight_decay: 0
+  lr_schedule: cosine
+  gradient_accumulation_steps: 8
+  final_lr: 0.0
+  early_stopping:
+    enabled: true
+    patience: 10
+    min_delta: 1.0e-4
+    min_epochs: 15
+    restore_best_weights: true
+
+validation:
+  eval_every: 1
+
+logging:
+  write_tag: linear_probe
+  auto_folder: true