diff --git a/README.md b/README.md index 8d83767..376989d 100644 --- a/README.md +++ b/README.md @@ -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//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 diff --git a/src/train_supervised.py b/src/train_supervised.py index 18cd5a9..a998694 100644 --- a/src/train_supervised.py +++ b/src/train_supervised.py @@ -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") diff --git a/tools/aggregate_seeds.py b/tools/aggregate_seeds.py index c4793ae..143cf36 100644 --- a/tools/aggregate_seeds.py +++ b/tools/aggregate_seeds.py @@ -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}//summary.json")