Spaces:
Sleeping
Sleeping
File size: 4,147 Bytes
cabc6bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Model-family benchmark: YOLOv8n / v8s / 11n / 11s on identical config (plan 3.1).
Trains each candidate on the same data, epochs, batch, image size, and seed as
the original yolov8s run, evaluates on the val split, and writes one
accuracy-vs-latency-vs-size table. Meant for a Colab GPU session:
!uv run python scripts/benchmark_models.py # full run (hours)
!uv run python scripts/benchmark_models.py --models yolov8n,yolo11n --epochs 10 # quick pass
!uv run python scripts/benchmark_models.py --skip-train # re-tabulate existing runs
Each model trains into runs/detect/bench_<model>/; --skip-train reuses those
weights, so an interrupted session can resume without retraining finished
models. Output table: docs/benchmark_models.md.
"""
from __future__ import annotations
import argparse
from pathlib import Path
from ultralytics import YOLO
DATA = "work/data.yaml"
# YOLO12 needs ultralytics >= 8.3.78; this repo pins 8.3.40, so the ladder tops
# out at YOLO11. The v8s row doubles as a sanity check against the original run.
MODELS = ["yolov8n", "yolov8s", "yolo11n", "yolo11s"]
def run_name(model: str) -> str:
return f"bench_{model}"
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--models", default=",".join(MODELS), help="comma-separated Ultralytics model names")
p.add_argument("--data", default=DATA)
# Defaults mirror the original training run (see runs/detect/yolov8n_v1_train/args.yaml).
p.add_argument("--epochs", type=int, default=35)
p.add_argument("--batch", type=int, default=32)
p.add_argument("--imgsz", type=int, default=640)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--device", default=None, help="e.g. '0' on Colab GPU; omit to auto-pick")
p.add_argument("--skip-train", action="store_true", help="only evaluate existing runs/detect/bench_* weights")
args = p.parse_args()
models = [m.strip() for m in args.models.split(",")]
rows: list[tuple[str, float, float, float, float, float, float]] = []
for name in models:
best = Path("runs/detect") / run_name(name) / "weights" / "best.pt"
if not args.skip_train and not best.exists():
print(f"\n=== Training {name} ===")
YOLO(f"{name}.pt").train(
data=args.data,
epochs=args.epochs,
batch=args.batch,
imgsz=args.imgsz,
seed=args.seed,
device=args.device,
name=run_name(name),
exist_ok=True,
)
if not best.exists():
print(f"skip {name}: {best} not found")
continue
print(f"\n=== Evaluating {name} ===")
model = YOLO(best)
metrics = model.val(data=args.data, split="val", imgsz=args.imgsz, device=args.device)
n_params = sum(x.numel() for x in model.model.parameters()) / 1e6
size_mb = best.stat().st_size / 1e6
latency = metrics.speed["inference"] # ms per image on this device
rows.append((name, n_params, size_mb, metrics.box.map50, metrics.box.map, latency, 1000 / latency))
lines = [
"# Model-family benchmark (plan 3.1)",
"",
f"_Identical config per model: {args.epochs} epochs, batch {args.batch}, imgsz {args.imgsz},"
f" seed {args.seed}, val split of `{args.data}`. Latency = Ultralytics val-time forward pass"
" on the training device — comparable within this table only._",
"",
"| Model | Params (M) | Size (MB) | mAP@0.5 | mAP@0.5:0.95 | Latency (ms) | FPS |",
"|---|---|---|---|---|---|---|",
]
for name, params, size, map50, map5095, ms, fps in rows:
lines.append(f"| {name} | {params:.1f} | {size:.1f} | {map50:.3f} | {map5095:.3f} | {ms:.1f} | {fps:.0f} |")
lines += ["", "TODO(human): one paragraph — which model would you deploy, and why?", ""]
out = Path("docs/benchmark_models.md")
out.write_text("\n".join(lines))
print(f"\nWrote {out}\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
|