bench.py

import time
import csv
import torch
from ultralytics import YOLO

# ==== SPEC DECISIONS ====
MODEL_NAME   = "yolo11n.pt"   # swap per rung: "yolo11n.pt" | "yolo11n.pt"
RUN_LABEL    = "tensorrt-fp16"    # swap per rung: "pytorch-cuda" | "tensorrt-fp32" | "tensorrt-fp16"
BATCH_SIZES  = [1, 4, 8]
WARMUP_RUNS  = 50
TIMED_RUNS   = 100
IMG_SIZE     = 640

def bench_one_config(model, batch_size):
    x = torch.rand(batch_size, 3, IMG_SIZE, IMG_SIZE).cuda()

    # ---- WARMUP (discard startup anomalies: lazy init, autotuning, clock ramp) ----
    for _ in range(WARMUP_RUNS):
        model(x, verbose=False)

    torch.cuda.synchronize()

    # ---- TIMED RUNS ----
    latencies = []
    for _ in range(TIMED_RUNS):
        start = time.perf_counter()
        model(x, verbose=False)
        torch.cuda.synchronize()  # GPU calls are async; without this we'd time the handoff, not the work
        latencies.append((time.perf_counter() - start) * 1000)  # ms

    print(latencies)  # raw, pre-sort: arrival order is diagnostic (clustered-early = warmup residue, scattered = contention)

    latencies.sort()
    median = latencies[len(latencies)//2]
    p95    = latencies[int(len(latencies)*0.95)]
    fps    = (batch_size * 1000) / median
    mem_mb = torch.cuda.max_memory_allocated() / 1e6  # NOTE: PyTorch allocator only — blind to TensorRT's own allocations
    return median, p95, fps, mem_mb

def main():
    model = YOLO(MODEL_NAME)
    rows = []
    for bs in BATCH_SIZES:
        torch.cuda.reset_peak_memory_stats()
        med, p95, fps, mem = bench_one_config(model, bs)
        print(f"batch={bs}: median={med:.1f}ms p95={p95:.1f}ms fps={fps:.0f} mem={mem:.0f}MB")
        rows.append([RUN_LABEL, bs, med, p95, fps, mem])

    with open("results.csv", "a", newline="") as f:
        w = csv.writer(f)
        w.writerows(rows)

if __name__ == "__main__":
    main()