510 lines
20 KiB
Python
510 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""Run checkpointed CUDA train/evaluate iterations until gates pass or a batch yields."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def write_json(path: Path, value: dict[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(json.dumps(value, indent=2), encoding="utf-8")
|
|
temporary.replace(path)
|
|
|
|
|
|
def dataset_audit_failures(
|
|
audit: dict[str, Any],
|
|
train_quality_audit: dict[str, Any],
|
|
) -> list[str]:
|
|
"""Return automated corpus blockers while leaving final human review deferred."""
|
|
failures = [str(item) for item in audit.get("failures") or []]
|
|
status = audit.get("status")
|
|
if status not in {"ok", "needs_human_review"}:
|
|
failures.append(f"unsupported audit status: {status}")
|
|
if audit.get("manifest_immutable") is not True:
|
|
failures.append("corpus manifest is not immutable")
|
|
if audit.get("spatial_leakage_status") != "ok":
|
|
failures.append("spatial leakage audit is not ok")
|
|
if train_quality_audit.get("status") != "ok":
|
|
failures.append("train tile quality audit is not ok")
|
|
if int(train_quality_audit.get("label_stats", {}).get("invalid_label_count", -1)) != 0:
|
|
failures.append("train tile quality audit contains invalid labels")
|
|
if int(train_quality_audit.get("label_stats", {}).get("missing_label_file_count", -1)) != 0:
|
|
failures.append("train tile quality audit contains missing label files")
|
|
if int(train_quality_audit.get("low_variance_positive_tile_count", -1)) != 0:
|
|
failures.append("dataset contains blank/low-variance positive tiles")
|
|
return failures
|
|
|
|
|
|
def select_calibration_threshold(report: dict[str, Any]) -> dict[str, Any]:
|
|
"""Choose a threshold without consulting test or background evidence."""
|
|
eligible = [item for item in report["sweeps"] if item["pure_empty_false_positives"] == 0]
|
|
if not eligible:
|
|
eligible = report["sweeps"]
|
|
return max(
|
|
eligible,
|
|
key=lambda item: (
|
|
min(region["f1"] for region in item["regions"].values()),
|
|
item["aggregate"]["f1"],
|
|
-item["pure_empty_false_positives"],
|
|
),
|
|
)
|
|
|
|
|
|
def calibration_failures(
|
|
chosen: dict[str, Any],
|
|
*,
|
|
min_aggregate_f1: float,
|
|
min_region_f1: float,
|
|
min_region_precision: float,
|
|
min_region_recall: float,
|
|
max_pure_empty_fp: int,
|
|
) -> list[str]:
|
|
failures: list[str] = []
|
|
if chosen["aggregate"]["f1"] < min_aggregate_f1:
|
|
failures.append("calibration_aggregate_f1_below_gate")
|
|
for region, values in chosen["regions"].items():
|
|
if values["f1"] < min_region_f1:
|
|
failures.append(f"calibration_{region}_f1_below_gate")
|
|
if values["precision"] < min_region_precision:
|
|
failures.append(f"calibration_{region}_precision_below_gate")
|
|
if values["recall"] < min_region_recall:
|
|
failures.append(f"calibration_{region}_recall_below_gate")
|
|
if chosen["pure_empty_false_positives"] > max_pure_empty_fp:
|
|
failures.append("calibration_pure_empty_false_positive_gate_failed")
|
|
return failures
|
|
|
|
|
|
def rejected_candidate_score(assessment: dict[str, Any]) -> tuple[float, ...]:
|
|
"""Rank rejected candidates by the weakest normalized release gate first."""
|
|
calibration = assessment["calibration"]
|
|
gates = assessment["gates"]
|
|
normalized: list[float] = [
|
|
calibration["aggregate"]["f1"] / gates["min_aggregate_f1"]
|
|
]
|
|
for metrics in calibration["regions"].values():
|
|
normalized.extend(
|
|
(
|
|
metrics["f1"] / gates["min_region_f1"],
|
|
metrics["precision"] / gates["min_region_precision"],
|
|
metrics["recall"] / gates["min_region_recall"],
|
|
)
|
|
)
|
|
normalized.sort()
|
|
return tuple(normalized)
|
|
|
|
|
|
def training_command(
|
|
yolo: str,
|
|
*,
|
|
model: Path,
|
|
data: Path,
|
|
project: Path,
|
|
name: str,
|
|
epochs: int,
|
|
patience: int = 18,
|
|
seed: int,
|
|
batch: int,
|
|
workers: int,
|
|
max_det: int = 1000,
|
|
imgsz: int = 640,
|
|
optimizer: str = "auto",
|
|
lr0: float | None = None,
|
|
mosaic: float = 1.0,
|
|
scale: float = 0.5,
|
|
translate: float = 0.1,
|
|
degrees: float = 0.0,
|
|
flipud: float = 0.0,
|
|
fliplr: float = 0.5,
|
|
warmup_epochs: float = 1.0,
|
|
warmup_bias_lr: float = 0.01,
|
|
hsv_h: float = 0.01,
|
|
hsv_s: float = 0.2,
|
|
hsv_v: float = 0.15,
|
|
) -> list[str]:
|
|
command = [
|
|
yolo,
|
|
"train",
|
|
f"model={model}",
|
|
f"data={data}",
|
|
f"epochs={epochs}",
|
|
f"imgsz={imgsz}",
|
|
f"batch={batch}",
|
|
"device=0",
|
|
f"workers={workers}",
|
|
f"patience={patience}",
|
|
"cache=disk",
|
|
"close_mosaic=20",
|
|
f"max_det={max_det}",
|
|
f"optimizer={optimizer}",
|
|
f"mosaic={mosaic}",
|
|
f"scale={scale}",
|
|
f"translate={translate}",
|
|
f"degrees={degrees}",
|
|
f"flipud={flipud}",
|
|
f"fliplr={fliplr}",
|
|
f"warmup_epochs={warmup_epochs}",
|
|
f"warmup_bias_lr={warmup_bias_lr}",
|
|
f"hsv_h={hsv_h}",
|
|
f"hsv_s={hsv_s}",
|
|
f"hsv_v={hsv_v}",
|
|
f"seed={seed}",
|
|
"deterministic=True",
|
|
f"project={project}",
|
|
f"name={name}",
|
|
"exist_ok=True",
|
|
]
|
|
if lr0 is not None:
|
|
command.append(f"lr0={lr0}")
|
|
return command
|
|
|
|
|
|
def failure_sampling_command(
|
|
*,
|
|
scripts_dir: Path,
|
|
train_summary: Path,
|
|
corpus_manifest: Path,
|
|
assessment: Path,
|
|
output_dir: Path,
|
|
sampling_round: int = 0,
|
|
) -> list[str]:
|
|
return [
|
|
sys.executable,
|
|
str(scripts_dir / "build_failure_driven_yolo_sampling.py"),
|
|
"--summary", str(train_summary),
|
|
"--corpus-manifest", str(corpus_manifest),
|
|
"--assessment", str(assessment),
|
|
"--output-dir", str(output_dir),
|
|
"--sampling-round", str(sampling_round),
|
|
]
|
|
|
|
|
|
def resumable_training_command(yolo: str, checkpoint: Path) -> list[str]:
|
|
return [yolo, "train", f"resume={checkpoint}", "device=0"]
|
|
|
|
|
|
def run(command: list[str], log_path: Path | None = None, *, allowed: set[int] = {0}) -> int:
|
|
if log_path:
|
|
log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with log_path.open("a", encoding="utf-8") as log:
|
|
completed = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=False)
|
|
else:
|
|
completed = subprocess.run(command, check=False)
|
|
if completed.returncode not in allowed:
|
|
raise RuntimeError(f"Command failed ({completed.returncode}): {' '.join(command)}")
|
|
return completed.returncode
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--initial-model", type=Path, required=True)
|
|
parser.add_argument("--train-yaml", type=Path, required=True)
|
|
parser.add_argument("--train-summary", type=Path, required=True)
|
|
parser.add_argument("--dataset-audit", type=Path, required=True)
|
|
parser.add_argument("--train-quality-audit", type=Path, required=True)
|
|
parser.add_argument("--calibration-summary", type=Path, required=True)
|
|
parser.add_argument("--test-summary", type=Path, required=True)
|
|
parser.add_argument("--background-summary", type=Path, required=True)
|
|
parser.add_argument("--corpus-manifest", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--epochs", type=int, default=160)
|
|
parser.add_argument("--patience", type=int, default=18)
|
|
parser.add_argument("--batch", type=int, default=2)
|
|
parser.add_argument("--workers", type=int, default=4)
|
|
parser.add_argument("--max-det", type=int, default=1000)
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
parser.add_argument(
|
|
"--thresholds",
|
|
type=float,
|
|
nargs="+",
|
|
default=[0.05, 0.075, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4],
|
|
help="Calibration-only confidence sweep used for threshold selection.",
|
|
)
|
|
parser.add_argument("--optimizer", default="auto")
|
|
parser.add_argument("--lr0", type=float)
|
|
parser.add_argument("--mosaic", type=float, default=1.0)
|
|
parser.add_argument("--scale", type=float, default=0.5)
|
|
parser.add_argument("--translate", type=float, default=0.1)
|
|
parser.add_argument("--degrees", type=float, default=0.0)
|
|
parser.add_argument("--flipud", type=float, default=0.0)
|
|
parser.add_argument("--fliplr", type=float, default=0.5)
|
|
parser.add_argument("--warmup-epochs", type=float, default=1.0)
|
|
parser.add_argument("--warmup-bias-lr", type=float, default=0.01)
|
|
parser.add_argument("--hsv-h", type=float, default=0.01)
|
|
parser.add_argument("--hsv-s", type=float, default=0.2)
|
|
parser.add_argument("--hsv-v", type=float, default=0.15)
|
|
parser.add_argument("--seed", type=int, default=20260731)
|
|
parser.add_argument("--yolo", default="yolo")
|
|
parser.add_argument("--min-aggregate-f1", type=float, default=0.55)
|
|
parser.add_argument("--min-region-f1", type=float, default=0.45)
|
|
parser.add_argument("--min-region-precision", type=float, default=0.5)
|
|
parser.add_argument("--min-region-recall", type=float, default=0.4)
|
|
parser.add_argument("--max-pure-empty-fp", type=int, default=0)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument(
|
|
"--evaluate-initial-model",
|
|
action="store_true",
|
|
help="Gate an already trained initial checkpoint before starting the next training iteration.",
|
|
)
|
|
args = parser.parse_args()
|
|
if args.iterations < 1:
|
|
raise SystemExit("--iterations must be positive")
|
|
dataset_audit = json.loads(args.dataset_audit.read_text(encoding="utf-8"))
|
|
train_quality_audit = json.loads(args.train_quality_audit.read_text(encoding="utf-8"))
|
|
audit_failures = dataset_audit_failures(dataset_audit, train_quality_audit)
|
|
if audit_failures:
|
|
raise SystemExit(f"Dataset audit is not eligible for training: {audit_failures}")
|
|
|
|
state_path = args.output_dir / "training-loop-state.json"
|
|
state: dict[str, Any] = {
|
|
"schema_version": 1,
|
|
"status": "running",
|
|
"started_at": datetime.now(UTC).isoformat(),
|
|
"initial_model": str(args.initial_model),
|
|
"train_yaml": str(args.train_yaml),
|
|
"dataset_audit": str(args.dataset_audit),
|
|
"dataset_audit_sha256": sha256(args.dataset_audit),
|
|
"train_quality_audit": str(args.train_quality_audit),
|
|
"train_quality_audit_sha256": sha256(args.train_quality_audit),
|
|
"corpus_manifest": str(args.corpus_manifest),
|
|
"iterations": [],
|
|
}
|
|
if state_path.is_file():
|
|
state = json.loads(state_path.read_text(encoding="utf-8"))
|
|
state["status"] = "running"
|
|
model = Path(state.get("next_model") or args.initial_model)
|
|
train_yaml = Path(state.get("next_train_yaml") or args.train_yaml)
|
|
first_index = len(state["iterations"]) + 1
|
|
scripts_dir = Path(__file__).resolve().parent
|
|
|
|
for offset in range(args.iterations):
|
|
index = first_index + offset
|
|
name = f"iteration-{index:03d}"
|
|
iteration_dir = args.output_dir / name
|
|
iteration_dir.mkdir(parents=True, exist_ok=True)
|
|
train_run = args.output_dir / "runs" / name
|
|
evaluate_existing = args.evaluate_initial_model and offset == 0 and not state["iterations"]
|
|
partial_checkpoint = train_run / "weights" / "last.pt"
|
|
resume_partial = not evaluate_existing and partial_checkpoint.is_file()
|
|
command = None if evaluate_existing else (
|
|
resumable_training_command(args.yolo, partial_checkpoint)
|
|
if resume_partial else training_command(
|
|
args.yolo,
|
|
model=model,
|
|
data=train_yaml,
|
|
project=args.output_dir / "runs",
|
|
name=name,
|
|
epochs=args.epochs,
|
|
patience=args.patience,
|
|
seed=args.seed + index,
|
|
batch=args.batch,
|
|
workers=args.workers,
|
|
max_det=args.max_det,
|
|
imgsz=args.imgsz,
|
|
optimizer=args.optimizer,
|
|
lr0=args.lr0,
|
|
mosaic=args.mosaic,
|
|
scale=args.scale,
|
|
translate=args.translate,
|
|
degrees=args.degrees,
|
|
flipud=args.flipud,
|
|
fliplr=args.fliplr,
|
|
warmup_epochs=args.warmup_epochs,
|
|
warmup_bias_lr=args.warmup_bias_lr,
|
|
hsv_h=args.hsv_h,
|
|
hsv_s=args.hsv_s,
|
|
hsv_v=args.hsv_v,
|
|
))
|
|
if args.dry_run:
|
|
print(json.dumps({
|
|
"training_command": command,
|
|
"evaluate_existing": evaluate_existing,
|
|
"resume_partial": resume_partial,
|
|
}, indent=2))
|
|
return 0
|
|
if evaluate_existing:
|
|
best = model
|
|
if not best.is_file():
|
|
raise RuntimeError(f"Initial checkpoint does not exist: {best}")
|
|
else:
|
|
assert command is not None
|
|
run(command, iteration_dir / "training.log")
|
|
best = train_run / "weights" / "best.pt"
|
|
if not best.is_file():
|
|
raise RuntimeError(f"Training produced no best checkpoint: {best}")
|
|
candidate = iteration_dir / "candidate.pt"
|
|
shutil.copy2(best, candidate)
|
|
|
|
reports: dict[str, Path] = {}
|
|
for role, summary in (("calibration", args.calibration_summary),):
|
|
report = iteration_dir / f"{role}.json"
|
|
reports[role] = report
|
|
run(
|
|
[
|
|
sys.executable,
|
|
str(scripts_dir / "evaluate_belgium_building_candidate.py"),
|
|
"--model",
|
|
str(candidate),
|
|
"--summary",
|
|
str(summary),
|
|
"--corpus-manifest",
|
|
str(args.corpus_manifest),
|
|
"--output",
|
|
str(report),
|
|
"--device",
|
|
"cuda:0",
|
|
"--max-det",
|
|
str(args.max_det),
|
|
"--imgsz",
|
|
str(args.imgsz),
|
|
"--thresholds",
|
|
*map(str, args.thresholds),
|
|
],
|
|
iteration_dir / f"{role}.log",
|
|
)
|
|
assessment = iteration_dir / "assessment.json"
|
|
calibration = json.loads(reports["calibration"].read_text(encoding="utf-8"))
|
|
chosen = select_calibration_threshold(calibration)
|
|
failures = calibration_failures(
|
|
chosen,
|
|
min_aggregate_f1=args.min_aggregate_f1,
|
|
min_region_f1=args.min_region_f1,
|
|
min_region_precision=args.min_region_precision,
|
|
min_region_recall=args.min_region_recall,
|
|
max_pure_empty_fp=args.max_pure_empty_fp,
|
|
)
|
|
if failures:
|
|
write_json(
|
|
assessment,
|
|
{
|
|
"schema_version": 1,
|
|
"status": "continue_training_loop",
|
|
"phase": "calibration_rejected",
|
|
"threshold_selection_source": "calibration_only",
|
|
"selected_threshold": chosen["threshold"],
|
|
"gates": {
|
|
"min_aggregate_f1": args.min_aggregate_f1,
|
|
"min_region_f1": args.min_region_f1,
|
|
"min_region_precision": args.min_region_precision,
|
|
"min_region_recall": args.min_region_recall,
|
|
"max_pure_empty_false_positives": args.max_pure_empty_fp,
|
|
},
|
|
"calibration": chosen,
|
|
"test": None,
|
|
"background": None,
|
|
"failures": failures,
|
|
},
|
|
)
|
|
else:
|
|
for role, summary in (
|
|
("test", args.test_summary),
|
|
("background", args.background_summary),
|
|
):
|
|
report = iteration_dir / f"{role}.json"
|
|
reports[role] = report
|
|
run(
|
|
[
|
|
sys.executable,
|
|
str(scripts_dir / "evaluate_belgium_building_candidate.py"),
|
|
"--model", str(candidate),
|
|
"--summary", str(summary),
|
|
"--corpus-manifest", str(args.corpus_manifest),
|
|
"--output", str(report),
|
|
"--device", "cuda:0",
|
|
"--max-det", str(args.max_det),
|
|
"--imgsz", str(args.imgsz),
|
|
"--thresholds", *map(str, args.thresholds),
|
|
],
|
|
iteration_dir / f"{role}.log",
|
|
)
|
|
run(
|
|
[
|
|
sys.executable,
|
|
str(scripts_dir / "assess_belgium_building_training_iteration.py"),
|
|
"--calibration", str(reports["calibration"]),
|
|
"--test", str(reports["test"]),
|
|
"--background", str(reports["background"]),
|
|
"--output", str(assessment),
|
|
],
|
|
iteration_dir / "assessment.log",
|
|
allowed={0, 2},
|
|
)
|
|
decision = json.loads(assessment.read_text(encoding="utf-8"))
|
|
record = {
|
|
"iteration": index,
|
|
"candidate": str(candidate),
|
|
"candidate_sha256": sha256(candidate),
|
|
"training_skipped_for_existing_checkpoint": evaluate_existing,
|
|
"training_resumed_from_partial_checkpoint": resume_partial,
|
|
"assessment": str(assessment),
|
|
"status": decision["status"],
|
|
"failures": decision["failures"],
|
|
}
|
|
state["iterations"].append(record)
|
|
score = rejected_candidate_score(decision) if decision["status"] != "training_complete" else ()
|
|
incumbent_score = tuple(state.get("incumbent_rejected_score", ()))
|
|
if not incumbent_score or score > incumbent_score:
|
|
state["incumbent_rejected_model"] = str(candidate)
|
|
state["incumbent_rejected_score"] = list(score)
|
|
record["promoted_to_training_incumbent"] = True
|
|
else:
|
|
record["promoted_to_training_incumbent"] = False
|
|
state["next_model"] = state.get("incumbent_rejected_model", str(candidate))
|
|
if decision["status"] == "training_complete":
|
|
state["status"] = "training_complete"
|
|
state["completed_at"] = datetime.now(UTC).isoformat()
|
|
write_json(state_path, state)
|
|
print(json.dumps(state, indent=2))
|
|
return 0
|
|
sampling_dir = iteration_dir / "failure-driven-training"
|
|
run(
|
|
failure_sampling_command(
|
|
scripts_dir=scripts_dir,
|
|
train_summary=args.train_summary,
|
|
corpus_manifest=args.corpus_manifest,
|
|
assessment=assessment,
|
|
output_dir=sampling_dir,
|
|
sampling_round=index,
|
|
),
|
|
iteration_dir / "failure-driven-sampling.log",
|
|
)
|
|
sampling_evidence = sampling_dir / "failure-driven-sampling.json"
|
|
next_train_yaml = sampling_dir / "dataset.yaml"
|
|
if not sampling_evidence.is_file() or not next_train_yaml.is_file():
|
|
raise RuntimeError("Failure-driven sampling produced incomplete evidence")
|
|
record["failure_driven_sampling"] = str(sampling_evidence)
|
|
record["failure_driven_sampling_sha256"] = sha256(sampling_evidence)
|
|
record["next_train_yaml"] = str(next_train_yaml)
|
|
state["next_train_yaml"] = str(next_train_yaml)
|
|
model = Path(state["next_model"])
|
|
train_yaml = next_train_yaml
|
|
write_json(state_path, state)
|
|
|
|
state["status"] = "continue_training_loop"
|
|
state["yielded_at"] = datetime.now(UTC).isoformat()
|
|
write_json(state_path, state)
|
|
print(json.dumps(state, indent=2))
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|