From 75fd385a88aaf2f092e38d4142faf6a0c1afe4d6 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Sun, 5 Apr 2026 06:57:15 +0200 Subject: [PATCH] feat(hpc): tensorboard analysis --- configs/hpc/smoke_test.yaml | 2 +- env/hpc/requirements.txt | 2 +- pyproject.toml | 5 +- scripts/analysis/README.md | 27 +++++ scripts/analysis/explore_tensorboard.py | 135 ++++++++++++++++++++++++ scripts/export_hpc_requirements.py | 13 ++- 6 files changed, 180 insertions(+), 4 deletions(-) create mode 100644 scripts/analysis/README.md create mode 100644 scripts/analysis/explore_tensorboard.py diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml index 72c6653..a00b800 100644 --- a/configs/hpc/smoke_test.yaml +++ b/configs/hpc/smoke_test.yaml @@ -6,6 +6,6 @@ track: true # Test WandB integration capture_video: false # No rendering for smoke test save_model: true # Test the end-of-training save routine num_envs: 4 -total_timesteps: 10000 +total_timesteps: 500 num_steps: 64 cuda: true diff --git a/env/hpc/requirements.txt b/env/hpc/requirements.txt index 5314edb..819ea32 100644 --- a/env/hpc/requirements.txt +++ b/env/hpc/requirements.txt @@ -5,7 +5,6 @@ flax>=0.12.2 gymnasium>=1.2.3 ipykernel==7.2.0 jax[cuda13]==0.9.0.1 -torch numpy>=2.0.0 protobuf>=5.0.0 warp-lang @@ -17,3 +16,4 @@ pyopengl>=3.1.10 pyopengl-accelerate>=3.1.10 tyro>=1.0.10 wandb==0.24.2 +torch>=2.4.0 diff --git a/pyproject.toml b/pyproject.toml index d1d468d..171bae1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,10 @@ dependencies = [ [project.optional-dependencies] cuda = [ - "jax[cuda12]==0.9.0.1", + "jax[cuda13]==0.9.0.1", +] +analysis = [ + "tensorboard", ] [dependency-groups] diff --git a/scripts/analysis/README.md b/scripts/analysis/README.md new file mode 100644 index 0000000..a99f952 --- /dev/null +++ b/scripts/analysis/README.md @@ -0,0 +1,27 @@ +# Experiment Analysis Tools + +This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights. + +## Scripts + +### 1. `explore_tensorboard.py` +A CLI tool to summarize TensorBoard `tfevents` files without a GUI. + +**Key Features:** +- Displays last values, min, max, and step counts for all scalar metrics. +- Calculates total run duration and estimated completion percentage. +- Exports granular scalar data to CSV for analysis in Excel/Pandas. + +**Usage:** +```bash +# General usage +python explore_tensorboard.py + +# Exporting data +python explore_tensorboard.py --csv data.csv +``` + +**Requirements:** +- `pandas` +- `tensorboard` +- `tensorflow-cpu` (or `tensorflow`) diff --git a/scripts/analysis/explore_tensorboard.py b/scripts/analysis/explore_tensorboard.py new file mode 100644 index 0000000..c9e9e5b --- /dev/null +++ b/scripts/analysis/explore_tensorboard.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +""" +Reproducible CLI tool to explore TensorBoard logs. +Designed for both local development and HPC diagnostics. + +Requirements: + pip install tensorboard + +Usage: + python explore_tensorboard.py [--csv output.csv] +""" + +import argparse +import os +import sys +import csv + +try: + from tensorboard.backend.event_processing import event_accumulator +except ImportError: + print("Error: Missing dependency. Please run: pip install tensorboard") + sys.exit(1) + +def explore_run(log_dir): + """ + Extracts and displays a summary of scalar metrics from a TensorBoard log directory. + """ + print(f"\n{'='*20} Exploring Run {'='*20}") + print(f"Directory: {log_dir}") + print(f"{'='*55}\n") + + if not os.path.exists(log_dir): + print(f"Error: Directory '{log_dir}' does not exist.") + return None + + # Initialize EventAccumulator + # size_guidance=0 loads all data points for each tag. + ea = event_accumulator.EventAccumulator(log_dir, size_guidance={ + event_accumulator.SCALARS: 0, + event_accumulator.TENSORS: 0, + }) + + print("Loading event files (this may take a moment for large runs)...") + ea.Reload() + + tags = ea.Tags() + scalar_tags = tags.get('scalars', []) + + if not scalar_tags: + print("No scalar metrics found in this directory.") + return None + + print(f"Found {len(scalar_tags)} scalar metrics.\n") + + data = {} + summary = [] + + # Process scalar values + for tag in scalar_tags: + events = ea.Scalars(tag) + if not events: + continue + + values = [e.value for e in events] + last_event = events[-1] + data[tag] = values + + summary.append({ + "Metric": tag, + "Steps": len(events), + "Last Value": f"{last_event.value:.4f}", + "Max": f"{max(values):.4f}", + "Min": f"{min(values):.4f}" + }) + + # Display summary table formatted manually + summary = sorted(summary, key=lambda x: x['Metric']) + print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}") + print("-" * 80) + for row in summary: + print(f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} {row['Max']:>12} {row['Min']:>12}") + + # Calculate and display global metadata + if 'charts/SPS' in data: + sps_events = ea.Scalars('charts/SPS') + if len(sps_events) > 1: + total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600 + print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours") + + # Estimate completion if total_timesteps is available in hyperparameters + try: + hp_tags = [t for t in tags.get('tensors', []) if 'hyperparameters' in t] + if hp_tags: + hp_event = ea.Tensors(hp_tags[0])[0] + hp_text = hp_event.tensor_proto.string_val[0].decode('utf-8') + if 'total_timesteps' in hp_text: + for line in hp_text.split('\n'): + if 'total_timesteps' in line: + target = int(line.split('|')[2].strip()) + current = ea.Scalars(scalar_tags[0])[-1].step + percent = (current / target) * 100 + print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)") + except Exception: + pass + + return data + +def main(): + parser = argparse.ArgumentParser(description="Clean, reproducible TensorBoard exploration tool.") + parser.add_argument("log_dir", help="Path to the TensorBoard run directory.") + parser.add_argument("--csv", help="Optional: Path to save all scalar data as a CSV.", default=None) + + args = parser.parse_args() + + scalar_data = explore_run(args.log_dir) + + if args.csv and scalar_data: + # Reloading for wall_time and steps + ea = event_accumulator.EventAccumulator(args.log_dir).Reload() + with open(args.csv, mode='w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"]) + writer.writeheader() + for tag in scalar_data.keys(): + for e in ea.Scalars(tag): + writer.writerow({ + "tag": tag, + "step": e.step, + "value": e.value, + "wall_time": e.wall_time + }) + + print(f"\nData exported to: {args.csv}") + +if __name__ == "__main__": + main() diff --git a/scripts/export_hpc_requirements.py b/scripts/export_hpc_requirements.py index fcc9850..078d54e 100644 --- a/scripts/export_hpc_requirements.py +++ b/scripts/export_hpc_requirements.py @@ -49,7 +49,18 @@ def main() -> None: with pyproject_path.open("rb") as f: data = tomllib.load(f) - deps: list[str] = data.get("project", {}).get("dependencies", []) + # Collect all dependencies, merging 'cuda' extras into base dependencies + dep_dict: dict[str, str] = {} + for dep in data.get("project", {}).get("dependencies", []): + dep_dict[normalise(pkg_name(dep))] = dep + + # Add cuda extras (takes precedence for HPC) + optional_deps = data.get("project", {}).get("optional-dependencies", {}) + for group in ["cuda"]: + for dep in optional_deps.get(group, []): + dep_dict[normalise(pkg_name(dep))] = dep + + deps = list(dep_dict.values()) final_deps: list[str] = [] print(f"Checking dependencies against HPC module list...", file=sys.stderr)