File size: 3,026 Bytes
25e8396 | 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""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
# Allow `pytest` from the repo root without an explicit install.
_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))
from hamer import build_reference, build_from_reference, sample_input # noqa: E402
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:
# CPU fallback — chip unavailable.
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)
|