import argparse import copy import csv import json import logging import subprocess import sys from datetime import datetime from pathlib import Path import yaml REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from train.train import compute_split_baselines, list_split_batches logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") def _load_yaml(path): with open(path, "r", encoding="utf-8") as handle: return yaml.safe_load(handle) def _write_yaml(path, payload): path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(payload, handle, sort_keys=False) def _deep_update(target, updates): for key, value in updates.items(): if isinstance(value, dict) and isinstance(target.get(key), dict): _deep_update(target[key], value) else: target[key] = value def _slugify(name): return "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in name) def _to_jsonable(value): if isinstance(value, dict): return {key: _to_jsonable(item) for key, item in value.items()} if isinstance(value, (list, tuple)): return [_to_jsonable(item) for item in value] if hasattr(value, "tolist"): return value.tolist() return value def _torchrun_path(): candidate = Path(sys.executable).with_name("torchrun") if candidate.exists(): return str(candidate) return "torchrun" def _run_command(command, log_path, cwd): log_path.parent.mkdir(parents=True, exist_ok=True) logging.info("Running: %s", " ".join(command)) with open(log_path, "w", encoding="utf-8") as handle: process = subprocess.run(command, cwd=cwd, stdout=handle, stderr=subprocess.STDOUT, text=True) return process.returncode def _find_run_dir(prefix_path): parent = prefix_path.parent pattern = prefix_path.name + "_*" candidates = [path for path in parent.glob(pattern) if path.is_dir()] if not candidates: raise FileNotFoundError(f"No run directory found for prefix {str(prefix_path)}") candidates.sort(key=lambda path: path.stat().st_mtime, reverse=True) return candidates[0] def _write_summary_csv(path, rows): path.parent.mkdir(parents=True, exist_ok=True) fieldnames = [ "experiment", "seed", "status", "run_dir", "weights_path", "model_miou", "model_mf1", "majority_miou", "majority_mf1", "anyview_miou", "anyview_mf1", "metrics_path", "log_train", "log_test", ] with open(path, "w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() for row in rows: writer.writerow({key: row.get(key, "") for key in fieldnames}) def build_experiment_config(base_config, experiment, seed, output_root, defaults=None): cfg = copy.deepcopy(base_config) _deep_update(cfg, experiment.get("overrides", {})) cfg["training"]["seed"] = int(seed) cfg.setdefault("data", {}) cfg["data"].setdefault("split_roots", {}) defaults = defaults or {} for split_name in ("train", "val", "test"): split_key = f"{split_name}_batches_root" split_root = experiment.get(split_key, defaults.get(split_key)) if split_root is None: split_root = base_config.get("data", {}).get("split_roots", {}).get(split_name) if split_root: cfg["data"]["split_roots"][split_name] = str(split_root) # Weekend sweeps already compute split baselines once globally and again at test time. # Disabling per-train baseline recomputation avoids redundant DDP startup work. cfg.setdefault("training", {}) cfg["training"]["compute_validation_baselines"] = bool( experiment.get("compute_validation_baselines", False) ) run_prefix = output_root / "runs" / f"{_slugify(experiment['name'])}_seed{seed}" cfg["training"]["output_dir"] = str(run_prefix) return cfg, run_prefix def build_train_command(config_path, nproc_per_node): if int(nproc_per_node) > 1: return [ _torchrun_path(), "--standalone", f"--nproc_per_node={int(nproc_per_node)}", "main.py", "--config", str(config_path), "--mode", "train", ] return [sys.executable, "main.py", "--config", str(config_path), "--mode", "train"] def build_test_command(config_path, weights_path, split_name, metrics_output, limit_files=None): command = [ sys.executable, "main.py", "--config", str(config_path), "--mode", "test", "--weights_path", str(weights_path), "--split", split_name, "--metrics_output", str(metrics_output), ] if limit_files is not None: command.extend(["--limit_files", str(int(limit_files))]) return command def main(): parser = argparse.ArgumentParser() parser.add_argument("--base-config", default="configs/config_deepchoice_base.yaml") parser.add_argument("--plan", default="experiments/weekend_plan.yaml") parser.add_argument("--output-root", default=None) parser.add_argument("--continue-on-error", action="store_true") args = parser.parse_args() repo_root = REPO_ROOT base_config = _load_yaml(repo_root / args.base_config) if not Path(args.base_config).is_absolute() else _load_yaml(args.base_config) plan = _load_yaml(repo_root / args.plan) if not Path(args.plan).is_absolute() else _load_yaml(args.plan) output_root = Path(args.output_root) if args.output_root else repo_root / "artifacts" / "experiments" / datetime.now().strftime("weekend_%Y%m%d_%H%M%S") output_root.mkdir(parents=True, exist_ok=True) defaults = plan.get("defaults", {}) summary_rows = [] baselines = None baseline_cfg = copy.deepcopy(base_config) baseline_eval = plan.get("baseline_evaluation", {}) if baseline_eval.get("enabled", False): split_name = baseline_eval.get("split", defaults.get("split", "test")) try: _deep_update(baseline_cfg, baseline_eval.get("overrides", {})) baseline_cfg.setdefault("data", {}) baseline_cfg["data"].setdefault("split_roots", {}) for split_key in ("train", "val", "test"): default_root = defaults.get(f"{split_key}_batches_root") if default_root: baseline_cfg["data"]["split_roots"][split_key] = str(default_root) baseline_limit = baseline_eval.get("limit_files") baseline_paths = list_split_batches(baseline_cfg, split_name, limit=baseline_limit) baselines = compute_split_baselines( baseline_cfg, paths=baseline_paths, file_batch_size=baseline_cfg.get("test", {}).get("file_batch_size", baseline_cfg["training"].get("eval_file_batch_size", 1)), desc=f"Computing {split_name} baselines", ) baseline_path = output_root / "metrics" / f"baselines_{split_name}.json" baseline_path.parent.mkdir(parents=True, exist_ok=True) baseline_path.write_text(json.dumps(_to_jsonable(baselines), indent=2), encoding="utf-8") summary_rows.append( { "experiment": f"baseline_{split_name}", "seed": "", "status": "ok", "run_dir": "", "weights_path": "", "model_miou": "", "model_mf1": "", "majority_miou": baselines["majority"]["miou"], "majority_mf1": baselines["majority"]["mf1"], "anyview_miou": baselines["anyview"]["miou"], "anyview_mf1": baselines["anyview"]["mf1"], "metrics_path": str(baseline_path), "log_train": "", "log_test": "", } ) except Exception as exc: logging.exception("Baseline evaluation failed on split %s", split_name) summary_rows.append( { "experiment": f"baseline_{split_name}", "seed": "", "status": f"failed: {exc}", } ) if not args.continue_on_error: raise for experiment in plan.get("experiments", []): exp_name = experiment["name"] split_name = experiment.get("split", defaults.get("split", "test")) test_limit_files = experiment.get("limit_files", defaults.get("limit_files")) nproc = experiment.get("nproc_per_node", defaults.get("nproc_per_node", 1)) run_test_after_train = experiment.get("run_test_after_train", defaults.get("run_test_after_train", True)) for seed in experiment.get("seeds", [base_config["training"].get("seed", 42)]): cfg, run_prefix = build_experiment_config(base_config, experiment, seed, output_root, defaults=defaults) config_path = output_root / "configs" / f"{_slugify(exp_name)}__seed{seed}.yaml" train_log = output_root / "logs" / f"{_slugify(exp_name)}__seed{seed}__train.log" test_log = output_root / "logs" / f"{_slugify(exp_name)}__seed{seed}__test.log" metrics_path = output_root / "metrics" / f"{_slugify(exp_name)}__seed{seed}__test.json" _write_yaml(config_path, cfg) row = { "experiment": exp_name, "seed": seed, "status": "pending", "run_dir": "", "weights_path": "", "metrics_path": str(metrics_path), "log_train": str(train_log), "log_test": str(test_log), } try: if ( baselines is not None and baseline_eval.get("enabled", False) and split_name == baseline_eval.get("split", defaults.get("split", "test")) ): cfg.setdefault("training", {}) cfg["training"]["precomputed_validation_baselines"] = _to_jsonable(baselines) _write_yaml(config_path, cfg) train_cmd = build_train_command(config_path, nproc) train_rc = _run_command(train_cmd, train_log, repo_root) if train_rc != 0: raise RuntimeError(f"Training failed with return code {train_rc}") run_dir = _find_run_dir(run_prefix) weights_path = run_dir / "best_model.pt" if not weights_path.exists(): raise FileNotFoundError(f"Missing checkpoint {str(weights_path)}") row["run_dir"] = str(run_dir) row["weights_path"] = str(weights_path) if run_test_after_train: test_cmd = build_test_command(config_path, weights_path, split_name, metrics_path, limit_files=test_limit_files) test_rc = _run_command(test_cmd, test_log, repo_root) if test_rc != 0: raise RuntimeError(f"Test failed with return code {test_rc}") metrics = json.loads(metrics_path.read_text(encoding="utf-8")) row["model_miou"] = metrics["miou"] row["model_mf1"] = metrics["mf1"] row["majority_miou"] = metrics["baselines"]["majority"]["miou"] row["majority_mf1"] = metrics["baselines"]["majority"]["mf1"] row["anyview_miou"] = metrics["baselines"]["anyview"]["miou"] row["anyview_mf1"] = metrics["baselines"]["anyview"]["mf1"] row["status"] = "ok" except Exception as exc: logging.exception("Experiment failed: %s seed=%s", exp_name, seed) row["status"] = f"failed: {exc}" summary_rows.append(row) _write_summary_csv(output_root / "summary.csv", summary_rows) (output_root / "summary.json").write_text(json.dumps(_to_jsonable(summary_rows), indent=2), encoding="utf-8") if not args.continue_on_error: raise continue summary_rows.append(row) _write_summary_csv(output_root / "summary.csv", summary_rows) (output_root / "summary.json").write_text(json.dumps(_to_jsonable(summary_rows), indent=2), encoding="utf-8") logging.info("Weekend experiment run complete. Summary written under %s", output_root) if __name__ == "__main__": main()