| |
| |
| """ST2 per-layer probe: every stacked [W_hi | W_lo] conv / transposed conv vs its two separate term convs, inside one |
| eager forward of the real graph on the shipped sample (EPI on), bit for bit per layer. |
| |
| bin/devrun -t 900 -- python -u code/scripts/diag/st2_probe.py |
| |
| Each wrapped layer runs both forms on the same input; the stacked result continues the forward. Prints one line per |
| layer call (identical / number of differing elements / max abs diff) and the conv config of the layer. |
| """ |
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def main() -> int: |
| import ttnn |
|
|
| from tt_sceneseg.device import close_device, open_device |
| from tt_sceneseg.host.preprocess import prepare_image |
| from tt_sceneseg.reference import assets, config |
| from tt_sceneseg.reference.weights import load_weights |
| from tt_sceneseg.tt import layers as L |
| from tt_sceneseg.tt.model import INPUT_SHAPE, TtSceneSeg |
| from tt_sceneseg.ttaw.tensors import to_numpy |
|
|
| code = Path(__file__).resolve().parents[2] |
| wdir = assets.test_weights_dir() |
| onnx = assets.find_in(wdir, config.ONNX_FILE) if wdir is not None else None |
| if onnx is None: |
| from huggingface_hub import snapshot_download |
| from tt_sceneseg.api import SceneSeg |
|
|
| d = snapshot_download(SceneSeg.DEFAULT_REPO, revision=SceneSeg.DEFAULT_REVISION, |
| allow_patterns=SceneSeg.ALLOW_PATTERNS) |
| onnx = assets.find_in(Path(d), config.ONNX_FILE) |
| dev = open_device(l1_small_size=98304) |
| model = TtSceneSeg(load_weights(onnx), "bgr", epi=True, term_stack=True) |
| stats = [] |
|
|
| def cmp(name, a, b): |
| x, y = to_numpy(a), to_numpy(b) |
| same = x.shape == y.shape and np.array_equal(x.view(np.uint8), y.view(np.uint8)) |
| nd = 0 if same else int((x != y).sum()) |
| md = 0.0 if same else float(np.abs(x.astype(np.float64) - y.astype(np.float64)).max()) |
| stats.append((name, same)) |
| print(f"{'SAME' if same else 'DIFF'} {name} {tuple(x.shape)} diff {nd} max {md:.3g}", flush=True) |
|
|
| orig_conv = L.TermConv2d.__call__ |
|
|
| def conv_call(self, x, *, act_terms=1, output_dtype=None): |
| if self.stacked is None: |
| return orig_conv(self, x, act_terms=act_terms, output_dtype=output_dtype) |
| xin = x.with_tensor(ttnn.clone(x.tensor)) |
| y1 = orig_conv(self, x, act_terms=act_terms, output_dtype=output_dtype) |
| st, self.stacked = self.stacked, None |
| y2 = orig_conv(self, xin, act_terms=act_terms, output_dtype=output_dtype) |
| self.stacked = st |
| cmp(f"conv {self.name} {x.height}x{x.width} {self.in_channels}->{self.out_channels} at={act_terms}", |
| y1.tensor, y2.tensor) |
| ttnn.deallocate(y2.tensor) |
| ttnn.deallocate(xin.tensor) |
| return y1 |
|
|
| orig_views = L.TermConvTranspose2d.term_views |
|
|
| def views(self, x): |
| if self.stacked is None: |
| return orig_views(self, x) |
| parts, cols, first = orig_views(self, x) |
| st, self.stacked = self.stacked, None |
| p2, _, _ = orig_views(self, x) |
| self.stacked = st |
| c = self.out_channels |
| shp = [int(v) for v in parts[0].shape] |
| for i, off in enumerate((0, c)): |
| h = ttnn.slice(parts[0], [0, 0, 0, off], [shp[0], shp[1], shp[2], off + c]) |
| cmp(f"deconv {self.name} term {i} {x.height}x{x.width} {self.in_channels}->{c}", h, p2[i]) |
| ttnn.deallocate(h) |
| for t in p2: |
| ttnn.deallocate(t) |
| return parts, cols, first |
|
|
| L.TermConv2d.__call__ = conv_call |
| L.TermConvTranspose2d.term_views = views |
| canvas = prepare_image(str(code / "tt_sceneseg" / "samples" / "highway_normal_1.png")).device_input() |
| import torch |
|
|
| image = ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(canvas.reshape(INPUT_SHAPE))), dtype=ttnn.uint8, |
| layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| out = model.forward(image) |
| print("class map", np.bincount(to_numpy(out["class_map"]).astype(np.int64).ravel(), minlength=3).tolist()) |
| n_diff = sum(not s for _, s in stats) |
| print(f"{len(stats)} layer calls, {n_diff} differ") |
| model.release() |
| close_device(dev) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
