File size: 5,048 Bytes
fc49781 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """Benchmark RF-DETR inference on CPU and, when available, CUDA.
Example:
python benchmark.py plan.jpg --device both --repeats 20 --output benchmark.json
"""
from __future__ import annotations
import argparse
import json
import platform
import statistics
import time
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from inference import DEFAULT_REPO_ID, FloorPlanDetector
def cpu_name() -> str:
cpuinfo = Path("/proc/cpuinfo")
if cpuinfo.is_file():
for line in cpuinfo.read_text(encoding="utf-8").splitlines():
if line.startswith("model name"):
return line.partition(":")[2].strip()
return platform.processor() or "CPU"
def synchronize(device: str) -> None:
if device == "cuda":
torch.cuda.synchronize()
def run_device(
image: np.ndarray,
*,
device: str,
repo_id: str,
revision: str | None,
model_dir: Path | None,
threshold: float,
warmups: int,
repeats: int,
) -> dict:
start = time.perf_counter()
detector = FloorPlanDetector(repo_id, revision=revision, model_dir=model_dir, device=device)
load_seconds = time.perf_counter() - start
for _ in range(warmups):
detector.model.predict(image, threshold=threshold)
synchronize(device)
samples_ms = []
result = None
for _ in range(repeats):
synchronize(device)
start = time.perf_counter()
result = detector.model.predict(image, threshold=threshold)
synchronize(device)
samples_ms.append((time.perf_counter() - start) * 1000)
median_ms = statistics.median(samples_ms)
return {
"device": device,
"hardware": torch.cuda.get_device_name(0) if device == "cuda" else cpu_name(),
"precision": "float32",
"model_load_seconds": round(load_seconds, 3),
"median_ms_per_image": round(median_ms, 2),
"p95_ms_per_image": round(float(np.percentile(samples_ms, 95)), 2),
"images_per_second_at_median": round(1000 / median_ms, 2),
"detections": len(result.confidence),
"samples_ms": [round(value, 2) for value in samples_ms],
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("image", type=Path)
parser.add_argument("--repo-id", default=DEFAULT_REPO_ID)
parser.add_argument("--revision", help="Hub commit SHA or tag to pin the checkpoint")
parser.add_argument("--model-dir", type=Path, help="Use a locally downloaded model")
parser.add_argument("--device", choices=("cpu", "cuda", "both"), default="both")
parser.add_argument("--threshold", type=float, default=0.35)
parser.add_argument("--warmups", type=int, default=3)
parser.add_argument("--repeats", type=int, default=20)
parser.add_argument("--output", type=Path, help="Optional JSON report path")
args = parser.parse_args()
if not args.image.is_file():
parser.error(f"image does not exist: {args.image}")
if not 0 <= args.threshold <= 1:
parser.error("--threshold must be between 0 and 1")
if args.warmups < 1 or args.repeats < 2:
parser.error("use at least one warm-up and two timed runs")
if args.device == "cuda" and not torch.cuda.is_available():
parser.error("CUDA GPU is unavailable on this machine")
with Image.open(args.image) as source:
image = np.array(source.convert("RGB"))
devices = ("cpu", "cuda") if args.device == "both" else (args.device,)
report = {
"model": args.repo_id,
"revision": args.revision or "main",
"image": args.image.name,
"image_width": image.shape[1],
"image_height": image.shape[0],
"input_resolution": 576,
"threshold": args.threshold,
"warmups": args.warmups,
"repeats": args.repeats,
"torch_version": torch.__version__,
"torch_threads": torch.get_num_threads(),
"results": {},
}
for device in devices:
if device == "cuda" and not torch.cuda.is_available():
report["results"][device] = {"status": "unavailable"}
print("CUDA: unavailable on this machine")
continue
result = run_device(
image,
device=device,
repo_id=args.repo_id,
revision=args.revision,
model_dir=args.model_dir,
threshold=args.threshold,
warmups=args.warmups,
repeats=args.repeats,
)
report["results"][device] = result
print(f"{device.upper()}: {result['median_ms_per_image']} ms median, "
f"{result['p95_ms_per_image']} ms p95, "
f"{result['images_per_second_at_median']} images/s, "
f"{result['detections']} detections")
if args.output:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(f"Report: {args.output}")
if __name__ == "__main__":
main()
|