update README

.gitignore

@@ -1,3 +1,8 @@
 *.swp
 *.swo
 __pycache__
+.venv/
+experiment_logs/
+wilds_data/
+submitit_logs/
+test.py

CODE_OF_CONDUCT.md

@@ -1,80 +0,0 @@
-# Code of Conduct
-
-## Our Pledge
-
-In the interest of fostering an open and welcoming environment, we as
-contributors and maintainers pledge to make participation in our project and
-our community a harassment-free experience for everyone, regardless of age, body
-size, disability, ethnicity, sex characteristics, gender identity and expression,
-level of experience, education, socio-economic status, nationality, personal
-appearance, race, religion, or sexual identity and orientation.
-
-## Our Standards
-
-Examples of behavior that contributes to creating a positive environment
-include:
-
-* Using welcoming and inclusive language
-* Being respectful of differing viewpoints and experiences
-* Gracefully accepting constructive criticism
-* Focusing on what is best for the community
-* Showing empathy towards other community members
-
-Examples of unacceptable behavior by participants include:
-
-* The use of sexualized language or imagery and unwelcome sexual attention or
-advances
-* Trolling, insulting/derogatory comments, and personal or political attacks
-* Public or private harassment
-* Publishing others' private information, such as a physical or electronic
-address, without explicit permission
-* Other conduct which could reasonably be considered inappropriate in a
-professional setting
-
-## Our Responsibilities
-
-Project maintainers are responsible for clarifying the standards of acceptable
-behavior and are expected to take appropriate and fair corrective action in
-response to any instances of unacceptable behavior.
-
-Project maintainers have the right and responsibility to remove, edit, or
-reject comments, commits, code, wiki edits, issues, and other contributions
-that are not aligned to this Code of Conduct, or to ban temporarily or
-permanently any contributor for other behaviors that they deem inappropriate,
-threatening, offensive, or harmful.
-
-## Scope
-
-This Code of Conduct applies within all project spaces, and it also applies when
-an individual is representing the project or its community in public spaces.
-Examples of representing a project or community include using an official
-project e-mail address, posting via an official social media account, or acting
-as an appointed representative at an online or offline event. Representation of
-a project may be further defined and clarified by project maintainers.
-
-This Code of Conduct also applies outside the project spaces when there is a
-reasonable belief that an individual's behavior may have a negative impact on
-the project or its community.
-
-## Enforcement
-
-Instances of abusive, harassing, or otherwise unacceptable behavior may be
-reported by contacting the project team at <opensource-conduct@meta.com>. All
-complaints will be reviewed and investigated and will result in a response that
-is deemed necessary and appropriate to the circumstances. The project team is
-obligated to maintain confidentiality with regard to the reporter of an incident.
-Further details of specific enforcement policies may be posted separately.
-
-Project maintainers who do not follow or enforce the Code of Conduct in good
-faith may face temporary or permanent repercussions as determined by other
-members of the project's leadership.
-
-## Attribution
-
-This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
-available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
-
-[homepage]: https://www.contributor-covenant.org
-
-For answers to common questions about this code of conduct, see
-https://www.contributor-covenant.org/faq

CONTRIBUTING.md

