133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""CUDA-train an isolated, non-promotable YOLO candidate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
EXPERIMENTAL_ROOT = Path("/app/storage/training/experimental")
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _inside(path: Path, root: Path) -> bool:
|
|
try:
|
|
path.resolve().relative_to(root.resolve())
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def validate_paths(dataset_dir: Path, output_dir: Path) -> dict:
|
|
if not _inside(dataset_dir, EXPERIMENTAL_ROOT) or not _inside(output_dir, EXPERIMENTAL_ROOT):
|
|
raise ValueError(f"dataset and output must remain below {EXPERIMENTAL_ROOT}")
|
|
marker_path = dataset_dir / "EXPERIMENTAL_ONLY.json"
|
|
yaml_path = dataset_dir / "dataset.yaml"
|
|
if not marker_path.is_file() or not yaml_path.is_file():
|
|
raise ValueError("experimental marker or dataset.yaml is missing")
|
|
marker = json.loads(marker_path.read_text(encoding="utf-8"))
|
|
if marker.get("promotion_allowed") is not False or marker.get("release_claim_allowed") is not False:
|
|
raise ValueError("experimental marker does not prohibit promotion and release claims")
|
|
if marker.get("dataset_yaml_sha256") != sha256(yaml_path):
|
|
raise ValueError("experimental dataset.yaml checksum mismatch")
|
|
return marker
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--dataset-dir", type=Path, required=True)
|
|
parser.add_argument("--base-model", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--run-name", default="v73-flanders-remediation")
|
|
parser.add_argument("--epochs", type=int, default=20)
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
parser.add_argument("--batch", type=int, default=8)
|
|
parser.add_argument("--learning-rate", type=float, default=0.0001)
|
|
parser.add_argument("--freeze", type=int, default=10)
|
|
parser.add_argument("--workers", type=int, default=0)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
validate_paths(args.dataset_dir, args.output_dir)
|
|
if not args.base_model.is_file():
|
|
raise SystemExit(f"base model missing: {args.base_model}")
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("CUDA is required for GeoIntel experimental training")
|
|
model = YOLO(str(args.base_model))
|
|
result = model.train(
|
|
data=str(args.dataset_dir / "dataset.yaml"),
|
|
epochs=args.epochs,
|
|
imgsz=args.imgsz,
|
|
batch=args.batch,
|
|
workers=args.workers,
|
|
device=0,
|
|
project=str(args.output_dir),
|
|
name=args.run_name,
|
|
exist_ok=False,
|
|
pretrained=True,
|
|
deterministic=True,
|
|
seed=73,
|
|
amp=False,
|
|
val=False,
|
|
optimizer="AdamW",
|
|
lr0=args.learning_rate,
|
|
lrf=0.1,
|
|
freeze=args.freeze,
|
|
mosaic=0.0,
|
|
translate=0.05,
|
|
scale=0.1,
|
|
plots=False,
|
|
verbose=True,
|
|
)
|
|
run_dir = Path(result.save_dir)
|
|
best = run_dir / "weights" / "best.pt"
|
|
if not best.is_file():
|
|
best = run_dir / "weights" / "last.pt"
|
|
candidate = run_dir / "candidate.experimental.pt"
|
|
shutil.copy2(best, candidate)
|
|
summary = {
|
|
"schema_version": 1,
|
|
"status": "trained_experimental_only",
|
|
"promotion_allowed": False,
|
|
"release_claim_allowed": False,
|
|
"human_review_pending": True,
|
|
"dataset_yaml_sha256": sha256(args.dataset_dir / "dataset.yaml"),
|
|
"base_model": str(args.base_model),
|
|
"base_model_sha256": sha256(args.base_model),
|
|
"candidate_model": str(candidate),
|
|
"candidate_model_sha256": sha256(candidate),
|
|
"epochs": args.epochs,
|
|
"imgsz": args.imgsz,
|
|
"batch": args.batch,
|
|
"learning_rate": args.learning_rate,
|
|
"optimizer": "AdamW",
|
|
"freeze": args.freeze,
|
|
"workers": args.workers,
|
|
"device": torch.cuda.get_device_name(0),
|
|
"required_next_gate": "independent frozen V72 calibration evaluation",
|
|
}
|
|
(run_dir / "experimental_training_summary.json").write_text(
|
|
json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
|
)
|
|
print(json.dumps(summary, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|