YannAhlgrim/wilds-ijepa
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add ram logging

Browse Source
This commit is contained in:
YannAhlgrim
2026-07-06 13:59:19 +02:00
parent 56f15dd856
commit 79a7dbb36d
3 changed files with 91 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+14 -5
View File
@@ -76,9 +76,17 @@ 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.
- 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`).
@@ -107,8 +115,9 @@ 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)
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`)
## License
src/train_supervised.py
+43 -1
View File
@@ -1,5 +1,6 @@
import os
import random
import resource
import shutil
import sys
import json
@@ -53,6 +54,33 @@ def _format_hms(seconds):
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
@@ -446,6 +474,9 @@ def main(args, resume_preempt=False):
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
@@ -535,12 +566,17 @@ def main(args, resume_preempt=False):
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"({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:
@@ -615,6 +651,9 @@ def main(args, resume_preempt=False):
"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
@@ -650,6 +689,9 @@ def main(args, resume_preempt=False):
"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")
tools/aggregate_seeds.py
+34 -2
View File
@@ -192,6 +192,12 @@ def main():
"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()):
@@ -228,6 +234,9 @@ def main():
"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"),
}
)
@@ -245,8 +254,14 @@ def main():
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, _ = _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)
# Headline (leaderboard) numbers.
id_f1_mean, id_f1_std, _ = _mean_std(metric_values("id_metrics", args.primary))
@@ -275,6 +290,14 @@ def main():
"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,
@@ -303,6 +326,12 @@ def main():
"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,
}
)
@@ -318,7 +347,8 @@ def main():
header = (
f"{'model':<40} {'seeds':>5} "
f"{'ID MacroF1':>18} {'ID AvgAcc':>18} "
f"{'OOD MacroF1':>18} {'OOD AvgAcc':>18} {'gap':>8}"
f"{'OOD MacroF1':>18} {'OOD AvgAcc':>18} {'gap':>8} "
f"{'peakRAM_GB':>14} {'peakVRAM_GB':>14}"
)
print(header)
print("-" * len(header))
@@ -332,7 +362,9 @@ def main():
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"{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")