162 lines
6.8 KiB
Python
162 lines
6.8 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
os.environ["DEV"] = "CPU"
|
|
|
|
from iqpilot.selfdrive.iqmodeld.egpu_policy import PolicyRunner, make_run_policy, packed_layout, queue_shapes
|
|
from iqpilot.selfdrive.iqmodeld.temporal_state import TemporalInputState
|
|
|
|
SPEC = {
|
|
"img": ((1, 12, 8, 16), "uint8"),
|
|
"big_img": ((1, 12, 8, 16), "uint8"),
|
|
"desire_pulse": ((1, 25, 8), "float32"),
|
|
"traffic_convention": ((1, 2), "float32"),
|
|
"action_t": ((1, 2), "float32"),
|
|
"features_buffer": ((1, 24, 512), "float32"),
|
|
}
|
|
FS = 4
|
|
OUT_LEN = 2580
|
|
HIDDEN = slice(1064, 1576)
|
|
|
|
|
|
def _pack(inputs):
|
|
from tinygrad.tensor import Tensor
|
|
parts = [inputs[k].cast("float32").reshape(-1) for k in ("img", "big_img", "features_buffer", "desire_pulse", "traffic_convention", "action_t")]
|
|
flat = Tensor.cat(*parts)
|
|
hidden = (flat[:512] * 0.001).reshape(1, 512)
|
|
return flat, hidden
|
|
|
|
|
|
def _fake_model(inputs):
|
|
from tinygrad.tensor import Tensor
|
|
flat, hidden = _pack(inputs)
|
|
n = flat.shape[0]
|
|
head = flat[:min(n, HIDDEN.start)]
|
|
out = Tensor.cat(head.pad((0, HIDDEN.start - head.shape[0])), hidden.reshape(-1), Tensor.zeros(OUT_LEN - HIDDEN.stop, device="CPU"))
|
|
return {"outputs": out.reshape(1, -1)}
|
|
|
|
|
|
class _Reference:
|
|
def __init__(self):
|
|
self.state = TemporalInputState(FS, SPEC)
|
|
|
|
def run(self, warped, desire, traffic, action_t):
|
|
inputs = self.state.push_and_materialize(warped, desire, traffic, action_t)
|
|
from tinygrad.tensor import Tensor
|
|
t = {k: Tensor(np.ascontiguousarray(v), device="CPU") for k, v in inputs.items()}
|
|
out = _fake_model(t)["outputs"].numpy().reshape(-1)
|
|
self.state.note_hidden_state(out, HIDDEN)
|
|
return out
|
|
|
|
|
|
def test_policy_queues_match_temporal_state():
|
|
from tinygrad.engine.jit import TinyJit
|
|
jit = TinyJit(make_run_policy(_fake_model, SPEC, FS, "CPU"), prune=True)
|
|
runner = PolicyRunner(jit, SPEC, FS, HIDDEN, "CPU")
|
|
ref = _Reference()
|
|
rng = np.random.default_rng(3)
|
|
desire = np.zeros(8, dtype=np.float32)
|
|
for i in range(14):
|
|
warped = rng.integers(0, 256, (2, 6, 8, 16), dtype=np.int64).astype(np.uint8)
|
|
if i in (2, 3, 9):
|
|
desire[:] = 0
|
|
desire[1 + (i % 3)] = 1
|
|
elif i == 5:
|
|
desire[:] = 0
|
|
traffic = np.array([1.0, 0.0], dtype=np.float32) if i % 2 else np.array([0.0, 1.0], dtype=np.float32)
|
|
action_t = np.array([0.1 * i, 0.2], dtype=np.float32)
|
|
got = runner.run(warped, desire, traffic, action_t)
|
|
want = ref.run(warped, desire, traffic, action_t)
|
|
np.testing.assert_array_equal(got, want, err_msg=f"frame {i}")
|
|
|
|
|
|
def test_layouts():
|
|
shapes, sizes = packed_layout(SPEC)
|
|
assert list(shapes) == ["desire", "traffic_convention", "action_t", "prev_feat"]
|
|
assert sum(sizes) == 8 + 2 + 2 + 512
|
|
q = queue_shapes(SPEC, FS)
|
|
assert q["img_q"][0] == (5, 6, 8, 16) and q["feat_q"][0] == (96, 1, 512) and q["desire_q"][0] == (100, 1, 8)
|
|
|
|
|
|
CAM = (64, 48)
|
|
|
|
|
|
def _nv12(cam_w, cam_h):
|
|
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
|
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
|
return (cam_w, cam_h, stride, y_height, uv_height)
|
|
|
|
|
|
def _numpy_warp_plane(src, m, w_dst, h_dst):
|
|
h_src, w_src = src.shape
|
|
x = np.tile(np.arange(w_dst, dtype=np.float32), h_dst)
|
|
y = np.repeat(np.arange(h_dst, dtype=np.float32), w_dst)
|
|
sx = (m[0, 0] * x + m[0, 1] * y + m[0, 2]) / (m[2, 0] * x + m[2, 1] * y + m[2, 2])
|
|
sy = (m[1, 0] * x + m[1, 1] * y + m[1, 2]) / (m[2, 0] * x + m[2, 1] * y + m[2, 2])
|
|
xi = np.clip(np.round(sx), 0, w_src - 1).astype(np.int64)
|
|
yi = np.clip(np.round(sy), 0, h_src - 1).astype(np.int64)
|
|
return src[yi, xi].reshape(h_dst, w_dst)
|
|
|
|
|
|
def _numpy_frame_prepare(frame, m, nv12, model_w, model_h):
|
|
cam_w, cam_h, stride, y_height, uv_height = nv12
|
|
m = m.astype(np.float32)
|
|
y_src = frame[:cam_h * stride].reshape(cam_h, stride)
|
|
uv = frame[stride * y_height:stride * y_height + uv_height * stride].reshape(uv_height, stride)
|
|
m_uv = m * np.array([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], dtype=np.float32)
|
|
y = _numpy_warp_plane(y_src, m, model_w, model_h)
|
|
u = _numpy_warp_plane(uv[:cam_h // 2, :cam_w:2], m_uv, model_w // 2, model_h // 2)
|
|
v = _numpy_warp_plane(uv[:cam_h // 2, 1:cam_w:2], m_uv, model_w // 2, model_h // 2)
|
|
f = np.concatenate([y.ravel(), u.ravel(), v.ravel()]).reshape(model_h * 3 // 2, model_w)
|
|
H, W = model_h, model_w
|
|
return np.stack([f[0:H:2, 0::2], f[1:H:2, 0::2], f[0:H:2, 1::2], f[1:H:2, 1::2],
|
|
f[H:H + H // 4].reshape(H // 2, W // 2), f[H + H // 4:H + H // 2].reshape(H // 2, W // 2)])
|
|
|
|
|
|
def _jittered_scale(rng, cam, model_w, model_h):
|
|
m = np.array([[cam[0] / model_w, 0.0, 0.0], [0.0, cam[1] / model_h, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
|
m += (0.05 * rng.standard_normal((3, 3))).astype(np.float32) * np.array([[1, 1, 1], [1, 1, 1], [0.01, 0.01, 0.1]], dtype=np.float32)
|
|
return m
|
|
|
|
|
|
def test_frame_layout():
|
|
from iqpilot.selfdrive.iqmodeld.egpu_policy import frame_layout, model_size, nv12_copy_size
|
|
shapes, sizes, npy_bytes = frame_layout(SPEC)
|
|
assert list(shapes) == ["tfm", "big_tfm", "desire", "traffic_convention", "action_t", "prev_feat"]
|
|
assert npy_bytes == (18 + 8 + 2 + 2 + 512) * 4
|
|
assert model_size(SPEC) == (32, 16)
|
|
assert nv12_copy_size(128, 64, 32) == 128 * 96
|
|
|
|
|
|
def test_model_runner_matches_device_warp():
|
|
from tinygrad.engine.jit import TinyJit
|
|
|
|
from iqpilot.selfdrive.iqmodeld.egpu_policy import ModelRunner, make_run_model, make_warp, model_size, nv12_copy_size
|
|
nv12 = _nv12(*CAM)
|
|
fcs = nv12_copy_size(nv12[2], nv12[3], nv12[4])
|
|
model_w, model_h = model_size(SPEC)
|
|
run_policy = make_run_policy(_fake_model, SPEC, FS, "CPU")
|
|
jit = TinyJit(make_run_model(make_warp(nv12, model_w, model_h, "CPU"), run_policy, SPEC, fcs, "CPU"), prune=True)
|
|
runner = ModelRunner(jit, SPEC, FS, HIDDEN, "CPU", fcs)
|
|
ref = PolicyRunner(TinyJit(make_run_policy(_fake_model, SPEC, FS, "CPU"), prune=True), SPEC, FS, HIDDEN, "CPU")
|
|
rng = np.random.default_rng(7)
|
|
desire = np.zeros(8, dtype=np.float32)
|
|
for i in range(10):
|
|
main = rng.integers(0, 256, fcs, dtype=np.int64).astype(np.uint8)
|
|
extra = rng.integers(0, 256, fcs, dtype=np.int64).astype(np.uint8)
|
|
tfm = _jittered_scale(rng, CAM, model_w, model_h)
|
|
big_tfm = _jittered_scale(rng, CAM, model_w, model_h)
|
|
if i in (2, 6):
|
|
desire[:] = 0
|
|
desire[1 + i % 3] = 1
|
|
traffic = np.array([1.0, 0.0], dtype=np.float32) if i % 2 else np.array([0.0, 1.0], dtype=np.float32)
|
|
action_t = np.array([0.1 * i, 0.2], dtype=np.float32)
|
|
got = runner.run(main, extra, tfm, big_tfm, desire, traffic, action_t)
|
|
warped = np.stack([_numpy_frame_prepare(main, tfm, nv12, model_w, model_h), _numpy_frame_prepare(extra, big_tfm, nv12, model_w, model_h)])
|
|
want = ref.run(warped, desire, traffic, action_t)
|
|
np.testing.assert_array_equal(got, want, err_msg=f"frame {i}")
|