File size: 7,226 Bytes
aec3f16 e0adbab aec3f16 e0adbab 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of sceneseg-p150: one eager forward and one traced replay of the whole graph between Tracy signposts.
ROOT=/home/ubuntu/experiments/tt-models
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --op-support-count 6000 \\
-o $ROOT/generated/profiler/sceneseg_baseline code/scripts/profile_ops.py
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace --end-signpost trace_end --tracing-mode
python code/scripts/profile_summary.py <ops_perf_results_*.csv> summary.json --md summary.md
Sections (signposts): ``eager`` .. ``eager_end``, one eager forward after the warm-up (``SceneSegRunner.run_eager``:
the same graph and programs as the trace, nothing compiled), and ``trace`` .. ``trace_end``, one replay of the
captured trace (the served device path; the uploaded input is the squashed sample). The device profiler buffer is
flushed (``ttnn.ReadDeviceProfiler``) before and after each section. Inside the eager section every block of the
graph opens a signpost (``m:<block>``: ``stem``, ``s<stage>.b<block>``, ``context``, ``SceneNeck.<up>``,
``SceneSegHead.<up>``, ``head_tail``, ``class_decision``) and every conv / linear / transposed conv layer opens one
before (``c:<key>``) and after (``c:<key>|post``: the glue that follows it, e.g. the SE gate multiply after ``fc2``
or the residual add after a project conv) its call; ``--no-layer-signposts`` turns them off. Through the identical
op order, the ops of the replay are attributed to the same layers. One frame issues more programs than the
profiler's default 1000-program buffer, so run with ``--op-support-count`` above the program count of the largest
section. The output folder must be absolute and outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations
import argparse
import functools
import json
import time
from pathlib import Path
from typing import Callable, List
from tt_sceneseg import SceneSeg
from tt_sceneseg.reference import assets
from tt_sceneseg.ttaw.profiling import read_device_profiler, signpost, signposted
SAMPLE = Path(__file__).resolve().parents[1] / "tt_sceneseg" / "samples" / "highway_normal_1.png"
class _Signposted:
"""A callable stand-in for a layer that emits ``c:<name>`` before and ``c:<name>|post`` after each call (and
around ``terms``, the un-summed call of a transposed conv)."""
def __init__(self, inner, name: str):
self.inner, self._name = inner, name
def __call__(self, *args, **kwargs):
signpost(f"c:{self._name}")
out = self.inner(*args, **kwargs)
signpost(f"c:{self._name}|post")
return out
def terms(self, *args, **kwargs):
signpost(f"c:{self._name}")
out = self.inner.terms(*args, **kwargs)
signpost(f"c:{self._name}|post")
return out
def __getattr__(self, name):
return getattr(self.inner, name)
def layer_signposts(m) -> Callable[[], None]:
"""Shadow the block methods and the layers of a ``tt.model.TtSceneSeg`` instance with versions that emit Tracy
signposts; returns the function that removes them again."""
names: List[str] = []
def wrap(attr: str, label: Callable) -> None:
fn = getattr(m, attr)
@functools.wraps(fn)
def inner(*args, **kwargs):
signpost(label(*args))
return fn(*args, **kwargs)
setattr(m, attr, inner)
names.append(attr)
wrap("stem", lambda *a: "m:stem")
wrap("mbconv", lambda s, b, *a: f"m:s{s}.b{b}")
wrap("se", lambda key, i, *a: f"c:{key}.{i}.se_sum")
wrap("head_conv", lambda *a: "m:head_conv")
wrap("context", lambda *a: "m:context")
wrap("up_block", lambda mod, block, *a: f"m:{mod}.{block[0]}")
wrap("head_tail", lambda *a: "m:head_tail")
wrap("channel_major", lambda *a: "m:class_decision")
saved = {}
for attr in ("convs", "linears", "deconvs"):
d = getattr(m, attr)
saved[attr] = dict(d)
for k in list(d):
d[k] = _Signposted(d[k], k)
for attr in ("stem_conv", "margin_conv", "logit_conv"):
layer = getattr(m, attr)
saved[attr] = layer
setattr(m, attr, _Signposted(layer, layer.name))
# S1 sub-pixel blocks: the Z conv(s) and the composed skip conv; the epilogue follows as "<last conv>|post"
s1_saved = []
for blk in getattr(m, "s1", {}).values():
s1_saved.append((blk, list(blk.zconvs), blk.sconv))
blk.zconvs = [_Signposted(zc, zc.name) for zc in blk.zconvs]
if blk.sconv is not None:
blk.sconv = _Signposted(blk.sconv, blk.sconv.name)
def undo() -> None:
for blk, zcs, sc in s1_saved:
blk.zconvs, blk.sconv = zcs, sc
for attr in names:
delattr(m, attr)
for attr, v in saved.items():
if isinstance(v, dict):
getattr(m, attr).clear()
getattr(m, attr).update(v)
else:
setattr(m, attr, v)
return undo
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("--input", default=str(SAMPLE))
ap.add_argument("--variant", default=None)
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
ap.add_argument("--replays", type=int, default=1, help="replays inside the trace section")
ap.add_argument("--no-eager", dest="eager", action="store_false")
ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
a = ap.parse_args()
import ttnn
weights = assets.find_weights_dir()
t0 = time.perf_counter()
with SceneSeg.from_pretrained(variant=a.variant, dispatch=a.dispatch, num_command_queues=a.cqs,
weights_dir=str(weights) if weights else None) as model:
sr, dev = model.runner, model.device
tr, var = sr.runner, sr.variant
print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info, default=str),
"knobs", json.dumps(model.knobs.as_dict()), flush=True)
canvas = model._prepare(a.input)["prepared"].device_input()
sr(canvas) # one served frame: upload + replay + readback
ttnn.synchronize_device(dev)
read_device_profiler(dev) # warm-up / capture / first frame out of the buffer
if a.eager:
undo = layer_signposts(sr.model) if a.layers else (lambda: None)
try:
with signposted("eager"):
sr.run_eager(canvas)
ttnn.synchronize_device(dev)
finally:
undo()
read_device_profiler(dev)
tr.upload({"image": sr._upload_value(canvas)})
ttnn.synchronize_device(dev)
read_device_profiler(dev)
with signposted("trace"):
tr.replay(var, n=a.replays)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
print("trace:", json.dumps(tr.describe(), default=str), flush=True)
if __name__ == "__main__":
main()
|