| |
| |
| """CPU emulation of the TT graph's numerics (no device): the real ``tt_sceneseg.tt`` graph executed by the fake ttnn |
| of ``common/tests/host`` (bf16-rounded operands and outputs, float64 products and sums), one and two weight terms, on |
| a few frames, against the fp32 CPU reference goldens: class-map agreement, foreground IoU, logits PCC. The evidence |
| behind the weight-terms default (PORT_LOG.md M1), next to the device runs. |
| |
| source bin/tt-env.sh; cd bundles/sceneseg-p150 |
| PYTHONPATH=$PWD/code:$PYTHONPATH TT_VISIBLE_DEVICES=none OMP_NUM_THREADS=4 python code/scripts/diag/ss_precision_emulation.py \ |
| --frames tut_highway_rain_2,ps019_f40_front_camera --terms 1,2 --json ../../logs/sceneseg/m1_precision_emulation.json |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| BUNDLE = Path(__file__).resolve().parents[3] |
| WS = BUNDLE.parents[1] |
| sys.path.insert(0, str(BUNDLE / "code")) |
|
|
|
|
| def install_fake(): |
| """The fake ttnn (+ its CNN ops and a row sum) as ``ttnn``; never the real one.""" |
| host = WS / "common" / "tests" / "host" |
| spec = importlib.util.spec_from_file_location("fake_ttnn", host / "fake_ttnn.py") |
| fake = importlib.util.module_from_spec(spec) |
| sys.modules["fake_ttnn"] = fake |
| spec.loader.exec_module(fake) |
| ttnn = fake.install() |
| spec = importlib.util.spec_from_file_location("fake_ttnn_cnn", host / "fake_ttnn_cnn.py") |
| cnn = importlib.util.module_from_spec(spec) |
| sys.modules["fake_ttnn_cnn"] = cnn |
| spec.loader.exec_module(cnn) |
| cnn.install(ttnn) |
|
|
| def sum_(t, dim=None, keepdim=False, **_): |
| return ttnn._op("sum", [t], lambda x: x.sum(axis=dim, keepdims=keepdim), t.dtype, t.layout, (dim, keepdim)) |
|
|
| ttnn.sum = sum_ |
| return ttnn |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| ap.add_argument("--frames", default="tut_highway_rain_2,ps019_f40_front_camera") |
| ap.add_argument("--terms", default="1,2") |
| ap.add_argument("--order", default="bgr") |
| ap.add_argument("--json", default="") |
| a = ap.parse_args() |
| ttnn = install_fake() |
| if not hasattr(ttnn, "subtract"): |
| ttnn.subtract = lambda x, y, **_: ttnn._op("subtract", [x, y], lambda u, v: u - v, x.dtype, x.layout) |
| import torch |
|
|
| torch.set_num_threads(4) |
| from tt_sceneseg.host.preprocess import prepare_bgr |
| from tt_sceneseg.reference import assets |
| from tt_sceneseg.reference.weights import load_weights |
| from tt_sceneseg.tests import sceneseg_eval as ev |
| from tt_sceneseg.tt.model import TtSceneSeg |
| from tt_sceneseg.ttaw.device import open_device |
| from tt_sceneseg.ttaw.metrics import pcc |
|
|
| dev = open_device(0, dispatch="eth", num_command_queues=1) |
| weights = load_weights(assets.onnx_path()) |
| gold = WS / "research" / "sceneseg" / "goldens" / "public" |
| inputs = WS / "research" / "sceneseg" / "public_data" / "inputs" |
| results = {} |
| for terms in (int(t) for t in a.terms.split(",")): |
| model = TtSceneSeg(weights, a.order, weight_terms=terms) |
| for fid in a.frames.split(","): |
| t0 = time.perf_counter() |
| bgr = np.ascontiguousarray(ev.load_rgb(inputs / fid / "source_bgr.png")[:, :, ::-1]) |
| canvas = prepare_bgr(bgr, channel_order=a.order).device_input() |
| image = ttnn.from_torch(torch.from_numpy(canvas), dtype=ttnn.uint8, layout=ttnn.ROW_MAJOR_LAYOUT, |
| device=dev) |
| out = model.forward(image, logits=True) |
| cmap = ttnn.to_torch(out["class_map"]).numpy().reshape(320, 640) |
| logits = ttnn.to_torch(out["logits"]).float().numpy() |
| with np.load(gold / f"{fid}.npz") as z: |
| ref_logits, ref_cls = z[f"logits_{a.order}"], z[f"cls_{a.order}"] |
| rec = {"logits_pcc": pcc(logits, ref_logits), "agreement": ev.argmax_agreement(cmap, ref_cls), |
| "foreground_iou": ev.foreground_iou(cmap, ref_cls), |
| "fg_fraction": float((cmap == 1).mean()), "fg_fraction_fp32": float((ref_cls == 1).mean()), |
| "seconds": round(time.perf_counter() - t0, 1)} |
| results[f"terms{terms}/{fid}"] = rec |
| print(f"terms={terms} {fid}: {json.dumps(rec)}", flush=True) |
| model.release() |
| if a.json: |
| Path(a.json).write_text(json.dumps({"order": a.order, "results": results}, indent=1) + "\n") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
