Files
scene-actor-extraction/scripts/optimizer/optimize.py
T
dtourolle d113c83189 feat(optimizer): sweep expansion bands + presence mode; tolerate scattered dropped votes
Make the flood-fill and expansion knobs reachable from the DE sweep:

- kpn_bindings: read expand_band_lo/hi and presence_mode from the replay
  cfg dict (presence_mode accepts "flood"/"track_extent" or a numeric
  >=0.5 toggle), and carry is_scene_boundary onto the replayed frame.
- replay.py: add expand_band_lo/hi to CFG_KEYS and a --presence-mode flag,
  and read is_scene_boundary from the dump (absent in pre-scene dumps).
- optimize.py: map the continuous presence_flood knob (0..1, >=0.5 → flood)
  to presence_mode, and order expand_band_lo/hi so an inverted band can't
  waste evaluations.

Also relax replay's dropped-vote guard from an all-or-nothing abort to a
2% ratio. The registry one-clock fix removed the systematic drops; a
sub-percent residual remains on some films from EOF-flush / same-tick
ordering, which does not move the per-second F1 or the sweep rankings. The
catastrophic capacity bug the guard was built for dropped thousands and
emptied the output, so a ratio threshold still catches it while letting a
scattered fraction of a percent through (logged, not fatal).
2026-08-09 10:21:45 +02:00

