seed runs

README.md

@@ -28,6 +28,9 @@ Reference: official I-JEPA README https://github.com/facebookresearch/ijepa/blob
 - `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 -->
@@ -63,6 +66,50 @@ Evaluation metrics are written to `experiment_logs/eval-wilds-vith14/iwildcam_te
 
 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**, and the number of
+  **epochs run** (accounting for early stopping) into the per-split metrics JSON
+  and into `params.yaml` in the eval folder.
+
+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, and ID-OOD generalization gap)
+- `experiment_logs/seed-runs/summary_all.csv` (one row per model, paper-ready)
+
 ## License
 
 See the `LICENSE` file for details about the license under which this code is made available.

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

src/train_supervised.py

@@ -1,7 +1,9 @@
 import os
+import random
 import shutil
 import sys
 import json
+import time
 import yaml
 import logging
 
@@ -35,6 +37,22 @@ 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 strip_module_prefix(state_dict):
     if not any(k.startswith("module.") for k in state_dict.keys()):
         return state_dict
@@ -213,6 +231,10 @@ def main(args, resume_preempt=False):
     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)
 
@@ -424,9 +446,13 @@ def main(args, resume_preempt=False):
             if is_best:
                 torch.save(save_dict, best_path)
 
+    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"]:
@@ -501,12 +527,22 @@ def main(args, resume_preempt=False):
             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
+    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)"
+        )
+
     if early_stopper.enabled and early_stopper.restore_best_weights:
         if rank == 0:
             logger.info("Restoring best model weights before exit")
@@ -523,9 +559,18 @@ def main(args, resume_preempt=False):
             )
 
     if rank == 0:
-        eval_args = {
+        # 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": m_args.get("seed", _GLOBAL_SEED),
+                    "seed": seed,
                     "model_name": m_args["model_name"],
                     "embed_dim": m_args["embed_dim"],
                     "num_classes": m_args["num_classes"],
@@ -542,54 +587,79 @@ def main(args, resume_preempt=False):
                     "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": "test",
+                    "split": split,
                     "download": True,
                 },
                 "logging": {
-                "write_tag": "iwildcam_test",
+                    "write_tag": write_tag,
                     "auto_folder": True,
                 },
             }
-        eval_result = eval_wilds_main(args=eval_args)
 
+        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,
+            "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),
+        }
+
+        # 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
-                metric_key = m_args.get("selection_metric", "macro_f1")
-                eval_root = os.path.join("experiment_logs", "eval-wilds")
-                ranking = {
-                    "metric_key": metric_key,
-                    "this_run_name": os.path.basename(os.path.normpath(eval_folder)),
-                    "this_is_best": None,
-                    "this_metric": None,
-                    "best_run_name": None,
-                    "best_metric": None,
-                }
-                if os.path.isdir(eval_root):
-                    rows = _collect_eval_rows(eval_root, metric_key)
-                    if rows:
-                        rows.sort(key=lambda r: r[0], reverse=True)
-                        best_value, best_run_name, _ = rows[0]
-                        this_run_name = ranking["this_run_name"]
-                        this_value = None
-                        for value, run_name, _ in rows:
-                            if run_name == this_run_name:
-                                this_value = value
-                                break
-                        ranking["this_metric"] = this_value
-                        ranking["best_run_name"] = best_run_name
-                        ranking["best_metric"] = best_value
-                        if this_value is not None:
-                            ranking["this_is_best"] = this_run_name == best_run_name
                 params_out["results"] = {
                     "best_val_loss": float(early_stopper.best_metric),
                     "best_epoch": int(early_stopper.best_epoch),
                     "best_checkpoint": best_path,
-                    "eval_metrics": eval_result.get("metrics") if eval_result else None,
+                    "seed": seed,
-                    "ranking": ranking,
+                    "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),
+                    "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)

src/utils/logging.py

@@ -88,6 +88,7 @@ def build_run_name(args):
     add("ipe", opt_args.get("ipe_scale"))
     add("cs", data_args.get("crop_scale"))
     add("eval", val_args.get("eval_every"))
+    add("seed", meta_args.get("seed"))
 
     return "-".join([p for p in parts if p])
 

tools/aggregate_seeds.py

@@ -0,0 +1,343 @@
+#!/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
+"""
+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 _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)",
+    )
+    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 = {}
+        for m in (ood_metrics, id_metrics):
+            if isinstance(m, dict) and isinstance(m.get("run_info"), dict):
+                run_info = m["run_info"]
+                break
+
+        group = _strip_seed(entry)
+        groups.setdefault(group, []).append(
+            {
+                "run_name": entry,
+                "seed": run_info.get("seed"),
+                "id_metrics": id_metrics,
+                "ood_metrics": ood_metrics,
+                "run_info": run_info,
+            }
+        )
+
+    if not groups:
+        print(f"No metrics found 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",
+    ]
+
+    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"),
+                }
+            )
+
+        # 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]
+        time_mean, time_std, _ = _mean_std(time_vals)
+        epoch_mean, epoch_std, _ = _mean_std(epoch_vals)
+
+        # 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},
+            },
+            "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,
+            }
+        )
+
+    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}"
+    )
+    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}"
+        )
+
+    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_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"