File size: 5,407 Bytes
aec3f16 | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager forward of the served graph (default numerics, ETH dispatch,
the device grid). Never use it for timing: graph capture and detailed buffer tracing synchronise after every op.
ROOT=/home/ubuntu/experiments/tt-models; cd $ROOT/bundles/sceneseg-p150
$ROOT/bin/devrun -t 2400 -- env PYTHONPATH=$PWD/code:$PYTHONPATH HF_HUB_OFFLINE=1 \
python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/sceneseg_baseline
ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/sceneseg_baseline \
--performance-path $ROOT/generated/profiler/sceneseg_baseline/reports/<ts>
What the pytest fixture of ttnn does (ttnn/graph_report.py ``_run_pytest_graph_report_fixture``), in a script:
warm-up forward (JIT, conv weight preparation), then the second forward inside ``ttnn.graph.full_graph_capture``
(Python I/O, stack traces, detailed buffer tracing), ``ttnn.graph_report.import_report`` -> ``db.sqlite``, and
``config.json`` (the settings the capture ran with). The intermediate graph JSON is large and is deleted after the
import; ``--no-stack-traces`` leaves the Python stack traces out of the capture (a smaller JSON).
"""
from __future__ import annotations
import json
import sys
import time
from pathlib import Path
def capture(json_path: Path, stack_traces: bool):
"""``ttnn.graph.full_graph_capture`` (Python I/O, stack traces, detailed buffer tracing, slow dispatch), or the
same without Python stack traces (``--no-stack-traces``: a much smaller graph JSON on a full disk)."""
import contextlib
import ttnn
if stack_traces:
return ttnn.graph.full_graph_capture(str(json_path))
@contextlib.contextmanager
def light():
g = ttnn.graph
prev_buf, prev_fast = g.is_detailed_buffer_tracing_enabled(), ttnn.CONFIG.enable_fast_runtime_mode
g.enable_detailed_buffer_tracing()
ttnn.CONFIG.enable_fast_runtime_mode = False
try:
g.begin_graph_capture(g.RunMode.NORMAL)
yield
g.end_graph_capture_to_file(str(json_path))
finally:
if g.is_graph_capture_active():
try:
g._cpp_end_graph_capture()
except Exception:
pass
if not prev_buf:
g.disable_detailed_buffer_tracing()
ttnn.CONFIG.enable_fast_runtime_mode = prev_fast
return light()
def main() -> None:
stack_traces = "--no-stack-traces" not in sys.argv
out_dir = Path([a for a in sys.argv[1:] if not a.startswith("--")][0]).resolve()
out_dir.mkdir(parents=True, exist_ok=True)
import torch
import ttnn
from tt_sceneseg.device import close_device, describe_device, open_device
from tt_sceneseg.host.preprocess import prepare_image
from tt_sceneseg.reference import config
from tt_sceneseg.reference.weights import load_weights
from tt_sceneseg.reference.assets import find_in, find_weights_dir
from tt_sceneseg.tt.model import INPUT_SHAPE, TtSceneSeg
sample = Path(__file__).resolve().parents[1] / "tt_sceneseg" / "samples" / "highway_normal_1.png"
canvas = prepare_image(str(sample)).device_input()
weights = load_weights(find_in(find_weights_dir(True), config.ONNX_FILE))
dev = open_device(allow_fallback=False)
try:
print("device:", json.dumps(describe_device(dev, 0)), flush=True)
model = TtSceneSeg(weights, config.DEFAULT_CHANNEL_ORDER)
x = ttnn.from_torch(torch.from_numpy(canvas.reshape(INPUT_SHAPE)), dtype=ttnn.uint8,
layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
def forward():
out = model.forward(x)
ttnn.synchronize_device(dev)
for t in [out["class_map"]]:
if t.is_allocated():
ttnn.deallocate(t)
t0 = time.perf_counter()
forward() # warm-up: JIT + weight preparation (host writes)
print(f"warm-up forward {time.perf_counter() - t0:.1f} s", flush=True)
json_path = out_dir / "graph_capture.json"
t0 = time.perf_counter()
with capture(json_path, stack_traces):
forward()
print(f"captured forward {time.perf_counter() - t0:.1f} s -> {json_path} "
f"({json_path.stat().st_size / 1e6:.1f} MB)", flush=True)
ttnn.deallocate(x)
model.release()
finally:
close_device(dev)
t0 = time.perf_counter()
db = ttnn.graph_report.import_report(json_path, out_dir)
ttnn.save_config_to_json_file(out_dir / "config.json")
cfg = json.loads((out_dir / "config.json").read_text()) # record what the capture ran with
cfg.update({"enable_fast_runtime_mode": False, "enable_graph_report": True, "enable_graph_python_stack_traces": stack_traces,
"enable_detailed_buffer_report": True, "root_report_path": str(out_dir.parent),
"report_name": out_dir.name})
(out_dir / "config.json").write_text(json.dumps(cfg, indent=4) + "\n")
print(f"imported in {time.perf_counter() - t0:.1f} s -> {db} ({Path(db).stat().st_size / 1e6:.1f} MB)", flush=True)
json_path.unlink(missing_ok=True)
if __name__ == "__main__":
main()
|