File size: 1,214 Bytes
58ba6df | 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 | # SPDX-License-Identifier: Apache-2.0
"""Device-profiler driver: 2 eager forwards (JIT + warm, profiled) then trace capture and
3 replays, real checkpoint, sample1 crop. Run under `python -m tracy -r -p -v`."""
import os, sys
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
os.environ.setdefault("TT_FUSED_HASH", "0")
import torch, ttnn
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
from hamer.preprocess import preprocess
hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
ref = Hamer().eval(); load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
x, _ = preprocess(_REPO.parent / "media/sample1.jpg", bbox_xyxy=(20, 170, 295, 260), is_right=True)
tt = build_from_reference(ref, device_id=0, strict=True)
with torch.no_grad():
tt(x) # eager, JIT
ttnn.ReadDeviceProfiler(tt.device)
tt(x) # eager warm -> trace captured on this call? (warmup_done) -- see model
ttnn.ReadDeviceProfiler(tt.device)
for _ in range(3):
tt(x)
ttnn.ReadDeviceProfiler(tt.device)
tt.close()
|