Delete raven/budget.py

raven/budget.py

@@ -1,176 +0,0 @@
-from __future__ import annotations
-
-import copy
-import json
-import math
-from pathlib import Path
-
-import yaml
-
-from .benchmark import benchmark_foundation
-from .config import load_config
-from .utils import read_jsonl
-
-
-def _counts(rows: list[dict]) -> dict:
-    out = {"train_real": 0, "train_ai": 0, "val": 0, "test": 0}
-    for r in rows:
-        split = str(r.get("split", "train"))
-        if split == "train":
-            out["train_ai" if int(r.get("label", 0)) == 1 else "train_real"] += 1
-        elif split == "val":
-            out["val"] += 1
-        elif split == "test":
-            out["test"] += 1
-    return out
-
-
-def _estimate(cfg: dict, counts: dict, ips: float) -> dict:
-    fc = cfg["foundation"]
-    rc = cfg["real_specialists"]
-    tc = cfg["training"]
-    cap = int(fc.get("train_cap_per_class", 0))
-    tr0 = counts["train_real"] if cap <= 0 else min(counts["train_real"], cap)
-    tr1 = counts["train_ai"] if cap <= 0 else min(counts["train_ai"], cap)
-    eval_n = counts["val"] + counts["test"]
-    base_n = tr0 + tr1 + eval_n
-    aug_per_class = int(fc.get("augment_max_per_class", 0))
-    aug_n = 0
-    if not bool(fc.get("reuse_semantic_for_augmented_views", True)):
-        aug_n = min(tr0, aug_per_class) + min(tr1, aug_per_class)
-    foundation_n = base_n + aug_n
-    foundation_h = foundation_n / max(ips, 1e-6) / 3600.0
-
-
-    runtime = str(fc.get("_planned_runtime", "unknown"))
-    overhead = {"directml": 1.65, "openvino": 1.80, "torch": 2.15}.get(runtime, 1.90)
-    cache_h = foundation_h * overhead
-
-    specialist_limit = int(rc.get("max_real_images", 0))
-    specialist_n = counts["train_real"] if specialist_limit <= 0 else min(counts["train_real"], specialist_limit)
-    specialist_epochs = int(rc.get("epochs", 12))
-    specialist_h = 0.05 + specialist_n * (0.030 + 0.0015 * specialist_epochs) / 3600.0
-
-
-    members = int(tc.get("ensemble_members", 5))
-    epochs = int(tc.get("epochs", 60))
-    refine = int(tc.get("hard_refine_epochs", 6))
-    selected_train = tr0 + tr1
-    fusion_h = 0.04 + (selected_train / 10000.0) * members * (epochs + refine) * 0.0018
-    calibration_h = 0.08
-    total = cache_h + specialist_h + fusion_h + calibration_h
-    return {
-        "selected_train_real": tr0,
-        "selected_train_ai": tr1,
-        "eval_images": eval_n,
-        "foundation_passes": foundation_n,
-        "foundation_hours_raw": foundation_h,
-        "cache_hours_est": cache_h,
-        "specialist_hours_est": specialist_h,
-        "fusion_hours_est": fusion_h,
-        "calibration_hours_est": calibration_h,
-        "total_hours_est": total,
-    }
-
-
-def plan_budget(
-    manifest: str,
-    base_config: str,
-    out_config: str,
-    target_hours: float = 4.5,
-    min_hours: float = 3.0,
-    max_hours: float = 6.0,
-    accelerator: str = "auto",
-    loops: int = 4,
-) -> dict:
-    rows = read_jsonl(manifest)
-    counts = _counts(rows)
-    cfg = copy.deepcopy(load_config(base_config))
-    bench = benchmark_foundation(cfg, accelerator, loops)
-    ips = float(bench["foundation_images_per_second"])
-    runtime = str(bench["selected_runtime"])
-    fc, rc, tc = cfg["foundation"], cfg["real_specialists"], cfg["training"]
-    fc["_planned_runtime"] = runtime
-
-
-    tc["ensemble_members"] = 5
-    tc["epochs"] = 60
-    tc["early_stop_patience"] = 7
-    tc["hard_refine_epochs"] = 6
-    tc["hard_refine_lr_multiplier"] = 0.22
-    tc["hard_replay_boost"] = 4.0
-    tc["contrastive_weight"] = 0.028
-    tc["branch_dropout"] = 0.12
-    tc["source_dropout_per_member"] = 1
-    rc["epochs"] = 12
-    rc["early_stop_patience"] = 3
-
-    rc["max_real_images"] = int(counts["train_real"])
-    fc["reuse_semantic_for_augmented_views"] = True
-    fc["augment_max_per_class"] = 6000
-    fc["train_coreset"] = True
-    fc["train_cap_per_class"] = 0
-
-
-    target = min(max(float(target_hours), float(min_hours)), float(max_hours))
-    runtime_factor = {"directml": 1.65, "openvino": 1.80, "torch": 2.15}.get(runtime, 1.90)
-    foundation_budget_h = target * 0.62
-    affordable_total = int(foundation_budget_h * 3600.0 * ips / runtime_factor)
-    eval_n = counts["val"] + counts["test"]
-    train_budget = max(2000, affordable_total - eval_n)
-
-
-    fc["train_cap_per_class"] = 0
-
-    est = _estimate(cfg, counts, ips)
-
-
-    if fc["train_cap_per_class"] == 0 and est["total_hours_est"] < min(target * 0.78, max_hours - 0.5):
-        remaining_h = max(0.0, target - est["total_hours_est"])
-        extra_passes = int(remaining_h * 3600.0 * ips / runtime_factor * 0.72)
-        per_cls_extra = min(counts["train_real"], counts["train_ai"], max(0, extra_passes // 2))
-        if per_cls_extra >= 500:
-            fc["reuse_semantic_for_augmented_views"] = False
-            fc["augment_max_per_class"] = int(min(6000, per_cls_extra))
-            est = _estimate(cfg, counts, ips)
-
-
-    if est["total_hours_est"] > max_hours:
-        fc["reuse_semantic_for_augmented_views"] = True
-        fc["augment_max_per_class"] = 0
-        tc["epochs"] = min(int(tc.get("epochs",60)), 50)
-        tc["hard_refine_epochs"] = min(int(tc.get("hard_refine_epochs",6)), 4)
-        fc["train_cap_per_class"] = 0
-        rc["max_real_images"] = int(counts["train_real"])
-        est = _estimate(cfg, counts, ips)
-
-    cfg["project"]["version"] = "5.1-full-data-debiased"
-    cfg["budget"] = {
-        "target_hours": target,
-        "minimum_hours": float(min_hours),
-        "maximum_hours": float(max_hours),
-        "benchmark_runtime": runtime,
-        "foundation_images_per_second": ips,
-        "planner_note": "Estimate is hardware-measured but still approximate; disk/image decode speed can shift wall time.",
-    }
-    fc.pop("_planned_runtime", None)
-    out = Path(out_config)
-    out.parent.mkdir(parents=True, exist_ok=True)
-    out.write_text(yaml.safe_dump(cfg, allow_unicode=True, sort_keys=False), encoding="utf-8")
-    result = {
-        "counts": counts,
-        "benchmark": bench,
-        "selected": {
-            "train_cap_per_class": int(fc.get("train_cap_per_class", 0)),
-            "reuse_semantic_for_augmented_views": bool(fc.get("reuse_semantic_for_augmented_views", True)),
-            "augment_max_per_class": int(fc.get("augment_max_per_class", 0)),
-            "real_specialist_max": int(rc.get("max_real_images", 0)),
-            "ensemble_members": int(tc.get("ensemble_members", 5)),
-            "hard_refine_epochs": int(tc.get("hard_refine_epochs", 0)),
-            "all_original_train_images": True,
-        },
-        "estimate": est,
-        "output_config": str(out),
-    }
-    Path(str(out) + ".plan.json").write_text(json.dumps(result, indent=2, ensure_ascii=False), encoding="utf-8")
-    return result