mirror of
https://github.com/infiniteCable2/openpilot.git
synced 2026-07-25 19:32:03 +08:00
modeld: single jit (#37758)
* compile_modeld.py * update estimates * missing image=2? * Revert "missing image=2?" This reverts commit 2f5952eb63ba1e3f24cbf5769e6b5e9170d7f0a6. * Revert "update estimates" This reverts commit 1f72feef2ffdec6126e3c941e899b46ace7b4b65. * Revert "compile_modeld.py" This reverts commit f10541502efca02725f368deda2a21d1f786f57d. * load warp in ModelState init * dead code * prep * compile modeld * update SConscript * tmp save plot locally * Revert "tmp save plot locally" This reverts commit ec22f15161ad3b0241a097546b35860f989219f5. * openpilot hacks? * no float16 * tmp more chunks * Revert "tmp more chunks" This reverts commit 9e1d9b4d0dc36ff530d2a70b565fbfabd7afb00d. * Revert "no float16" This reverts commit 6204956e98e3c0818ed1985ede8eeccb810f63e3. * realize boundaries * Revert "realize boundaries" This reverts commit ffaa19259eba70944e7793e8f51a0f87089531b3. * prune=False? * Reapply "tmp more chunks" This reverts commit 2599c41cea93b4a6b4e946cdffc6a617663a7d23. * tg bug? * load first? * Revert "load first?" This reverts commit f643d082d76a424b23295e254179eb111e936e61. * revert * Reapply "tmp save plot locally" This reverts commit 1b95b82ee58654bd908b1cb04ab0ddbcd1a5955d. * 0 tol pc * warp -> modeld * rename * bypass chunking? * dont chunk * Revert "dont chunk" This reverts commit cc97fc67b3203456e123f02babe5c83b87c7e264. * dont chunk * debug * Revert "debug" This reverts commit b3c2f2e7a095fd32f8d8562a68fd1cca42357eac. * Revert "dont chunk" This reverts commit 42bd9b6f6ad0722c50348ba11ba7e2a64fdf997d. * Revert "bypass chunking?" This reverts commit ad5422a93483ffd8a59ba62e5fb72ced3b5d04d0. * corrupt model outputs * Revert "corrupt model outputs" This reverts commit 245feb94480e02f83a20b65a9488652bcbfc88b0. * image=0 for warp, match master * dedupe enqueue * pass traffic convention * tg buffer for desire * dedupe buffer creation * compile_modeld: nuke stale cached pkl before compiling The UNSAFE CI checkout keeps gitignored files (.pkl, .sconsign.dblite), so stale pkl files from previous commits can persist and be reused instead of being recompiled. Delete them explicitly before compiling. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test vs compile * all outputs need to be different on different inputs * randomize numpy inputs * randomize on every step * SConscript: nuke stale pkl+chunks before compile_modeld Move the stale artifact cleanup from compile_modeld.py into the SConscript build command. This ensures stale gitignored pkl and chunk files are deleted even if scons decides to skip the compile step (due to a stale .sconsign.dblite from UNSAFE CI checkout). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * compile_modeld: restore Context(IMAGE=0) for warp The warp operations must run under IMAGE=0 to avoid QCOM image texture optimizations that corrupt the output buffer after ~33 frames. This was accidentally commented out in a855173. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * modeld: create SubMaster before model loading Move PubMaster/SubMaster creation before the model loading step. During model loading (3.5s+), process_replay may send liveCalibration. If SubMaster doesn't exist yet, the message is dropped and the warp transform stays as zeros, producing garbage warped images. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Revert "modeld: create SubMaster before model loading" This reverts commit 968c987c2fbb3fce141c4e345d10ddea559b6c50. * stale metadata? * claude debug * Revert "claude debug" This reverts commit 49e754c6affa45a8ea8834588a00227b8090b17a. * Revert "stale metadata?" This reverts commit 870388513c0d4a67dcf970cd277b6db56cb2b478. * modeld: realize jit outputs before parsing * Update modeld.py * modeld: fix NameError by removing redundant MODELS_DIR definition * test buffers in test vs. compile * 2x inputs before running * fixup 2x inputs test * realize onnx weights? * Revert "realize onnx weights?" This reverts commit 49c8b9a505db38ff22f342db011a3a6b6526d398. * move openpilot_hacks flag to sconscript * stricter test vs compile * correct timings * more run more fail? * Revert "more run more fail?" This reverts commit 9e94bb63940751ec29e81b634c42449113e1f2e5. * numpy shenanigans * correct shapes * dont assert timings for now * Revert "correct shapes" This reverts commit 5b9ff6c84c0022327d21801d179e9e51c39e8f78. * Revert "numpy shenanigans" This reverts commit b4f6fb3078d7e9b09698895b88728fd8eea8c8a8. * no need to nuke * comment unused * don't use NPY device * copy instead of from_blob * to device before jit * Revert "to device before jit" This reverts commit 7a59ed9b1ac88657b5a3917986b6ff92e59a2ee3. * Revert "copy instead of from_blob" This reverts commit 196c4892a06ffba89ef631876372cecf137cc1b4. * Revert "don't use NPY device" This reverts commit 18abf43bbac46ad47a60c03dd8d1ef40b3f59227. * 3 runs is enough * no_memory_planner=1 * lint * restore model_replay.py * on policy -> policy * unused * prepare only enqueues full images * warp with image=2? * unused args * test vs compile, check different inputs different outputs * avoid uop cache collision * dont need realize here * misc * input queues diverged * strict zip * monkey patch for now * memory planner * prev desire correct order * dedupe pkl paths / compile targets * don't change behavior, warp and enqueue frames when skipping model eval * actually prepare only * warm up warp jit * correct path * oops * explicit warmup * need continuous + can't have dupplicate jit inputs * whitespace * bufs -> input_queues * master tg * /N_RUNS * bump tg, remove uop cache patch * more readable * Revert "bump tg, remove uop cache patch" This reverts commit 499acca2591becd389de4025943f9e776a5b337c. * missing dep --------- Co-authored-by: Bruce Wayne <harald.the.engineer@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Armand du Parc Locmaria
committed by
GitHub
GitHub
parent
c001f3c9b4
commit
d81d66193f
+20
-14
@@ -16,7 +16,6 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "
|
||||
def estimate_pickle_max_size(onnx_size):
|
||||
return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty
|
||||
|
||||
# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689
|
||||
# get fastest TG config
|
||||
available = set(Device.get_available_devices())
|
||||
if 'CUDA' in available:
|
||||
@@ -24,9 +23,10 @@ if 'CUDA' in available:
|
||||
tg_flags = f'DEV={tg_backend}'
|
||||
elif 'QCOM' in available:
|
||||
tg_backend = 'QCOM'
|
||||
tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0'
|
||||
tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
|
||||
else:
|
||||
tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM'
|
||||
# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689
|
||||
tg_flags = f'DEV={tg_backend} THREADS=0'
|
||||
|
||||
def write_tg_compiled_flags(target, source, env):
|
||||
@@ -53,14 +53,23 @@ for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
|
||||
image_flag = {
|
||||
'larch64': 'IMAGE=2',
|
||||
}.get(arch, 'IMAGE=0')
|
||||
script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
|
||||
compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
|
||||
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
||||
warp_targets = []
|
||||
for cam in [_ar_ox_fisheye, _os_fisheye]:
|
||||
w, h = cam.width, cam.height
|
||||
warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath]
|
||||
lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd)
|
||||
script_files = [File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)]
|
||||
compile_modeld_cmd = f'{tg_flags} {mac_brew_string} {image_flag} python3 {Dir("#selfdrive/modeld").abspath}/compile_modeld.py '
|
||||
driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']]
|
||||
driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']]
|
||||
|
||||
from openpilot.selfdrive.modeld.compile_modeld import MODELD_CONFIGS, DM_WARP_CONFIGS
|
||||
policy_pkls = [File(cfg.pkl_path).abspath for cfg in MODELD_CONFIGS]
|
||||
modeld_targets = policy_pkls + [File(cfg.pkl_path).abspath for cfg in DM_WARP_CONFIGS]
|
||||
compile_node = lenv.Command(modeld_targets, tinygrad_files + script_files + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], compile_modeld_cmd)
|
||||
|
||||
# chunk the combined policy pkls
|
||||
for policy_pkl in policy_pkls:
|
||||
onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
|
||||
chunk_targets = get_chunk_paths(policy_pkl, estimate_pickle_max_size(onnx_sizes_sum))
|
||||
def do_chunk(target, source, env, pkl=policy_pkl, chunks=chunk_targets):
|
||||
chunk_file(pkl, chunks)
|
||||
lenv.Command(chunk_targets, compile_node, do_chunk)
|
||||
|
||||
def tg_compile(flags, model_name):
|
||||
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
|
||||
@@ -81,7 +90,4 @@ def tg_compile(flags, model_name):
|
||||
do_chunk,
|
||||
)
|
||||
|
||||
# Compile small models
|
||||
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
|
||||
tg_compile(tg_flags, model_name)
|
||||
|
||||
tg_compile(tg_flags, 'dmonitoring_model')
|
||||
|
||||
Executable
+289
@@ -0,0 +1,289 @@
|
||||
#!/usr/bin/env python3
|
||||
import time
|
||||
import pickle
|
||||
from dataclasses import dataclass
|
||||
from itertools import product
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
|
||||
from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
|
||||
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
||||
|
||||
# https://github.com/tinygrad/tinygrad/issues/15682
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
_orig = UOp.__reduce__
|
||||
UOp.__reduce__ = lambda self: (UOp.unique, ()) if self.op is Ops.UNIQUE else _orig(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompileConfig:
|
||||
cam_w: int
|
||||
cam_h: int
|
||||
prepare_only: bool
|
||||
prefix: str
|
||||
|
||||
@property
|
||||
def pkl_path(self):
|
||||
return str(MODELS_DIR / f'{self.prefix}{"warp_" if self.prepare_only else ""}{self.cam_w}x{self.cam_h}_tinygrad.pkl')
|
||||
|
||||
CAMERA_CONFIGS = [
|
||||
(_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208
|
||||
(_os_fisheye.width, _os_fisheye.height), # mici: 1344x760
|
||||
]
|
||||
MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_') for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])]
|
||||
DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS]
|
||||
|
||||
|
||||
UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
|
||||
UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
|
||||
|
||||
|
||||
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad):
|
||||
w_dst, h_dst = dst_shape
|
||||
h_src, w_src = src_shape
|
||||
|
||||
x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
|
||||
y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
|
||||
|
||||
# inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
|
||||
src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
|
||||
src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2]
|
||||
src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2]
|
||||
|
||||
src_x = src_x / src_w
|
||||
src_y = src_y / src_w
|
||||
|
||||
x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int')
|
||||
y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int')
|
||||
idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
|
||||
|
||||
return src_flat[idx]
|
||||
|
||||
|
||||
def frames_to_tensor(frames):
|
||||
H = (frames.shape[0] * 2) // 3
|
||||
W = frames.shape[1]
|
||||
in_img1 = Tensor.cat(frames[0:H:2, 0::2],
|
||||
frames[1:H:2, 0::2],
|
||||
frames[0:H:2, 1::2],
|
||||
frames[1:H:2, 1::2],
|
||||
frames[H:H+H//4].reshape((H//2, W//2)),
|
||||
frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2))
|
||||
return in_img1
|
||||
|
||||
|
||||
def make_frame_prepare(cam_w, cam_h, model_w, model_h):
|
||||
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
||||
uv_offset = stride * y_height
|
||||
stride_pad = stride - cam_w
|
||||
|
||||
def frame_prepare_tinygrad(input_frame, M_inv):
|
||||
# UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
|
||||
M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]])
|
||||
# deinterleave NV12 UV plane (UVUV... -> separate U, V)
|
||||
uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
y = warp_perspective_tinygrad(input_frame[:cam_h*stride],
|
||||
M_inv, (model_w, model_h),
|
||||
(cam_h, cam_w), stride_pad).realize()
|
||||
u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w))
|
||||
tensor = frames_to_tensor(yuv)
|
||||
return tensor
|
||||
return frame_prepare_tinygrad
|
||||
|
||||
|
||||
def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip):
|
||||
img = vision_input_shapes['img'] # (1, 12, 128, 256)
|
||||
n_frames = img[1] // 6
|
||||
img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
|
||||
|
||||
fb = policy_input_shapes['features_buffer'] # (1, 25, 512)
|
||||
dp = policy_input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
tc = policy_input_shapes['traffic_convention'] # (1, 2)
|
||||
|
||||
npy = {
|
||||
'desire': np.zeros(dp[2], dtype=np.float32),
|
||||
'traffic_convention': np.zeros(tc, dtype=np.float32),
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
input_queues = {
|
||||
'img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(),
|
||||
'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(),
|
||||
'feat_q': Tensor.zeros(frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]).contiguous().realize(),
|
||||
'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2]).contiguous().realize(),
|
||||
**{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
|
||||
}
|
||||
return input_queues, npy
|
||||
|
||||
|
||||
def shift_and_sample(buf, new_val, sample_fn):
|
||||
buf.assign(buf[1:].cat(new_val, dim=0).contiguous())
|
||||
return sample_fn(buf)
|
||||
|
||||
|
||||
def sample_skip(buf, frame_skip):
|
||||
return buf[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def sample_desire(buf, frame_skip):
|
||||
return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def make_warp_dm(cam_w, cam_h, dm_w, dm_h):
|
||||
stride, y_height, _, _ = get_nv12_info(cam_w, cam_h)
|
||||
stride_pad = stride - cam_w
|
||||
|
||||
def warp_dm(input_frame, M_inv):
|
||||
M_inv = M_inv.to(Device.DEFAULT)
|
||||
result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w)
|
||||
return result
|
||||
return warp_dm
|
||||
|
||||
|
||||
def make_run_policy(vision_runner, policy_runner, cam_w, cam_h,
|
||||
vision_features_slice, frame_skip, prepare_only=False):
|
||||
model_w, model_h = MEDMODEL_INPUT_SIZE
|
||||
frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
||||
|
||||
def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame):
|
||||
img = shift_and_sample(img_q, frame_prepare(frame, tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn)
|
||||
|
||||
if prepare_only:
|
||||
return img, big_img
|
||||
|
||||
vision_out = next(iter(vision_runner({'img': img, 'big_img': big_img}).values())).cast('float32')
|
||||
|
||||
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
|
||||
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
|
||||
desire_buf = shift_and_sample(desire_q, desire.to(Device.DEFAULT).reshape(1, 1, -1), sample_desire_fn)
|
||||
|
||||
inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention.to(Device.DEFAULT)}
|
||||
policy_out = next(iter(policy_runner(inputs).values())).cast('float32')
|
||||
|
||||
return vision_out, policy_out
|
||||
return run_policy
|
||||
|
||||
|
||||
def compile_modeld(cam_w, cam_h, prepare_only, pkl_path):
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from openpilot.selfdrive.modeld.constants import ModelConstants
|
||||
|
||||
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
||||
print(f"Compiling combined policy JIT for {cam_w}x{cam_h}...")
|
||||
|
||||
vision_runner = OnnxRunner(MODELS_DIR / 'driving_vision.onnx')
|
||||
policy_runner = OnnxRunner(MODELS_DIR / 'driving_policy.onnx')
|
||||
|
||||
with open(MODELS_DIR / 'driving_vision_metadata.pkl', 'rb') as f:
|
||||
vision_metadata = pickle.load(f)
|
||||
vision_features_slice = vision_metadata['output_slices']['hidden_state']
|
||||
vision_input_shapes = vision_metadata['input_shapes']
|
||||
with open(MODELS_DIR / 'driving_policy_metadata.pkl', 'rb') as f:
|
||||
policy_input_shapes = pickle.load(f)['input_shapes']
|
||||
|
||||
frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
|
||||
_run = make_run_policy(vision_runner, policy_runner,
|
||||
cam_w, cam_h, vision_features_slice, frame_skip, prepare_only)
|
||||
run_policy_jit = TinyJit(_run, prune=True)
|
||||
|
||||
N_RUNS = 3
|
||||
SEED = 42
|
||||
|
||||
def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True):
|
||||
input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip)
|
||||
np.random.seed(seed)
|
||||
|
||||
for i in range(N_RUNS):
|
||||
frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize()
|
||||
big_frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize()
|
||||
for v in npy.values():
|
||||
v[:] = np.random.randn(*v.shape).astype(v.dtype)
|
||||
Device.default.synchronize()
|
||||
st = time.perf_counter()
|
||||
outs = fn(**input_queues, frame=frame, big_frame=big_frame)
|
||||
mt = time.perf_counter()
|
||||
for o in outs:
|
||||
# .realize() not needed once jitted, but needed for unjitted fn
|
||||
o.realize()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/{N_RUNS}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
val = [np.copy(v.numpy()) for v in outs]
|
||||
buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()]
|
||||
|
||||
if test_val is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True))
|
||||
assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
if test_buffers is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True))
|
||||
assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
return fn, val, buffers
|
||||
|
||||
print('run unjitted')
|
||||
_, test_val, test_buffers = random_inputs_run_fn(_run, seed=SEED)
|
||||
print('capture + replay')
|
||||
run_policy_jit, _, _ = random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers)
|
||||
|
||||
print('pickle round trip')
|
||||
with open(pkl_path, "wb") as f:
|
||||
pickle.dump(run_policy_jit, f)
|
||||
print(f" Saved to {pkl_path}")
|
||||
with open(pkl_path, "rb") as f:
|
||||
run_policy_jit = pickle.load(f)
|
||||
random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True)
|
||||
random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False)
|
||||
|
||||
|
||||
def compile_dm_warp(cam_w, cam_h, pkl_path):
|
||||
dm_w, dm_h = DM_INPUT_SIZE
|
||||
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
||||
|
||||
print(f"Compiling DM warp for {cam_w}x{cam_h}...")
|
||||
|
||||
warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h)
|
||||
warp_dm_jit = TinyJit(warp_dm, prune=True)
|
||||
|
||||
for i in range(10):
|
||||
inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
Device.default.synchronize()
|
||||
st = time.perf_counter()
|
||||
warp_dm_jit(*inputs)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
with open(pkl_path, "wb") as f:
|
||||
pickle.dump(warp_dm_jit, f)
|
||||
print(f" Saved to {pkl_path}")
|
||||
|
||||
|
||||
def run_and_save_pickle():
|
||||
for cfg in MODELD_CONFIGS:
|
||||
compile_modeld(cfg.cam_w, cfg.cam_h, cfg.prepare_only, cfg.pkl_path)
|
||||
for cfg in DM_WARP_CONFIGS:
|
||||
compile_dm_warp(cfg.cam_w, cfg.cam_h, cfg.pkl_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_and_save_pickle()
|
||||
@@ -1,198 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
import time
|
||||
import pickle
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
|
||||
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
||||
|
||||
MODELS_DIR = Path(__file__).parent / 'models'
|
||||
|
||||
CAMERA_CONFIGS = [
|
||||
(_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208
|
||||
(_os_fisheye.width, _os_fisheye.height), # mici: 1344x760
|
||||
]
|
||||
|
||||
UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
|
||||
UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
|
||||
|
||||
IMG_BUFFER_SHAPE = (30, MEDMODEL_INPUT_SIZE[1] // 2, MEDMODEL_INPUT_SIZE[0] // 2)
|
||||
|
||||
|
||||
def warp_pkl_path(w, h):
|
||||
return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
|
||||
|
||||
|
||||
def dm_warp_pkl_path(w, h):
|
||||
return MODELS_DIR / f'dm_warp_{w}x{h}_tinygrad.pkl'
|
||||
|
||||
|
||||
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad):
|
||||
w_dst, h_dst = dst_shape
|
||||
h_src, w_src = src_shape
|
||||
|
||||
x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
|
||||
y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
|
||||
|
||||
# inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
|
||||
src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
|
||||
src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2]
|
||||
src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2]
|
||||
|
||||
src_x = src_x / src_w
|
||||
src_y = src_y / src_w
|
||||
|
||||
x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int')
|
||||
y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int')
|
||||
idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
|
||||
|
||||
return src_flat[idx]
|
||||
|
||||
|
||||
def frames_to_tensor(frames, model_w, model_h):
|
||||
H = (frames.shape[0] * 2) // 3
|
||||
W = frames.shape[1]
|
||||
in_img1 = Tensor.cat(frames[0:H:2, 0::2],
|
||||
frames[1:H:2, 0::2],
|
||||
frames[0:H:2, 1::2],
|
||||
frames[1:H:2, 1::2],
|
||||
frames[H:H+H//4].reshape((H//2, W//2)),
|
||||
frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2))
|
||||
return in_img1
|
||||
|
||||
|
||||
def make_frame_prepare(cam_w, cam_h, model_w, model_h):
|
||||
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
||||
uv_offset = stride * y_height
|
||||
stride_pad = stride - cam_w
|
||||
|
||||
def frame_prepare_tinygrad(input_frame, M_inv):
|
||||
# UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
|
||||
M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]])
|
||||
# deinterleave NV12 UV plane (UVUV... -> separate U, V)
|
||||
uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
y = warp_perspective_tinygrad(input_frame[:cam_h*stride],
|
||||
M_inv, (model_w, model_h),
|
||||
(cam_h, cam_w), stride_pad).realize()
|
||||
u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w))
|
||||
tensor = frames_to_tensor(yuv, model_w, model_h)
|
||||
return tensor
|
||||
return frame_prepare_tinygrad
|
||||
|
||||
|
||||
def make_update_img_input(frame_prepare, model_w, model_h):
|
||||
def update_img_input_tinygrad(frame_buffer, frame, M_inv):
|
||||
M_inv = M_inv.to(Device.DEFAULT)
|
||||
new_img = frame_prepare(frame, M_inv)
|
||||
frame_buffer.assign(frame_buffer[6:].cat(new_img, dim=0).contiguous())
|
||||
return Tensor.cat(frame_buffer[:6], frame_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2)
|
||||
return update_img_input_tinygrad
|
||||
|
||||
|
||||
def make_update_both_imgs(frame_prepare, model_w, model_h):
|
||||
update_img = make_update_img_input(frame_prepare, model_w, model_h)
|
||||
|
||||
def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
|
||||
calib_big_img_buffer, new_big_img, M_inv_big):
|
||||
calib_img_pair = update_img(calib_img_buffer, new_img, M_inv)
|
||||
calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big)
|
||||
return calib_img_pair, calib_big_img_pair
|
||||
return update_both_imgs_tinygrad
|
||||
|
||||
|
||||
def make_warp_dm(cam_w, cam_h, dm_w, dm_h):
|
||||
stride, y_height, _, _ = get_nv12_info(cam_w, cam_h)
|
||||
stride_pad = stride - cam_w
|
||||
|
||||
def warp_dm(input_frame, M_inv):
|
||||
M_inv = M_inv.to(Device.DEFAULT)
|
||||
result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w)
|
||||
return result
|
||||
return warp_dm
|
||||
|
||||
|
||||
def compile_modeld_warp(cam_w, cam_h):
|
||||
model_w, model_h = MEDMODEL_INPUT_SIZE
|
||||
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
||||
|
||||
print(f"Compiling modeld warp for {cam_w}x{cam_h}...")
|
||||
|
||||
frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h)
|
||||
update_both_imgs = make_update_both_imgs(frame_prepare, model_w, model_h)
|
||||
update_img_jit = TinyJit(update_both_imgs, prune=True)
|
||||
|
||||
full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize()
|
||||
big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize()
|
||||
for i in range(10):
|
||||
img_inputs = [full_buffer,
|
||||
Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
big_img_inputs = [big_full_buffer,
|
||||
Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
inputs = img_inputs + big_img_inputs
|
||||
Device.default.synchronize()
|
||||
|
||||
st = time.perf_counter()
|
||||
_ = update_img_jit(*inputs)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
pkl_path = warp_pkl_path(cam_w, cam_h)
|
||||
with open(pkl_path, "wb") as f:
|
||||
pickle.dump(update_img_jit, f)
|
||||
print(f" Saved to {pkl_path}")
|
||||
|
||||
jit = pickle.load(open(pkl_path, "rb"))
|
||||
jit(*inputs)
|
||||
|
||||
|
||||
def compile_dm_warp(cam_w, cam_h):
|
||||
dm_w, dm_h = DM_INPUT_SIZE
|
||||
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
||||
|
||||
print(f"Compiling DM warp for {cam_w}x{cam_h}...")
|
||||
|
||||
warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h)
|
||||
warp_dm_jit = TinyJit(warp_dm, prune=True)
|
||||
|
||||
for i in range(10):
|
||||
inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(),
|
||||
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
||||
Device.default.synchronize()
|
||||
st = time.perf_counter()
|
||||
warp_dm_jit(*inputs)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
pkl_path = dm_warp_pkl_path(cam_w, cam_h)
|
||||
with open(pkl_path, "wb") as f:
|
||||
pickle.dump(warp_dm_jit, f)
|
||||
print(f" Saved to {pkl_path}")
|
||||
|
||||
|
||||
def run_and_save_pickle():
|
||||
for cam_w, cam_h in CAMERA_CONFIGS:
|
||||
compile_modeld_warp(cam_w, cam_h)
|
||||
compile_dm_warp(cam_w, cam_h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_and_save_pickle()
|
||||
+33
-117
@@ -26,6 +26,7 @@ from openpilot.common.transformations.model import get_warp_matrix
|
||||
from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
|
||||
from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan
|
||||
from openpilot.selfdrive.modeld.parse_model_outputs import Parser
|
||||
from openpilot.selfdrive.modeld.compile_modeld import CompileConfig, make_input_queues
|
||||
from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
|
||||
from openpilot.common.file_chunker import read_file_chunked
|
||||
from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
|
||||
@@ -34,17 +35,13 @@ from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
|
||||
PROCESS_NAME = "selfdrive.modeld.modeld"
|
||||
SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
|
||||
|
||||
VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl'
|
||||
VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl'
|
||||
POLICY_PKL_PATH = MODELS_DIR / 'driving_policy_tinygrad.pkl'
|
||||
POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl'
|
||||
|
||||
LAT_SMOOTH_SECONDS = 0.0
|
||||
LONG_SMOOTH_SECONDS = 0.3
|
||||
MIN_LAT_CONTROL_SPEED = 0.3
|
||||
|
||||
IMG_QUEUE_SHAPE = (6*(ModelConstants.MODEL_RUN_FREQ//ModelConstants.MODEL_CONTEXT_FREQ + 1), 128, 256)
|
||||
assert IMG_QUEUE_SHAPE[0] == 30
|
||||
|
||||
|
||||
def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
|
||||
@@ -79,106 +76,36 @@ class FrameMeta:
|
||||
if vipc is not None:
|
||||
self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
|
||||
|
||||
class InputQueues:
|
||||
def __init__ (self, model_fps, env_fps, n_frames_input):
|
||||
assert env_fps % model_fps == 0
|
||||
assert env_fps >= model_fps
|
||||
self.model_fps = model_fps
|
||||
self.env_fps = env_fps
|
||||
self.n_frames_input = n_frames_input
|
||||
|
||||
self.dtypes = {}
|
||||
self.shapes = {}
|
||||
self.q = {}
|
||||
|
||||
def update_dtypes_and_shapes(self, input_dtypes, input_shapes) -> None:
|
||||
self.dtypes.update(input_dtypes)
|
||||
if self.env_fps == self.model_fps:
|
||||
self.shapes.update(input_shapes)
|
||||
else:
|
||||
for k in input_shapes:
|
||||
shape = list(input_shapes[k])
|
||||
if 'img' in k:
|
||||
n_channels = shape[1] // self.n_frames_input
|
||||
shape[1] = (self.env_fps // self.model_fps + (self.n_frames_input - 1)) * n_channels
|
||||
else:
|
||||
shape[1] = (self.env_fps // self.model_fps) * shape[1]
|
||||
self.shapes[k] = tuple(shape)
|
||||
|
||||
def reset(self) -> None:
|
||||
self.q = {k: np.zeros(self.shapes[k], dtype=self.dtypes[k]) for k in self.dtypes.keys()}
|
||||
|
||||
def enqueue(self, inputs:dict[str, np.ndarray]) -> None:
|
||||
for k in inputs.keys():
|
||||
if inputs[k].dtype != self.dtypes[k]:
|
||||
raise ValueError(f'supplied input <{k}({inputs[k].dtype})> has wrong dtype, expected {self.dtypes[k]}')
|
||||
input_shape = list(self.shapes[k])
|
||||
input_shape[1] = -1
|
||||
single_input = inputs[k].reshape(tuple(input_shape))
|
||||
sz = single_input.shape[1]
|
||||
self.q[k][:,:-sz] = self.q[k][:,sz:]
|
||||
self.q[k][:,-sz:] = single_input
|
||||
|
||||
def get(self, *names) -> dict[str, np.ndarray]:
|
||||
if self.env_fps == self.model_fps:
|
||||
return {k: self.q[k] for k in names}
|
||||
else:
|
||||
out = {}
|
||||
for k in names:
|
||||
shape = self.shapes[k]
|
||||
if 'img' in k:
|
||||
n_channels = shape[1] // (self.env_fps // self.model_fps + (self.n_frames_input - 1))
|
||||
out[k] = np.concatenate([self.q[k][:, s:s+n_channels] for s in np.linspace(0, shape[1] - n_channels, self.n_frames_input, dtype=int)], axis=1)
|
||||
elif 'pulse' in k:
|
||||
# any pulse within interval counts
|
||||
out[k] = self.q[k].reshape((shape[0], shape[1] * self.model_fps // self.env_fps, self.env_fps // self.model_fps, -1)).max(axis=2)
|
||||
else:
|
||||
idxs = np.arange(-1, -shape[1], -self.env_fps // self.model_fps)[::-1]
|
||||
out[k] = self.q[k][:, idxs]
|
||||
return out
|
||||
|
||||
class ModelState:
|
||||
inputs: dict[str, np.ndarray]
|
||||
output: np.ndarray
|
||||
prev_desire: np.ndarray # for tracking the rising edge of the pulse
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, cam_w: int, cam_h: int):
|
||||
with open(VISION_METADATA_PATH, 'rb') as f:
|
||||
vision_metadata = pickle.load(f)
|
||||
self.vision_input_shapes = vision_metadata['input_shapes']
|
||||
self.vision_input_names = list(self.vision_input_shapes.keys())
|
||||
self.vision_output_slices = vision_metadata['output_slices']
|
||||
vision_output_size = vision_metadata['output_shapes']['outputs'][1]
|
||||
|
||||
with open(POLICY_METADATA_PATH, 'rb') as f:
|
||||
policy_metadata = pickle.load(f)
|
||||
self.policy_input_shapes = policy_metadata['input_shapes']
|
||||
self.policy_output_slices = policy_metadata['output_slices']
|
||||
policy_output_size = policy_metadata['output_shapes']['outputs'][1]
|
||||
|
||||
self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
|
||||
|
||||
# policy inputs
|
||||
self.numpy_inputs = {k: np.zeros(self.policy_input_shapes[k], dtype=np.float32) for k in self.policy_input_shapes}
|
||||
self.full_input_queues = InputQueues(ModelConstants.MODEL_CONTEXT_FREQ, ModelConstants.MODEL_RUN_FREQ, ModelConstants.N_FRAMES)
|
||||
for k in ['desire_pulse', 'features_buffer']:
|
||||
self.full_input_queues.update_dtypes_and_shapes({k: self.numpy_inputs[k].dtype}, {k: self.numpy_inputs[k].shape})
|
||||
self.full_input_queues.reset()
|
||||
|
||||
self.img_queues = {'img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize(),
|
||||
'big_img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize()}
|
||||
self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip)
|
||||
self.full_frames : dict[str, Tensor] = {}
|
||||
self._blob_cache : dict[int, Tensor] = {}
|
||||
self.transforms_np = {k: np.zeros((3,3), dtype=np.float32) for k in self.img_queues}
|
||||
self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()}
|
||||
self.vision_output = np.zeros(vision_output_size, dtype=np.float32)
|
||||
self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
|
||||
self.policy_output = np.zeros(policy_output_size, dtype=np.float32)
|
||||
self.parser = Parser()
|
||||
self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {}
|
||||
self.update_imgs = None
|
||||
self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH)))
|
||||
self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH)))
|
||||
self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
|
||||
self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path))
|
||||
self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path))
|
||||
self.warp_enqueue(
|
||||
**self.input_queues,
|
||||
frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(),
|
||||
big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize())
|
||||
|
||||
def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
|
||||
parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
|
||||
@@ -186,18 +113,6 @@ class ModelState:
|
||||
|
||||
def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
|
||||
inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None:
|
||||
# Model decides when action is completed, so desire input is just a pulse triggered on rising edge
|
||||
inputs['desire_pulse'][0] = 0
|
||||
new_desire = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0)
|
||||
self.prev_desire[:] = inputs['desire_pulse']
|
||||
if self.update_imgs is None:
|
||||
for key in bufs.keys():
|
||||
w, h = bufs[key].width, bufs[key].height
|
||||
self.frame_buf_params[key] = get_nv12_info(w, h)
|
||||
warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
|
||||
with open(warp_path, "rb") as f:
|
||||
self.update_imgs = pickle.load(f)
|
||||
|
||||
for key in bufs.keys():
|
||||
ptr = bufs[key].data.ctypes.data
|
||||
yuv_size = self.frame_buf_params[key][3]
|
||||
@@ -206,30 +121,31 @@ class ModelState:
|
||||
if cache_key not in self._blob_cache:
|
||||
self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8')
|
||||
self.full_frames[key] = self._blob_cache[cache_key]
|
||||
for key in bufs.keys():
|
||||
self.transforms_np[key][:,:] = transforms[key][:,:]
|
||||
|
||||
out = self.update_imgs(self.img_queues['img'], self.full_frames['img'], self.transforms['img'],
|
||||
self.img_queues['big_img'], self.full_frames['big_img'], self.transforms['big_img'])
|
||||
vision_inputs = {'img': out[0], 'big_img': out[1]}
|
||||
# Model decides when action is completed, so desire input is just a pulse triggered on rising edge
|
||||
inputs['desire_pulse'][0] = 0
|
||||
self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0)
|
||||
self.prev_desire[:] = inputs['desire_pulse']
|
||||
self.npy['traffic_convention'][:] = inputs['traffic_convention']
|
||||
self.npy['tfm'][:,:] = transforms['img'][:,:]
|
||||
self.npy['big_tfm'][:,:] = transforms['big_img'][:,:]
|
||||
|
||||
if prepare_only:
|
||||
self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'])
|
||||
return None
|
||||
|
||||
self.vision_output = self.vision_run(**vision_inputs).contiguous().realize().uop.base.buffer.numpy().flatten()
|
||||
vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(self.vision_output, self.vision_output_slices))
|
||||
vision_output, policy_output = self.run_policy(
|
||||
**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']
|
||||
)
|
||||
|
||||
self.full_input_queues.enqueue({'features_buffer': vision_outputs_dict['hidden_state'], 'desire_pulse': new_desire})
|
||||
for k in ['desire_pulse', 'features_buffer']:
|
||||
self.numpy_inputs[k][:] = self.full_input_queues.get(k)[k]
|
||||
self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention']
|
||||
|
||||
self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten()
|
||||
policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices))
|
||||
vision_output = vision_output.numpy().flatten()
|
||||
policy_output = policy_output.numpy().flatten()
|
||||
vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices))
|
||||
policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices))
|
||||
combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict}
|
||||
if SEND_RAW_PRED:
|
||||
combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()])
|
||||
|
||||
if SEND_RAW_PRED:
|
||||
combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()])
|
||||
return combined_outputs_dict
|
||||
|
||||
|
||||
@@ -241,11 +157,6 @@ def main(demo=False):
|
||||
# also need to move the aux USB interrupts for good timings
|
||||
config_realtime_process(7, 54)
|
||||
|
||||
st = time.monotonic()
|
||||
cloudlog.warning("loading model")
|
||||
model = ModelState()
|
||||
cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
|
||||
|
||||
# visionipc clients
|
||||
while True:
|
||||
available_streams = VisionIpcClient.available_streams("camerad", block=False)
|
||||
@@ -269,6 +180,11 @@ def main(demo=False):
|
||||
if use_extra_client:
|
||||
cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})")
|
||||
|
||||
st = time.monotonic()
|
||||
cloudlog.warning("loading model")
|
||||
model = ModelState(vipc_client_main.width, vipc_client_main.height)
|
||||
cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
|
||||
|
||||
# messaging
|
||||
pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"])
|
||||
sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"])
|
||||
|
||||
@@ -2,7 +2,7 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
MODELS_DIR = Path(__file__).parent / 'models'
|
||||
MODELS_DIR = Path(__file__).resolve().parent / 'models'
|
||||
COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json'
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user