File size: 4,311 Bytes
d160c10 | 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 | # SPDX-License-Identifier: Apache-2.0
"""Bring-up check of ``tt.l1res`` (OPT round 4): the real encoder weights, the round-3 encoder vs the L1-resident one,
per bottleneck of layers 3-4, every intermediate (conv1, conv2, conv3, identity, output) compared bit for bit.
Both encoders get the same layer-2 output (the round-3 encoder's).
bin/devrun -t 900 -- python -u code/scripts/debug_l1res.py
"""
from __future__ import annotations
import sys
import numpy as np
def main():
import torch
import ttnn
from tt_bevdet.device import close_device, open_device
from tt_bevdet.reference.config import default_config
from tt_bevdet.reference.weights import BEVDetWeights
from tt_bevdet.tests import paths
from tt_bevdet.tt.conv_tune import apply_tune
from tt_bevdet.tt.image_encoder import TtImageEncoder
from tt_bevdet.tt.l1res import L1ResBlock, block_mc, ensure
from tt_bevdet.tt.params import DEFAULT_PRECISION, BEVDetParams
from tt_bevdet.ttaw.ops.conv import is_dram
def bits(t):
if t.memory_config().is_sharded():
t = ttnn.to_memory_config(t, ttnn.DRAM_MEMORY_CONFIG)
return ttnn.to_torch(ttnn.from_device(t)).to(torch.bfloat16).view(torch.int16).numpy().reshape(
-1, int(t.shape[-1]))
def cmp(tag, a, b):
x, y = bits(a), bits(b)
print(f" {tag}: {int(np.count_nonzero(x != y))} differ; l1 mc "
f"{'dram' if is_dram(b) else b.memory_config().shard_spec.shape if b.memory_config().is_sharded() else 'l1'}",
flush=True)
cfg = default_config()
params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg)
dev = open_device()
try:
kw = dict(precision=DEFAULT_PRECISION, stem_in="kin", conv_l1="all", backbone_res="next12", mm1x1=True)
ref = TtImageEncoder(params, **kw)
l1 = TtImageEncoder(params, l1res=True, **kw)
apply_tune(ref.convs())
apply_tune(l1.convs())
n, (h, w) = cfg.n_cams, cfg.input_hw
rng = np.random.default_rng(7)
crop = rng.integers(0, 256, (n, h, w, 3), dtype=np.uint8)
din = ttnn.from_torch(torch.from_numpy(crop.reshape(1, 1, n * h, w * 3)), dtype=ttnn.uint8,
layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
x = ref.kin_stem(din)
for blk in ref.stages.stages[0] + ref.stages.stages[1]:
x = blk(x)
xr = xl = x
gy = int(dev.compute_with_storage_grid_size().y)
for s in (2, 3):
for rb, lb in zip(ref.stages.stages[s], l1.stages.stages[s]):
b = rb
lbb: L1ResBlock = lb
print(f"{b.name}: in ref {'dram' if is_dram(xr.tensor) else 'l1'} / l1 "
f"{'dram' if is_dram(xl.tensor) else 'l1'}", flush=True)
y_r = b.conv1(xr)
if is_dram(xl.tensor):
y_l = lbb.blk.conv1(xl)
else:
y_l = lbb._mm(ensure(xl, block_mc(xl.rows, xl.channels, 8, gy))[0], lbb.w1,
b.conv3.in_channels, True, gy)
cmp("conv1", y_r.tensor, y_l.tensor)
z_r = b.conv2(y_r)
getattr(lbb.blk.conv2, "conv").slicing = "l1" if is_dram(y_l.tensor) else "auto"
z_l = lbb.blk.conv2(y_l)
cmp("conv2", z_r.tensor, z_l.tensor)
zs, mv = ensure(z_l, block_mc(z_l.rows, z_l.channels, 8, gy))
print(" conv2 moved:", mv, flush=True)
y3_r = b.conv3(z_r)
y3_l = lbb._mm(zs, lbb.w3, b.conv3.out_channels, False, gy)
cmp("conv3", y3_r.tensor, y3_l.tensor)
if b.downsample is not None:
d_r = b.downsample(xr)
xd = xl if is_dram(xl.tensor) else xl.with_tensor(
ttnn.to_memory_config(xl.tensor, ttnn.DRAM_MEMORY_CONFIG)) # as L1ResBlock does
d_l = lbb.blk.downsample(xd)
cmp("downsample", d_r.tensor, d_l.tensor)
xr = b(xr)
xl = lb(xl)
cmp("out", xr.tensor, xl.tensor)
finally:
close_device(dev)
return 0
if __name__ == "__main__":
sys.exit(main())
|