File size: 6,418 Bytes
fe51a63 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device diagnostic: per-module error of the device graph on one catalogued frame, next to the CPU emulation.
cd bundles/frnet-p150 && bin/devrun -t 1200 -- env PYTHONPATH=code python code/scripts/diag_frame_taps.py \\
--frame nuscenes_scene-0916_40.ot128 --out /home/ubuntu/experiments/tt-models/logs/frnet/m1b/diag_0916_40.json
The exact fp32 TT dataflow (``reference/rewrites.py FRNetPlanTorch``) is the reference of every tap; the CPU
emulation (``--emulation``, default the ``plan_exact_points`` row of ``code/scripts/precision_study.py``: bf16 conv
operands, gathers and segment tables, exact point layers) and the device ``taps`` variant (one replay; knobs from the
``FRNET_*`` environment) are compared with it tap by tap: where the device's ``1 - PCC`` grows faster than the
emulation's is where the device loses accuracy that the emulation does not model. Also reported: ``pred_probs`` PCC
and raw-point argmax agreement vs ONNX Runtime for the device and the emulation. The CPU work runs first (in the
device lock: ~1-2 min at OT128).
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))
sys.path.insert(0, str(HERE.parent))
TAPS = ("encoder.pf3", "encoder.voxel_feats", "backbone.stem", "backbone.point_stem", "backbone.fusion_stem",
*(f"backbone.{k}{i}" for i in range(1, 5) for k in ("layer", "point_fusion", "pixel_fusion", "attention",
"stage")),
"decoder.fuse1", "decoder.fuse2", "decoder.point_fuse1", "decoder.point_fuse2",
*(f"head.mlp{i}" for i in range(4)), "head.logits", "pred_probs")
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--frame", required=True)
ap.add_argument("--out", type=Path, required=True)
ap.add_argument("--emulation", default="plan_exact_points")
ap.add_argument("--threads", type=int, default=4)
ap.add_argument("--cpu-only", action="store_true", help="stop after the CPU part (a dry run without a device)")
args = ap.parse_args()
import torch
torch.set_num_threads(args.threads)
from precision_study import CONFIGS
from tt_frnet.host.preprocess import build_plan
from tt_frnet.reference.config import FRNetConfig
from tt_frnet.reference.rewrites import Emulation, FRNetPlanTorch
from tt_frnet.reference.weights import load_weights
from tt_frnet.tests import _data
from tt_frnet.ttaw.golden import TapRegistry
from tt_frnet.ttaw.metrics import argmax_agreement, pcc
variant = args.frame.rpartition(".")[2]
wdir = _data.weights_dir()
cfg = FRNetConfig.from_model_dir(wdir, variant)
w = load_weights(wdir, variant)
inputs, ort_probs, _ = _data.load_frame(args.frame)
plan = build_plan(inputs, cfg, pre_norm=(w.pre_norm_scale, w.pre_norm_shift))
n, n_raw = plan.n, inputs.num_points_raw
H, W = plan.frustum_hw
t0 = time.time()
exact = TapRegistry(include=list(TAPS))
FRNetPlanTorch(w)(plan, taps=exact)
emu = TapRegistry(include=list(TAPS))
emu_probs = plan.unsort(FRNetPlanTorch(w, emulation=Emulation.parse(CONFIGS[args.emulation]))(plan, taps=emu))
cpu_s = time.time() - t0
def rows(name, a):
"""point taps: the first N sorted rows; voxel taps: the first M rows."""
a = np.asarray(a)
if a.shape[0] == plan.n_cap:
return a[:n]
return a[: plan.m] if a.shape[0] in (cfg.m_cap, cfg.m_cap + 1) else a
if args.cpu_only:
ort = np.asarray(ort_probs, np.float32)
print(f"cpu {cpu_s:.0f} s; emulation pred_probs pcc {pcc(emu_probs, ort):.6f} argmax "
f"{argmax_agreement(emu_probs[:n_raw], ort[:n_raw]):.5f}")
return 0
from tt_frnet.device import device_session
from tt_frnet.tt.model import TtFRNet, build_runner, frame_arrays
report = {"frame": args.frame, "N": n, "M": plan.m, "emulation": args.emulation, "cpu_s": round(cpu_s, 1),
"taps": {}}
with device_session() as device:
graph = TtFRNet(device, w, cfg)
arrays = frame_arrays(plan, cfg.m_cap)
runner = build_runner(graph, variants=("taps",), name="frnet_diag", init=arrays)
try:
runner.capture()
dev = runner("taps", inputs=arrays)
report["knobs"] = graph.knobs.as_dict()
for name in TAPS:
if name not in dev or name not in exact:
continue
ref = np.asarray(exact[name], np.float32)
d = np.asarray(dev[name], np.float32)
if ref.ndim == 4:
_, c, h, ww = ref.shape
d = d.reshape(-1, d.shape[-1])[: h * ww, :c].reshape(1, h, ww, c).transpose(0, 3, 1, 2)
e = np.asarray(emu[name], np.float32)
else:
ref = rows(name, ref)[:, : ref.shape[-1]]
d = rows(name, d.reshape(-1, d.shape[-1]))[:, : ref.shape[-1]]
e = rows(name, np.asarray(emu[name], np.float32))[:, : ref.shape[-1]]
p_dev, p_emu = float(pcc(d, ref)), float(pcc(e, ref))
ratio = (1 - p_dev) / max(1 - p_emu, 1e-12)
report["taps"][name] = {"pcc_dev": p_dev, "pcc_emu": p_emu, "err_ratio": ratio}
print(f"{name:26s} dev {p_dev:.6f} emu {p_emu:.6f} (1-pcc) ratio {ratio:6.2f}", flush=True)
probs = plan.unsort(np.asarray(dev["pred_probs"], np.float32).reshape(-1, 32)[:, :27])
finally:
runner.release()
graph.release()
ort = np.asarray(ort_probs, np.float32)
report["device"] = {"pred_probs_pcc": float(pcc(probs, ort)),
"argmax_agreement_raw": float(argmax_agreement(probs[:n_raw], ort[:n_raw]))}
report["emulation_result"] = {"pred_probs_pcc": float(pcc(emu_probs, ort)),
"argmax_agreement_raw": float(argmax_agreement(emu_probs[:n_raw], ort[:n_raw]))}
print("device", report["device"], "emulation", report["emulation_result"])
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(json.dumps(report, indent=1))
return 0
if __name__ == "__main__":
sys.exit(main())
|