@@ -1,36 +0,0 @@
-# Contributing to ijepa
-We want to make contributing to this project as easy and transparent as
-possible.
-
-## Pull Requests
-We actively welcome your pull requests.
-
-1. Fork the repo and create your branch from `main`.
-2. If you've added code that should be tested, add tests.
-3. If you've changed APIs, update the documentation.
-4. Ensure the test suite passes.
-5. Make sure your code lints.
-6. If you haven't already, complete the Contributor License Agreement ("CLA").
-
-## Contributor License Agreement ("CLA")
-In order to accept your pull request, we need you to submit a CLA. You only need
-to do this once to work on any of Meta's open source projects.
-
-Complete your CLA here: <https://code.facebook.com/cla>
-
-## Issues
-We use GitHub issues to track public bugs. Please ensure your description is
-clear and has sufficient instructions to be able to reproduce the issue.
-
-Meta has a [bounty program](https://www.facebook.com/whitehat/) for the safe
-disclosure of security bugs. In those cases, please go through the process
-outlined on that page and do not file a public issue.
-
-## Coding Style  
-* 4 spaces for indentation rather than tabs
-* 80 character line length
-* PEP8 formatting
-
-## License
-By contributing to this repository, you agree that your contributions will be licensed
-under the LICENSE file in the root directory of this source tree.

README.md

@@ -1,3 +1,210 @@
 # WILDS-IJEPA
 
-This is a fork of the official IJEPA github repo. The learning method will be applied to wild animal images (unconstrained environments) to improve the generalization capabilities of wildlife image classification models across different locations while addressing the challenge of limited data.
+Fork of the official I-JEPA repo, adapted for WILDS-iWildCam.
+
+Reference: official I-JEPA README https://github.com/facebookresearch/ijepa/blob/main/README.md
+
+- SSL pretraining on WILDS-iWildCam unlabeled dataset: https://arxiv.org/abs/2112.05090 (Extending the WILDS Benchmark for Unsupervised Adaptation)
+- Supervised training on WILDS-iWildCam labeled dataset: https://arxiv.org/abs/2012.07421 (WILDS: A Benchmark of in-the-Wild Distribution Shifts)
+- Supervised learning supports full fine-tuning or freezing the encoder
+
+<p align="center">
+  <img src="assets/ijepa_masking_turkey.jpg" width="640" alt="I-JEPA masking on an iWildCam camera-trap image">
+</p>
+
+<p align="center">
+  <em>I-JEPA masking on a real iWildCam camera-trap image: the context block (blue) is
+  encoded to predict the representations of several target blocks (green/red/orange).</em>
+</p>
+
+## Models
+
+- ViT-H, 14x14 patches, 224x224 resolution (trained)
+- 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
+
+## Results
+
+Linear-probing evaluation on iWildCam2020-WILDS with a **frozen** I-JEPA encoder and a
+single trained linear head. All values are the mean over 5 seeds (±std). The headline
+metric is F1-Macro on the Target Out-of-Distribution (OOD) split; the gap
+Δ = F1(ID) − F1(OOD) measures the drop under distribution shift. Only the
+ImageNet-pretrained checkpoints are reported here.
+
+| Model | ID F1-Macro | OOD F1-Macro | Gap ΔF1 |
+|---|---|---|---|
+| ViT-H/14 (224, IN-1K) | 0.338 ±0.015 | 0.214 ±0.006 | 0.124 |
+| ViT-H/16 (448, IN-1K) | 0.385 ±0.005 | 0.247 ±0.003 | 0.138 |
+| **ViT-H/14 (224, IN-22K)** | 0.348 ±0.003 | **0.260 ±0.003** | 0.088 |
+| ViT-g/16 (224, IN-22K) | 0.371 ±0.002 | 0.255 ±0.003 | 0.116 |
+
+Key takeaways:
+
+- The best frozen probe (**ViT-H/14, IN-22K**) reaches **0.260 OOD F1-Macro**, ranking
+  **#16** on the iWildCam2020-WILDS leaderboard — despite training only a linear head
+  rather than fine-tuning the full backbone.
+- Its ID→OOD generalization gap (ΔF1 = 0.088) is comparable to the full-fine-tuning
+  CLIP leaders on the leaderboard (FLYP ΔF1 = 0.139, AutoFT ΔF1 = 0.115).
+- Absolute F1-Macro scales with pretraining data (IN-22K > IN-1K) and input resolution
+  (the higher-resolution ViT-H/16 448 is the strongest IN-1K checkpoint).
+
+### Label efficiency
+
+Because I-JEPA pretrains without labels, the representations stay useful when labeled
+data is scarce. OOD Target F1-Macro when the linear probe trains on 1%, 10%, 50%, and
+100% of the labeled Source split (mean over 5 seeds):
+
+| Model | 1% | 10% | 50% | 100% |
+|---|---|---|---|---|
+| ViT-H/14 (224, IN-22K) | 0.193 ±0.009 | 0.236 ±0.011 | 0.261 ±0.009 | 0.260 ±0.003 |
+| ViT-g/16 (224, IN-22K) | 0.204 ±0.009 | 0.243 ±0.007 | 0.259 ±0.014 | 0.255 ±0.003 |
+| ViT-H/14 (224, IN-1K) | 0.120 ±0.036 | 0.190 ±0.010 | 0.220 ±0.012 | 0.214 ±0.006 |
+| ViT-H/16 (448, IN-1K) | 0.146 ±0.022 | 0.214 ±0.012 | 0.230 ±0.011 | 0.247 ±0.003 |
+
+A small labeled subset already recovers most of the full-data performance
+(diminishing returns), with the IN-22K backbones degrading most gracefully — a
+practical advantage for wildlife monitoring where labeled camera-trap data is expensive.
+
+These results are from the accompanying master thesis evaluating I-JEPA on iWildCam2020-WILDS.
+
+## Repo layout
+
+- `src/`: core model, masks, and training utilities
+- `src/train.py`: SSL training loop
+- `src/train_supervised.py`: supervised training loop
+- `configs/`: training configs
+- `configs/wilds_vith14_ep300.yaml`: SSL 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
+- `configs/grids/seeds/`: per-model seed grids for multi-seed paper runs
+- `tools/run_seed_sweep.sh`: launch each model across all seeds (one by one)
+- `tools/aggregate_seeds.py`: aggregate seed runs into mean +/- std (ID + OOD)
+- `requirements.txt`: dependencies
+
+<!-- Optional: add a sample iWildCam image grid here -->
+
+## Requirements
+
+- Python 3.8+ (compatible and newer)
+- PyTorch (CUDA 12.1 wheel index): https://download.pytorch.org/whl/cu121
+- Key deps: torchvision, submitit, wilds, PyYAML, numpy
+- Full list: `requirements.txt`
+
+## SLURM commands
+
+SSL pretraining:
+
+```
+python3 main_distributed.py --fname configs/wilds_vith14_ep300.yaml --folder $submitit_folder --partition $slurm_partition --nodes $nodes --tasks-per-node $tasks_per_node --time $time
+```
+
+Supervised fine-tuning:
+
+```
+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:
+
+```
+python3 main_eval_wilds.py --fname configs/eval_wilds_vith14.yaml --folder $submitit_folder --partition $slurm_partition --nodes $nodes --tasks-per-node $tasks_per_node --time $time
+```
+
+Evaluation metrics are written to `experiment_logs/eval-wilds-vith14/iwildcam_test_metrics.json` by default.
+
+Variable hints: set `$submitit_folder`, `$slurm_partition`, `$nodes`, `$tasks_per_node`, and `$time` to match your SLURM cluster.
+
+## Multi-seed runs (paper results)
+
+To report mean +/- std over seeds, each supervised model is trained across 5
+seeds (0-4). Seeding is config-driven via `meta.seed` (applied in
+`src/train_supervised.py`), and the run folder name includes `-seed{N}` so seeds
+do not collide.
+
+Each run automatically:
+- evaluates on **both** WILDS splits: `id_test` (ID) and `test` (OOD), so the
+  generalization gap can be measured;
+- records the WILDS metrics, the wall-clock **training time**, the number of
+  **epochs run** (accounting for early stopping), and the **effective memory
+  usage** into the per-split metrics JSON and into `params.yaml` in the eval
+  folder.
+
+Effective memory is captured as a high-water mark during training:
+- `peak_host_ram_gb`: peak process RSS (`resource.getrusage`), to compare
+  against the SLURM `mem_per_gpu` request (e.g. 180G) and right-size future jobs.
+  With `tasks_per_node: 1` this reflects the whole training worker.
+- `peak_gpu_alloc_gb` / `peak_gpu_reserved_gb`: peak GPU VRAM
+  (`torch.cuda.max_memory_allocated` / `max_memory_reserved`).
+
+The four leaderboard columns are: Test ID Macro F1, Test ID Avg Acc,
+Test OOD Macro F1, Test OOD Avg Acc (headline metric: `F1-macro_all`).
+
+Launch all models, one at a time, each across all seeds (SLURM/submitit):
+
+```
+bash tools/run_seed_sweep.sh --partition $slurm_partition --time $time
+```
+
+Run a subset of models:
+
+```
+bash tools/run_seed_sweep.sh --partition $slurm_partition --models "vith14_224 vith16_448"
+```
+
+Per-model seed grids live in `configs/grids/seeds/` (each sets
+`meta.seed: [0, 1, 2, 3, 4]` over the corresponding `configs/supervised_*.yaml`
+base config). They are launched via `tools/run_grid.py`.
+
+Aggregate mean +/- std across seeds after the jobs finish:
+
+```
+python3 tools/aggregate_seeds.py --root experiment_logs/eval-wilds
+```
+
+Outputs:
+- `experiment_logs/seed-runs/<model>/summary.json` (per-seed rows + mean/std for
+  all metrics, training time, epochs, peak memory, and ID-OOD generalization gap)
+- `experiment_logs/seed-runs/summary_all.csv` (one row per model, paper-ready;
+  includes `peak_host_ram_gb_mean/std` and `peak_gpu_alloc_gb_mean/std`)
+
+## Label-efficiency experiments
+
+To measure how well the frozen representations work with fewer labels, train
+linear probes on 1%, 10%, 50%, and 100% of the labeled Source split. The subset
+is stratified by class and deterministic per seed (so every class is represented
+even at 1%).
+
+Grids for all supervised models are generated under `configs/grids/label_efficiency/`.
+Launch the full sweep:
+
+```
+bash tools/run_label_efficiency.sh --partition $slurm_partition --time $time
+```
+
+Run a subset of models or fractions:
+
+```
+bash tools/run_label_efficiency.sh --partition $slurm_partition \
+  --models "vith14_224_in22k vitg16_224_in22k" \
+  --fractions "0.01 0.10 0.50"
+```
+
+Each grid submits one submitit job per seed (5 seeds per fraction). After the
+jobs finish, aggregate into a paper-style Table 4 CSV:
+
+```
+python3 tools/aggregate_label_efficiency.py --root experiment_logs/eval-wilds
+```
+
+Outputs:
+- `experiment_logs/label-efficiency/summary.csv` (columns: 1%, 10%, 50%, 100% OOD F1-Macro)
+- `experiment_logs/label-efficiency/<model>/summary.json`
+
+## License
+
+See the `LICENSE` file for details about the license under which this code is made available.
+
+## Citation
+
+To be defined.

configs/eval_wilds_vith14.yaml

@@ -0,0 +1,36 @@
+meta:
+  seed: 0
+  model_name: vit_huge
+  embed_dim: 1280
+  num_classes: 182
+  patch_size: 14
+  crop_size: 224
+  use_bfloat16: true
+  representation_type: last_avgpool
+  head_type: linear
+  checkpoint_folder: experiment_logs/supervised/
+  read_checkpoint: linear_probe-best.pth.tar
+
+data:
+  batch_size: 128
+  root_path: ./wilds_data
+  num_workers: 10
+  pin_mem: true
+  split: test
+  download: true
+
+optimization:
+  lr: 0.01
+  weight_decay: 5.0e-4
+  optimizer: adamw
+  epochs: 300
+  freeze_weights: true
+  lr_schedule: cosine
+  warmup: 0
+  momentum: 0.99
+  start_lr: 0.0
+  final_lr: 0.0
+
+logging:
+  write_tag: iwildcam_test
+  auto_folder: true

configs/grids/label_efficiency/vitb16_448_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vitb16_448
+# Generated from configs/grids/seeds/vitb16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitb16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitb16_448_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vitb16_448
+# Generated from configs/grids/seeds/vitb16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitb16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitb16_448_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vitb16_448
+# Generated from configs/grids/seeds/vitb16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitb16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitb16_448_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vitb16_448
+# Generated from configs/grids/seeds/vitb16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitb16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitg16_224_in22k_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vitg16_224_in22k
+# Generated from configs/grids/seeds/vitg16_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitg16_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitg16_224_in22k_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vitg16_224_in22k
+# Generated from configs/grids/seeds/vitg16_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitg16_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitg16_224_in22k_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vitg16_224_in22k
+# Generated from configs/grids/seeds/vitg16_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitg16_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vitg16_224_in22k_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vitg16_224_in22k
+# Generated from configs/grids/seeds/vitg16_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitg16_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vith14_224
+# Generated from configs/grids/seeds/vith14_224.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vith14_224
+# Generated from configs/grids/seeds/vith14_224.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vith14_224
+# Generated from configs/grids/seeds/vith14_224.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vith14_224
+# Generated from configs/grids/seeds/vith14_224.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in1k_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vith14_224_in1k
+# Generated from configs/grids/seeds/vith14_224_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in1k_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vith14_224_in1k
+# Generated from configs/grids/seeds/vith14_224_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in1k_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vith14_224_in1k
+# Generated from configs/grids/seeds/vith14_224_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in1k_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vith14_224_in1k
+# Generated from configs/grids/seeds/vith14_224_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in22k_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vith14_224_in22k
+# Generated from configs/grids/seeds/vith14_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in22k_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vith14_224_in22k
+# Generated from configs/grids/seeds/vith14_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in22k_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vith14_224_in22k
+# Generated from configs/grids/seeds/vith14_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith14_224_in22k_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vith14_224_in22k
+# Generated from configs/grids/seeds/vith14_224_in22k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in22k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vith16_448
+# Generated from configs/grids/seeds/vith16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vith16_448
+# Generated from configs/grids/seeds/vith16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vith16_448
+# Generated from configs/grids/seeds/vith16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vith16_448
+# Generated from configs/grids/seeds/vith16_448.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_in1k_frac0.01.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (1%) for vith16_448_in1k
+# Generated from configs/grids/seeds/vith16_448_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.01
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_in1k_frac0.10.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (10%) for vith16_448_in1k
+# Generated from configs/grids/seeds/vith16_448_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.1
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_in1k_frac0.50.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (50%) for vith16_448_in1k
+# Generated from configs/grids/seeds/vith16_448_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 0.5
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/label_efficiency/vith16_448_in1k_frac1.00.yaml

@@ -0,0 +1,20 @@
+# Label-efficiency sweep (100%) for vith16_448_in1k
+# Generated from configs/grids/seeds/vith16_448_in1k.yaml
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448_in1k.yaml
+constants:
+  logging.write_tag: linear_probe
+  data.label_fraction: 1.0
+grid:
+  meta.seed:
+  - 0
+  - 1
+  - 2
+  - 3
+  - 4
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/lp_grid.yaml

@@ -0,0 +1,22 @@
+base_config: configs/supervised_wilds_vith16-448_ep300-lp.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  optimization.optimizer: [adamw]
+  optimization.lr: [0.01, 0.05, 0.001]
+  optimization.weight_decay: [5.0e-4, 0]
+  optimization.momentum: [0.9]
+  optimization.lr_schedule: [cosine]
+  optimization.gradient_accumulation_steps: [1]
+  data.batch_size: [128, 256, 512, 1024]
+  meta.representation_type: [last_avgpool, last4_avgpool_concat]
+  meta.head_type: [linear, bn_linear]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 180

configs/grids/seeds/vitb16_448.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vitb16_448
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitb16_448.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vitg16_224_in22k.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vitg16_224_in22k
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vitg16_224_in22k.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vith14_224.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vith14_224
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vith14_224_in1k.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vith14_224_in1k
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in1k.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vith14_224_in22k.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vith14_224_in22k
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith14_224_in22k.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vith16_448.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vith16_448
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/grids/seeds/vith16_448_in1k.yaml

@@ -0,0 +1,17 @@
+# Seed sweep for vith16_448_in1k
+# Runs the model across 5 seeds for mean +/- std reporting.
+# Each seed becomes a separate submitit job via tools/run_grid.py.
+base_config: configs/supervised_vith16_448_in1k.yaml
+
+constants:
+  logging.write_tag: linear_probe
+
+grid:
+  meta.seed: [0, 1, 2, 3, 4]
+
+launch:
+  folder: submitit_logs/
+  partition: gpu1
+  nodes: 1
+  tasks_per_node: 1
+  time: 4300

configs/in22k_vitg16_ep44.yaml

@@ -1,56 +0,0 @@
-# NOTE: ImageNet-22K (IN22k) dataloader is not implemented
-# please implement IN22k data loader based on your data
-# storage format, and update the paths in your config
-# to load from your IN22k dataset.
-data:
-  batch_size: 16
-  color_jitter_strength: 0.0
-  crop_scale:
-  - 0.3
-  - 1.0
-  crop_size: 224
-  image_folder: imagenet_full_size/061417/
-  num_workers: 10
-  pin_mem: true
-  root_path: $replace_this_with_absolute_path_to_your_datasets_directory
-  use_color_distortion: false
-  use_gaussian_blur: false
-  use_horizontal_flip: false
-logging:
-  folder: $replace_this_with_path_for_experiment_logs/in2kk_vith16.224-bs.2048-ep.44/
-  write_tag: jepa
-mask:
-  allow_overlap: false
-  aspect_ratio:
-  - 0.75
-  - 1.5
-  enc_mask_scale:
-  - 0.85
-  - 1.0
-  min_keep: 10
-  num_enc_masks: 1
-  num_pred_masks: 4
-  patch_size: 16
-  pred_mask_scale:
-  - 0.15
-  - 0.2
-meta:
-  copy_data: false
-  load_checkpoint: false
-  model_name: vit_giant
-  pred_depth: 16
-  pred_emb_dim: 384
-  read_checkpoint: null
-  use_bfloat16: true
-optimization:
-  ema:
-  - 0.996
-  - 1.0
-  epochs: 44
-  final_lr: 1.0e-06
-  final_weight_decay: 0.4
-  ipe_scale: 1.0
-  lr: 0.001
-  start_lr: 0.0002
-  warmup: 3
-  weight_decay: 0.04

configs/in22k_vith14_ep66.yaml

@@ -1,56 +0,0 @@
-# NOTE: ImageNet-22K (IN22k) dataloader is not implemented
-# please implement IN22k data loader based on your data
-# storage format, and update the paths in your config
-# to load from your IN22k dataset.
-data:
-  batch_size: 32
-  color_jitter_strength: 0.0
-  crop_scale:
-  - 0.3
-  - 1.0
-  crop_size: 224
-  image_folder: imagenet_full_size/061417/
-  num_workers: 10
-  pin_mem: true
-  root_path: $replace_this_with_absolute_path_to_your_datasets_directory
-  use_color_distortion: false
-  use_gaussian_blur: false
-  use_horizontal_flip: false
-logging:
-  folder: $replace_this_with_path_for_experiment_logs/in2kk_vith14.224-bs.2048-ep.66/
-  write_tag: jepa
-mask:
-  allow_overlap: false
-  aspect_ratio:
-  - 0.75
-  - 1.5
-  enc_mask_scale:
-  - 0.85
-  - 1.0
-  min_keep: 10
-  num_enc_masks: 1
-  num_pred_masks: 4
-  patch_size: 14
-  pred_mask_scale:
-  - 0.15
-  - 0.2
-meta:
-  copy_data: false
-  load_checkpoint: false
-  model_name: vit_huge
-  pred_depth: 12
-  pred_emb_dim: 384
-  read_checkpoint: null
-  use_bfloat16: true
-optimization:
-  ema:
-  - 0.996
-  - 1.0
-  epochs: 66
-  final_lr: 1.0e-06
-  final_weight_decay: 0.4
-  ipe_scale: 1.0
-  lr: 0.001
-  start_lr: 0.0002
-  warmup: 3
-  weight_decay: 0.04

configs/supervised_vitb16_448.yaml

@@ -0,0 +1,48 @@
+meta:
+  model_name: vit_base
+  embed_dim: 768
+  load_checkpoint: true
+  checkpoint_folder: experiment_logs/vitb16.448-bs.16-ep.600/
+  read_checkpoint: jepa-ep600.pth.tar
+  use_bfloat16: true
+  num_classes: 182
+  representation_type: last4_avgpool_concat
+  head_type: linear
+
+data:
+  batch_size: 16
+  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: 1
+  lr: 0.01
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 8
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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_vitg16_224_in22k.yaml

@@ -0,0 +1,50 @@
+meta:
+  model_name: vit_giant
+  tag: in22k
+  embed_dim: 1408
+  load_checkpoint: true
+  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: bn_linear
+  use_gradient_checkpointing: true
+
+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: 16
+
+optimization:
+  optimizer: adamw
+  freeze_weights: true
+  epochs: 1
+  lr: 0.001
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 16
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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.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: 1
+  lr: 0.001
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 32
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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: 1
+  lr: 0.01
+  weight_decay: 0
+  lr_schedule: cosine
+  gradient_accumulation_steps: 64
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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

@@ -0,0 +1,49 @@
+meta:
+  model_name: vit_huge
+  tag: in22k
+  embed_dim: 1280
+  load_checkpoint: true
+  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: 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: 1
+  lr: 0.001
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 32
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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_vith16_448.yaml

@@ -0,0 +1,48 @@
+meta:
+  model_name: vit_huge
+  embed_dim: 1280
+  load_checkpoint: true
+  checkpoint_folder: experiment_logs/vith16-bs.2048-448/
+  read_checkpoint: jepa-ep50.pth.tar
+  use_bfloat16: true
+  num_classes: 182
+  representation_type: last_avgpool
+  head_type: bn_linear
+
+data:
+  batch_size: 16
+  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: 1
+  lr: 0.01
+  weight_decay: 5.0e-4
+  lr_schedule: cosine
+  gradient_accumulation_steps: 64
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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_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: 16
+  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: 1
+  lr: 0.001
+  weight_decay: 0
+  lr_schedule: cosine
+  gradient_accumulation_steps: 32
+  final_lr: 0.0
+  early_stopping:
+    enabled: false
+    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/wilds_vitb16-448_ep600.yaml

@@ -0,0 +1,54 @@
+data:
+  batch_size: 256
+  color_jitter_strength: 0.0
+  crop_scale:
+    - 1.0
+    - 1.0
+  crop_size: 448
+  num_workers: 10
+  pin_mem: true
+  root_path: ./wilds_data
+  use_color_distortion: false
+  use_gaussian_blur: false
+  use_horizontal_flip: false
+  use_random_resized_crop: false
+logging:
+  folder: experiment_logs/vitb16.448-bs.16-ep.600-paper
+  write_tag: jepa
+mask:
+  allow_overlap: false
+  aspect_ratio:
+    - 0.75
+    - 1.5
+  enc_mask_scale:
+    - 0.85
+    - 1.0
+  min_keep: 10
+  num_enc_masks: 1
+  num_pred_masks: 4
+  patch_size: 16
+  pred_mask_scale:
+    - 0.15
+    - 0.2
+meta:
+  copy_data: false
+  load_checkpoint: false
+  model_name: vit_base
+  embed_dim: 768
+  pred_depth: 6
+  pred_emb_dim: 384
+  read_checkpoint: null
+  use_bfloat16: true
+optimization:
+  ema:
+    - 0.996
+    - 1.0
+  epochs: 600
+  final_lr: 1.0e-06
+  final_weight_decay: 0.4
+  ipe_scale: 1.0
+  lr: 0.001
+  start_lr: 0.0001
+  warmup: 15
+  wd_schedule: linear
+  weight_decay: 0.04

configs/wilds_vith14_ep300.yaml

@@ -1,38 +1,37 @@
 data:
-  batch_size: 128
+  batch_size: 256
   color_jitter_strength: 0.0
   crop_scale:
-  - 0.3
-  - 1.0
+    - 0.3
+    - 1.0
   crop_size: 224
-  image_folder: imagenet_full_size/061417/
   num_workers: 10
   pin_mem: true
-  root_path: $replace_this_with_absolute_path_to_your_datasets_directory
+  root_path: ./wilds_data
   use_color_distortion: false
   use_gaussian_blur: false
   use_horizontal_flip: false
 logging:
-  folder: $replace_this_with_path_for_experiment_logs/vith14.224-bs.2048-ep.300/
+  folder: experiment_logs/vith14.224-bs.128-ep.300/
   write_tag: jepa
 mask:
   allow_overlap: false
   aspect_ratio:
-  - 0.75
-  - 1.5
+    - 0.75
+    - 1.5
   enc_mask_scale:
-  - 0.85
-  - 1.0
+    - 0.85
+    - 1.0
   min_keep: 10
   num_enc_masks: 1
   num_pred_masks: 4
   patch_size: 14
   pred_mask_scale:
-  - 0.15
-  - 0.2
+    - 0.15
+    - 0.2
 meta:
   copy_data: false
-  load_checkpoint: false
+  load_checkpoint: true
   model_name: vit_huge
   pred_depth: 12
   pred_emb_dim: 384
@@ -40,13 +39,14 @@ meta:
   use_bfloat16: true
 optimization:
   ema:
-  - 0.996
-  - 1.0
+    - 0.996
+    - 1.0
   epochs: 300
   final_lr: 1.0e-06
   final_weight_decay: 0.4
   ipe_scale: 1.0
   lr: 0.001
-  start_lr: 0.0002
-  warmup: 40
+  start_lr: 0.0001
+  warmup: 15
+  wd_schedule: linear
   weight_decay: 0.04

configs/wilds_vith16-448_ep300.yaml

@@ -1,35 +1,35 @@
 data:
-  batch_size: 16
+  batch_size: 128
   color_jitter_strength: 0.0
   crop_scale:
-  - 0.3
-  - 1.0
+    - 1.0
+    - 1.0
   crop_size: 448
-  image_folder: imagenet_full_size/061417/
   num_workers: 10
   pin_mem: true
-  root_path: $replace_this_with_absolute_path_to_your_datasets_directory
+  root_path: ./wilds_data
   use_color_distortion: false
   use_gaussian_blur: false
   use_horizontal_flip: false
+  use_random_resized_crop: false
 logging:
-  folder: $replace_this_with_path_for_experiment_logs/vith16.448-bs.2048-ep.300/
+  folder: experiment_logs/vith16.448-bs.128-ep.300/
   write_tag: jepa
 mask:
   allow_overlap: false
   aspect_ratio:
-  - 0.75
-  - 1.5
+    - 0.75
+    - 1.5
   enc_mask_scale:
-  - 0.85
-  - 1.0
+    - 0.85
+    - 1.0
   min_keep: 10
   num_enc_masks: 1
   num_pred_masks: 4
   patch_size: 16
   pred_mask_scale:
-  - 0.15
-  - 0.2
+    - 0.15
+    - 0.2
 meta:
   copy_data: false
   load_checkpoint: false
@@ -38,15 +38,18 @@ meta:
   pred_emb_dim: 384
   read_checkpoint: null
   use_bfloat16: true
+  use_gradient_checkpointing: true
 optimization:
+  gradient_accumulation_steps: 2
   ema:
-  - 0.996
-  - 1.0
+    - 0.996
+    - 1.0
   epochs: 300
   final_lr: 1.0e-06
   final_weight_decay: 0.4
   ipe_scale: 1.0
   lr: 0.001
-  start_lr: 0.0002
-  warmup: 40
+  start_lr: 0.0001
+  warmup: 15
+  wd_schedule: linear
   weight_decay: 0.04

main_distributed.py

@@ -21,46 +21,43 @@ logger = logging.getLogger()
 
 
 parser = argparse.ArgumentParser()
+parser.add_argument("--folder", type=str, help="location to save submitit logs")
 parser.add_argument(
-    '--folder', type=str,
-    help='location to save submitit logs')
+    "--batch-launch",
+    action="store_true",
+    help="whether fname points to a file to batch-lauch several config files",
+)
 parser.add_argument(
-    '--batch-launch', action='store_true',
-    help='whether fname points to a file to batch-lauch several config files')
+    "--fname",
+    type=str,
+    help="yaml file containing config file names to launch",
+    default="configs.yaml",
+)
+parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
 parser.add_argument(
-    '--fname', type=str,
-    help='yaml file containing config file names to launch',
-    default='configs.yaml')
+    "--nodes", type=int, default=1, help="num. nodes to request for job"
+)
 parser.add_argument(
-    '--partition', type=str,
-    help='cluster partition to submit jobs on')
-parser.add_argument(
-    '--nodes', type=int, default=1,
-    help='num. nodes to request for job')
-parser.add_argument(
-    '--tasks-per-node', type=int, default=1,
-    help='num. procs to per node')
-parser.add_argument(
-    '--time', type=int, default=4300,
-    help='time in minutes to run job')
+    "--tasks-per-node", type=int, default=1, help="num. procs to per node"
+)
+parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
 
 
 class Trainer:
-
-    def __init__(self, fname='configs.yaml', load_model=None):
+    def __init__(self, fname="configs.yaml", load_model=None):
         self.fname = fname
         self.load_model = load_model
 
     def __call__(self):
         fname = self.fname
         load_model = self.load_model
-        logger.info(f'called-params {fname}')
+        logger.info(f"called-params {fname}")
 
         # -- load script params
         params = None
-        with open(fname, 'r') as y_file:
+        with open(fname, "r") as y_file:
             params = yaml.load(y_file, Loader=yaml.FullLoader)
-            logger.info('loaded params...')
+            logger.info("loaded params...")
             pp = pprint.PrettyPrinter(indent=4)
             pp.pprint(params)
 
@@ -69,21 +66,24 @@ class Trainer:
 
     def checkpoint(self):
         fb_trainer = Trainer(self.fname, True)
-        return submitit.helpers.DelayedSubmission(fb_trainer,)
+        return submitit.helpers.DelayedSubmission(
+            fb_trainer,
+        )
 
 
 def launch():
-    executor = submitit.AutoExecutor(
+    executor = submitit.SlurmExecutor(
         folder=os.path.join(args.folder, 'job_%j'),
-        slurm_max_num_timeout=20)
+        max_num_timeout=20)
     executor.update_parameters(
-        slurm_partition=args.partition,
-        slurm_mem_per_gpu='55G',
-        timeout_min=args.time,
+        partition=args.partition,
+        mem_per_gpu='180G',
+        time=args.time,
         nodes=args.nodes,
-        tasks_per_node=args.tasks_per_node,
+        ntasks_per_node=args.tasks_per_node,
         cpus_per_task=10,
-        gpus_per_node=args.tasks_per_node)
+        gpus_per_node=args.tasks_per_node,
+    )
 
     config_fnames = [args.fname]
 
@@ -91,7 +91,9 @@ def launch():
     with executor.batch():
         for cf in config_fnames:
             fb_trainer = Trainer(cf)
-            job = executor.submit(fb_trainer,)
+            job = executor.submit(
+                fb_trainer,
+            )
             trainers.append(fb_trainer)
             jobs.append(job)
 
@@ -99,6 +101,6 @@ def launch():
         print(job.job_id)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     args = parser.parse_args()
     launch()

main_distributed_supervised.py

@@ -0,0 +1,106 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import argparse
+import logging
+import os
+import pprint
+import sys
+import yaml
+
+import submitit
+
+from src.train_supervised import main as app_main
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger()
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--folder", type=str, help="location to save submitit logs")
+parser.add_argument(
+    "--batch-launch",
+    action="store_true",
+    help="whether fname points to a file to batch-lauch several config files",
+)
+parser.add_argument(
+    "--fname",
+    type=str,
+    help="yaml file containing config file names to launch",
+    default="configs.yaml",
+)
+parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
+parser.add_argument(
+    "--nodes", type=int, default=1, help="num. nodes to request for job"
+)
+parser.add_argument(
+    "--tasks-per-node", type=int, default=1, help="num. procs to per node"
+)
+parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
+
+
+class Trainer:
+    def __init__(self, fname="configs.yaml", load_model=None):
+        self.fname = fname
+        self.load_model = load_model
+
+    def __call__(self):
+        fname = self.fname
+        load_model = self.load_model
+        logger.info(f"called-params {fname}")
+
+        # -- load script params
+        params = None
+        with open(fname, "r") as y_file:
+            params = yaml.load(y_file, Loader=yaml.FullLoader)
+            logger.info("loaded params...")
+            pp = pprint.PrettyPrinter(indent=4)
+            pp.pprint(params)
+
+        resume_preempt = False if load_model is None else load_model
+        app_main(args=params, resume_preempt=resume_preempt)
+
+    def checkpoint(self):
+        fb_trainer = Trainer(self.fname, True)
+        return submitit.helpers.DelayedSubmission(
+            fb_trainer,
+        )
+
+
+def launch():
+    executor = submitit.SlurmExecutor(
+        folder=os.path.join(args.folder, "job_%j"), max_num_timeout=20
+    )
+    executor.update_parameters(
+        partition=args.partition,
+        mem_per_gpu="180G",
+        time=args.time,
+        nodes=args.nodes,
+        ntasks_per_node=args.tasks_per_node,
+        cpus_per_task=10,
+        gpus_per_node=args.tasks_per_node,
+    )
+
+    config_fnames = [args.fname]
+
+    jobs, trainers = [], []
+    with executor.batch():
+        for cf in config_fnames:
+            fb_trainer = Trainer(cf)
+            job = executor.submit(
+                fb_trainer,
+            )
+            trainers.append(fb_trainer)
+            jobs.append(job)
+
+    for job in jobs:
+        print(job.job_id)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    launch()

main_eval_wilds.py

@@ -0,0 +1,102 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import argparse
+import logging
+import os
+import pprint
+import sys
+import yaml
+
+import submitit
+
+from src.eval_wilds import main as app_main
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger()
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--folder", type=str, help="location to save submitit logs")
+parser.add_argument(
+    "--batch-launch",
+    action="store_true",
+    help="whether fname points to a file to batch-lauch several config files",
+)
+parser.add_argument(
+    "--fname",
+    type=str,
+    help="yaml file containing config file names to launch",
+    default="configs.yaml",
+)
+parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
+parser.add_argument(
+    "--nodes", type=int, default=1, help="num. nodes to request for job"
+)
+parser.add_argument(
+    "--tasks-per-node", type=int, default=1, help="num. procs to per node"
+)
+parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
+
+
+class Trainer:
+    def __init__(self, fname="configs.yaml"):
+        self.fname = fname
+
+    def __call__(self):
+        fname = self.fname
+        logger.info(f"called-params {fname}")
+
+        params = None
+        with open(fname, "r") as y_file:
+            params = yaml.load(y_file, Loader=yaml.FullLoader)
+            logger.info("loaded params...")
+            pp = pprint.PrettyPrinter(indent=4)
+            pp.pprint(params)
+
+        app_main(args=params)
+
+    def checkpoint(self):
+        fb_trainer = Trainer(self.fname)
+        return submitit.helpers.DelayedSubmission(
+            fb_trainer,
+        )
+
+
+def launch():
+    executor = submitit.SlurmExecutor(
+        folder=os.path.join(args.folder, "job_%j"), max_num_timeout=20
+    )
+    executor.update_parameters(
+        partition=args.partition,
+        mem_per_gpu="55G",
+        time=args.time,
+        nodes=args.nodes,
+        ntasks_per_node=args.tasks_per_node,
+        cpus_per_task=10,
+        gpus_per_node=args.tasks_per_node,
+    )
+
+    config_fnames = [args.fname]
+
+    jobs, trainers = [], []
+    with executor.batch():
+        for cf in config_fnames:
+            fb_trainer = Trainer(cf)
+            job = executor.submit(
+                fb_trainer,
+            )
+            trainers.append(fb_trainer)
+            jobs.append(job)
+
+    for job in jobs:
+        print(job.job_id)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    launch()

requirements.txt

@@ -0,0 +1,49 @@
+certifi==2026.2.25
+charset-normalizer==3.4.6
+cloudpickle==3.1.2
+cmake==4.3.1
+filelock==3.19.1
+idna==3.11
+Jinja2==3.1.6
+joblib==1.5.3
+lit==18.1.8
+littleutils==0.2.4
+MarkupSafe==3.0.3
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.4
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+ogb==1.3.6
+outdated==0.2.2
+packaging==26.0
+pandas==2.3.3
+pillow==11.3.0
+python-dateutil==2.9.0.post0
+pytz==2026.1.post1
+PyYAML==6.0.3
+requests==2.32.5
+scikit-learn==1.6.1
+scipy==1.13.1
+six==1.17.0
+submitit==1.5.0
+sympy==1.14.0
+threadpoolctl==3.6.0
+torch==2.0.1
+torchaudio==2.0.2
+torchvision==0.15.2
+tqdm==4.67.3
+triton==2.0.0
+typing_extensions==4.15.0
+tzdata==2025.3
+urllib3==2.6.3
+wilds==2.0.0

results/job_73125_full_history.csv

@@ -0,0 +1,81 @@
+JobIDRaw,JobID,JobName,Partition,Submit,Start,End,State,ExitCode,ElapsedRaw,Elapsed,AllocCPUS,AllocNodes,ReqMem,TotalCPU,CPUTimeRAW,MaxRSS,MaxVMSize,AveRSS,AveVMSize
+73125,73125,submitit,gpu8,2026-04-18T15:04:34,2026-04-18T15:04:34,2026-04-18T20:02:51,REQUEUED,0:0,17897,04:58:17,80,1,440G,1-22:05:51,1431760,,,,
+73125.batch,73125.batch,batch,,2026-04-18T15:04:34,2026-04-18T15:04:34,2026-04-18T20:02:53,CANCELLED,0:15,17899,04:58:19,80,1,,00:00.005,1431920,6356K,0,6356K,0
+73125.extern,73125.extern,extern,,2026-04-18T15:04:34,2026-04-18T15:04:34,2026-04-18T20:03:17,COMPLETED,0:0,17923,04:58:43,80,1,,00:00.001,1433840,428K,0,428K,0
+73125.0,73125.0,python3,,2026-04-18T15:04:39,2026-04-18T15:04:39,2026-04-18T20:03:30,CANCELLED,0:9,17931,04:58:51,80,1,,1-22:05:51,1434480,10548560K,0,10474636.50K,0
+73125,73125,submitit,gpu8,2026-04-18T20:03:39,2026-04-18T20:06:18,2026-04-19T01:04:48,REQUEUED,0:0,17910,04:58:30,80,1,440G,2-00:15:14,1432800,,,,
+73125.batch,73125.batch,batch,,2026-04-18T20:06:18,2026-04-18T20:06:18,2026-04-19T01:04:49,CANCELLED,0:15,17911,04:58:31,80,1,,00:00.005,1432880,6364K,0,6364K,0
+73125.extern,73125.extern,extern,,2026-04-18T20:06:18,2026-04-18T20:06:18,2026-04-19T01:05:16,COMPLETED,0:0,17938,04:58:58,80,1,,00:00.001,1435040,440K,0,440K,0
+73125.0,73125.0,python3,,2026-04-18T20:06:23,2026-04-18T20:06:23,2026-04-19T01:05:26,FAILED,127:0,17943,04:59:03,80,1,,2-00:15:14,1435440,12046456K,0,11749292.50K,0
+73125,73125,submitit,gpu8,2026-04-19T01:05:35,2026-04-19T01:08:18,2026-04-19T06:06:48,REQUEUED,0:0,17910,04:58:30,80,1,440G,2-09:24:18,1432800,,,,
+73125.batch,73125.batch,batch,,2026-04-19T01:08:18,2026-04-19T01:08:18,2026-04-19T06:06:49,CANCELLED,0:15,17911,04:58:31,80,1,,00:00.005,1432880,6544K,0,6544K,0
+73125.extern,73125.extern,extern,,2026-04-19T01:08:18,2026-04-19T01:08:18,2026-04-19T06:07:17,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,436K,0,436K,0
+73125.0,73125.0,python3,,2026-04-19T01:08:23,2026-04-19T01:08:23,2026-04-19T06:07:26,FAILED,127:0,17943,04:59:03,80,1,,2-09:24:18,1435440,12048284K,0,11797508K,0
+73125,73125,submitit,gpu8,2026-04-19T06:07:35,2026-04-19T06:10:18,2026-04-19T11:09:00,REQUEUED,0:0,17922,04:58:42,80,1,440G,2-11:37:35,1433760,,,,
+73125.batch,73125.batch,batch,,2026-04-19T06:10:18,2026-04-19T06:10:18,2026-04-19T11:09:01,CANCELLED,0:15,17923,04:58:43,80,1,,00:00.005,1433840,6284K,0,6284K,0
+73125.extern,73125.extern,extern,,2026-04-19T06:10:18,2026-04-19T06:10:18,2026-04-19T11:09:17,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,100K,0,100K,0
+73125.0,73125.0,python3,,2026-04-19T06:10:24,2026-04-19T06:10:24,2026-04-19T11:09:38,FAILED,127:0,17954,04:59:14,80,1,,2-11:37:35,1436320,12254580K,0,11977551K,0
+73125,73125,submitit,gpu8,2026-04-19T11:09:46,2026-04-19T11:11:51,2026-04-19T16:10:18,REQUEUED,0:0,17907,04:58:27,80,1,440G,1-20:14:17,1432560,,,,
+73125.batch,73125.batch,batch,,2026-04-19T11:11:51,2026-04-19T11:11:51,2026-04-19T16:10:19,CANCELLED,0:15,17908,04:58:28,80,1,,00:00.005,1432640,6364K,0,6364K,0
+73125.extern,73125.extern,extern,,2026-04-19T11:11:51,2026-04-19T11:11:51,2026-04-19T16:10:47,COMPLETED,0:0,17936,04:58:56,80,1,,00:00.001,1434880,428K,0,428K,0
+73125.0,73125.0,python3,,2026-04-19T11:11:56,2026-04-19T11:11:56,2026-04-19T16:10:59,FAILED,127:0,17943,04:59:03,80,1,,1-20:14:17,1435440,12147412K,0,11909438K,0
+73125,73125,submitit,gpu8,2026-04-19T16:11:17,2026-04-19T16:13:18,2026-04-19T21:11:20,REQUEUED,0:0,17882,04:58:02,80,1,440G,1-23:13:33,1430560,,,,
+73125.batch,73125.batch,batch,,2026-04-19T16:13:18,2026-04-19T16:13:18,2026-04-19T21:11:21,CANCELLED,0:15,17883,04:58:03,80,1,,00:00.005,1430640,6284K,0,6284K,0
+73125.extern,73125.extern,extern,,2026-04-19T16:13:18,2026-04-19T16:13:18,2026-04-19T21:11:48,COMPLETED,0:0,17910,04:58:30,80,1,,00:00.001,1432800,100K,0,100K,0
+73125.0,73125.0,python3,,2026-04-19T16:13:22,2026-04-19T16:13:22,2026-04-19T21:11:57,FAILED,127:0,17915,04:58:35,80,1,,1-23:13:33,1433200,12068060K,0,11925202K,0
+73125,73125,submitit,gpu8,2026-04-19T21:12:07,2026-04-19T21:14:19,2026-04-20T02:12:49,REQUEUED,0:0,17910,04:58:30,80,1,440G,1-22:37:38,1432800,,,,
+73125.batch,73125.batch,batch,,2026-04-19T21:14:19,2026-04-19T21:14:19,2026-04-20T02:12:50,CANCELLED,0:15,17911,04:58:31,80,1,,00:00.005,1432880,6328K,0,6328K,0
+73125.extern,73125.extern,extern,,2026-04-19T21:14:19,2026-04-19T21:14:19,2026-04-20T02:13:18,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,92K,0,92K,0
+73125.0,73125.0,python3,,2026-04-19T21:14:25,2026-04-19T21:14:25,2026-04-20T02:13:27,FAILED,127:0,17942,04:59:02,80,1,,1-22:37:38,1435360,11812372K,0,11514502K,0
+73125,73125,submitit,gpu8,2026-04-20T02:13:36,2026-04-20T02:16:19,2026-04-20T07:14:49,REQUEUED,0:0,17910,04:58:30,80,1,440G,1-22:04:59,1432800,,,,
+73125.batch,73125.batch,batch,,2026-04-20T02:16:19,2026-04-20T02:16:19,2026-04-20T07:14:51,CANCELLED,0:15,17912,04:58:32,80,1,,00:00.005,1432960,6280K,0,6280K,0
+73125.extern,73125.extern,extern,,2026-04-20T02:16:19,2026-04-20T02:16:19,2026-04-20T07:15:18,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,120K,0,120K,0
+73125.0,73125.0,python3,,2026-04-20T02:16:24,2026-04-20T02:16:24,2026-04-20T07:15:28,FAILED,127:0,17944,04:59:04,80,1,,1-22:04:59,1435520,12213376K,0,11796571.50K,0
+73125,73125,submitit,gpu8,2026-04-20T07:15:36,2026-04-20T07:18:19,2026-04-20T12:16:58,REQUEUED,0:0,17919,04:58:39,80,1,440G,1-19:54:07,1433520,,,,
+73125.batch,73125.batch,batch,,2026-04-20T07:18:19,2026-04-20T07:18:19,2026-04-20T12:17:00,CANCELLED,0:15,17921,04:58:41,80,1,,00:00.005,1433680,6360K,0,6360K,0
+73125.extern,73125.extern,extern,,2026-04-20T07:18:19,2026-04-20T07:18:19,2026-04-20T12:17:18,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,428K,0,428K,0
+73125.0,73125.0,python3,,2026-04-20T07:18:25,2026-04-20T07:18:25,2026-04-20T12:17:37,FAILED,127:0,17952,04:59:12,80,1,,1-19:54:07,1436160,12116592K,0,11885582K,0
+73125,73125,submitit,gpu8,2026-04-20T12:17:45,2026-04-20T12:19:57,2026-04-20T17:18:25,REQUEUED,0:0,17908,04:58:28,80,1,440G,1-20:45:44,1432640,,,,
+73125.batch,73125.batch,batch,,2026-04-20T12:19:57,2026-04-20T12:19:57,2026-04-20T17:18:26,CANCELLED,0:15,17909,04:58:29,80,1,,00:00.005,1432720,6360K,0,6360K,0
+73125.extern,73125.extern,extern,,2026-04-20T12:19:57,2026-04-20T12:19:57,2026-04-20T17:18:48,COMPLETED,0:0,17931,04:58:51,80,1,,00:00.001,1434480,460K,0,460K,0
+73125.0,73125.0,python3,,2026-04-20T12:20:03,2026-04-20T12:20:03,2026-04-20T17:19:02,FAILED,127:0,17939,04:58:59,80,1,,1-20:45:44,1435120,12113352K,0,11940565.50K,0
+73125,73125,submitit,gpu8,2026-04-20T17:19:11,2026-04-20T17:21:20,2026-04-20T22:19:49,REQUEUED,0:0,17909,04:58:29,80,1,440G,1-23:24:41,1432720,,,,
+73125.batch,73125.batch,batch,,2026-04-20T17:21:20,2026-04-20T17:21:20,2026-04-20T22:19:51,CANCELLED,0:15,17911,04:58:31,80,1,,00:00.005,1432880,6572K,0,6572K,0
+73125.extern,73125.extern,extern,,2026-04-20T17:21:20,2026-04-20T17:21:20,2026-04-20T22:20:18,COMPLETED,0:0,17938,04:58:58,80,1,,00:00.002,1435040,764K,0,764K,0
+73125.0,73125.0,python3,,2026-04-20T17:21:24,2026-04-20T17:21:24,2026-04-20T22:20:27,CANCELLED,0:9,17943,04:59:03,80,1,,1-23:24:41,1435440,12095104K,0,11970123.50K,0
+73125,73125,submitit,gpu8,2026-04-20T22:20:36,2026-04-20T22:22:45,2026-04-21T03:20:50,REQUEUED,0:0,17885,04:58:05,80,1,440G,1-22:11:12,1430800,,,,
+73125.batch,73125.batch,batch,,2026-04-20T22:22:45,2026-04-20T22:22:45,2026-04-21T03:20:52,CANCELLED,0:15,17887,04:58:07,80,1,,00:00.005,1430960,6564K,0,6564K,0
+73125.extern,73125.extern,extern,,2026-04-20T22:22:45,2026-04-20T22:22:45,2026-04-21T03:21:19,COMPLETED,0:0,17914,04:58:34,80,1,,00:00.002,1433120,488K,0,488K,0
+73125.0,73125.0,python3,,2026-04-20T22:22:50,2026-04-20T22:22:50,2026-04-21T03:21:27,FAILED,127:0,17917,04:58:37,80,1,,1-22:11:12,1433360,12250596K,0,11970431K,0
+73125,73125,submitit,gpu8,2026-04-21T03:21:36,2026-04-21T03:24:20,2026-04-21T08:23:04,REQUEUED,0:0,17924,04:58:44,80,1,440G,1-22:50:58,1433920,,,,
+73125.batch,73125.batch,batch,,2026-04-21T03:24:20,2026-04-21T03:24:20,2026-04-21T08:23:05,CANCELLED,0:15,17925,04:58:45,80,1,,00:00.005,1434000,6360K,0,6360K,0
+73125.extern,73125.extern,extern,,2026-04-21T03:24:20,2026-04-21T03:24:20,2026-04-21T08:23:19,COMPLETED,0:0,17939,04:58:59,80,1,,00:00.001,1435120,380K,0,380K,0
+73125.0,73125.0,python3,,2026-04-21T03:24:24,2026-04-21T03:24:24,2026-04-21T08:23:42,CANCELLED,0:9,17958,04:59:18,80,1,,1-22:50:58,1436640,12143984K,0,11908241.50K,0
+73125,73125,submitit,gpu8,2026-04-21T08:23:52,2026-04-21T08:25:56,2026-04-21T13:24:19,REQUEUED,0:0,17903,04:58:23,80,1,440G,1-20:02:26,1432240,,,,
+73125.batch,73125.batch,batch,,2026-04-21T08:25:56,2026-04-21T08:25:56,2026-04-21T13:24:21,CANCELLED,0:15,17905,04:58:25,80,1,,00:00.005,1432400,6364K,0,6364K,0
+73125.extern,73125.extern,extern,,2026-04-21T08:25:56,2026-04-21T08:25:56,2026-04-21T13:24:48,COMPLETED,0:0,17932,04:58:52,80,1,,00:00.002,1434560,764K,0,764K,0
+73125.0,73125.0,python3,,2026-04-21T08:26:01,2026-04-21T08:26:01,2026-04-21T13:24:57,FAILED,127:0,17936,04:58:56,80,1,,1-20:02:26,1434880,12256488K,0,11947183K,0
+73125,73125,submitit,gpu8,2026-04-21T13:25:06,2026-04-21T15:03:17,2026-04-21T20:01:24,REQUEUED,0:0,17887,04:58:07,80,1,440G,1-19:18:22,1430960,,,,
+73125.batch,73125.batch,batch,,2026-04-21T15:03:17,2026-04-21T15:03:17,2026-04-21T20:01:25,CANCELLED,0:15,17888,04:58:08,80,1,,00:00.005,1431040,6364K,0,6364K,0
+73125.extern,73125.extern,extern,,2026-04-21T15:03:17,2026-04-21T15:03:17,2026-04-21T20:01:49,COMPLETED,0:0,17912,04:58:32,80,1,,00:00.001,1432960,512K,0,512K,0
+73125.0,73125.0,python3,,2026-04-21T15:03:22,2026-04-21T15:03:22,2026-04-21T20:02:02,CANCELLED,0:9,17920,04:58:40,80,1,,1-19:18:22,1433600,11718344K,0,11577435.50K,0
+73125,73125,submitit,gpu8,2026-04-21T20:02:11,2026-04-21T20:24:49,2026-04-22T01:22:50,REQUEUED,0:0,17881,04:58:01,80,1,440G,1-22:29:18,1430480,,,,
+73125.batch,73125.batch,batch,,2026-04-21T20:24:49,2026-04-21T20:24:49,2026-04-22T01:22:52,CANCELLED,0:15,17883,04:58:03,80,1,,00:00.006,1430640,6360K,0,6360K,0
+73125.extern,73125.extern,extern,,2026-04-21T20:24:49,2026-04-21T20:24:49,2026-04-22T01:23:19,COMPLETED,0:0,17910,04:58:30,80,1,,00:00.001,1432800,380K,0,380K,0
+73125.0,73125.0,python3,,2026-04-21T20:24:53,2026-04-21T20:24:53,2026-04-22T01:23:28,FAILED,127:0,17915,04:58:35,80,1,,1-22:29:18,1433200,12041376K,0,11856561.50K,0
+73125,73125,submitit,gpu8,2026-04-22T01:23:37,2026-04-24T01:24:15,2026-04-24T06:22:27,REQUEUED,0:0,17892,04:58:12,80,1,440G,2-03:34:14,1431360,,,,
+73125.batch,73125.batch,batch,,2026-04-24T01:24:15,2026-04-24T01:24:15,2026-04-24T06:22:28,CANCELLED,0:15,17893,04:58:13,80,1,,00:00.006,1431440,6360K,0,6360K,0
+73125.extern,73125.extern,extern,,2026-04-24T01:24:15,2026-04-24T01:24:15,2026-04-24T06:22:50,COMPLETED,0:0,17915,04:58:35,80,1,,00:00.001,1433200,764K,0,764K,0
+73125.0,73125.0,python3,,2026-04-24T01:24:20,2026-04-24T01:24:20,2026-04-24T06:23:03,FAILED,127:0,17923,04:58:43,80,1,,2-03:34:14,1433840,12068204K,0,11851032.50K,0
+73125,73125,submitit,gpu8,2026-04-24T06:23:12,2026-04-26T06:25:56,2026-04-26T11:24:30,REQUEUED,0:0,17914,04:58:34,80,1,440G,2-09:23:29,1433120,,,,
+73125.batch,73125.batch,batch,,2026-04-26T06:25:56,2026-04-26T06:25:56,2026-04-26T11:24:31,CANCELLED,0:15,17915,04:58:35,80,1,,00:00.006,1433200,6372K,0,6372K,0
+73125.extern,73125.extern,extern,,2026-04-26T06:25:56,2026-04-26T06:25:56,2026-04-26T11:24:48,COMPLETED,0:0,17932,04:58:52,80,1,,00:00.001,1434560,376K,0,376K,0
+73125.0,73125.0,python3,,2026-04-26T06:26:00,2026-04-26T06:26:00,2026-04-26T11:25:08,FAILED,127:0,17948,04:59:08,80,1,,2-09:23:29,1435840,12064528K,0,11898462K,0
+73125,73125,submitit,gpu8,2026-04-26T11:25:17,2026-04-26T11:27:20,2026-04-26T16:25:49,REQUEUED,0:0,17909,04:58:29,80,1,440G,2-17:21:04,1432720,,,,
+73125.batch,73125.batch,batch,,2026-04-26T11:27:20,2026-04-26T11:27:20,2026-04-26T16:25:50,CANCELLED,0:15,17910,04:58:30,80,1,,00:00.005,1432800,6372K,0,6372K,0
+73125.extern,73125.extern,extern,,2026-04-26T11:27:20,2026-04-26T11:27:20,2026-04-26T16:26:18,COMPLETED,0:0,17938,04:58:58,80,1,,00:00.001,1435040,440K,0,440K,0
+73125.0,73125.0,python3,,2026-04-26T11:27:26,2026-04-26T11:27:26,2026-04-26T16:26:26,FAILED,127:0,17940,04:59:00,80,1,,2-17:21:04,1435200,12117404K,0,12036375.50K,0
+73125,73125,submitit,gpu8,2026-04-26T16:26:35,2026-04-26T16:29:21,2026-04-26T20:52:30,COMPLETED,0:0,15789,04:23:09,80,1,440G,2-11:34:11,1263120,,,,
+73125.batch,73125.batch,batch,,2026-04-26T16:29:21,2026-04-26T16:29:21,2026-04-26T20:52:30,COMPLETED,0:0,15789,04:23:09,80,1,,00:00.015,1263120,6364K,0,6364K,0
+73125.extern,73125.extern,extern,,2026-04-26T16:29:21,2026-04-26T16:29:21,2026-04-26T20:52:30,COMPLETED,0:0,15789,04:23:09,80,1,,00:00.001,1263120,420K,0,420K,0
+73125.0,73125.0,python3,,2026-04-26T16:29:26,2026-04-26T16:29:26,2026-04-26T20:52:30,COMPLETED,0:0,15784,04:23:04,80,1,,2-11:34:11,1262720,11900144K,0,11689215.50K,0

src/datasets/wilds.py

@@ -0,0 +1,158 @@
+from logging import getLogger
+
+import numpy as np
+import torch
+from wilds import get_dataset
+from wilds.common.data_loaders import get_eval_loader
+
+logger = getLogger()
+
+
+def _stratified_subset_indices(labels, fraction, seed):
+    """Return a deterministic, seed-dependent stratified subset of indices.
+
+    For each class, keep ceil(count * fraction) samples, guaranteeing every
+    class is represented even for very small fractions.
+
+    Parameters
+    ----------
+    labels : array-like of int
+        Per-sample class labels.
+    fraction : float
+        Fraction of samples to keep (0 < fraction <= 1).
+    seed : int
+        Seed for NumPy's RNG; different seeds give different subsets.
+
+    Returns
+    -------
+    indices : np.ndarray
+        Sorted array of selected global indices.
+    """
+    if not (0 < fraction <= 1.0):
+        raise ValueError(f"label_fraction must be in (0, 1], got {fraction}")
+
+    rng = np.random.default_rng(seed)
+    labels = np.asarray(labels).reshape(-1)
+    classes = np.unique(labels)
+
+    selected = []
+    for cls in classes:
+        cls_idx = np.nonzero(labels == cls)[0]
+        n_keep = max(1, int(np.ceil(len(cls_idx) * fraction)))
+        n_keep = min(n_keep, len(cls_idx))
+        selected.extend(rng.choice(cls_idx, size=n_keep, replace=False).tolist())
+
+    indices = np.array(sorted(selected), dtype=np.int64)
+    return indices
+
+
+def make_iwildcam(
+    transform,
+    batch_size,
+    collator=None,
+    split="extra_unlabeled",
+    num_workers=8,
+    world_size=1,
+    rank=0,
+    root_path="./wilds_data",
+    download=True,
+    pin_mem=True,
+    drop_last=True,
+    label_fraction=None,
+    seed=0,
+):
+    unlabeled = True if split == "extra_unlabeled" else False
+    shuffle = True if split == "extra_unlabeled" or split == "train" else False
+
+    full_dataset = get_dataset(
+        dataset="iwildcam", download=download, root_dir=root_path, unlabeled=unlabeled
+    )
+
+    dataset = full_dataset.get_subset(split, transform=transform)
+
+    # Subset the labeled training data for label-efficiency experiments.
+    if label_fraction is not None and split == "train":
+        labels = dataset.y_array
+        indices = _stratified_subset_indices(
+            labels=labels, fraction=float(label_fraction), seed=int(seed)
+        )
+        dataset = torch.utils.data.Subset(dataset, indices)
+        logger.info(
+            f"iWildCam {split} subset created with {len(dataset)} samples "
+            f"(label_fraction={label_fraction}, seed={seed})"
+        )
+
+    # Use the unified wrapper that mimics the ImageNet structure
+    dataset = WildsToTorchWrapper(dataset, is_unlabeled=unlabeled)
+
+    logger.info(f"iWildCam {split} dataset created with {len(dataset)} samples")
+
+    dist_sampler = torch.utils.data.distributed.DistributedSampler(
+        dataset=dataset, num_replicas=world_size, rank=rank, shuffle=shuffle
+    )
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=collator,
+        sampler=dist_sampler,
+        batch_size=batch_size,
+        drop_last=drop_last,
+        pin_memory=pin_mem,
+        num_workers=num_workers,
+        persistent_workers=(num_workers > 0),
+    )
+
+    return dataset, data_loader, dist_sampler
+
+
+def make_iwildcam_eval(
+    transform,
+    batch_size,
+    split="test",
+    num_workers=8,
+    root_path="./wilds_data",
+    download=True,
+    pin_mem=True,
+):
+    full_dataset = get_dataset(
+        dataset="iwildcam", download=download, root_dir=root_path
+    )
+
+    dataset = full_dataset.get_subset(split, transform=transform)
+
+    logger.info(f"iWildCam {split} eval dataset created with {len(dataset)} samples")
+
+    data_loader = get_eval_loader(
+        "standard",
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_mem,
+    )
+
+    return full_dataset, dataset, data_loader
+
+
+class WildsToTorchWrapper(torch.utils.data.Dataset):
+    """
+    Mimics the ImageNet wrapper by always returning (image, target).
+    """
+
+    def __init__(self, wilds_subset, is_unlabeled=False):
+        self.dataset = wilds_subset
+        self.is_unlabeled = is_unlabeled
+
+    def __getitem__(self, i):
+        if self.is_unlabeled:
+            # metadata is at index 1 for unlabeled WILDS
+            x, _ = self.dataset[i]
+            target = -1  # Dummy target to match (img, target) signature
+        else:
+            # image, target, metadata
+            x, y, _ = self.dataset[i]
+            target = y
+
+        return x, target
+
+    def __len__(self):
+        return len(self.dataset)

src/eval_wilds.py

@@ -0,0 +1,273 @@
+import os
+import sys
+import json
+import resource
+import time
+import yaml
+import logging
+
+import numpy as np
+
+import torch
+import torch.distributed as dist
+
+from src.datasets.wilds import make_iwildcam_eval
+from src.helper import init_model
+from src.models.head import ViTClassifier
+from src.transforms import make_transform_eval
+from src.utils.distributed import init_distributed
+from src.utils.logging import resolve_log_dir
+
+
+_GLOBAL_SEED = 0
+np.random.seed(_GLOBAL_SEED)
+torch.manual_seed(_GLOBAL_SEED)
+torch.backends.cudnn.benchmark = True
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger()
+
+
+def _peak_host_ram_gb():
+    """Peak resident set size (RSS) of this process, in GB.
+
+    Uses resource.getrusage(RUSAGE_SELF).ru_maxrss (kilobytes on Linux). This is
+    the process high-water mark, folded into run_info so the seed aggregator can
+    report peak_host_ram_gb for eval-only runs.
+    """
+    try:
+        maxrss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        return float(maxrss_kb) / (1024.0 * 1024.0)
+    except (ValueError, OSError):
+        return None
+
+
+def _peak_gpu_mem_gb(device):
+    """Peak allocated and reserved GPU memory (GB) since the last reset."""
+    if not torch.cuda.is_available():
+        return None, None
+    try:
+        alloc = torch.cuda.max_memory_allocated(device) / 1e9
+        reserved = torch.cuda.max_memory_reserved(device) / 1e9
+        return float(alloc), float(reserved)
+    except (RuntimeError, ValueError):
+        return None, None
+
+
+def strip_module_prefix(state_dict):
+    if not any(k.startswith("module.") for k in state_dict.keys()):
+        return state_dict
+    return {
+        k[len("module.") :] if k.startswith("module.") else k: v
+        for k, v in state_dict.items()
+    }
+
+
+def _load_yaml(path):
+    with open(path, "r") as f:
+        return yaml.load(f, Loader=yaml.FullLoader)
+
+
+def _get_seed(args):
+    return int(args.get("meta", {}).get("seed", _GLOBAL_SEED))
+
+
+def _set_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _resolve_checkpoint_path(meta_args):
+    if meta_args.get("checkpoint_path"):
+        return meta_args["checkpoint_path"]
+    folder = meta_args.get("checkpoint_folder")
+    fname = meta_args.get("read_checkpoint")
+    if folder is None or fname is None:
+        return None
+    return fname if os.path.isabs(fname) else os.path.join(folder, fname)
+
+
+def _load_model_state(model, checkpoint_path, device):
+    if checkpoint_path is None:
+        raise ValueError("checkpoint_path or checkpoint_folder/read_checkpoint required")
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    if "model" not in checkpoint:
+        raise KeyError(f"No model weights found in checkpoint: {checkpoint_path}")
+    state = strip_module_prefix(checkpoint["model"])
+    msg = model.load_state_dict(state, strict=True)
+    model.to(device)
+    logger.info(f"Loaded model from {checkpoint_path} with msg: {msg}")
+
+
+def main(args):
+    force_single = bool(args.get("meta", {}).get("force_single_process", False))
+    if force_single:
+        world_size, rank = 1, 0
+    else:
+        world_size, rank = init_distributed()
+
+    if (
+        not force_single
+        and dist.is_available()
+        and dist.is_initialized()
+        and world_size > 1
+        and rank != 0
+    ):
+        dist.barrier()
+        dist.destroy_process_group()
+        return
+
+    seed = _get_seed(args)
+    _set_seed(seed)
+
+    meta_args = args.get("meta", {})
+    data_args = args.get("data", {})
+    log_args = args.get("logging", {})
+
+    if not torch.cuda.is_available():
+        device = torch.device("cpu")
+    else:
+        device = torch.device(f"cuda:{torch.cuda.current_device()}")
+
+    folder = resolve_log_dir(args, stage="eval")
+    tag = log_args.get("write_tag", "wilds_eval")
+
+    params_path = os.path.join(folder, "params-eval.yaml")
+    with open(params_path, "w") as f:
+        yaml.dump(args, f)
+
+    model_name = meta_args["model_name"]
+    embed_dim = meta_args["embed_dim"]
+    num_classes = meta_args["num_classes"]
+    patch_size = meta_args.get("patch_size", 16)
+    crop_size = meta_args.get("crop_size", 224)
+    use_bfloat16 = bool(meta_args.get("use_bfloat16", True))
+    use_autocast = use_bfloat16 and torch.cuda.is_available()
+
+    encoder, _ = init_model(
+        device=device,
+        patch_size=patch_size,
+        crop_size=crop_size,
+        model_name=model_name,
+    )
+    model = ViTClassifier(
+        encoder,
+        num_classes,
+        embed_dim,
+        representation_type=meta_args.get("representation_type", "last_avgpool"),
+        head_type=meta_args.get("head_type", "linear"),
+    ).to(device)
+
+    checkpoint_path = _resolve_checkpoint_path(meta_args)
+    _load_model_state(model, checkpoint_path, device)
+    model.eval()
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+    eval_start = time.time()
+
+    eval_transform = make_transform_eval(crop_size=crop_size)
+    split = data_args.get("split", "test")
+    root_path = data_args.get("root_path", "./wilds_data")
+
+    full_dataset, eval_dataset, eval_loader = make_iwildcam_eval(
+        transform=eval_transform,
+        split=split,
+        batch_size=data_args.get("batch_size", 128),
+        root_path=root_path,
+        num_workers=data_args.get("num_workers", 8),
+        pin_mem=bool(data_args.get("pin_mem", True)),
+        download=bool(data_args.get("download", True)),
+    )
+
+    all_y_pred = []
+    all_y_true = []
+    all_metadata = []
+
+    with torch.no_grad():
+        for batch in eval_loader:
+            if len(batch) == 3:
+                imgs, y_true, metadata = batch
+            else:
+                raise ValueError("Expected eval loader to return (x, y, metadata)")
+            imgs = imgs.to(device, non_blocking=True)
+            with torch.cuda.amp.autocast(enabled=use_autocast, dtype=torch.bfloat16):
+                logits = model(imgs)
+            preds = logits.argmax(dim=1).cpu()
+            all_y_pred.append(preds)
+            all_y_true.append(y_true.cpu())
+            all_metadata.append(metadata.cpu())
+
+    if all_y_pred:
+        all_y_pred = torch.cat(all_y_pred, dim=0)
+        all_y_true = torch.cat(all_y_true, dim=0)
+    all_metadata = torch.cat(all_metadata, dim=0) if all_metadata else torch.empty(0)
+
+    if int(all_metadata.shape[0]) != int(all_y_pred.shape[0]):
+        raise ValueError(
+            "Metadata length mismatch with predictions: "
+            f"{int(all_metadata.shape[0])} vs {int(all_y_pred.shape[0])}"
+        )
+
+    metrics = full_dataset.eval(all_y_pred, all_y_true, all_metadata)
+
+    eval_time_seconds = time.time() - eval_start
+    peak_host_ram_gb = _peak_host_ram_gb()
+    peak_gpu_alloc_gb, peak_gpu_reserved_gb = _peak_gpu_mem_gb(device)
+    run_info = {
+        "seed": seed,
+        "eval_time_seconds": float(eval_time_seconds),
+        "peak_host_ram_gb": peak_host_ram_gb,
+        "peak_gpu_alloc_gb": peak_gpu_alloc_gb,
+        "peak_gpu_reserved_gb": peak_gpu_reserved_gb,
+    }
+
+    # WILDS `full_dataset.eval` returns a (metrics_dict, summary_str) tuple/list.
+    # Wrap it in a dict so run_info + split live alongside the metrics, matching
+    # what tools/aggregate_seeds.py expects (it recurses for the metric keys).
+    if isinstance(metrics, (list, tuple)):
+        metrics_dict = metrics[0] if metrics and isinstance(metrics[0], dict) else {}
+        summary_str = metrics[1] if len(metrics) > 1 else None
+        metrics_out = dict(metrics_dict)
+        if summary_str is not None:
+            metrics_out["summary"] = summary_str
+    elif isinstance(metrics, dict):
+        metrics_out = dict(metrics)
+    else:
+        metrics_out = {"value": metrics}
+    metrics_out["run_info"] = run_info
+    metrics_out["split"] = split
+
+    metrics_path = os.path.join(folder, f"{tag}_metrics.json")
+    with open(metrics_path, "w") as f:
+        json.dump(metrics_out, f, indent=2, sort_keys=True)
+    logger.info(
+        f"Eval metrics saved to {metrics_path} "
+        f"(peak_host_ram_gb={peak_host_ram_gb}, "
+        f"peak_gpu_alloc_gb={peak_gpu_alloc_gb})"
+    )
+    logger.info(f"Eval metrics: {metrics}")
+
+    if (
+        not force_single
+        and dist.is_available()
+        and dist.is_initialized()
+        and world_size > 1
+    ):
+        dist.barrier()
+        dist.destroy_process_group()
+
+    return {
+        "metrics": metrics,
+        "metrics_path": metrics_path,
+        "folder": folder,
+    }
+
+
+if __name__ == "__main__":
+    raise RuntimeError(
+        "Use main_eval_wilds.py to launch this script with a config file."
+    )

src/helper.py

@@ -13,7 +13,8 @@ import torch
 import src.models.vision_transformer as vit
 from src.utils.schedulers import (
     WarmupCosineSchedule,
-    CosineWDSchedule)
+    CosineWDSchedule,
+    LinearWDSchedule)
 from src.utils.tensors import trunc_normal_
 
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
@@ -71,17 +72,20 @@ def init_model(
     model_name='vit_base',
     crop_size=224,
     pred_depth=6,
-    pred_emb_dim=384
+    pred_emb_dim=384,
+    use_gradient_checkpointing=False
 ):
     encoder = vit.__dict__[model_name](
         img_size=[crop_size],
-        patch_size=patch_size)
+        patch_size=patch_size,
+        use_checkpoint=use_gradient_checkpointing)
     predictor = vit.__dict__['vit_predictor'](
         num_patches=encoder.patch_embed.num_patches,
         embed_dim=encoder.embed_dim,
         predictor_embed_dim=pred_emb_dim,
         depth=pred_depth,
-        num_heads=encoder.num_heads)
+        num_heads=encoder.num_heads,
+        use_checkpoint=use_gradient_checkpointing)
 
     def init_weights(m):
         if isinstance(m, torch.nn.Linear):
@@ -116,7 +120,8 @@ def init_opt(
     final_wd=1e-6,
     final_lr=0.0,
     use_bfloat16=False,
-    ipe_scale=1.25
+    ipe_scale=1.25,
+    wd_schedule='cosine'
 ):
     param_groups = [
         {
@@ -147,10 +152,20 @@ def init_opt(
         ref_lr=ref_lr,
         final_lr=final_lr,
         T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
-    wd_scheduler = CosineWDSchedule(
-        optimizer,
-        ref_wd=wd,
-        final_wd=final_wd,
-        T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
-    scaler = torch.cuda.amp.GradScaler() if use_bfloat16 else None
+    wd_schedule = str(wd_schedule).lower()
+    if wd_schedule == 'linear':
+        logger.info('Using linear weight-decay schedule')
+        wd_scheduler = LinearWDSchedule(
+            optimizer,
+            ref_wd=wd,
+            final_wd=final_wd,
+            T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
+    else:
+        logger.info('Using cosine weight-decay schedule')
+        wd_scheduler = CosineWDSchedule(
+            optimizer,
+            ref_wd=wd,
+            final_wd=final_wd,
+            T_max=int(ipe_scale*num_epochs*iterations_per_epoch))
+    scaler = None
     return optimizer, scaler, scheduler, wd_scheduler

src/models/head.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+
+
+class ViTClassifier(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        num_classes,
+        embed_dim,
+        representation_type="last_avgpool",
+        head_type="linear",
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.representation_type = str(representation_type).lower()
+
+        if self.representation_type == "last_avgpool":
+            in_dim = embed_dim
+        elif self.representation_type == "last4_avgpool_concat":
+            in_dim = 4 * embed_dim
+        else:
+            raise ValueError(f"Unknown representation_type: {representation_type}")
+
+        head_type = str(head_type).lower()
+        if head_type == "linear":
+            self.head = nn.Linear(in_dim, num_classes)
+        elif head_type == "bn_linear":
+            self.head = nn.Sequential(
+                nn.BatchNorm1d(in_dim, affine=False, eps=1e-6),
+                nn.Linear(in_dim, num_classes),
+            )
+        else:
+            raise ValueError(f"Unknown head_type: {head_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 _extract_representation(self, x):
+        if self.representation_type == "last_avgpool":
+            features = self.encoder(x)
+            return features.mean(dim=1)
+
+        _, layer_outputs = self.encoder(
+            x, return_layer_outputs=True, num_last_layers=4
+        )
+        pooled = [layer_tokens.mean(dim=1) for layer_tokens in layer_outputs]
+        return torch.cat(pooled, dim=-1)
+
+    def forward(self, x):
+        if any(p.requires_grad for p in self.encoder.parameters()):
+            representation = self._extract_representation(x)
+        else:
+            with torch.no_grad():
+                representation = self._extract_representation(x)
+
+        logits = self.head(representation)
+        return logits

src/models/vision_transformer.py

@@ -11,6 +11,7 @@ import numpy as np
 
 import torch
 import torch.nn as nn
+from torch.utils.checkpoint import checkpoint as _checkpoint
 
 from src.utils.tensors import (
     trunc_normal_,
@@ -82,6 +83,12 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
     return emb
 
 
+def _maybe_checkpoint(blk, x, use_checkpoint, training):
+    if use_checkpoint and training and x.requires_grad:
+        return _checkpoint(blk, x, use_reentrant=False)
+    return blk(x)
+
+
 def drop_path(x, drop_prob: float = 0., training: bool = False):
     if drop_prob == 0. or not training:
         return x
@@ -234,9 +241,11 @@ class VisionTransformerPredictor(nn.Module):
         drop_path_rate=0.0,
         norm_layer=nn.LayerNorm,
         init_std=0.02,
+        use_checkpoint=False,
         **kwargs
     ):
         super().__init__()
+        self.use_checkpoint = use_checkpoint
         self.predictor_embed = nn.Linear(embed_dim, predictor_embed_dim, bias=True)
         self.mask_token = nn.Parameter(torch.zeros(1, 1, predictor_embed_dim))
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
@@ -316,7 +325,7 @@ class VisionTransformerPredictor(nn.Module):
 
         # -- fwd prop
         for blk in self.predictor_blocks:
-            x = blk(x)
+            x = _maybe_checkpoint(blk, x, self.use_checkpoint, self.training)
         x = self.predictor_norm(x)
 
         # -- return preds for mask tokens
@@ -346,11 +355,13 @@ class VisionTransformer(nn.Module):
         drop_path_rate=0.0,
         norm_layer=nn.LayerNorm,
         init_std=0.02,
+        use_checkpoint=False,
         **kwargs
     ):
         super().__init__()
         self.num_features = self.embed_dim = embed_dim
         self.num_heads = num_heads
+        self.use_checkpoint = use_checkpoint
         # --
         self.patch_embed = PatchEmbed(
             img_size=img_size[0],
@@ -398,7 +409,13 @@ class VisionTransformer(nn.Module):
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, masks=None):
+    def forward(
+        self,
+        x,
+        masks=None,
+        return_layer_outputs=False,
+        num_last_layers=4,
+    ):
         if masks is not None:
             if not isinstance(masks, list):
                 masks = [masks]
@@ -416,11 +433,20 @@ class VisionTransformer(nn.Module):
             x = apply_masks(x, masks)
 
         # -- fwd prop
+        layer_outputs = []
         for i, blk in enumerate(self.blocks):
-            x = blk(x)
+            x = _maybe_checkpoint(blk, x, self.use_checkpoint, self.training)
+            if return_layer_outputs:
+                layer_outputs.append(x)
 
         if self.norm is not None:
             x = self.norm(x)
+            if return_layer_outputs:
+                layer_outputs = [self.norm(t) for t in layer_outputs]
+
+        if return_layer_outputs:
+            k = min(int(num_last_layers), len(layer_outputs))
+            return x, layer_outputs[-k:]
 
         return x
 

src/train.py

@@ -13,7 +13,7 @@ try:
     # --          SURE TO UPDATE THIS TO GET LOCAL-RANK ON NODE, OR ENSURE
     # --          THAT YOUR JOBS ARE LAUNCHED WITH ONLY 1 DEVICE VISIBLE
     # --          TO EACH PROCESS
-    os.environ['CUDA_VISIBLE_DEVICES'] = os.environ['SLURM_LOCALID']
+    os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["SLURM_LOCALID"]
 except Exception:
     pass
 
@@ -25,28 +25,19 @@ import yaml
 import numpy as np
 
 import torch
+import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn.functional as F
 from torch.nn.parallel import DistributedDataParallel
 
 from src.masks.multiblock import MaskCollator as MBMaskCollator
 from src.masks.utils import apply_masks
-from src.utils.distributed import (
-    init_distributed,
-    AllReduce
-)
-from src.utils.logging import (
-    CSVLogger,
-    gpu_timer,
-    grad_logger,
-    AverageMeter)
+from src.utils.distributed import init_distributed, AllReduce
+from src.utils.logging import CSVLogger, gpu_timer, grad_logger, AverageMeter
 from src.utils.tensors import repeat_interleave_batch
-from src.datasets.imagenet1k import make_imagenet1k
+from src.datasets.wilds import make_iwildcam
 
-from src.helper import (
-    load_checkpoint,
-    init_model,
-    init_opt)
+from src.helper import load_checkpoint, init_model, init_opt
 from src.transforms import make_transforms
 
 # --
@@ -65,98 +56,104 @@ logger = logging.getLogger()
 
 
 def main(args, resume_preempt=False):
-
     # ----------------------------------------------------------------------- #
     #  PASSED IN PARAMS FROM CONFIG FILE
     # ----------------------------------------------------------------------- #
 
     # -- META
-    use_bfloat16 = args['meta']['use_bfloat16']
-    model_name = args['meta']['model_name']
-    load_model = args['meta']['load_checkpoint'] or resume_preempt
-    r_file = args['meta']['read_checkpoint']
-    copy_data = args['meta']['copy_data']
-    pred_depth = args['meta']['pred_depth']
-    pred_emb_dim = args['meta']['pred_emb_dim']
+    use_bfloat16 = args["meta"]["use_bfloat16"]
+    model_name = args["meta"]["model_name"]
+    use_gradient_checkpointing = args["meta"].get("use_gradient_checkpointing", False)
+    load_model = args["meta"]["load_checkpoint"] or resume_preempt
+    r_file = args["meta"]["read_checkpoint"]
+    copy_data = args["meta"]["copy_data"]
+    pred_depth = args["meta"]["pred_depth"]
+    pred_emb_dim = args["meta"]["pred_emb_dim"]
     if not torch.cuda.is_available():
-        device = torch.device('cpu')
+        device = torch.device("cpu")
     else:
-        device = torch.device('cuda:0')
+        device = torch.device("cuda:0")
         torch.cuda.set_device(device)
 
     # -- DATA
-    use_gaussian_blur = args['data']['use_gaussian_blur']
-    use_horizontal_flip = args['data']['use_horizontal_flip']
-    use_color_distortion = args['data']['use_color_distortion']
-    color_jitter = args['data']['color_jitter_strength']
+    use_gaussian_blur = args["data"]["use_gaussian_blur"]
+    use_horizontal_flip = args["data"]["use_horizontal_flip"]
+    use_color_distortion = args["data"]["use_color_distortion"]
+    color_jitter = args["data"]["color_jitter_strength"]
+    use_random_resized_crop = args["data"].get("use_random_resized_crop", True)
     # --
-    batch_size = args['data']['batch_size']
-    pin_mem = args['data']['pin_mem']
-    num_workers = args['data']['num_workers']
-    root_path = args['data']['root_path']
-    image_folder = args['data']['image_folder']
-    crop_size = args['data']['crop_size']
-    crop_scale = args['data']['crop_scale']
+    batch_size = args["data"]["batch_size"]
+    pin_mem = args["data"]["pin_mem"]
+    num_workers = args["data"]["num_workers"]
+    root_path = args["data"]["root_path"]
+    crop_size = args["data"]["crop_size"]
+    crop_scale = args["data"]["crop_scale"]
     # --
 
     # -- MASK
-    allow_overlap = args['mask']['allow_overlap']  # whether to allow overlap b/w context and target blocks
-    patch_size = args['mask']['patch_size']  # patch-size for model training
-    num_enc_masks = args['mask']['num_enc_masks']  # number of context blocks
-    min_keep = args['mask']['min_keep']  # min number of patches in context block
-    enc_mask_scale = args['mask']['enc_mask_scale']  # scale of context blocks
-    num_pred_masks = args['mask']['num_pred_masks']  # number of target blocks
-    pred_mask_scale = args['mask']['pred_mask_scale']  # scale of target blocks
-    aspect_ratio = args['mask']['aspect_ratio']  # aspect ratio of target blocks
+    allow_overlap = args["mask"][
+        "allow_overlap"
+    ]  # whether to allow overlap b/w context and target blocks
+    patch_size = args["mask"]["patch_size"]  # patch-size for model training
+    num_enc_masks = args["mask"]["num_enc_masks"]  # number of context blocks
+    min_keep = args["mask"]["min_keep"]  # min number of patches in context block
+    enc_mask_scale = args["mask"]["enc_mask_scale"]  # scale of context blocks
+    num_pred_masks = args["mask"]["num_pred_masks"]  # number of target blocks
+    pred_mask_scale = args["mask"]["pred_mask_scale"]  # scale of target blocks
+    aspect_ratio = args["mask"]["aspect_ratio"]  # aspect ratio of target blocks
     # --
 
     # -- OPTIMIZATION
-    ema = args['optimization']['ema']
-    ipe_scale = args['optimization']['ipe_scale']  # scheduler scale factor (def: 1.0)
-    wd = float(args['optimization']['weight_decay'])
-    final_wd = float(args['optimization']['final_weight_decay'])
-    num_epochs = args['optimization']['epochs']
-    warmup = args['optimization']['warmup']
-    start_lr = args['optimization']['start_lr']
-    lr = args['optimization']['lr']
-    final_lr = args['optimization']['final_lr']
+    ema = args["optimization"]["ema"]
+    ipe_scale = args["optimization"]["ipe_scale"]  # scheduler scale factor (def: 1.0)
+    wd = float(args["optimization"]["weight_decay"])
+    final_wd = float(args["optimization"]["final_weight_decay"])
+    wd_schedule = args["optimization"].get("wd_schedule", "cosine")
+    num_epochs = args["optimization"]["epochs"]
+    warmup = args["optimization"]["warmup"]
+    start_lr = args["optimization"]["start_lr"]
+    lr = args["optimization"]["lr"]
+    final_lr = args["optimization"]["final_lr"]
+    accum_steps = args["optimization"].get("gradient_accumulation_steps", 1)
 
     # -- LOGGING
-    folder = args['logging']['folder']
-    tag = args['logging']['write_tag']
+    folder = args["logging"]["folder"]
+    tag = args["logging"]["write_tag"]
 
-    dump = os.path.join(folder, 'params-ijepa.yaml')
-    with open(dump, 'w') as f:
+    dump = os.path.join(folder, "params-ijepa.yaml")
+    with open(dump, "w") as f:
         yaml.dump(args, f)
     # ----------------------------------------------------------------------- #
 
     try:
-        mp.set_start_method('spawn')
+        mp.set_start_method("spawn")
     except Exception:
         pass
 
     # -- init torch distributed backend
     world_size, rank = init_distributed()
-    logger.info(f'Initialized (rank/world-size) {rank}/{world_size}')
+    logger.info(f"Initialized (rank/world-size) {rank}/{world_size}")
     if rank > 0:
         logger.setLevel(logging.ERROR)
 
     # -- log/checkpointing paths
-    log_file = os.path.join(folder, f'{tag}_r{rank}.csv')
-    save_path = os.path.join(folder, f'{tag}' + '-ep{epoch}.pth.tar')
-    latest_path = os.path.join(folder, f'{tag}-latest.pth.tar')
+    log_file = os.path.join(folder, f"{tag}_r{rank}.csv")
+    save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar")
+    latest_path = os.path.join(folder, f"{tag}-latest.pth.tar")
     load_path = None
     if load_model:
         load_path = os.path.join(folder, r_file) if r_file is not None else latest_path
 
     # -- make csv_logger
-    csv_logger = CSVLogger(log_file,
-                           ('%d', 'epoch'),
-                           ('%d', 'itr'),
-                           ('%.5f', 'loss'),
-                           ('%.5f', 'mask-A'),
-                           ('%.5f', 'mask-B'),
-                           ('%d', 'time (ms)'))
+    csv_logger = CSVLogger(
+        log_file,
+        ("%d", "epoch"),
+        ("%d", "itr"),
+        ("%.5f", "loss"),
+        ("%.5f", "mask-A"),
+        ("%.5f", "mask-B"),
+        ("%d", "time (ms)"),
+    )
 
     # -- init model
     encoder, predictor = init_model(
@@ -165,7 +162,9 @@ def main(args, resume_preempt=False):
         crop_size=crop_size,
         pred_depth=pred_depth,
         pred_emb_dim=pred_emb_dim,
-        model_name=model_name)
+        model_name=model_name,
+        use_gradient_checkpointing=use_gradient_checkpointing,
+    )
     target_encoder = copy.deepcopy(encoder)
 
     # -- make data transforms
@@ -178,7 +177,8 @@ def main(args, resume_preempt=False):
         nenc=num_enc_masks,
         npred=num_pred_masks,
         allow_overlap=allow_overlap,
-        min_keep=min_keep)
+        min_keep=min_keep,
+    )
 
     transform = make_transforms(
         crop_size=crop_size,
@@ -186,22 +186,22 @@ def main(args, resume_preempt=False):
         gaussian_blur=use_gaussian_blur,
         horizontal_flip=use_horizontal_flip,
         color_distortion=use_color_distortion,
-        color_jitter=color_jitter)
+        color_jitter=color_jitter,
+        use_random_resized_crop=use_random_resized_crop,
+    )
 
     # -- init data-loaders/samplers
-    _, unsupervised_loader, unsupervised_sampler = make_imagenet1k(
-            transform=transform,
-            batch_size=batch_size,
-            collator=mask_collator,
-            pin_mem=pin_mem,
-            training=True,
-            num_workers=num_workers,
-            world_size=world_size,
-            rank=rank,
-            root_path=root_path,
-            image_folder=image_folder,
-            copy_data=copy_data,
-            drop_last=True)
+    _, unsupervised_loader, unsupervised_sampler = make_iwildcam(
+        transform=transform,
+        batch_size=batch_size,
+        collator=mask_collator,
+        pin_mem=pin_mem,
+        num_workers=num_workers,
+        world_size=world_size,
+        rank=rank,
+        root_path=root_path,
+        drop_last=True,
+    )
     ipe = len(unsupervised_loader)
 
     # -- init optimizer and scheduler
@@ -217,29 +217,41 @@ def main(args, resume_preempt=False):
         warmup=warmup,
         num_epochs=num_epochs,
         ipe_scale=ipe_scale,
-        use_bfloat16=use_bfloat16)
-    encoder = DistributedDataParallel(encoder, static_graph=True)
-    predictor = DistributedDataParallel(predictor, static_graph=True)
-    target_encoder = DistributedDataParallel(target_encoder)
-    for p in target_encoder.parameters():
-        p.requires_grad = False
+        use_bfloat16=use_bfloat16,
+        wd_schedule=wd_schedule,
+    )
+    if dist.is_available() and dist.is_initialized() and world_size > 1:
+        encoder = DistributedDataParallel(encoder, static_graph=True)
+        predictor = DistributedDataParallel(predictor, static_graph=True)
+        target_encoder = DistributedDataParallel(target_encoder)
+        for p in target_encoder.parameters():
+            p.requires_grad = False
+    else:
+        logger.info("Distributed process group not initialized; running without DDP")
+        for p in target_encoder.parameters():
+            p.requires_grad = False
 
     # -- momentum schedule
-    momentum_scheduler = (ema[0] + i*(ema[1]-ema[0])/(ipe*num_epochs*ipe_scale)
-                          for i in range(int(ipe*num_epochs*ipe_scale)+1))
+    momentum_scheduler = (
+        ema[0] + i * (ema[1] - ema[0]) / (ipe * num_epochs * ipe_scale)
+        for i in range(int(ipe * num_epochs * ipe_scale) + 1)
+    )
 
     start_epoch = 0
     # -- load training checkpoint
     if load_model:
-        encoder, predictor, target_encoder, optimizer, scaler, start_epoch = load_checkpoint(
-            device=device,
-            r_path=load_path,
-            encoder=encoder,
-            predictor=predictor,
-            target_encoder=target_encoder,
-            opt=optimizer,
-            scaler=scaler)
-        for _ in range(start_epoch*ipe):
+        encoder, predictor, target_encoder, optimizer, scaler, start_epoch = (
+            load_checkpoint(
+                device=device,
+                r_path=load_path,
+                encoder=encoder,
+                predictor=predictor,
+                target_encoder=target_encoder,
+                opt=optimizer,
+                scaler=scaler,
+            )
+        )
+        for _ in range(start_epoch * ipe):
             scheduler.step()
             wd_scheduler.step()
             next(momentum_scheduler)
@@ -247,25 +259,26 @@ def main(args, resume_preempt=False):
 
     def save_checkpoint(epoch):
         save_dict = {
-            'encoder': encoder.state_dict(),
-            'predictor': predictor.state_dict(),
-            'target_encoder': target_encoder.state_dict(),
-            'opt': optimizer.state_dict(),
-            'scaler': None if scaler is None else scaler.state_dict(),
-            'epoch': epoch,
-            'loss': loss_meter.avg,
-            'batch_size': batch_size,
-            'world_size': world_size,
-            'lr': lr
+            "encoder": encoder.state_dict(),
+            "predictor": predictor.state_dict(),
+            "target_encoder": target_encoder.state_dict(),
+            "opt": optimizer.state_dict(),
+            "scaler": None if scaler is None else scaler.state_dict(),
+            "epoch": epoch,
+            "loss": loss_meter.avg,
+            "batch_size": batch_size,
+            "world_size": world_size,
+            "lr": lr,
         }
         if rank == 0:
             torch.save(save_dict, latest_path)
             if (epoch + 1) % checkpoint_freq == 0:
-                torch.save(save_dict, save_path.format(epoch=f'{epoch + 1}'))
+                torch.save(save_dict, save_path.format(epoch=f"{epoch + 1}"))
 
     # -- TRAINING LOOP
+    optimizer.zero_grad()
     for epoch in range(start_epoch, num_epochs):
-        logger.info('Epoch %d' % (epoch + 1))
+        logger.info("Epoch %d" % (epoch + 1))
 
         # -- update distributed-data-loader epoch
         unsupervised_sampler.set_epoch(epoch)
@@ -283,15 +296,12 @@ def main(args, resume_preempt=False):
                 masks_1 = [u.to(device, non_blocking=True) for u in masks_enc]
                 masks_2 = [u.to(device, non_blocking=True) for u in masks_pred]
                 return (imgs, masks_1, masks_2)
+
             imgs, masks_enc, masks_pred = load_imgs()
             maskA_meter.update(len(masks_enc[0][0]))
             maskB_meter.update(len(masks_pred[0][0]))
 
             def train_step():
-                _new_lr = scheduler.step()
-                _new_wd = wd_scheduler.step()
-                # --
-
                 def forward_target():
                     with torch.no_grad():
                         h = target_encoder(imgs)
@@ -313,66 +323,89 @@ def main(args, resume_preempt=False):
                     return loss
 
                 # Step 1. Forward
-                with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=use_bfloat16):
+                with torch.cuda.amp.autocast(
+                    dtype=torch.bfloat16, enabled=use_bfloat16
+                ):
                     h = forward_target()
                     z = forward_context()
                     loss = loss_fn(z, h)
 
-                #  Step 2. Backward & step
-                if use_bfloat16:
-                    scaler.scale(loss).backward()
-                    scaler.step(optimizer)
-                    scaler.update()
-                else:
-                    loss.backward()
-                    optimizer.step()
-                grad_stats = grad_logger(encoder.named_parameters())
-                optimizer.zero_grad()
+                # Step 2. Backward (accumulate gradients)
+                (loss / accum_steps).backward()
 
-                # Step 3. momentum update of target encoder
-                with torch.no_grad():
-                    m = next(momentum_scheduler)
-                    for param_q, param_k in zip(encoder.parameters(), target_encoder.parameters()):
-                        param_k.data.mul_(m).add_((1.-m) * param_q.detach().data)
+                # Step 3. Optimizer / scheduler / momentum (every accum_steps)
+                is_accumulated = ((itr + 1) % accum_steps == 0)
+                if is_accumulated:
+                    _new_lr = scheduler.step()
+                    _new_wd = wd_scheduler.step()
+                    optimizer.step()
+                    optimizer.zero_grad(set_to_none=True)
+                    grad_stats = grad_logger(encoder.named_parameters())
+
+                    with torch.no_grad():
+                        m = next(momentum_scheduler)
+                        for param_q, param_k in zip(
+                            encoder.parameters(), target_encoder.parameters()
+                        ):
+                            param_k.data.mul_(m).add_((1.0 - m) * param_q.detach().data)
+                else:
+                    _new_lr = None
+                    _new_wd = None
+                    grad_stats = None
 
                 return (float(loss), _new_lr, _new_wd, grad_stats)
+
             (loss, _new_lr, _new_wd, grad_stats), etime = gpu_timer(train_step)
             loss_meter.update(loss)
             time_meter.update(etime)
 
             # -- Logging
             def log_stats():
-                csv_logger.log(epoch + 1, itr, loss, maskA_meter.val, maskB_meter.val, etime)
+                csv_logger.log(
+                    epoch + 1, itr, loss, maskA_meter.val, maskB_meter.val, etime
+                )
                 if (itr % log_freq == 0) or np.isnan(loss) or np.isinf(loss):
-                    logger.info('[%d, %5d] loss: %.3f '
-                                'masks: %.1f %.1f '
-                                '[wd: %.2e] [lr: %.2e] '
-                                '[mem: %.2e] '
-                                '(%.1f ms)'
-                                % (epoch + 1, itr,
-                                   loss_meter.avg,
-                                   maskA_meter.avg,
-                                   maskB_meter.avg,
-                                   _new_wd,
-                                   _new_lr,
-                                   torch.cuda.max_memory_allocated() / 1024.**2,
-                                   time_meter.avg))
+                    lr_to_log = _new_lr if _new_lr is not None else 0.
+                    wd_to_log = _new_wd if _new_wd is not None else 0.
+                    logger.info(
+                        "[%d, %5d] loss: %.3f "
+                        "masks: %.1f %.1f "
+                        "[wd: %.2e] [lr: %.2e] "
+                        "[mem: %.2e] "
+                        "(%.1f ms)"
+                        % (
+                            epoch + 1,
+                            itr,
+                            loss_meter.avg,
+                            maskA_meter.avg,
+                            maskB_meter.avg,
+                            wd_to_log,
+                            lr_to_log,
+                            torch.cuda.max_memory_allocated() / 1024.0**2,
+                            time_meter.avg,
+                        )
+                    )
 
                     if grad_stats is not None:
-                        logger.info('[%d, %5d] grad_stats: [%.2e %.2e] (%.2e, %.2e)'
-                                    % (epoch + 1, itr,
-                                       grad_stats.first_layer,
-                                       grad_stats.last_layer,
-                                       grad_stats.min,
-                                       grad_stats.max))
+                        logger.info(
+                            "[%d, %5d] grad_stats: [%.2e %.2e] (%.2e, %.2e)"
+                            % (
+                                epoch + 1,
+                                itr,
+                                grad_stats.first_layer,
+                                grad_stats.last_layer,
+                                grad_stats.min,
+                                grad_stats.max,
+                            )
+                        )
 
             log_stats()
 
-            assert not np.isnan(loss), 'loss is nan'
+            assert not np.isnan(loss), "loss is nan"
 
         # -- Save Checkpoint after every epoch
-        logger.info('avg. loss %.3f' % loss_meter.avg)
-        save_checkpoint(epoch+1)
+        logger.info("avg. loss %.3f" % loss_meter.avg)
+        save_checkpoint(epoch + 1)
 
 
 if __name__ == "__main__":

src/train_supervised.py

@@ -0,0 +1,731 @@
+import os
+import random
+import resource
+import shutil
+import sys
+import json
+import time
+import yaml
+import logging
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+from src.datasets.wilds import make_iwildcam
+from src.helper import init_model
+from src.models.head import ViTClassifier
+from src.transforms import make_transforms, make_transform_eval
+from src.utils.distributed import init_distributed
+from src.utils.logging import CSVLogger, AverageMeter, resolve_log_dir
+from src.utils.optimizers import LARS
+from src.eval_wilds import main as eval_wilds_main
+
+# --
+log_freq = 10
+checkpoint_freq = 50
+# --
+
+_GLOBAL_SEED = 0
+np.random.seed(_GLOBAL_SEED)
+torch.manual_seed(_GLOBAL_SEED)
+torch.backends.cudnn.benchmark = True
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger()
+
+
+def _set_seed(seed):
+    """Seed all RNGs so multi-seed runs are reproducible and distinct."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _format_hms(seconds):
+    seconds = int(round(seconds))
+    h, rem = divmod(seconds, 3600)
+    m, s = divmod(rem, 60)
+    return f"{h:02d}:{m:02d}:{s:02d}"
+
+
+def _peak_host_ram_gb():
+    """Peak resident set size (RSS) of this process, in GB.
+
+    Uses resource.getrusage(RUSAGE_SELF).ru_maxrss, which on Linux is reported
+    in kilobytes. This is the process high-water mark; with tasks_per_node=1 it
+    reflects the whole training worker. Compare against the SLURM mem request
+    (e.g. 180G) to right-size future jobs.
+    """
+    try:
+        maxrss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        return float(maxrss_kb) / (1024.0 * 1024.0)
+    except (ValueError, OSError):
+        return None
+
+
+def _peak_gpu_mem_gb(device):
+    """Peak allocated and reserved GPU memory (GB) since the last reset."""
+    if not torch.cuda.is_available():
+        return None, None
+    try:
+        alloc = torch.cuda.max_memory_allocated(device) / 1e9
+        reserved = torch.cuda.max_memory_reserved(device) / 1e9
+        return float(alloc), float(reserved)
+    except (RuntimeError, ValueError):
+        return None, None
+
+
+def strip_module_prefix(state_dict):
+    if not any(k.startswith("module.") for k in state_dict.keys()):
+        return state_dict
+    return {
+        k[len("module.") :] if k.startswith("module.") else k: v
+        for k, v in state_dict.items()
+    }
+
+
+def distributed_average(value, device):
+    tensor = torch.tensor([value], device=device, dtype=torch.float64)
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+        tensor /= dist.get_world_size()
+    return float(tensor.item())
+
+
+def distributed_sum(value, device):
+    tensor = torch.tensor([value], device=device, dtype=torch.float64)
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    return float(tensor.item())
+
+
+def _find_metric(metrics, key):
+    if isinstance(metrics, dict):
+        if key in metrics:
+            return metrics[key]
+        for value in metrics.values():
+            found = _find_metric(value, key)
+            if found is not None:
+                return found
+    if isinstance(metrics, list):
+        for item in metrics:
+            found = _find_metric(item, key)
+            if found is not None:
+                return found
+    return None
+
+
+def _collect_eval_rows(eval_root, metric_key):
+    rows = []
+    for dirpath, _, filenames in os.walk(eval_root):
+        for fname in filenames:
+            if not fname.endswith("_metrics.json"):
+                continue
+            path = os.path.join(dirpath, fname)
+            try:
+                with open(path, "r") as f:
+                    metrics = json.load(f)
+            except (OSError, json.JSONDecodeError):
+                continue
+            value = _find_metric(metrics, metric_key)
+            if value is None:
+                continue
+            run_name = os.path.basename(os.path.dirname(path))
+            rows.append((float(value), run_name, path))
+    return rows
+
+
+class EarlyStopping:
+    def __init__(
+        self,
+        enabled=True,
+        patience=10,
+        min_delta=0.0,
+        min_epochs=0,
+        restore_best_weights=True,
+    ):
+        self.enabled = enabled
+        self.patience = patience
+        self.min_delta = min_delta
+        self.min_epochs = min_epochs
+        self.restore_best_weights = restore_best_weights
+        self.best_metric = float("inf")
+        self.best_epoch = -1
+        self.bad_epochs = 0
+        self.best_state = None
+
+    def state_dict(self):
+        return {
+            "enabled": self.enabled,
+            "patience": self.patience,
+            "min_delta": self.min_delta,
+            "min_epochs": self.min_epochs,
+            "restore_best_weights": self.restore_best_weights,
+            "best_metric": self.best_metric,
+            "best_epoch": self.best_epoch,
+            "bad_epochs": self.bad_epochs,
+        }
+
+    def load_state_dict(self, state):
+        if not state:
+            return
+        self.enabled = state.get("enabled", self.enabled)
+        self.patience = state.get("patience", self.patience)
+        self.min_delta = state.get("min_delta", self.min_delta)
+        self.min_epochs = state.get("min_epochs", self.min_epochs)
+        self.restore_best_weights = state.get(
+            "restore_best_weights", self.restore_best_weights
+        )
+        self.best_metric = state.get("best_metric", self.best_metric)
+        self.best_epoch = state.get("best_epoch", self.best_epoch)
+        self.bad_epochs = state.get("bad_epochs", self.bad_epochs)
+
+    def step(self, epoch, metric, model_module):
+        if not self.enabled:
+            return False, False
+
+        improved = metric < (self.best_metric - self.min_delta)
+        if improved:
+            self.best_metric = metric
+            self.best_epoch = epoch
+            self.bad_epochs = 0
+            if self.restore_best_weights:
+                self.best_state = {
+                    k: v.detach().cpu().clone()
+                    for k, v in model_module.state_dict().items()
+                }
+            return True, False
+
+        self.bad_epochs += 1
+        should_stop = (
+            epoch + 1
+        ) >= self.min_epochs and self.bad_epochs >= self.patience
+        return False, should_stop
+
+    def restore(self, model_module, device):
+        if self.restore_best_weights and self.best_state is not None:
+            model_module.load_state_dict(self.best_state)
+            model_module.to(device)
+
+
+def evaluate(model, loader, criterion, device, use_bfloat16):
+    model.eval()
+    loss_sum = 0.0
+    n_correct = 0.0
+    n_total = 0.0
+
+    with torch.no_grad():
+        for imgs, labels in loader:
+            imgs = imgs.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            with torch.cuda.amp.autocast(enabled=use_bfloat16, dtype=torch.bfloat16):
+                outputs = model(imgs)
+                loss = criterion(outputs, labels)
+
+            batch_size = labels.size(0)
+            preds = outputs.argmax(dim=1)
+            n_correct += float((preds == labels).sum().item())
+            n_total += float(batch_size)
+            loss_sum += float(loss.item()) * float(batch_size)
+
+    global_loss_sum = distributed_sum(loss_sum, device)
+    global_correct = distributed_sum(n_correct, device)
+    global_total = distributed_sum(n_total, device)
+
+    val_loss = global_loss_sum / max(global_total, 1.0)
+    val_acc = global_correct / max(global_total, 1.0)
+    return val_loss, val_acc
+
+
+def main(args, resume_preempt=False):
+    del resume_preempt
+
+    world_size, rank = init_distributed()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for supervised distributed training")
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
+
+    m_args = args["meta"]
+    o_args = args["optimization"]
+    d_args = args["data"]
+    mk_args = args["mask"]
+    l_args = args["logging"]
+    v_args = args["validation"]
+    es_args = o_args["early_stopping"]
+
+    seed = int(m_args.get("seed", _GLOBAL_SEED))
+    _set_seed(seed)
+    logger.info(f"Using seed {seed}")
+    accum_steps = o_args.get("gradient_accumulation_steps", 1)
+    use_gradient_checkpointing = m_args.get("use_gradient_checkpointing", True)
+
+    folder = resolve_log_dir(args, stage="train")
+    tag = l_args["write_tag"]
+
+    with open(os.path.join(folder, "params-supervised.yaml"), "w") as f:
+        yaml.dump(args, f)
+    with open(os.path.join(folder, "params.yaml"), "w") as f:
+        yaml.dump(args, f)
+
+    save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar")
+    latest_path = os.path.join(folder, f"{tag}-latest.pth.tar")
+    best_path = os.path.join(folder, f"{tag}-best.pth.tar")
+    log_file = os.path.join(folder, f"{tag}_r{rank}.csv")
+
+    csv_logger = CSVLogger(
+        log_file,
+        ("%d", "epoch"),
+        ("%.6f", "train_loss"),
+        ("%.6f", "val_loss"),
+        ("%.6f", "val_acc"),
+        ("%.6e", "lr"),
+        ("%.6f", "best_val_loss"),
+        ("%d", "best_epoch"),
+        ("%d", "early_stop"),
+    )
+
+    encoder, _ = init_model(
+        device=device,
+        patch_size=mk_args["patch_size"],
+        crop_size=d_args["crop_size"],
+        model_name=m_args["model_name"],
+        use_gradient_checkpointing=use_gradient_checkpointing,
+    )
+
+    representation_type = m_args.get("representation_type", "last_avgpool")
+    head_type = m_args.get("head_type", "linear")
+    model = ViTClassifier(
+        encoder,
+        m_args["num_classes"],
+        m_args["embed_dim"],
+        representation_type=representation_type,
+        head_type=head_type,
+    ).to(device)
+
+    if o_args["freeze_weights"]:
+        logger.info("Freezing encoder weights (Linear Probing mode)")
+        for param in model.encoder.parameters():
+            param.requires_grad = False
+        model.encoder.eval()
+    else:
+        logger.info("Training full model (Fine-tuning mode)")
+
+    params = [p for p in model.parameters() if p.requires_grad]
+    optimizer_name = o_args["optimizer"].lower()
+    if optimizer_name == "adamw":
+        optimizer = torch.optim.AdamW(
+            params, lr=o_args["lr"], weight_decay=o_args["weight_decay"]
+        )
+    elif optimizer_name == "lars":
+        optimizer = LARS(
+            params,
+            lr=o_args["lr"],
+            weight_decay=o_args["weight_decay"],
+            momentum=o_args.get("momentum", 0.9),
+            eta=o_args.get("lars_eta", 0.001),
+            eps=o_args.get("lars_eps", 1e-8),
+            exclude_bias_and_norm=o_args.get("lars_exclude_bias_and_norm", True),
+        )
+    else:
+        optimizer = torch.optim.SGD(
+            params,
+            lr=o_args["lr"],
+            momentum=o_args.get("momentum", 0.9),
+            weight_decay=o_args["weight_decay"],
+        )
+
+    scheduler = None
+    lr_schedule = o_args.get("lr_schedule", "cosine").lower()
+    if lr_schedule == "step":
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            milestones=o_args.get("step_milestones", [15, 30, 45]),
+            gamma=o_args.get("step_gamma", 0.1),
+        )
+    elif lr_schedule == "cosine":
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"]
+        )
+
+    train_transform = make_transforms(
+        crop_size=d_args["crop_size"],
+        horizontal_flip=d_args["use_horizontal_flip"],
+        color_distortion=d_args["use_color_distortion"],
+        color_jitter=d_args["color_jitter_strength"],
+        gaussian_blur=d_args["use_gaussian_blur"],
+        use_random_resized_crop=d_args.get("use_random_resized_crop", True),
+    )
+    val_transform = make_transform_eval(
+        crop_size=d_args["crop_size"],
+    )
+
+    label_fraction = d_args.get("label_fraction", None)
+    if label_fraction is not None:
+        logger.info(
+            f"Using label_fraction={label_fraction} for the training split (seed={seed})"
+        )
+
+    _, train_loader, train_sampler = make_iwildcam(
+        transform=train_transform,
+        split="train",
+        batch_size=d_args["batch_size"],
+        root_path=d_args["root_path"],
+        rank=rank,
+        world_size=world_size,
+        collator=None,
+        num_workers=d_args["num_workers"],
+        pin_mem=d_args["pin_mem"],
+        drop_last=True,
+        label_fraction=label_fraction,
+        seed=seed,
+    )
+
+    _, val_loader, val_sampler = make_iwildcam(
+        transform=val_transform,
+        split="val",
+        batch_size=d_args["batch_size"],
+        root_path=d_args["root_path"],
+        rank=rank,
+        world_size=world_size,
+        collator=None,
+        num_workers=d_args["num_workers"],
+        pin_mem=d_args["pin_mem"],
+        drop_last=False,
+    )
+
+    criterion = nn.CrossEntropyLoss().to(device)
+    model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])
+
+    eval_every = int(v_args["eval_every"])
+
+    early_stopper = EarlyStopping(
+        enabled=es_args["enabled"],
+        patience=es_args["patience"],
+        min_delta=es_args["min_delta"],
+        min_epochs=es_args["min_epochs"],
+        restore_best_weights=es_args["restore_best_weights"],
+    )
+
+    start_epoch = 0
+
+    checkpoint_to_load = None
+    resuming_interrupted = False
+    if os.path.exists(latest_path):
+        checkpoint_to_load = latest_path
+        resuming_interrupted = True
+    elif m_args["load_checkpoint"]:
+        r_file = m_args["read_checkpoint"]
+        checkpoint_folder = m_args["checkpoint_folder"]
+        if os.path.isabs(r_file):
+            checkpoint_to_load = r_file
+        else:
+            checkpoint_to_load = os.path.join(checkpoint_folder, r_file)
+
+    if checkpoint_to_load is not None and not os.path.exists(checkpoint_to_load):
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_to_load}")
+
+    if checkpoint_to_load and os.path.exists(checkpoint_to_load):
+        checkpoint = torch.load(checkpoint_to_load, map_location="cpu")
+        if resuming_interrupted and "model" in checkpoint:
+            model.module.load_state_dict(checkpoint["model"])
+            if "opt" in checkpoint:
+                optimizer.load_state_dict(checkpoint["opt"])
+            if scheduler is not None and "scheduler" in checkpoint:
+                scheduler.load_state_dict(checkpoint["scheduler"])
+            start_epoch = int(checkpoint.get("epoch", 0))
+            early_stopper.load_state_dict(checkpoint.get("early_stopping", {}))
+            logger.info(
+                f"Resuming training from {checkpoint_to_load} at epoch {start_epoch}"
+            )
+        else:
+            encoder_state = checkpoint.get("encoder")
+            if encoder_state is None and "model" in checkpoint:
+                encoder_state = {
+                    k.replace("encoder.", "", 1): v
+                    for k, v in checkpoint["model"].items()
+                    if k.startswith("encoder.")
+                }
+            if encoder_state is None:
+                raise KeyError(
+                    f"No encoder weights found in checkpoint: {checkpoint_to_load}"
+                )
+            encoder_state = strip_module_prefix(encoder_state)
+            msg = model.module.encoder.load_state_dict(encoder_state, strict=False)
+            logger.info(
+                f"Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}"
+            )
+
+    def save_checkpoint(epoch, train_loss, val_loss, val_acc, is_best=False):
+        save_dict = {
+            "model": model.module.state_dict(),
+            "opt": optimizer.state_dict(),
+            "scheduler": None if scheduler is None else scheduler.state_dict(),
+            "epoch": epoch,
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "val_acc": val_acc,
+            "args": args,
+            "early_stopping": early_stopper.state_dict(),
+        }
+        if rank == 0:
+            torch.save(save_dict, latest_path)
+            if epoch % checkpoint_freq == 0:
+                torch.save(save_dict, save_path.format(epoch=epoch))
+            if is_best:
+                torch.save(save_dict, best_path)
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+
+    train_start_time = time.perf_counter()
+    epochs_run = 0
+    early_stopped = False
+    for epoch in range(start_epoch, o_args["epochs"]):
+        train_sampler.set_epoch(epoch)
+        val_sampler.set_epoch(epoch)
+        epochs_run = epoch + 1
+
+        model.train()
+        if o_args["freeze_weights"]:
+            model.module.encoder.eval()
+
+        loss_meter = AverageMeter()
+
+        optimizer.zero_grad()
+        for itr, (imgs, labels) in enumerate(train_loader):
+            imgs = imgs.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            with torch.cuda.amp.autocast(
+                enabled=m_args["use_bfloat16"], dtype=torch.bfloat16
+            ):
+                outputs = model(imgs)
+                loss = criterion(outputs, labels)
+
+            (loss / accum_steps).backward()
+
+            if (itr + 1) % accum_steps == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+
+            loss_meter.update(loss.item(), n=labels.size(0))
+
+            if itr % log_freq == 0 and rank == 0:
+                logger.info(
+                    f"Epoch {epoch + 1} [{itr}/{len(train_loader)}] Train Loss: {loss_meter.avg:.4f}"
+                )
+
+        train_loss = distributed_average(loss_meter.avg, device)
+
+        do_eval = ((epoch + 1) % eval_every == 0) or (epoch + 1 == o_args["epochs"])
+        if do_eval:
+            val_loss, val_acc = evaluate(
+                model=model,
+                loader=val_loader,
+                criterion=criterion,
+                device=device,
+                use_bfloat16=m_args["use_bfloat16"],
+            )
+            is_best, should_stop = early_stopper.step(epoch + 1, val_loss, model.module)
+        else:
+            val_loss = float("nan")
+            val_acc = float("nan")
+            is_best, should_stop = False, False
+
+        if scheduler is not None:
+            scheduler.step()
+
+        if rank == 0:
+            logger.info(
+                f"Epoch {epoch + 1} done | train_loss={train_loss:.6f} val_loss={val_loss:.6f} val_acc={val_acc:.6f} best_val_loss={early_stopper.best_metric:.6f}"
+            )
+
+        csv_logger.log(
+            epoch + 1,
+            train_loss,
+            val_loss,
+            val_acc,
+            optimizer.param_groups[0]["lr"],
+            early_stopper.best_metric,
+            early_stopper.best_epoch,
+            int(should_stop),
+        )
+
+        save_checkpoint(epoch + 1, train_loss, val_loss, val_acc, is_best=is_best)
+
+        stop_tensor = torch.tensor([int(should_stop)], device=device)
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            dist.broadcast(stop_tensor, src=0)
+
+        if bool(stop_tensor.item()):
+            early_stopped = True
+            if rank == 0:
+                logger.info(
+                    f"Early stopping at epoch {epoch + 1}. Best val_loss={early_stopper.best_metric:.6f} @ epoch {early_stopper.best_epoch}"
+                )
+            break
+
+    train_time_seconds = time.perf_counter() - train_start_time
+    peak_host_ram_gb = _peak_host_ram_gb()
+    peak_gpu_alloc_gb, peak_gpu_reserved_gb = _peak_gpu_mem_gb(device)
+    if rank == 0:
+        logger.info(
+            f"Training finished: epochs_run={epochs_run} "
+            f"early_stopped={early_stopped} "
+            f"train_time={_format_hms(train_time_seconds)} "
+            f"({train_time_seconds:.1f}s) "
+            f"peak_host_ram_gb={peak_host_ram_gb} "
+            f"peak_gpu_alloc_gb={peak_gpu_alloc_gb} "
+            f"peak_gpu_reserved_gb={peak_gpu_reserved_gb}"
+        )
+
+    if early_stopper.enabled and early_stopper.restore_best_weights:
+        if rank == 0:
+            logger.info("Restoring best model weights before exit")
+        early_stopper.restore(model.module, device)
+        if rank == 0:
+            torch.save(
+                {
+                    "model": model.module.state_dict(),
+                    "epoch": early_stopper.best_epoch,
+                    "val_loss": early_stopper.best_metric,
+                    "args": args,
+                },
+                best_path,
+            )
+
+    if rank == 0:
+        # Evaluate on both the in-distribution (id_test) and out-of-distribution
+        # (test) splits so the generalization gap can be measured.
+        # WILDS-iWildCam: "id_test" == Test ID, "test" == Test OOD.
+        eval_splits = [
+            ("id_test", "iwildcam_id_test"),
+            ("test", "iwildcam_test"),
+        ]
+
+        def _make_eval_args(split, write_tag):
+            return {
+                "meta": {
+                    "seed": seed,
+                    "model_name": m_args["model_name"],
+                    "embed_dim": m_args["embed_dim"],
+                    "num_classes": m_args["num_classes"],
+                    "patch_size": mk_args.get("patch_size", m_args.get("patch_size", 16)),
+                    "crop_size": d_args.get("crop_size", m_args.get("crop_size", 224)),
+                    "use_bfloat16": m_args.get("use_bfloat16", True),
+                    "representation_type": representation_type,
+                    "head_type": head_type,
+                    "checkpoint_path": best_path,
+                    "force_single_process": True,
+                },
+                "data": {
+                    "batch_size": d_args.get("batch_size", 128),
+                    "root_path": d_args.get("root_path", "./wilds_data"),
+                    "num_workers": d_args.get("num_workers", 8),
+                    "pin_mem": d_args.get("pin_mem", True),
+                    "split": split,
+                    "download": True,
+                },
+                "logging": {
+                    "write_tag": write_tag,
+                    "auto_folder": True,
+                },
+            }
+
+        eval_results = {}
+        eval_folder = None
+        for split, write_tag in eval_splits:
+            eval_args = _make_eval_args(split, write_tag)
+            result = eval_wilds_main(args=eval_args)
+            eval_results[split] = result
+            if eval_folder is None:
+                eval_folder = eval_args.get("logging", {}).get("folder")
+
+        # Common run-level info folded into every metrics JSON + the params summary.
+        run_info = {
+            "seed": seed,
+            "label_fraction": label_fraction if label_fraction is not None else 1.0,
+            "train_time_seconds": float(train_time_seconds),
+            "train_time_hms": _format_hms(train_time_seconds),
+            "epochs_run": int(epochs_run),
+            "configured_epochs": int(o_args["epochs"]),
+            "best_epoch": int(early_stopper.best_epoch),
+            "early_stopped": bool(early_stopped),
+            "best_val_loss": float(early_stopper.best_metric),
+            "peak_host_ram_gb": peak_host_ram_gb,
+            "peak_gpu_alloc_gb": peak_gpu_alloc_gb,
+            "peak_gpu_reserved_gb": peak_gpu_reserved_gb,
+        }
+
+        # Fold run_info into each split's metrics JSON so an aggregator can read
+        # metrics + seed + timing + epochs from a single file per split.
+        for split, result in eval_results.items():
+            if not result:
+                continue
+            metrics_path = result.get("metrics_path")
+            if not metrics_path or not os.path.exists(metrics_path):
+                continue
+            try:
+                with open(metrics_path, "r") as f:
+                    metrics_obj = json.load(f)
+                metrics_obj["run_info"] = run_info
+                metrics_obj["split"] = split
+                with open(metrics_path, "w") as f:
+                    json.dump(metrics_obj, f, indent=2, sort_keys=True)
+            except (OSError, json.JSONDecodeError):
+                logger.warning(f"Could not augment metrics JSON for split {split}")
+
+        if eval_folder:
+            try:
+                params_out = yaml.safe_load(yaml.dump(args))
+                params_out.setdefault("meta", {})["representation_type"] = representation_type
+                params_out.setdefault("meta", {})["head_type"] = head_type
+                params_out["results"] = {
+                    "best_val_loss": float(early_stopper.best_metric),
+                    "best_epoch": int(early_stopper.best_epoch),
+                    "best_checkpoint": best_path,
+                    "seed": seed,
+                    "train_time_seconds": float(train_time_seconds),
+                    "train_time_hms": _format_hms(train_time_seconds),
+                    "epochs_run": int(epochs_run),
+                    "configured_epochs": int(o_args["epochs"]),
+                    "early_stopped": bool(early_stopped),
+                    "peak_host_ram_gb": peak_host_ram_gb,
+                    "peak_gpu_alloc_gb": peak_gpu_alloc_gb,
+                    "peak_gpu_reserved_gb": peak_gpu_reserved_gb,
+                    "eval_metrics_id_test": (
+                        eval_results.get("id_test", {}).get("metrics")
+                        if eval_results.get("id_test")
+                        else None
+                    ),
+                    "eval_metrics_test": (
+                        eval_results.get("test", {}).get("metrics")
+                        if eval_results.get("test")
+                        else None
+                    ),
+                }
+                with open(os.path.join(eval_folder, "params-supervised.yaml"), "w") as f:
+                    yaml.dump(params_out, f)
+                with open(os.path.join(eval_folder, "params.yaml"), "w") as f:
+                    yaml.dump(params_out, f)
+            except OSError:
+                logger.warning("Could not write supervised params to eval folder")
+
+        if os.path.exists(folder):
+            try:
+                shutil.rmtree(folder)
+            except OSError:
+                logger.warning("Could not remove supervised run folder")
+
+
+if __name__ == "__main__":
+    raise RuntimeError(
+        "Use main_distributed_supervised.py to launch this script with a config file."
+    )

src/transforms.py

@@ -23,23 +23,28 @@ def make_transforms(
     horizontal_flip=False,
     color_distortion=False,
     gaussian_blur=False,
-    normalization=((0.485, 0.456, 0.406),
-                   (0.229, 0.224, 0.225))
+    use_random_resized_crop=True,
+    normalization=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
 ):
-    logger.info('making imagenet data transforms')
+    logger.info("making imagenet data transforms")
 
     def get_color_distortion(s=1.0):
         # s is the strength of color distortion.
-        color_jitter = transforms.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)
+        color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
         rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)
         rnd_gray = transforms.RandomGrayscale(p=0.2)
-        color_distort = transforms.Compose([
-            rnd_color_jitter,
-            rnd_gray])
+        color_distort = transforms.Compose([rnd_color_jitter, rnd_gray])
         return color_distort
 
     transform_list = []
-    transform_list += [transforms.RandomResizedCrop(crop_size, scale=crop_scale)]
+    if use_random_resized_crop:
+        transform_list += [transforms.RandomResizedCrop(crop_size, scale=crop_scale)]
+    else:
+        resize_size = int(crop_size * 256 / 224)
+        transform_list += [
+            transforms.Resize(resize_size),
+            transforms.CenterCrop(crop_size),
+        ]
     if horizontal_flip:
         transform_list += [transforms.RandomHorizontalFlip()]
     if color_distortion:
@@ -53,8 +58,25 @@ def make_transforms(
     return transform
 
 
+def make_transform_eval(
+    crop_size=224, normalization=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+):
+    logger.info("making imagenet evaluation transforms")
+
+    resize_size = int(crop_size * 256 / 224)
+    transform = transforms.Compose(
+        [
+            transforms.Resize(resize_size),
+            transforms.CenterCrop(crop_size),
+            transforms.ToTensor(),
+            transforms.Normalize(normalization[0], normalization[1]),
+        ]
+    )
+    return transform
+
+
 class GaussianBlur(object):
-    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
+    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.0):
         self.prob = p
         self.radius_min = radius_min
         self.radius_max = radius_max

src/utils/distributed.py

@@ -6,6 +6,7 @@
 #
 
 import os
+import socket
 
 import torch
 import torch.distributed as dist
@@ -15,33 +16,40 @@ from logging import getLogger
 logger = getLogger()
 
 
-def init_distributed(port=40112, rank_and_world_size=(None, None)):
+def _find_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(('', 0))
+        return s.getsockname()[1]
+
+
+def init_distributed(port=None, rank_and_world_size=(None, None)):
 
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size(), dist.get_rank()
 
     rank, world_size = rank_and_world_size
-    os.environ['MASTER_ADDR'] = 'localhost'
 
     if (rank is None) or (world_size is None):
         try:
             world_size = int(os.environ['SLURM_NTASKS'])
             rank = int(os.environ['SLURM_PROCID'])
-            os.environ['MASTER_ADDR'] = os.environ['HOSTNAME']
         except Exception:
             logger.info('SLURM vars not set (distributed training not available)')
-            world_size, rank = 1, 0
-            return world_size, rank
+            return 1, 0
+
+    os.environ['MASTER_ADDR'] = '127.0.0.1'
+    os.environ['MASTER_PORT'] = str(port if port is not None else _find_free_port())
+    os.environ['NCCL_SOCKET_IFNAME'] = 'lo'
+    os.environ['NCCL_NET_GIB_EXTRA_IFS'] = 'lo'
 
     try:
-        os.environ['MASTER_PORT'] = str(port)
         torch.distributed.init_process_group(
             backend='nccl',
             world_size=world_size,
             rank=rank)
     except Exception as e:
-        world_size, rank = 1, 0
-        logger.info(f'distributed training not available {e}')
+        logger.warning(f'NCCL init failed ({e}); falling back to single-process')
+        return 1, 0
 
     return world_size, rank
 

src/utils/logging.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 #
 
+import os
 import torch
 
 
@@ -28,6 +29,102 @@ def gpu_timer(closure, log_timings=True):
     return result, elapsed_time
 
 
+def _format_value(value):
+    if value is None:
+        return None
+    if isinstance(value, float):
+        return f"{value:g}"
+    if isinstance(value, (list, tuple)):
+        return "-".join(_format_value(v) for v in value)
+    if isinstance(value, bool):
+        return "1" if value else "0"
+    return str(value)
+
+
+def build_run_name(args):
+    meta_args = args.get("meta", {})
+    data_args = args.get("data", {})
+    opt_args = args.get("optimization", {})
+    mask_args = args.get("mask", {})
+    val_args = args.get("validation", {})
+
+    model_name = meta_args.get("model_name", "model")
+    tag = meta_args.get("tag")
+    representation_type = meta_args.get("representation_type")
+    head_type = meta_args.get("head_type")
+    patch_size = mask_args.get("patch_size", meta_args.get("patch_size"))
+    crop_size = data_args.get("crop_size", meta_args.get("crop_size"))
+    batch_size = data_args.get("batch_size")
+    optimizer = opt_args.get("optimizer", "opt")
+    lr = opt_args.get("lr")
+    weight_decay = opt_args.get("weight_decay")
+    epochs = opt_args.get("epochs")
+
+    parts = [model_name]
+
+    def add(prefix, value):
+        formatted = _format_value(value)
+        if formatted is not None:
+            parts.append(f"{prefix}{formatted}")
+
+    add("p", patch_size)
+    add("c", crop_size)
+    if tag:
+        parts.append(str(tag))
+    add("bs", batch_size)
+    add("rep", representation_type)
+    add("head", head_type)
+    parts.append(str(optimizer).lower())
+    add("lr", lr)
+    add("wd", weight_decay)
+    add("ep", epochs)
+    add("sched", opt_args.get("lr_schedule"))
+    add("ms", opt_args.get("step_milestones"))
+    add("sg", opt_args.get("step_gamma"))
+    add("wu", opt_args.get("warmup"))
+    add("mom", opt_args.get("momentum"))
+    add("leta", opt_args.get("lars_eta"))
+    add("leps", opt_args.get("lars_eps"))
+    add("ipe", opt_args.get("ipe_scale"))
+    add("cs", data_args.get("crop_scale"))
+    add("lf", data_args.get("label_fraction"))
+    add("eval", val_args.get("eval_every"))
+    add("seed", meta_args.get("seed"))
+
+    return "-".join([p for p in parts if p])
+
+
+def _extract_run_name_from_checkpoint(meta_args):
+    checkpoint_path = meta_args.get("checkpoint_path")
+    if checkpoint_path:
+        folder = os.path.dirname(checkpoint_path)
+    else:
+        folder = meta_args.get("checkpoint_folder")
+    if not folder:
+        return None
+    return os.path.basename(os.path.normpath(folder))
+
+
+def resolve_log_dir(args, stage="train"):
+    log_args = args.setdefault("logging", {})
+    auto_folder = log_args.get("auto_folder", True)
+    if auto_folder or not log_args.get("folder"):
+        run_name = log_args.get("run_name")
+        if stage == "eval" and not run_name:
+            meta_args = args.get("meta", {})
+            run_name = _extract_run_name_from_checkpoint(meta_args)
+        if not run_name:
+            run_name = build_run_name(args)
+        base_dir = "experiment_logs"
+        if stage == "eval":
+            folder = os.path.join(base_dir, "eval-wilds", run_name)
+        else:
+            folder = os.path.join(base_dir, run_name)
+        log_args["folder"] = folder
+    os.makedirs(log_args["folder"], exist_ok=True)
+    return log_args["folder"]
+
+
 class CSVLogger(object):
 
     def __init__(self, fname, *argv):

src/utils/optimizers.py

@@ -0,0 +1,84 @@
+import math
+
+import torch
+
+
+class LARS(torch.optim.Optimizer):
+    def __init__(
+        self,
+        params,
+        lr,
+        weight_decay=0.0,
+        momentum=0.9,
+        eta=0.001,
+        eps=1e-8,
+        exclude_bias_and_norm=True,
+    ):
+        if lr <= 0.0:
+            raise ValueError(f"Invalid lr: {lr}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay: {weight_decay}")
+        if momentum < 0.0:
+            raise ValueError(f"Invalid momentum: {momentum}")
+        if eta <= 0.0:
+            raise ValueError(f"Invalid eta: {eta}")
+        if eps <= 0.0:
+            raise ValueError(f"Invalid eps: {eps}")
+
+        defaults = dict(
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            eta=eta,
+            eps=eps,
+            exclude_bias_and_norm=exclude_bias_and_norm,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            eta = group["eta"]
+            eps = group["eps"]
+            exclude_bias_and_norm = group["exclude_bias_and_norm"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("LARS does not support sparse gradients")
+
+                param_norm = torch.norm(p)
+                grad_norm = torch.norm(grad)
+
+                lars_lr = 1.0
+                if not exclude_bias_and_norm or p.ndim > 1:
+                    if param_norm > 0.0 and grad_norm > 0.0:
+                        lars_lr = eta * param_norm / (grad_norm + weight_decay * param_norm + eps)
+
+                d_p = grad
+                if weight_decay != 0.0 and (not exclude_bias_and_norm or p.ndim > 1):
+                    d_p = d_p.add(p, alpha=weight_decay)
+
+                if momentum != 0.0:
+                    param_state = self.state.setdefault(p, {})
+                    if "momentum_buffer" not in param_state:
+                        buf = param_state["momentum_buffer"] = torch.clone(d_p).detach()
+                    else:
+                        buf = param_state["momentum_buffer"]
+                        buf.mul_(momentum).add_(d_p)
+                    d_p = buf
+
+                p.add_(d_p, alpha=-lr * lars_lr)
+
+        return loss

src/utils/schedulers.py

@@ -74,3 +74,30 @@ class CosineWDSchedule(object):
             if ('WD_exclude' not in group) or not group['WD_exclude']:
                 group['weight_decay'] = new_wd
         return new_wd
+
+
+class LinearWDSchedule(object):
+
+    def __init__(
+        self,
+        optimizer,
+        ref_wd,
+        T_max,
+        final_wd=0.
+    ):
+        self.optimizer = optimizer
+        self.ref_wd = ref_wd
+        self.final_wd = final_wd
+        self.T_max = T_max
+        self._step = 0.
+
+    def step(self):
+        self._step += 1
+        progress = self._step / self.T_max
+        progress = min(max(progress, 0.0), 1.0)
+        new_wd = self.ref_wd + progress * (self.final_wd - self.ref_wd)
+
+        for group in self.optimizer.param_groups:
+            if ('WD_exclude' not in group) or not group['WD_exclude']:
+                group['weight_decay'] = new_wd
+        return new_wd

tools/aggregate_label_efficiency.py

@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""
+Aggregate label-efficiency supervised runs into a paper-style Table 4 CSV.
+
+For each model, this collects the OOD F1-Macro across seeds and label
+fractions (1%, 10%, 50%, 100%) and outputs mean +/- std in a wide CSV table.
+
+It expects the run_info written by train_supervised.py to contain
+`seed` and `data.label_fraction` (or the eval folder name to contain
+`-lf<FRACTION>-seed<SEED>`). The 100% runs are the full-data runs already
+configured for the paper.
+
+Outputs:
+  experiment_logs/label-efficiency/summary.csv
+  experiment_logs/label-efficiency/<model>/summary.json
+
+Usage:
+  python3 tools/aggregate_label_efficiency.py --root experiment_logs/eval-wilds
+  python3 tools/aggregate_label_efficiency.py --root experiment_logs/eval-wilds --metric F1-macro_all
+"""
+import argparse
+import csv
+import json
+import math
+import os
+import re
+from collections import defaultdict
+
+import yaml
+
+OOD_METRICS_FILE = "iwildcam_test_metrics.json"
+
+# Match trailing -lf0.01-seed0 or -seed0.
+_LF_SEED_SUFFIX_RE = re.compile(r"-lf(?P<frac>0\.\d+|[1-9]\d*\.?\d*)-seed(?P<seed>\d+)$")
+_SEED_SUFFIX_RE = re.compile(r"-seed(?P<seed>\d+)$")
+
+
+FRACTIONS = [0.01, 0.10, 0.50, 1.00]
+
+
+def _find_metric(metrics, key):
+    if isinstance(metrics, dict):
+        if key in metrics:
+            return metrics[key]
+        for value in metrics.values():
+            found = _find_metric(value, key)
+            if found is not None:
+                return found
+    elif isinstance(metrics, list):
+        for item in metrics:
+            found = _find_metric(item, key)
+            if found is not None:
+                return found
+    return None
+
+
+def _extract_run_info(metrics_obj):
+    def _search(obj):
+        if isinstance(obj, dict):
+            ri = obj.get("run_info")
+            if isinstance(ri, dict):
+                return ri
+            for v in obj.values():
+                found = _search(v)
+                if found is not None:
+                    return found
+        elif isinstance(obj, list):
+            for item in obj:
+                found = _search(item)
+                if found is not None:
+                    return found
+        return None
+
+    if not metrics_obj:
+        return {}
+    return _search(metrics_obj) or {}
+
+
+def _load_json(path):
+    try:
+        with open(path, "r") as f:
+            return json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def _load_yaml(path):
+    try:
+        with open(path, "r") as f:
+            return yaml.load(f, Loader=yaml.FullLoader)
+    except (OSError, yaml.YAMLError):
+        return None
+
+
+def _parse_fraction_and_seed(run_name, run_info, run_dir):
+    """Return (fraction, seed) from run_info, folder name, or params.yaml."""
+    fraction = run_info.get("label_fraction")
+    seed = run_info.get("seed")
+
+    # Try folder name first for both values.
+    m = _LF_SEED_SUFFIX_RE.search(run_name)
+    if m:
+        if fraction is None:
+            try:
+                fraction = float(m.group("frac"))
+            except ValueError:
+                pass
+        if seed is None:
+            try:
+                seed = int(m.group("seed"))
+            except ValueError:
+                pass
+    else:
+        m = _SEED_SUFFIX_RE.search(run_name)
+        if m and seed is None:
+            try:
+                seed = int(m.group("seed"))
+            except ValueError:
+                pass
+
+    # Fallback to params.yaml.
+    if fraction is None or seed is None:
+        params = _load_yaml(os.path.join(run_dir, "params.yaml"))
+        if params:
+            if fraction is None:
+                fraction = _find_metric(params, "label_fraction")
+            if seed is None:
+                seed = _find_metric(params, "seed")
+
+    return fraction, seed
+
+
+def _model_key(run_name):
+    """Strip the label-fraction and seed suffix to obtain a model group key."""
+    key = _LF_SEED_SUFFIX_RE.sub("", run_name)
+    key = _SEED_SUFFIX_RE.sub("", key)
+    return key
+
+
+def _mean_std(values):
+    vals = [v for v in values if v is not None and not _is_nan(v)]
+    if not vals:
+        return None, None, 0
+    n = len(vals)
+    mean = sum(vals) / n
+    if n > 1:
+        std = math.sqrt(sum((v - mean) ** 2 for v in vals) / (n - 1))
+    else:
+        std = 0.0
+    return mean, std, n
+
+
+def _is_nan(v):
+    try:
+        return math.isnan(float(v))
+    except (TypeError, ValueError):
+        return False
+
+
+def _fmt(mean, std):
+    if mean is None:
+        return ""
+    if std is None or std == 0.0:
+        return f"{mean:.4f}"
+    return f"{mean:.4f} +/- {std:.4f}"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description=__doc__,
+    )
+    parser.add_argument(
+        "--root",
+        default="experiment_logs/eval-wilds",
+        help="root folder holding per-run eval subfolders (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--out",
+        default="experiment_logs/label-efficiency",
+        help="output folder for summaries (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--metric",
+        default="F1-macro_all",
+        help="OOD metric to aggregate (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--min-seeds",
+        type=int,
+        default=1,
+        help="only report fractions with at least this many seeds (default: %(default)s)",
+    )
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.root):
+        print(f"Root folder not found: {args.root}")
+        return
+
+    # model_key -> {fraction: [records]}
+    groups = defaultdict(lambda: defaultdict(list))
+    skipped = []
+
+    for entry in sorted(os.listdir(args.root)):
+        run_dir = os.path.join(args.root, entry)
+        if not os.path.isdir(run_dir):
+            continue
+
+        ood_metrics = _load_json(os.path.join(run_dir, OOD_METRICS_FILE))
+        if ood_metrics is None:
+            continue
+
+        run_info = _extract_run_info(ood_metrics) or {}
+        fraction, seed = _parse_fraction_and_seed(entry, run_info, run_dir)
+        value = _find_metric(ood_metrics, args.metric)
+
+        if fraction is None:
+            skipped.append((entry, "no label_fraction"))
+            continue
+        if value is None:
+            skipped.append((entry, f"metric {args.metric} missing"))
+            continue
+
+        key = _model_key(entry)
+        groups[key][fraction].append(
+            {
+                "run_name": entry,
+                "seed": seed,
+                "value": float(value),
+                "run_info": run_info,
+            }
+        )
+
+    if not groups:
+        print(f"No usable metrics found under {args.root}")
+        return
+
+    os.makedirs(args.out, exist_ok=True)
+
+    # CSV fieldnames.
+    fieldnames = ["model"]
+    for frac in FRACTIONS:
+        frac_label = f"frac{frac:.2f}"
+        fieldnames.extend(
+            [
+                f"{frac_label}_mean",
+                f"{frac_label}_std",
+                f"{frac_label}_n",
+            ]
+        )
+    fieldnames.append("seeds")
+
+    csv_rows = []
+
+    for model_key in sorted(groups.keys()):
+        fractions = groups[model_key]
+        row = {"model": model_key}
+        per_frac = {}
+        all_seeds = set()
+
+        for frac in FRACTIONS:
+            records = fractions.get(frac, [])
+            values = [r["value"] for r in records]
+            mean, std, n = _mean_std(values)
+            frac_label = f"frac{frac:.2f}"
+            row[f"{frac_label}_mean"] = mean
+            row[f"{frac_label}_std"] = std
+            row[f"{frac_label}_n"] = n
+            per_frac[frac] = {
+                "mean": mean,
+                "std": std,
+                "n": n,
+                "seeds": [r["seed"] for r in records],
+                "values": values,
+                "run_names": [r["run_name"] for r in records],
+            }
+            for r in records:
+                if r["seed"] is not None:
+                    all_seeds.add(r["seed"])
+
+        row["seeds"] = " ".join(str(s) for s in sorted(all_seeds))
+        csv_rows.append(row)
+
+        # Write per-model JSON summary.
+        model_dir = os.path.join(args.out, model_key)
+        os.makedirs(model_dir, exist_ok=True)
+        with open(os.path.join(model_dir, "summary.json"), "w") as f:
+            json.dump(
+                {
+                    "model": model_key,
+                    "metric": args.metric,
+                    "fractions": {f"{k:.2f}": v for k, v in per_frac.items()},
+                },
+                f,
+                indent=2,
+                sort_keys=True,
+            )
+
+    csv_path = os.path.join(args.out, "summary.csv")
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in csv_rows:
+            writer.writerow(row)
+
+    # Terminal table.
+    print(f"\nAggregated {len(csv_rows)} model(s). Metric: {args.metric}\n")
+    header = f"{'model':<50}"
+    for frac in FRACTIONS:
+        header += f"  {f'{int(frac*100)}%':>18}"
+    print(header)
+    print("-" * len(header))
+    for row in csv_rows:
+        line = f"{row['model'][:50]:<50}"
+        for frac in FRACTIONS:
+            frac_label = f"frac{frac:.2f}"
+            mean = row[f"{frac_label}_mean"]
+            std = row[f"{frac_label}_std"]
+            line += f"  {_fmt(mean, std):>18}"
+        print(line)
+
+    if skipped:
+        print(f"\nSkipped {len(skipped)} run(s):")
+        for run_name, reason in skipped:
+            print(f"  {run_name}: {reason}")
+
+    print(f"\nWrote per-model summaries to: {args.out}/<model>/summary.json")
+    print(f"Wrote combined CSV to:        {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

tools/aggregate_seeds.py

@@ -0,0 +1,470 @@
+#!/usr/bin/env python3
+"""
+Aggregate multi-seed WILDS-iWildCam supervised runs into mean +/- std.
+
+For each model configuration (grouped across seeds), this reads the per-seed
+evaluation metrics for both splits:
+
+  - ID  (in-distribution):  iwildcam_id_test_metrics.json   (split "id_test")
+  - OOD (out-of-distribution): iwildcam_test_metrics.json    (split "test")
+
+and computes, across seeds, the mean and std of every WILDS metric plus the
+training time and epochs recorded during training. It also reports the
+generalization gap (ID - OOD) on the headline metric.
+
+Leaderboard columns reported:
+  Test ID Macro F1 | Test ID Avg Acc | Test OOD Macro F1 | Test OOD Avg Acc
+
+Outputs:
+  experiment_logs/seed-runs/<model>/summary.json   (per-seed rows + mean/std)
+  experiment_logs/seed-runs/summary_all.csv        (one row per model)
+
+Usage:
+  python3 tools/aggregate_seeds.py --root experiment_logs/eval-wilds
+  python3 tools/aggregate_seeds.py --primary F1-macro_all --acc acc_avg
+  python3 tools/aggregate_seeds.py --root experiment_logs/eval-wilds --min-seeds 5
+"""
+import argparse
+import csv
+import json
+import math
+import os
+import re
+
+
+ID_METRICS_FILE = "iwildcam_id_test_metrics.json"
+OOD_METRICS_FILE = "iwildcam_test_metrics.json"
+
+# Trailing "-seedN" (and any leftover separators) so runs of the same config
+# collapse into one group.
+_SEED_SUFFIX_RE = re.compile(r"-seed\d+$")
+
+
+def _find_metric(metrics, key):
+    """Recursively search a nested dict/list for `key`."""
+    if isinstance(metrics, dict):
+        if key in metrics:
+            return metrics[key]
+        for value in metrics.values():
+            found = _find_metric(value, key)
+            if found is not None:
+                return found
+    elif isinstance(metrics, list):
+        for item in metrics:
+            found = _find_metric(item, key)
+            if found is not None:
+                return found
+    return None
+
+
+def _strip_seed(run_name):
+    return _SEED_SUFFIX_RE.sub("", run_name)
+
+
+def _extract_run_info(*metrics_objs):
+    """Find a run_info dict inside any metrics object (dict- or list-form).
+
+    Eval-only runs write metrics as a list ([{metrics}, "summary"]) with no
+    run_info; training runs write a dict with a top-level run_info. Search both.
+    """
+
+    def _search(obj):
+        if isinstance(obj, dict):
+            ri = obj.get("run_info")
+            if isinstance(ri, dict):
+                return ri
+            for v in obj.values():
+                found = _search(v)
+                if found is not None:
+                    return found
+        elif isinstance(obj, list):
+            for item in obj:
+                found = _search(item)
+                if found is not None:
+                    return found
+        return None
+
+    for m in metrics_objs:
+        found = _search(m)
+        if found is not None:
+            return found
+    return {}
+
+
+def _seed_from_name(run_name):
+    """Parse a trailing -seedN from a run folder name; None if absent."""
+    m = re.search(r"-seed(\d+)$", run_name)
+    return int(m.group(1)) if m else None
+
+
+def _seed_from_params(run_dir):
+    """Read `seed:` from params-eval.yaml without requiring PyYAML."""
+    path = os.path.join(run_dir, "params-eval.yaml")
+    try:
+        with open(path, "r") as f:
+            for line in f:
+                m = re.match(r"\s*seed\s*:\s*(\d+)\s*$", line)
+                if m:
+                    return int(m.group(1))
+    except OSError:
+        pass
+    return None
+
+
+def _resolve_seed(run_info, run_name, run_dir):
+    """Seed detection chain: run_info -> folder name -> params-eval.yaml."""
+    seed = run_info.get("seed")
+    if seed is not None:
+        return seed
+    seed = _seed_from_name(run_name)
+    if seed is not None:
+        return seed
+    return _seed_from_params(run_dir)
+
+
+def _load_json(path):
+    try:
+        with open(path, "r") as f:
+            return json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def _mean_std(values):
+    vals = [v for v in values if v is not None and not _is_nan(v)]
+    if not vals:
+        return None, None, 0
+    n = len(vals)
+    mean = sum(vals) / n
+    if n > 1:
+        var = sum((v - mean) ** 2 for v in vals) / (n - 1)  # sample std
+        std = math.sqrt(var)
+    else:
+        std = 0.0
+    return mean, std, n
+
+
+def _is_nan(v):
+    try:
+        return math.isnan(float(v))
+    except (TypeError, ValueError):
+        return False
+
+
+def _collect_metric_keys(metrics_obj):
+    """All scalar metric keys in a WILDS metrics dict (excludes our extras)."""
+    keys = set()
+    if isinstance(metrics_obj, dict):
+        for k, v in metrics_obj.items():
+            if k in ("run_info", "split"):
+                continue
+            if isinstance(v, (int, float)) and not isinstance(v, bool):
+                keys.add(k)
+    return keys
+
+
+def _fmt(mean, std):
+    if mean is None:
+        return ""
+    if std is None:
+        return f"{mean:.4f}"
+    return f"{mean:.4f} +/- {std:.4f}"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description=__doc__,
+    )
+    parser.add_argument(
+        "--root",
+        default="experiment_logs/eval-wilds",
+        help="root folder holding per-run eval subfolders (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--out",
+        default="experiment_logs/seed-runs",
+        help="output folder for aggregated summaries (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--primary",
+        default="F1-macro_all",
+        help="headline metric key (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--acc",
+        default="acc_avg",
+        help="average-accuracy metric key (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--min-seeds",
+        type=int,
+        default=1,
+        help="only report models with at least this many seeds (default: %(default)s)",
+    )
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.root):
+        print(f"Root folder not found: {args.root}")
+        return
+
+    # group_name -> list of per-seed records
+    groups = {}
+    for entry in sorted(os.listdir(args.root)):
+        run_dir = os.path.join(args.root, entry)
+        if not os.path.isdir(run_dir):
+            continue
+        id_metrics = _load_json(os.path.join(run_dir, ID_METRICS_FILE))
+        ood_metrics = _load_json(os.path.join(run_dir, OOD_METRICS_FILE))
+        if id_metrics is None and ood_metrics is None:
+            continue
+
+        run_info = _extract_run_info(ood_metrics, id_metrics)
+
+        group = _strip_seed(entry)
+        groups.setdefault(group, []).append(
+            {
+                "run_name": entry,
+                "seed": _resolve_seed(run_info, entry, run_dir),
+                "id_metrics": id_metrics,
+                "ood_metrics": ood_metrics,
+                "run_info": run_info,
+            }
+        )
+
+    if not groups:
+        print(f"No metrics found under {args.root}")
+        return
+
+    if args.min_seeds > 1:
+        dropped = {
+            name: len(recs)
+            for name, recs in groups.items()
+            if len(recs) < args.min_seeds
+        }
+        groups = {
+            name: recs for name, recs in groups.items() if len(recs) >= args.min_seeds
+        }
+        if dropped:
+            print(
+                f"Skipping {len(dropped)} model(s) with fewer than "
+                f"{args.min_seeds} seed(s):"
+            )
+            for name in sorted(dropped):
+                print(f"  {name} ({dropped[name]} seed(s))")
+        if not groups:
+            print(f"No models have >= {args.min_seeds} seeds under {args.root}")
+            return
+
+    os.makedirs(args.out, exist_ok=True)
+
+    csv_rows = []
+    csv_fields = [
+        "model",
+        "num_seeds",
+        "seeds",
+        "id_macro_f1_mean",
+        "id_macro_f1_std",
+        "id_avg_acc_mean",
+        "id_avg_acc_std",
+        "ood_macro_f1_mean",
+        "ood_macro_f1_std",
+        "ood_avg_acc_mean",
+        "ood_avg_acc_std",
+        "gen_gap_macro_f1_mean",  # ID - OOD on primary metric
+        "train_time_seconds_mean",
+        "train_time_seconds_std",
+        "epochs_run_mean",
+        "epochs_run_std",
+        "peak_host_ram_gb_mean",
+        "peak_host_ram_gb_std",
+        "peak_gpu_alloc_gb_mean",
+        "peak_gpu_alloc_gb_std",
+        "peak_gpu_reserved_gb_mean",
+        "peak_gpu_reserved_gb_std",
+    ]
+
+    for group_name in sorted(groups.keys()):
+        records = sorted(groups[group_name], key=lambda r: (r["seed"] is None, r["seed"]))
+        seeds = [r["seed"] for r in records]
+
+        # Determine the union of metric keys present on either split.
+        all_metric_keys = set()
+        for r in records:
+            all_metric_keys |= _collect_metric_keys(r["id_metrics"])
+            all_metric_keys |= _collect_metric_keys(r["ood_metrics"])
+
+        def metric_values(split_key, mkey):
+            return [_find_metric(r[split_key], mkey) for r in records]
+
+        # Per-seed rows for the JSON summary.
+        per_seed = []
+        for r in records:
+            per_seed.append(
+                {
+                    "seed": r["seed"],
+                    "run_name": r["run_name"],
+                    "id": {
+                        "macro_f1": _find_metric(r["id_metrics"], args.primary),
+                        "avg_acc": _find_metric(r["id_metrics"], args.acc),
+                    },
+                    "ood": {
+                        "macro_f1": _find_metric(r["ood_metrics"], args.primary),
+                        "avg_acc": _find_metric(r["ood_metrics"], args.acc),
+                    },
+                    "train_time_seconds": r["run_info"].get("train_time_seconds"),
+                    "train_time_hms": r["run_info"].get("train_time_hms"),
+                    "epochs_run": r["run_info"].get("epochs_run"),
+                    "configured_epochs": r["run_info"].get("configured_epochs"),
+                    "best_epoch": r["run_info"].get("best_epoch"),
+                    "early_stopped": r["run_info"].get("early_stopped"),
+                    "peak_host_ram_gb": r["run_info"].get("peak_host_ram_gb"),
+                    "peak_gpu_alloc_gb": r["run_info"].get("peak_gpu_alloc_gb"),
+                    "peak_gpu_reserved_gb": r["run_info"].get("peak_gpu_reserved_gb"),
+                }
+            )
+
+        # Aggregate every metric for both splits.
+        def agg_all(split_key):
+            out = {}
+            for mkey in sorted(all_metric_keys):
+                mean, std, n = _mean_std(metric_values(split_key, mkey))
+                if n > 0:
+                    out[mkey] = {"mean": mean, "std": std, "n": n}
+            return out
+
+        id_agg = agg_all("id_metrics")
+        ood_agg = agg_all("ood_metrics")
+
+        time_vals = [r["run_info"].get("train_time_seconds") for r in records]
+        epoch_vals = [r["run_info"].get("epochs_run") for r in records]
+        host_ram_vals = [r["run_info"].get("peak_host_ram_gb") for r in records]
+        gpu_alloc_vals = [r["run_info"].get("peak_gpu_alloc_gb") for r in records]
+        gpu_reserved_vals = [r["run_info"].get("peak_gpu_reserved_gb") for r in records]
+        time_mean, time_std, _ = _mean_std(time_vals)
+        epoch_mean, epoch_std, _ = _mean_std(epoch_vals)
+        host_ram_mean, host_ram_std, host_ram_n = _mean_std(host_ram_vals)
+        gpu_alloc_mean, gpu_alloc_std, _ = _mean_std(gpu_alloc_vals)
+        gpu_reserved_mean, gpu_reserved_std, _ = _mean_std(gpu_reserved_vals)
+
+        if host_ram_n == 0:
+            # No usable peak_host_ram_gb in any seed's run_info. This usually means
+            # the metrics JSONs were (re)generated by a standalone eval run that did
+            # not fold in train_supervised.py's run_info, or run_info is absent.
+            missing_run_info = sum(1 for r in records if not r["run_info"])
+            print(
+                f"[warn] {group_name}: RAM was not recorded for these eval runs "
+                f"(no peak_host_ram_gb across {len(records)} seed(s); "
+                f"{missing_run_info} missing run_info entirely). RAM column stays blank."
+            )
+
+        # Headline (leaderboard) numbers.
+        id_f1_mean, id_f1_std, _ = _mean_std(metric_values("id_metrics", args.primary))
+        id_acc_mean, id_acc_std, _ = _mean_std(metric_values("id_metrics", args.acc))
+        ood_f1_mean, ood_f1_std, _ = _mean_std(metric_values("ood_metrics", args.primary))
+        ood_acc_mean, ood_acc_std, _ = _mean_std(metric_values("ood_metrics", args.acc))
+
+        gen_gap = None
+        if id_f1_mean is not None and ood_f1_mean is not None:
+            gen_gap = id_f1_mean - ood_f1_mean
+
+        summary = {
+            "model": group_name,
+            "primary_metric": args.primary,
+            "acc_metric": args.acc,
+            "num_seeds": len(records),
+            "seeds": seeds,
+            "leaderboard": {
+                "test_id_macro_f1": {"mean": id_f1_mean, "std": id_f1_std},
+                "test_id_avg_acc": {"mean": id_acc_mean, "std": id_acc_std},
+                "test_ood_macro_f1": {"mean": ood_f1_mean, "std": ood_f1_std},
+                "test_ood_avg_acc": {"mean": ood_acc_mean, "std": ood_acc_std},
+                "generalization_gap_macro_f1": gen_gap,
+            },
+            "training": {
+                "train_time_seconds": {"mean": time_mean, "std": time_std},
+                "epochs_run": {"mean": epoch_mean, "std": epoch_std},
+            },
+            "resources": {
+                "peak_host_ram_gb": {"mean": host_ram_mean, "std": host_ram_std},
+                "peak_gpu_alloc_gb": {"mean": gpu_alloc_mean, "std": gpu_alloc_std},
+                "peak_gpu_reserved_gb": {
+                    "mean": gpu_reserved_mean,
+                    "std": gpu_reserved_std,
+                },
+            },
+            "id_metrics_aggregated": id_agg,
+            "ood_metrics_aggregated": ood_agg,
+            "per_seed": per_seed,
+        }
+
+        model_out_dir = os.path.join(args.out, group_name)
+        os.makedirs(model_out_dir, exist_ok=True)
+        with open(os.path.join(model_out_dir, "summary.json"), "w") as f:
+            json.dump(summary, f, indent=2, sort_keys=True)
+
+        csv_rows.append(
+            {
+                "model": group_name,
+                "num_seeds": len(records),
+                "seeds": " ".join(str(s) for s in seeds),
+                "id_macro_f1_mean": id_f1_mean,
+                "id_macro_f1_std": id_f1_std,
+                "id_avg_acc_mean": id_acc_mean,
+                "id_avg_acc_std": id_acc_std,
+                "ood_macro_f1_mean": ood_f1_mean,
+                "ood_macro_f1_std": ood_f1_std,
+                "ood_avg_acc_mean": ood_acc_mean,
+                "ood_avg_acc_std": ood_acc_std,
+                "gen_gap_macro_f1_mean": gen_gap,
+                "train_time_seconds_mean": time_mean,
+                "train_time_seconds_std": time_std,
+                "epochs_run_mean": epoch_mean,
+                "epochs_run_std": epoch_std,
+                "peak_host_ram_gb_mean": host_ram_mean,
+                "peak_host_ram_gb_std": host_ram_std,
+                "peak_gpu_alloc_gb_mean": gpu_alloc_mean,
+                "peak_gpu_alloc_gb_std": gpu_alloc_std,
+                "peak_gpu_reserved_gb_mean": gpu_reserved_mean,
+                "peak_gpu_reserved_gb_std": gpu_reserved_std,
+            }
+        )
+
+    csv_path = os.path.join(args.out, "summary_all.csv")
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=csv_fields)
+        writer.writeheader()
+        for row in csv_rows:
+            writer.writerow(row)
+
+    # Terminal table.
+    print(f"\nAggregated {len(csv_rows)} model(s). Primary metric: {args.primary}\n")
+    header = (
+        f"{'model':<40} {'seeds':>5}  "
+        f"{'ID MacroF1':>18}  {'ID AvgAcc':>18}  "
+        f"{'OOD MacroF1':>18}  {'OOD AvgAcc':>18}  {'gap':>8}  "
+        f"{'peakRAM_GB':>14}  {'peakVRAM_GB':>14}"
+    )
+    print(header)
+    print("-" * len(header))
+    for row in csv_rows:
+        gap = row["gen_gap_macro_f1_mean"]
+        gap_str = "" if gap is None else f"{gap:.4f}"
+        model_str = row["model"][:40]
+        print(
+            f"{model_str:<40} {row['num_seeds']:>5}  "
+            f"{_fmt(row['id_macro_f1_mean'], row['id_macro_f1_std']):>18}  "
+            f"{_fmt(row['id_avg_acc_mean'], row['id_avg_acc_std']):>18}  "
+            f"{_fmt(row['ood_macro_f1_mean'], row['ood_macro_f1_std']):>18}  "
+            f"{_fmt(row['ood_avg_acc_mean'], row['ood_avg_acc_std']):>18}  "
+            f"{gap_str:>8}  "
+            f"{_fmt(row['peak_host_ram_gb_mean'], row['peak_host_ram_gb_std']):>14}  "
+            f"{_fmt(row['peak_gpu_alloc_gb_mean'], row['peak_gpu_alloc_gb_std']):>14}"
+        )
+
+    print(f"\nWrote per-model summaries to: {args.out}/<model>/summary.json")
+    print(f"Wrote combined CSV to:        {csv_path}")
+
+
+if __name__ == "__main__":
+    main()

tools/run_grid.py

@@ -0,0 +1,115 @@
+import argparse
+import itertools
+import os
+import sys
+import tempfile
+
+import yaml
+import submitit
+
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+from src.train_supervised import main as app_main
+
+
+def _set_by_dotted_key(config, dotted_key, value):
+    keys = dotted_key.split(".")
+    cur = config
+    for key in keys[:-1]:
+        if key not in cur or not isinstance(cur[key], dict):
+            cur[key] = {}
+        cur = cur[key]
+    cur[keys[-1]] = value
+
+
+def _deep_update(base, updates):
+    for key, value in updates.items():
+        _set_by_dotted_key(base, key, value)
+    return base
+
+
+def _load_yaml(path):
+    with open(path, "r") as f:
+        return yaml.load(f, Loader=yaml.FullLoader)
+
+
+def _expand_grid(grid_dict):
+    keys = list(grid_dict.keys())
+    values = [grid_dict[k] for k in keys]
+    for combo in itertools.product(*values):
+        yield dict(zip(keys, combo))
+
+
+class GridTrainer:
+    def __init__(self, fname):
+        self.fname = fname
+
+    def __call__(self):
+        with open(self.fname, "r") as y_file:
+            params = yaml.load(y_file, Loader=yaml.FullLoader)
+        app_main(args=params, resume_preempt=False)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--grid", required=True, help="grid yaml file")
+    parser.add_argument("--folder", type=str, help="location to save submitit logs")
+    parser.add_argument("--partition", type=str, help="cluster partition to submit jobs on")
+    parser.add_argument("--nodes", type=int, default=1, help="num. nodes to request for job")
+    parser.add_argument(
+        "--tasks-per-node", type=int, default=1, help="num. procs to per node"
+    )
+    parser.add_argument("--time", type=int, default=4300, help="time in minutes to run job")
+    args = parser.parse_args()
+
+    grid_cfg = _load_yaml(args.grid)
+    base_config_path = grid_cfg["base_config"]
+    base_params = _load_yaml(base_config_path)
+
+    constants = grid_cfg.get("constants", {})
+    grid = grid_cfg.get("grid", {})
+    launch = grid_cfg.get("launch", {})
+
+    params_list = []
+    for overrides in _expand_grid(grid):
+        params = yaml.safe_load(yaml.dump(base_params))
+        _deep_update(params, constants)
+        _deep_update(params, overrides)
+        params_list.append(params)
+
+    log_folder = args.folder or launch.get("folder")
+    if not log_folder:
+        raise ValueError("submitit folder required via --folder or grid launch.folder")
+
+    executor = submitit.SlurmExecutor(
+        folder=os.path.join(log_folder, "job_%j"), max_num_timeout=20
+    )
+    executor.update_parameters(
+        partition=args.partition or launch.get("partition"),
+        mem_per_gpu=launch.get("mem_per_gpu", "180G"),
+        time=args.time or int(launch.get("time", 4300)),
+        nodes=args.nodes or int(launch.get("nodes", 1)),
+        ntasks_per_node=args.tasks_per_node or int(launch.get("tasks_per_node", 1)),
+        cpus_per_task=int(launch.get("cpus_per_task", 10)),
+        gpus_per_node=args.tasks_per_node or int(launch.get("tasks_per_node", 1)),
+    )
+
+    temp_dir = tempfile.mkdtemp(prefix="grid_configs_", dir=log_folder)
+
+    jobs = []
+    with executor.batch():
+        for idx, params in enumerate(params_list):
+            tmp_path = os.path.join(temp_dir, f"grid_{idx}.yaml")
+            with open(tmp_path, "w") as f:
+                yaml.dump(params, f)
+            job = executor.submit(GridTrainer(tmp_path))
+            jobs.append(job)
+
+    for job in jobs:
+        print(job.job_id)
+
+
+if __name__ == "__main__":
+    main()

tools/run_label_efficiency.sh

@@ -0,0 +1,99 @@
+#!/usr/bin/env bash
+#
+# Run label-efficiency supervised experiments for all model grids.
+#
+# For every grid under configs/grids/label_efficiency/, this submits one
+# submitit job per seed (via tools/run_grid.py). Models are launched
+# sequentially so you can run them "one by one"; within a grid, the 5 seeds
+# are submitted together.
+#
+# Each run automatically:
+#   - uses a stratified subset of the Source split (data.label_fraction)
+#   - seeds training from meta.seed
+#   - evaluates on id_test (ID) and test (OOD) WILDS splits
+#   - records WILDS metrics + training time + epochs_run into the metrics JSON
+#
+# After all jobs finish, aggregate with:
+#   python3 tools/aggregate_label_efficiency.py --root experiment_logs/eval-wilds
+#
+# Usage:
+#   bash tools/run_label_efficiency.sh [--partition P] [--time MIN] [--folder DIR]
+#                                    [--models "vith14_224_in22k vitg16_224_in22k"]
+#                                    [--fractions "0.01 0.10 0.50"]
+#
+set -euo pipefail
+
+PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+GRID_DIR="${PROJECT_ROOT}/configs/grids/label_efficiency"
+
+PARTITION=""
+TIME=""
+FOLDER=""
+MODELS=""
+FRACTIONS=""
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --partition)  PARTITION="$2"; shift 2 ;;
+    --time)       TIME="$2"; shift 2 ;;
+    --folder)     FOLDER="$2"; shift 2 ;;
+    --models)     MODELS="$2"; shift 2 ;;
+    --fractions)  FRACTIONS="$2"; shift 2 ;;
+    *) echo "Unknown argument: $1" >&2; exit 1 ;;
+  esac
+done
+
+# Resolve the list of grid files to run.
+if [[ -n "${MODELS}" ]]; then
+  GRIDS=()
+  for m in ${MODELS}; do
+    if [[ -n "${FRACTIONS}" ]]; then
+      for frac in ${FRACTIONS}; do
+        g="${GRID_DIR}/${m}_frac${frac}.yaml"
+        if [[ ! -f "${g}" ]]; then
+          echo "Grid not found for model '${m}' fraction '${frac}': ${g}" >&2
+          exit 1
+        fi
+        GRIDS+=("${g}")
+      done
+    else
+      for g in "${GRID_DIR}/${m}"_frac*.yaml; do
+        if [[ -f "${g}" ]]; then
+          GRIDS+=("${g}")
+        fi
+      done
+    fi
+  done
+else
+  # All models and fractions, sorted.
+  GRIDS=()
+  while IFS= read -r g; do GRIDS+=("${g}"); done < <(ls "${GRID_DIR}"/*.yaml | sort)
+fi
+
+if [[ ${#GRIDS[@]} -eq 0 ]]; then
+  echo "No grid files found in ${GRID_DIR}" >&2
+  exit 1
+fi
+
+echo "Launching label-efficiency sweeps for ${#GRIDS[@]} grid(s):"
+for g in "${GRIDS[@]}"; do echo "  - $(basename "${g}")"; done
+echo
+
+for g in "${GRIDS[@]}"; do
+  echo "=================================================================="
+  echo "Model grid: $(basename "${g}")"
+  echo "=================================================================="
+
+  cmd=("${PROJECT_ROOT}/.venv/bin/python" "${PROJECT_ROOT}/tools/run_grid.py" --grid "${g}")
+  [[ -n "${PARTITION}" ]] && cmd+=(--partition "${PARTITION}")
+  [[ -n "${TIME}" ]]      && cmd+=(--time "${TIME}")
+  [[ -n "${FOLDER}" ]]    && cmd+=(--folder "${FOLDER}")
+
+  echo "+ ${cmd[*]}"
+  "${cmd[@]}"
+  echo
+done
+
+echo "All label-efficiency jobs submitted."
+echo "When they finish, aggregate results with:"
+echo "  ${PROJECT_ROOT}/.venv/bin/python tools/aggregate_label_efficiency.py --root experiment_logs/eval-wilds"

tools/run_seed_sweep.sh

@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+#
+# Run each supervised model through all seeds, one model at a time.
+#
+# For every model grid under configs/grids/seeds/, this submits one submitit
+# job per seed (via tools/run_grid.py). Models are launched sequentially so you
+# can run them "one by one"; within a model, the 5 seeds are submitted together.
+#
+# Each run automatically:
+#   - seeds training from meta.seed (config-driven, see src/train_supervised.py)
+#   - evaluates on id_test (ID) and test (OOD) WILDS splits
+#   - records WILDS metrics + training time + epochs_run into the metrics JSON
+#
+# After all jobs finish, aggregate with:
+#   python3 tools/aggregate_seeds.py --root experiment_logs/eval-wilds
+#
+# Usage:
+#   bash tools/run_seed_sweep.sh [--partition P] [--time MIN] [--folder DIR]
+#                                [--models "vith14_224 vith16_448"]
+#
+set -euo pipefail
+
+PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+GRID_DIR="${PROJECT_ROOT}/configs/grids/seeds"
+
+PARTITION=""
+TIME=""
+FOLDER=""
+MODELS=""
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --partition) PARTITION="$2"; shift 2 ;;
+    --time)      TIME="$2"; shift 2 ;;
+    --folder)    FOLDER="$2"; shift 2 ;;
+    --models)    MODELS="$2"; shift 2 ;;
+    *) echo "Unknown argument: $1" >&2; exit 1 ;;
+  esac
+done
+
+# Resolve the list of grid files to run.
+if [[ -n "${MODELS}" ]]; then
+  GRIDS=()
+  for m in ${MODELS}; do
+    g="${GRID_DIR}/${m}.yaml"
+    if [[ ! -f "${g}" ]]; then
+      echo "Grid not found for model '${m}': ${g}" >&2
+      exit 1
+    fi
+    GRIDS+=("${g}")
+  done
+else
+  # All models, sorted.
+  GRIDS=()
+  while IFS= read -r g; do GRIDS+=("${g}"); done < <(ls "${GRID_DIR}"/*.yaml | sort)
+fi
+
+echo "Launching seed sweeps for ${#GRIDS[@]} model(s):"
+for g in "${GRIDS[@]}"; do echo "  - $(basename "${g}")"; done
+echo
+
+for g in "${GRIDS[@]}"; do
+  echo "=================================================================="
+  echo "Model grid: $(basename "${g}")"
+  echo "=================================================================="
+
+  cmd=("${PROJECT_ROOT}/.venv/bin/python" "${PROJECT_ROOT}/tools/run_grid.py" --grid "${g}")
+  [[ -n "${PARTITION}" ]] && cmd+=(--partition "${PARTITION}")
+  [[ -n "${TIME}" ]]      && cmd+=(--time "${TIME}")
+  [[ -n "${FOLDER}" ]]    && cmd+=(--folder "${FOLDER}")
+
+  echo "+ ${cmd[*]}"
+  "${cmd[@]}"
+  echo
+done
+
+echo "All seed-sweep jobs submitted."
+echo "When they finish, aggregate results with:"
+echo "  ${PROJECT_ROOT}/.venv/bin/python tools/aggregate_seeds.py --root experiment_logs/eval-wilds"

tools/summarize_grid.py

@@ -0,0 +1,470 @@
+import argparse
+import json
+import os
+import re
+
+
+def _find_metric(metrics, key):
+    if isinstance(metrics, dict):
+        if key in metrics:
+            return metrics[key]
+        for value in metrics.values():
+            found = _find_metric(value, key)
+            if found is not None:
+                return found
+    if isinstance(metrics, list):
+        for item in metrics:
+            found = _find_metric(item, key)
+            if found is not None:
+                return found
+    return None
+
+
+def _parse_yaml_value(value):
+    value = value.strip()
+    if not value or value in ("null", "~"):
+        return None
+    if value == "true":
+        return True
+    if value == "false":
+        return False
+    if value.startswith("[") and value.endswith("]"):
+        try:
+            return json.loads(value)
+        except (json.JSONDecodeError, ValueError):
+            pass
+    try:
+        return int(value)
+    except ValueError:
+        pass
+    try:
+        return float(value)
+    except ValueError:
+        pass
+    if (value.startswith('"') and value.endswith('"')) or (
+        value.startswith("'") and value.endswith("'")
+    ):
+        return value[1:-1]
+    return value
+
+
+def _load_params_simple(dirpath):
+    path = os.path.join(dirpath, "params.yaml")
+    alt_path = os.path.join(dirpath, "params-eval.yaml")
+    for p in [path, alt_path]:
+        if os.path.isfile(p):
+            try:
+                with open(p, "r") as f:
+                    return _parse_simple_yaml(f.read())
+            except (OSError, ValueError):
+                continue
+    return None
+
+
+def _parse_simple_yaml(text):
+    lines = text.split("\n")
+
+    def _effective(raw, indent):
+        stripped = raw.lstrip(" ")
+        cur_indent = len(raw) - len(stripped)
+        if cur_indent != indent:
+            return None, cur_indent
+        if " #" in stripped:
+            effective = stripped[: stripped.index(" #")].rstrip()
+        else:
+            effective = stripped
+        return effective, cur_indent
+
+    def _peek_nonblank(idx):
+        while idx < len(lines):
+            raw = lines[idx].rstrip()
+            if raw.strip() and not raw.strip().startswith("#"):
+                return idx
+            idx += 1
+        return len(lines)
+
+    def _collect_list_items(start, indent):
+        items = []
+        i = start
+        while i < len(lines):
+            raw = lines[i].rstrip()
+            if not raw.strip() or raw.strip().startswith("#"):
+                i += 1
+                continue
+            effective, cur_indent = _effective(raw, indent)
+            if effective is None:
+                if cur_indent < indent:
+                    break
+                i += 1
+                continue
+            if effective.startswith("- "):
+                items.append(_parse_yaml_value(effective[2:]))
+                i += 1
+            else:
+                break
+        return items, i
+
+    def _parse_block(start, indent):
+        result = {}
+        list_items = []
+        is_list = False
+        i = start
+
+        while i < len(lines):
+            raw = lines[i].rstrip()
+            if not raw.strip() or raw.strip().startswith("#"):
+                i += 1
+                continue
+
+            stripped = raw.lstrip(" ")
+            cur_indent = len(raw) - len(stripped)
+
+            if cur_indent < indent:
+                break
+            if cur_indent > indent:
+                i += 1
+                continue
+
+            if " #" in stripped:
+                effective = stripped[: stripped.index(" #")].rstrip()
+            else:
+                effective = stripped
+
+            if not effective:
+                i += 1
+                continue
+
+            if effective.startswith("- "):
+                is_list = True
+                list_items.append(_parse_yaml_value(effective[2:]))
+                i += 1
+            elif effective.endswith(":"):
+                key = effective[:-1].strip()
+                nxt = _peek_nonblank(i + 1)
+                if nxt < len(lines):
+                    nxt_raw = lines[nxt].rstrip()
+                    nxt_stripped = nxt_raw.lstrip(" ")
+                    nxt_indent = len(nxt_raw) - len(nxt_stripped)
+                else:
+                    nxt_indent = -1
+
+                if nxt_indent > indent:
+                    sub_val, i = _parse_block(i + 1, nxt_indent)
+                    result[key] = sub_val
+                elif nxt_indent == indent and nxt_stripped.startswith("- "):
+                    lst, i = _collect_list_items(i + 1, indent)
+                    result[key] = lst
+                else:
+                    result[key] = None
+                    i += 1
+            elif ": " in effective:
+                key, _, val_str = effective.partition(": ")
+                result[key.strip()] = _parse_yaml_value(val_str)
+                i += 1
+            else:
+                i += 1
+
+        if is_list:
+            return list_items, i
+        return result, i
+
+    return _parse_block(0, 0)[0]
+
+
+def _get_in_params(params, path):
+    if params is None:
+        return None
+    keys = path.split(".")
+    current = params
+    for key in keys:
+        if isinstance(current, dict) and key in current:
+            current = current[key]
+        else:
+            return None
+    return current
+
+
+_MODEL_PREFIXES = {
+    "vith": "vit_huge",
+    "vitb": "vit_base",
+    "vitl": "vit_large",
+    "vitt": "vit_tiny",
+    "vits": "vit_small",
+    "vitg": "vit_giant",
+}
+
+
+def _infer_model_from_run_name(run_name):
+    lower = run_name.lower()
+    for model in sorted(set(_MODEL_PREFIXES.values()), key=len, reverse=True):
+        if lower.startswith(model):
+            return model
+    for prefix, model in sorted(_MODEL_PREFIXES.items(), key=lambda x: len(x[0]), reverse=True):
+        if lower.startswith(prefix):
+            return model
+    return None
+
+
+def _get_model_name(dirpath, run_name=None):
+    params = _load_params_simple(dirpath)
+    name = None
+    if params is not None:
+        name = _get_in_params(params, "meta.model_name")
+    if name is None and run_name:
+        name = _infer_model_from_run_name(run_name)
+    if name is None:
+        return "unknown"
+    base = name.split("-")[0]
+    for part in name.split("-")[1:]:
+        if part.startswith("in"):
+            return f"{base}-{part}"
+    if run_name:
+        for part in run_name.split("-"):
+            if part.lower().startswith("in"):
+                return f"{base}-{part}"
+            m = re.search(r"in\d+\w*", part)
+            if m:
+                return f"{base}-{m.group()}"
+    return base
+
+
+def _get_patch_size(params):
+    if params is None:
+        return "unknown"
+    ps = _get_in_params(params, "mask.patch_size")
+    if ps is None:
+        ps = _get_in_params(params, "meta.patch_size")
+    return ps if ps is not None else "unknown"
+
+
+def _get_crop_size(params):
+    if params is None:
+        return "unknown"
+    cs = _get_in_params(params, "data.crop_size")
+    if cs is None:
+        cs = _get_in_params(params, "meta.crop_size")
+    return cs if cs is not None else "unknown"
+
+
+def _collect_rows(root_dir, metric_key, col_paths):
+    rows = []
+    for dirpath, _, filenames in os.walk(root_dir):
+        for fname in filenames:
+            if not fname.endswith("_metrics.json"):
+                continue
+            path = os.path.join(dirpath, fname)
+            try:
+                with open(path, "r") as f:
+                    metrics = json.load(f)
+            except (OSError, json.JSONDecodeError):
+                continue
+            # Only consider out-of-distribution (OOD) metrics. WILDS eval files
+            # record their split at the top level; "test" is the OOD split.
+            if not isinstance(metrics, dict) or metrics.get("split") != "test":
+                continue
+            value = _find_metric(metrics, metric_key)
+            if value is None:
+                continue
+            run_name = os.path.basename(os.path.dirname(path))
+            model_name = _get_model_name(dirpath, run_name)
+            params = _load_params_simple(dirpath)
+            patch_size = _get_patch_size(params)
+            crop_size = _get_crop_size(params)
+            col_values = [
+                _get_in_params(params, cp) for cp in col_paths
+            ]
+            rows.append(
+                (float(value), model_name, patch_size, crop_size, col_values, run_name, path)
+            )
+    return rows
+
+
+def _format_value(val, width=None):
+    if val is None:
+        s = "\u2014"
+    elif isinstance(val, bool):
+        s = str(val)
+    elif isinstance(val, int):
+        s = str(val)
+    elif isinstance(val, float):
+        if abs(val) < 0.001 or abs(val) >= 10000:
+            s = f"{val:.2e}"
+        elif val == int(val):
+            s = f"{val:.1f}"
+        else:
+            s = f"{val:.6f}".rstrip("0").rstrip(".")
+    elif isinstance(val, list):
+        s = str(val)
+    else:
+        s = str(val)
+    if width is not None and len(s) > width:
+        s = s[: width - 1] + "\u2026"
+    return s
+
+
+def _col_width(header, values):
+    w = len(header)
+    for val in values:
+        s = _format_value(val)
+        w = max(w, len(s))
+    return w
+
+
+_BLOCK = "\u2588"
+_LIGHT_H = "\u2500"
+_LIGHT_V = "\u2502"
+_LIGHT_D = "\u253c"
+_HEAVY_H = "\u2501"
+_HEAVY_V = "\u2503"
+_HEAVY_D = "\u254b"
+
+
+def _print_separator(widths, heavy=False):
+    h = _HEAVY_H if heavy else _LIGHT_H
+    d = _HEAVY_D if heavy else _LIGHT_D
+    parts = []
+    for i, w in enumerate(widths):
+        if i > 0:
+            parts.append(d)
+        segments = max(w, 1) - 1
+        if segments <= 0:
+            parts.append(h)
+        else:
+            parts.append(h * segments)
+    line = h.join(parts) if len(parts) > 0 else ""
+    print(f"  {line}")
+
+
+def _print_row(values, widths, align_right=None):
+    if align_right is None:
+        align_right = [True] * len(values)
+    parts = []
+    for i, (val, w) in enumerate(zip(values, widths)):
+        s = _format_value(val, w)
+        if align_right[i]:
+            s = s.rjust(w)
+        else:
+            s = s.ljust(w)
+        if i > 0:
+            parts.append(f" {_LIGHT_V} ")
+        parts.append(s)
+    print("  " + "".join(parts))
+
+
+def _print_header(full_cols, widths):
+    align = [False] + [True] * (len(full_cols) - 2) + [False]
+    _print_row(full_cols, widths, align)
+    _print_separator(widths)
+
+
+def _print_results(model_type, rows, metric_key, col_headers, top, show_run_name):
+    if not rows:
+        return
+
+    display_rows = rows[:top]
+
+    n = len(display_rows)
+    rank_width = max(3, len(str(n)))
+    metric_vals = [r[0] for r in display_rows]
+    metric_width = max(
+        len(metric_key), max(len(_format_value(v)) for v in metric_vals)
+    )
+
+    col_widths = []
+    for i, ch in enumerate(col_headers):
+        col_vals = [r[4][i] for r in display_rows]
+        cw = _col_width(ch, col_vals)
+        col_widths.append(cw)
+
+    widths = [rank_width, metric_width] + col_widths
+    full_cols = ["#", metric_key] + col_headers
+
+    if show_run_name:
+        run_name_vals = [r[5] for r in display_rows]
+        run_name_width = max(len("run_name"), max(len(v) for v in run_name_vals))
+        widths.append(run_name_width)
+        full_cols.append("run_name")
+
+    title = f"{model_type} \u2014 Top {top} by {metric_key}"
+    sep_len = sum(w + 3 for w in widths) + 2
+    print(_HEAVY_H * sep_len)
+    print(f"  {title}")
+    print(_HEAVY_H * sep_len)
+    _print_header(full_cols, widths)
+
+    for idx, row in enumerate(display_rows, start=1):
+        metric_str = _format_value(row[0])
+        col_strs = [_format_value(row[4][i]) for i in range(len(col_headers))]
+        vals = [str(idx), metric_str] + col_strs
+        if show_run_name:
+            vals.append(row[5])
+        align = [False] + [True] * (len(full_cols) - 2) + [False]
+        _print_row(vals, widths, align)
+
+    print()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="""Summarize grid evaluation results grouped by model / patch size / resolution combination.
+
+Examples:
+  %(prog)s --top 10
+  %(prog)s --metric macro_f1 --cols data.batch_size optimization.lr --top 5
+""",
+    )
+    parser.add_argument(
+        "--root",
+        default="experiment_logs/eval-wilds",
+        help="root folder for eval logs (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--metric",
+        default="macro_f1",
+        help="metric key to rank by (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--top",
+        type=int,
+        default=10,
+        help="top N results per model type (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--cols",
+        nargs="*",
+        default=[],
+        help="extra columns from params.yaml, e.g. data.batch_size optimization.lr",
+    )
+    parser.add_argument(
+        "--no-run-name",
+        action="store_false",
+        dest="show_run_name",
+        help="hide the run_name column",
+    )
+    args = parser.parse_args()
+
+    rows = _collect_rows(args.root, args.metric, args.cols)
+
+    if not rows:
+        print("No metrics found.")
+        return
+
+    col_headers = [c.split(".")[-1] for c in args.cols]
+
+    groups = {}
+    for row in rows:
+        model = row[1]
+        patch = row[2]
+        crop = row[3]
+        key = f"{model}  patch={patch}  resolution={crop}"
+        groups.setdefault(key, []).append(row)
+
+    for group_name in sorted(groups.keys()):
+        group = groups[group_name]
+        group.sort(key=lambda r: r[0], reverse=True)
+        _print_results(group_name, group, args.metric, col_headers, args.top, args.show_run_name)
+
+
+if __name__ == "__main__":
+    main()