297 lines
14 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
optimize.py — Differential Evolution over pipeline thresholds, scored against X-Ray.
Replaces the coarse grid sweep with scipy's differential_evolution over the
continuous knob space. Each candidate config is a full-9-film replay (real KPN
nodes) scored against Amazon X-Ray presence, micro-averaged. The gallery is loaded
once per process (binding caches by path), so an evaluation is just N cheap replays.
Objective: **maximize micro-F1** (DE minimizes, so we return -F1). NOTE: X-Ray recall
is a face-vs-cast-in-scene ceiling (see [[xray-validation-results]]), so unconstrained
F1 tends to push prob_threshold DOWN to recover unreachable recall — trading real
precision for it. We therefore log precision/recall at every evaluation and print
them at the optimum so the trade-off is visible and you can pick another operating
point from the trajectory (--trajectory).
Usage:
python scripts/optimizer/optimize.py --manifest films.json \
--gallery gallery_arcface_w600k_r50.json \
--params prob_threshold:0.5:0.999 ownership_logodds:0.5:4 track_alpha:0:1 \
--popsize 20 --maxiter 25 --trajectory traj.json
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import numpy as np
from scipy.optimize import differential_evolution
REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts" / "optimizer"))
sys.path.insert(0, str(REPO / "scripts" / "validation"))
sys.path.insert(0, str(REPO / "scripts"))
import json as _json
import os
import subprocess
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor
# Concurrent per-eval replays. Each replay is an isolated subprocess, so parallelism
# is deadlock-safe; with 9 films/eval, 8 workers replays nearly all at once. Tune via
# REPLAY_WORKERS (8 is the measured sweet spot on this 24GB GPU).
REPLAY_WORKERS = int(os.environ.get("REPLAY_WORKERS", "8"))
# DE-level parallelism: how many population candidates get evaluated concurrently
# (each spawning its own REPLAY_WORKERS film subprocesses). Total concurrent GPU
# replay processes ≈ DE_WORKERS × min(REPLAY_WORKERS, n_films). Threads, not
# multiprocessing — each objective() call just waits on subprocess.run, so threads
# share the GIL fine and avoid pickling the objective/gallery-key cache.
DE_WORKERS = int(os.environ.get("DE_WORKERS", "1"))
from second_score import score_seconds # noqa: E402 uniform per-second TPI/FPI scoring
from sample_eval import load_gallery_keys # noqa: E402
from replay import dump_embedder_stamp # noqa: E402
from sae_stamp import EmbedderMismatch, verify_gallery_stamp # noqa: E402
_GAL_KEYS: dict = {} # gallery path → key set (fair-recall FN mask), loaded once
_REPLAY_TIMEOUT = 45 # seconds per film; a wedged replay is killed, not left to hang
REPLAY_CLI = str(Path(__file__).resolve().parent / "replay.py")
def _gallery_keys(path):
if path not in _GAL_KEYS:
_GAL_KEYS[path] = load_gallery_keys(path)
return _GAL_KEYS[path]
def _replay_subprocess(dump, gallery, cfg, build_dir):
"""Run one replay in a SUBPROCESS with a timeout, returning its presence JSON.
In-process replay intermittently DEADLOCKS at network teardown — a KPN worker
stuck mid-rocBLAS GEMM inside the ROCm driver makes ~PyNode's jthread.join() hang
forever (root-caused via gdb, 2026-07-15). Isolating each replay means a wedged
GPU thread only kills that subprocess; the sweep continues. Returns None on
timeout/failure (the caller drops that film from the average)."""
with tempfile.NamedTemporaryFile("r", suffix=".json", delete=False) as tf:
out = tf.name
argv = [sys.executable, REPLAY_CLI, "--dump", dump, "--gallery", gallery,
"--out", out, "--build-dir", build_dir]
for k, v in cfg.items():
if isinstance(v, bool): # store_true flags: pass the flag, not a value
if v:
argv.append(f"--{k.replace('_', '-')}")
else:
argv += [f"--{k.replace('_', '-')}", str(v)]
try:
subprocess.run(argv, timeout=_REPLAY_TIMEOUT, capture_output=True, check=True)
return _json.loads(Path(out).read_text())
except (subprocess.TimeoutExpired, subprocess.CalledProcessError,
FileNotFoundError, ValueError) as e:
print(f"[opt] replay failed for {Path(dump).name}: {type(e).__name__}",
file=sys.stderr)
return None
finally:
try:
Path(out).unlink()
except OSError:
pass
def evaluate(cfg, films, build_dir, step=None):
"""Objective = MACRO-mean over films of each film's duration-weighted per-scene F1.
Each film's replay runs in a subprocess (timeout-guarded) to survive the
intermittent ROCm teardown deadlock. If ANY film's replay times out, this
evaluation is scored f1=0.0 (see below) rather than averaging over the
survivors — a partial-coverage eval must never look better than a complete
one, or DE will converge onto configs that make the hardest film time out.
(An earlier version averaged over survivors, which silently rewarded
truncation; the rep4 `mbf_full_noexp` winner was one such corrupted eval.)
UNIFORM PER-SECOND scoring (second_score.py): every second of the film is sampled;
GT(t) = the cast of the X-Ray scene containing t, Pred(t) = actors whose presence
window covers t. Counts instances — TPI / FPI / FN — with FPI weighted 10× when the
named actor isn't in the film's cast at all (a real misID vs a timing slip). FN
counts only gallery-known actors (fair recall). Reports agreement_rate = mean
per-second Jaccard (the "% of on-screen actors we agree with X-Ray about, over
time"). Objective = macro-mean across films of the per-second weighted F1.
expand_gallery: controlled by env SAE_EXPAND (default on). Set SAE_EXPAND=0 to run
the no-expansion arm — the overnight matrix tests both to quantify what expansion buys.
The 9 films' replays run CONCURRENTLY (REPLAY_WORKERS) — each is an isolated
subprocess, so parallelism is safe (a wedged one only kills itself)."""
if os.environ.get("SAE_EXPAND", "1") == "1":
cfg = {**cfg, "expand_gallery": True}
def _one(film):
pj = _replay_subprocess(film["dump"], film.get("gallery"), cfg, build_dir)
if pj is None:
return None
return score_seconds(pj, film["xray"],
gallery_keys=_gallery_keys(film.get("gallery")))
with ThreadPoolExecutor(max_workers=REPLAY_WORKERS) as ex:
per_film = [m for m in ex.map(_one, films) if m is not None]
n = len(per_film)
n_expected = len(films)
# Incomplete coverage (a replay timed out) is scored as a failure, not
# averaged over survivors: dropping the hardest film would otherwise inflate
# the score and let DE reward exactly the configs that cause timeouts. We
# still record the real survivor counts so a truncated eval is diagnosable
# in the trajectory (f1=0.0, films_scored < films_expected).
if n < n_expected:
agg = {"precision": 0.0, "recall": 0.0, "f1": 0.0, "agreement": 0.0,
"TPI": sum(m["TPI"] for m in per_film),
"FPI": sum(m["FPI"] for m in per_film),
"FPI_misid": sum(m["FPI_misid"] for m in per_film),
"FN": sum(m["FN"] for m in per_film)}
agg["films_scored"] = n
agg["films_expected"] = n_expected
return agg
return {"precision": sum(m["precision"] for m in per_film) / n,
"recall": sum(m["recall"] for m in per_film) / n,
"f1": sum(m["f1"] for m in per_film) / n,
"agreement": sum(m["agreement_rate"] for m in per_film) / n,
"TPI": sum(m["TPI"] for m in per_film),
"FPI": sum(m["FPI"] for m in per_film),
"FPI_misid": sum(m["FPI_misid"] for m in per_film),
"FN": sum(m["FN"] for m in per_film),
"films_scored": n,
"films_expected": n_expected}
def main():
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--manifest", required=True)
p.add_argument("--gallery", help="default gallery if not per-film")
p.add_argument("--params", nargs="+", required=True,
help="knob:lo:hi (e.g. prob_threshold:0.5:0.999). Int knobs kept float, rounded in cfg.")
p.add_argument("--build-dir", default=str(REPO / "build"))
p.add_argument("--step", type=float, default=5.0)
p.add_argument("--popsize", type=int, default=20)
p.add_argument("--maxiter", type=int, default=25)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--trajectory", help="write every evaluation here (JSON lines)")
p.add_argument("--out", help="write best config + metrics")
# TRACES: GR-004 | SR-001
p.add_argument("--require-gallery-stamp", action="store_true",
help="unprovable gallery/dump model binding is a hard error, "
"not a warning (also via SAE_REQUIRE_GALLERY_STAMP=1)")
args = p.parse_args()
if args.require_gallery_stamp:
# Set the env var rather than threading a flag through cfg: replays run as
# subprocesses and inherit it, so strictness cannot be lost in the handoff.
os.environ["SAE_REQUIRE_GALLERY_STAMP"] = "1"
films = json.loads(Path(args.manifest).read_text())
for f in films:
f.setdefault("gallery", args.gallery)
if not Path(f["dump"]).exists():
sys.exit(f"[opt] missing dump for {f['name']}: {f['dump']}")
# TRACES: GR-004 | SR-001
# every (dump, gallery) pair is checked ONCE here,
# before the first evaluation. A DE sweep is thousands of replays; discovering
# a cross-model pair at the end (or never) means every number it produced was
# noise. Each replay subprocess re-checks its own pair anyway.
for f in films:
try:
verify_gallery_stamp(f["gallery"], stamp=dump_embedder_stamp(f["dump"]),
embedder_desc=f"embedding dump {Path(f['dump']).name}",
require_stamp=args.require_gallery_stamp)
except EmbedderMismatch as e:
sys.exit(f"[opt] {f['name']}: {e}")
names, bounds = [], []
int_knobs = {"track_max_frames_missing", "cut_inactive_max_frames"}
for spec in args.params:
k, lo, hi = spec.split(":")
names.append(k); bounds.append((float(lo), float(hi)))
print(f"[opt] DE over {names} bounds={bounds}", file=sys.stderr)
print(f"[opt] {len(films)} films, popsize={args.popsize}, maxiter={args.maxiter}", file=sys.stderr)
traj = []
evals = [0]
t0 = time.time()
traj_lock = threading.Lock()
def vec_to_cfg(x):
cfg = {}
for k, v in zip(names, x):
cfg[k] = int(round(v)) if k in int_knobs else float(v)
# The expansion band is [lo, hi]; independent DE bounds can invert it,
# and an inverted band admits nothing (track_gallery.hpp). Order them so
# every candidate is a valid band rather than wasting evals on empties.
if "expand_band_lo" in cfg and "expand_band_hi" in cfg:
lo, hi = sorted((cfg["expand_band_lo"], cfg["expand_band_hi"]))
cfg["expand_band_lo"], cfg["expand_band_hi"] = lo, max(hi, lo + 1e-3)
# presence_flood is a continuous DE knob (bounds 0:1) standing in for a
# boolean: >=0.5 selects flood-fill presence. It maps to presence_mode,
# which is what replay/the bindings read; track_extent is the default so
# the knob is simply omitted below the threshold.
if "presence_flood" in cfg:
flood = cfg.pop("presence_flood") >= 0.5
if flood:
cfg["presence_mode"] = "flood"
return cfg
def objective(x):
cfg = vec_to_cfg(x)
m = evaluate(cfg, films, args.build_dir, args.step)
with traj_lock:
evals[0] += 1
rec = {"eval": evals[0], "config": cfg, **m, "t": round(time.time() - t0, 1)}
traj.append(rec)
print(f"[opt] eval {evals[0]:3d} thr={cfg['prob_threshold']:.2f} "
f"own={cfg.get('ownership_logodds', float('nan')):.2f} → "
f"F1={m['f1']*100:.1f}% P={m['precision']*100:.1f}% R={m['recall']*100:.1f}% "
f"agree={m.get('agreement', 0)*100:.1f}% misID={m.get('FPI_misid', 0)}",
file=sys.stderr)
if args.trajectory:
with open(args.trajectory, "a") as tf:
tf.write(json.dumps(rec) + "\n")
return -m["f1"]
de_kwargs = dict(
popsize=args.popsize, maxiter=args.maxiter,
seed=args.seed, polish=False, tol=1e-4, mutation=(0.5, 1.0), recombination=0.7,
init="sobol")
if DE_WORKERS > 1:
pool = ThreadPoolExecutor(max_workers=DE_WORKERS)
de_kwargs["workers"] = pool.map
result = differential_evolution(objective, bounds, **de_kwargs)
best_cfg = vec_to_cfg(result.x)
best = evaluate(best_cfg, films, args.build_dir, args.step)
print("\n══ DE optimum (by F1) ═══════════════════════════")
print(f" config : {best_cfg}")
print(f" F1 : {best['f1']*100:.2f}%")
print(f" precision: {best['precision']*100:.2f}% recall: {best['recall']*100:.2f}%")
print(f" TP/FP/FN: {best['TP']}/{best['FP']}/{best['FN']}")
print(f" evaluations: {evals[0]} time: {time.time()-t0:.0f}s")
# Also surface the highest-precision config seen (the ship-safe operating point).
if traj:
hp = max(traj, key=lambda r: (r["precision"], r["recall"]))
print("\n── highest-precision config seen (ship-safe) ──")
print(f" config : {hp['config']}")
print(f" P={hp['precision']*100:.2f}% R={hp['recall']*100:.2f}% F1={hp['f1']*100:.2f}%")
if args.out:
Path(args.out).write_text(json.dumps(
{"best_by_f1": {"config": best_cfg, **best}, "n_evals": evals[0]}, indent=2))
if __name__ == "__main__":
main()